Repository: Lightning-AI/pytorch-lightning
Pull request: Update deepspeed model-parallel docs
## What does this PR do?

Updates docs for `configure_sharded_model` in DeepSpeed model-parallel docs. DeepSpeed recently changed how parameters get moved to the device in a breaking change (https://github.com/microsoft/DeepSpeed/pull/3611). From now on, all layers need to be defined in this hook.

Also fixes #17806

<!-- Does your PR introduce any breaking changes? If yes, please list them. -->

<details>
  <summary><b>Before submitting</b></summary>

- Was this **discussed/agreed** via a GitHub issue? (not for typos and docs)
- [ ] Did you read the [contributor guideline](https://github.com/Lightning-AI/lightning/blob/master/.github/CONTRIBUTING.md), **Pull Request** section?
- [ ] Did you make sure your **PR does only one thing**, instead of bundling different changes together?
- Did you make sure to **update the documentation** with your changes? (if necessary)
- Did you write any **new necessary tests**? (not for typos and docs)
- [ ] Did you verify new and **existing tests pass** locally with your changes?
- Did you list all the **breaking changes** introduced by this pull request?
- Did you **update the CHANGELOG**? (not for typos, docs, test updates, or minor internal changes/refactors)

<!-- In the CHANGELOG, separate each item in the unreleased section by a blank line to reduce collisions -->

</details>

## PR review

Anyone in the community is welcome to review the PR.
Before you start reviewing, make sure you have read the [review guidelines](https://github.com/Lightning-AI/lightning/wiki/Review-guidelines). In short, see the following bullet-list:

<details>
  <summary>Reviewer checklist</summary>

- [ ] Is this pull request ready for review? (if not, please submit in draft mode)
- [ ] Check that all items from **Before submitting** are resolved
- [ ] Make sure the title is self-explanatory and the description concisely explains the PR
- [ ] Add labels and milestones (and optionally projects) to the PR so it can be classified

</details>

<!--

Did you have fun?

Make sure you had fun coding 🙃

-->


cc @borda @awaelchli

Problem statement:
DeepSpeed doesn't move tensors to GPU in deepspeed 0.9.3 and above
### Bug description

Lightning version (pytorch-lightning 2.0.2)
deepspeed version (0.9.4)

Trained chatglm, but got weight and input on different devices exception. Weight is on cpu, while input is on gpu. It works if deepspeed 0.9.2 is used. I noticed a difference in DeepSpeedEngine._configure_distributed_model() between 0.9.2 and 0.9.3. 

## deepspeed 0.9.2
        if not self.dont_change_device:
            self.module.to(self.device)

## deepspeed 0.9.3 and above
        # zero.Init() handles device placement of model
        if not (self.dont_change_device or is_zero3_model):
            self.module.to(self.device)

In use stage 3, model won't move to GPU, which caused the following exception.  
```
----------------------------------------
Traceback (most recent call last):
 File "/home/zhanqing/.conda/envs/llm_5/lib/python3.9/runpy.py", line 197, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/home/zhanqing/.conda/envs/llm_5/lib/python3.9/runpy.py", line 87, in _run_code
    exec(code, run_globals)
  File "/home/zhanqing/.vscode-server/extensions/ms-python.python-2023.8.0/pythonFiles/lib/python/debugpy/adapter/../../debugpy/launcher/../../debugpy/__main__.py", line 39, in <module>
    cli.main()
  File "/home/zhanqing/.vscode-server/extensions/ms-python.python-2023.8.0/pythonFiles/lib/python/debugpy/adapter/../../debugpy/launcher/../../debugpy/../debugpy/server/cli.py", line 430, in main
    run()
  File "/home/zhanqing/.vscode-server/extensions/ms-python.python-2023.8.0/pythonFiles/lib/python/debugpy/adapter/../../debugpy/launcher/../../debugpy/../debugpy/server/cli.py", line 284, in run_file
    runpy.run_path(target, run_name="__main__")
  File "/home/zhanqing/.vscode-server/extensions/ms-python.python-2023.8.0/pythonFiles/lib/python/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_runpy.py", line 321, in run_path
    return _run_module_code(code, init_globals, run_name,
  File "/home/zhanqing/.vscode-server/extensions/ms-python.python-2023.8.0/pythonFiles/lib/python/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_runpy.py", line 135, in _run_module_code
    _run_code(code, mod_globals, init_globals,
  File "/home/zhanqing/.vscode-server/extensions/ms-python.python-2023.8.0/pythonFiles/lib/python/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_runpy.py", line 124, in _run_code
    exec(code, run_globals)
  File "/nfs_data/data/personal/zhanqing/projects/nlp_basedon_lightning/llm/src/training/trainer.py", line 75, in <module>
    main(lightning_cli)
  File "/nfs_data/data/personal/zhanqing/projects/nlp_basedon_lightning/llm/src/training/trainer.py", line 60, in main
    trainer.train(lightning_cli=lightning_cli)
  File "/nfs_data/data/personal/zhanqing/projects/nlp_basedon_lightning/llm/src/training/trainer.py", line 53, in train
    lightning_cli.trainer.fit(lightning_cli.model, lightning_cli.datamodule)  # 完成后self.model会进入cpu
  File "/home/zhanqing/.conda/envs/llm_5/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 520, in fit
    call._call_and_handle_interrupt(
  File "/home/zhanqing/.conda/envs/llm_5/lib/python3.9/site-packages/pytorch_lightning/trainer/call.py", line 42, in _call_and_handle_interrupt
    return trainer.strategy.launcher.launch(trainer_fn, *args, trainer=trainer, **kwargs)
  File "/home/zhanqing/.conda/envs/llm_5/lib/python3.9/site-packages/pytorch_lightning/strategies/launchers/subprocess_script.py", line 92, in launch
    return function(*args, **kwargs)
  File "/home/zhanqing/.conda/envs/llm_5/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 559, in _fit_impl
    self._run(model, ckpt_path=ckpt_path)
  File "/home/zhanqing/.conda/envs/llm_5/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 935, in _run
    results = self._run_stage()
  File "/home/zhanqing/.conda/envs/llm_5/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 976, in _run_stage
    self._run_sanity_check()
  File "/home/zhanqing/.conda/envs/llm_5/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 1005, in _run_sanity_check
    val_loop.run()
  File "/home/zhanqing/.conda/envs/llm_5/lib/python3.9/site-packages/pytorch_lightning/loops/utilities.py", line 177, in _decorator
    return loop_run(self, *args, **kwargs)
  File "/home/zhanqing/.conda/envs/llm_5/lib/python3.9/site-packages/pytorch_lightning/loops/evaluation_loop.py", line 115, in run
    self._evaluation_step(batch, batch_idx, dataloader_idx)
  File "/home/zhanqing/.conda/envs/llm_5/lib/python3.9/site-packages/pytorch_lightning/loops/evaluation_loop.py", line 375, in _evaluation_step
    output = call._call_strategy_hook(trainer, hook_name, *step_kwargs.values())
  File "/home/zhanqing/.conda/envs/llm_5/lib/python3.9/site-packages/pytorch_lightning/trainer/call.py", line 288, in _call_strategy_hook
    output = fn(*args, **kwargs)
  File "/home/zhanqing/.conda/envs/llm_5/lib/python3.9/site-packages/pytorch_lightning/strategies/deepspeed.py", line 906, in validation_step
    return self.model(*args, **kwargs)
  File "/home/zhanqing/.conda/envs/llm_5/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/zhanqing/.conda/envs/llm_5/lib/python3.9/site-packages/deepspeed/utils/nvtx.py", line 15, in wrapped_fn
    ret_val = func(*args, **kwargs)
  File "/home/zhanqing/.conda/envs/llm_5/lib/python3.9/site-packages/deepspeed/runtime/engine.py", line 1736, in forward
    loss = self.module(*inputs, **kwargs)
  File "/home/zhanqing/.conda/envs/llm_5/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1212, in _call_impl
    result = forward_call(*input, **kwargs)
  File "/home/zhanqing/.conda/envs/llm_5/lib/python3.9/site-packages/pytorch_lightning/overrides/base.py", line 102, in forward
    return self._forward_module.validation_step(*inputs, **kwargs)
  File "/home/zhanqing/.conda/envs/llm_5/lib/python3.9/site-packages/lightning_base/lightning_module/abstract_lightning_module.py", line 54, in validation_step
    loss = self._step_internal(batch)
  File "/nfs_data/data/personal/zhanqing/projects/nlp_basedon_lightning/llm/src/training/lightning_module.py", line 136, in _step_internal
    outputs = self.model(
  File "/home/zhanqing/.conda/envs/llm_5/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1212, in _call_impl
    result = forward_call(*input, **kwargs)
  File "/home/zhanqing/.conda/envs/llm_5/lib/python3.9/site-packages/peft/peft_model.py", line 678, in forward
    return self.base_model(
  File "/home/zhanqing/.conda/envs/llm_5/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1212, in _call_impl
    result = forward_call(*input, **kwargs)
  File "/home/zhanqing/.cache/huggingface/modules/transformers_modules/chatglm-6b/modeling_chatglm.py", line 1190, in forward
    transformer_outputs = self.transformer(
  File "/home/zhanqing/.conda/envs/llm_5/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1212, in _call_impl
    result = forward_call(*input, **kwargs)
  File "/home/zhanqing/.cache/huggingface/modules/transformers_modules/chatglm-6b/modeling_chatglm.py", line 996, in forward
    layer_ret = layer(
  File "/home/zhanqing/.conda/envs/llm_5/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1212, in _call_impl
    result = forward_call(*input, **kwargs)
  File "/home/zhanqing/.cache/huggingface/modules/transformers_modules/chatglm-6b/modeling_chatglm.py", line 627, in forward
    attention_outputs = self.attention(
  File "/home/zhanqing/.conda/envs/llm_5/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1212, in _call_impl
    result = forward_call(*input, **kwargs)
  File "/home/zhanqing/.cache/huggingface/modules/transformers_modules/chatglm-6b/modeling_chatglm.py", line 460, in forward
    cos, sin = self.rotary_emb(q1, seq_len=position_ids.max() + 1)
  File "/home/zhanqing/.conda/envs/llm_5/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1212, in _call_impl
    result = forward_call(*input, **kwargs)
  File "/home/zhanqing/.cache/huggingface/modules/transformers_modules/chatglm-6b/modeling_chatglm.py", line 203, in forward
    freqs = torch.einsum('i,j->ij', t, self.inv_freq)
  File "/home/zhanqing/.conda/envs/llm_5/lib/python3.9/site-packages/torch/functional.py", line 378, in einsum
    return _VF.einsum(equation, operands)  # type: ignore[attr-defined]
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!
```
### What version are you seeing the problem on?

v2.0

### How to reproduce the bug

_No response_

### Error messages and logs

```
# Error messages and logs here please
```


### Environment

<details>
  <summary>Current environment</summary>

```
#- Lightning Component (e.g. Trainer, LightningModule, LightningApp, LightningWork, LightningFlow):
#- PyTorch Lightning Version (e.g., 1.5.0):
#- Lightning App Version (e.g., 0.5.2):
#- PyTorch Version (e.g., 2.0):
#- Python version (e.g., 3.9):
#- OS (e.g., Linux):
#- CUDA/cuDNN version:
#- GPU models and configuration:
#- How you installed Lightning(`conda`, `pip`, source):
#- Running environment of LightningApp (e.g. local, cloud):
```

</details>


### More info

_No response_

cc @awaelchli

Hints from issue discussion:
@calvinzhan This was a change introduced in deepspeed (perhaps breaking change since it was done in a minor release) https://github.com/microsoft/DeepSpeed/pull/3611. We had to update our tests too, see here: #17748
As of deepspeed 0.9.3, all model layers/parameters must be created under the `deepspeed.Init()` context manager, otherwise they won't get moved to the GPU. In Lightning, this translates to defining the layers in the designated `configure_sharded_model` hook. 

Could you please report this to DeepSpeed to make sure this is intended? 
@awaelchli 

I created model in the following way, since I don't have enough GPU memory. Have to put the model on CPU. I don't know who should move weights to GPU later, lightning or deepspeed. I did report an issue to deepspeed.

    def configure_sharded_model(self):
        deepspeed.zero.Init(module=self.model, remote_device="cpu", pin_memory=True)

@calvinzhan You don't need to move the layers to the GPU yourself, nor do you need to call `deepspeed.zero.Init` anywhere in Lightning. You can just define the layers here:

```py
def configure_sharded_model(self):
    self.layer = nn.Linear(...)
    # etc.
```
(the way you normally define your layer modules in `__init__()`).

You can find further examples of how to do this in: https://lightning.ai/docs/pytorch/stable/advanced/model_parallel.html.

Commit message: update configure_model docs

Patch to review:
diff --git a/docs/source-pytorch/advanced/model_parallel.rst b/docs/source-pytorch/advanced/model_parallel.rst
index 2408ca937628a..4707c7eca1846 100644
--- a/docs/source-pytorch/advanced/model_parallel.rst
+++ b/docs/source-pytorch/advanced/model_parallel.rst
@@ -189,10 +189,11 @@ Here's an example using that uses ``wrap`` to create your model:
     class MyModel(pl.LightningModule):
         def __init__(self):
             super().__init__()
+
+        def configure_model(self):
             self.linear_layer = nn.Linear(32, 32)
             self.block = nn.Sequential(nn.Linear(32, 32), nn.Linear(32, 32))
 
-        def configure_model(self):
             # modules are sharded across processes
             # as soon as they are wrapped with `wrap`.
             # During the forward/backward passes, weights get synced across processes


Code review comments:
diff_hunk: @@ -189,10 +189,11 @@ Here's an example using that uses ``wrap`` to create your model:
     class MyModel(pl.LightningModule):
         def __init__(self):
             super().__init__()
path: docs/source-pytorch/advanced/model_parallel.rst
text: @user1:
This can be removed now