Registered buffers not moved to correct device when using DeepSpeed Stage 3

Nobody has claimed this yet.

### Bug description

Using the DeepSpeed `Strategy` configuration
```yaml
_target_: lightning.pytorch.strategies.DeepSpeedStrategy
zero_optimization: true
stage: 3
allgather_bucket_size: 2e8
reduce_bucket_size: 2e8
offload_optimizer: false
offload_parameters: false
partition_activations: false
cpu_checkpointing: false
contiguous_gradients: false
overlap_comm: false
```

I am experiencing an issue (specifically with DeepSpeed stage 3, not stages 1-2) where the tensors registered within sub-`nn.Modules` of my `LightningModule`'s main `lit_model.network` `nn.Module` are not moved by `register_buffer()` to the correct device upon training the `lit_module.network`. In particular, I am trying to register buffers as

```python
distance_bins_tensor = tensor([0.0, 1.0, 2.0, 3.0])
self.register_buffer("distance_bins", distance_bins_tensor)
```

within the various submodules of my `lit_module.network`. When my optimizer tries to perform a step, I get the error
```bash
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:6 and cpu!
```
when trying to use these registered buffers e.g., by multiplying them by feature tensors loaded onto (in this case) `cuda:6`.

### What version are you seeing the problem on?

v2.4

### How to reproduce the bug

_No response_

### Error messages and logs

```
# Error messages and logs here please
```

### Environment

Current environment

* CUDA:
- GPU:
- NVIDIA A100 80GB PCIe
- NVIDIA A100 80GB PCIe
- available: True
- version: 11.8
* Lightning:
- adam-atan2-pytorch: 0.0.10
- alphafold3-pytorch: 0.0.41
- alphafold3-pytorch-lightning-hydra: 0.1.111
- frame-averaging-pytorch: 0.0.19
- lightning: 2.4.0
- lightning-utilities: 0.11.6
- pytorch-lightning: 2.4.0
- rotary-embedding-torch: 0.6.1
- torch: 2.3.0+cu118
- torch-geometric: 2.5.3
- torchaudio: 2.3.0+cu118
- torchmetrics: 1.4.1
- torchtyping: 0.1.4
- torchvision: 0.18.0+cu118
* Packages:
- adam-atan2-pytorch: 0.0.10
- aiofiles: 23.2.1
- aiohttp: 3.9.5
- aiosignal: 1.3.1
- alembic: 1.13.1
- alphafold3-pytorch: 0.0.41
- alphafold3-pytorch-lightning-hydra: 0.1.111
- annotated-types: 0.7.0
- antlr4-python3-runtime: 4.9.3
- anyio: 4.4.0
- appdirs: 1.4.4
- argcomplete: 3.3.0
- asttokens: 2.4.1
- async-timeout: 4.0.3
- attrs: 23.2.0
- autopage: 0.5.2
- beartype: 0.18.5
- beautifulsoup4: 4.12.3
- biopandas: 0.5.1.dev0
- biopython: 1.83
- bioservices: 1.11.2
- cattrs: 23.2.3
- certifi: 2024.8.30
- cfgv: 3.4.0
- chardet: 5.2.0
- charset-normalizer: 3.3.2
- click: 8.1.7
- cliff: 4.7.0
- cmaes: 0.10.0
- cmd2: 2.4.3
- colorama: 0.4.6
- colorlog: 6.8.2
- colt5-attention: 0.11.0
- comm: 0.2.2
- contourpy: 1.2.1
- cycler: 0.12.1
- debugpy: 1.8.1
- decorator: 5.1.1
- deepdiff: 7.0.1
- deepspeed: 0.15.0
- distlib: 0.3.8
- docker-pycreds: 0.4.0
- easydev: 0.13.2
- einops: 0.8.0
- einx: 0.2.2
- environs: 11.0.0
- exceptiongroup: 1.2.1
- executing: 2.0.1
- fastapi: 0.112.2
- ffmpy: 0.4.0
- filelock: 3.13.1
- fonttools: 4.52.4
- frame-averaging-pytorch: 0.0.19
- freetype-py: 2.3.0
- frozendict: 2.4.4
- frozenlist: 1.4.1
- fsspec: 2024.2.0
- gemmi: 0.6.6
- gevent: 24.2.1
- gitdb: 4.0.11
- gitpython: 3.1.43
- gradio: 4.43.0
- gradio-client: 1.3.0
- gradio-molecule3d: 0.0.5
- graphein: 1.7.6
- greenlet: 3.0.3
- grequests: 0.7.0
- h11: 0.14.0
- hjson: 3.1.0
- httpcore: 1.0.5
- httpx: 0.27.2
- huggingface-hub: 0.23.4
- hydra-colorlog: 1.2.0
- hydra-core: 1.3.2
- hydra-optuna-sweeper: 1.2.0
- identify: 2.5.36
- idna: 3.7
- importlib-resources: 6.4.4
- iniconfig: 2.0.0
- ipykernel: 6.29.4
- ipython: 8.24.0
- jaxtyping: 0.2.28
- jedi: 0.19.1
- jinja2: 3.1.3
- joblib: 1.4.2
- jupyter-client: 8.6.2
- jupyter-core: 5.7.2
- kiwisolver: 1.4.5
- lightning: 2.4.0
- lightning-utilities: 0.11.6
- line-profiler: 4.1.3
- local-attention: 1.9.1
- loguru: 0.7.2
- looseversion: 1.1.2
- lxml: 5.2.2
- mako: 1.3.5
- markdown-it-py: 3.0.0
- markupsafe: 2.1.5
- marshmallow: 3.21.3
- matplotlib: 3.8.4
- matplotlib-inline: 0.1.7
- mdurl: 0.1.2
- mmtf-python: 1.1.3
- mpmath: 1.3.0
- msgpack: 1.0.8
- multidict: 6.0.5
- multipledispatch: 1.0.0
- munkres: 1.1.4
- nest-asyncio: 1.6.0
- networkx: 3.2.1
- ninja: 1.11.1.1
- nodeenv: 1.8.0
- numpy: 1.23.5
- nvidia-cublas-cu11: 11.11.3.6
- nvidia-cuda-cupti-cu11: 11.8.87
- nvidia-cuda-nvrtc-cu11: 11.8.89
- nvidia-cuda-runtime-cu11: 11.8.89
- nvidia-cudnn-cu11: 8.7.0.84
- nvidia-cufft-cu11: 10.9.0.58
- nvidia-curand-cu11: 10.3.0.86
- nvidia-cusolver-cu11: 11.4.1.48
- nvidia-cusparse-cu11: 11.7.5.86
- nvidia-ml-py: 12.560.30
- nvidia-nccl-cu11: 2.20.5
- nvidia-nvtx-cu11: 11.8.86
- omegaconf: 2.3.0
- optree: 0.11.0
- optuna: 2.10.1
- ordered-set: 4.1.0
- orjson: 3.10.7
- packaging: 24.0
- pandas: 1.5.3
- parso: 0.8.4
- pbr: 6.0.0
- pdbeccdutils: 0.8.5
- pexpect: 4.9.0
- pillow: 10.2.0
- pip: 24.0
- pipx: 1.5.0
- platformdirs: 4.2.2
- plotly: 5.22.0
- pluggy: 1.5.0
- polars: 1.3.0
- pre-commit: 3.7.1
- prettytable: 3.10.0
- prompt-toolkit: 3.0.45
- protobuf: 4.25.4
- psutil: 5.9.8
- ptyprocess: 0.7.0
- pure-eval: 0.2.2
- py-cpuinfo: 9.0.0
- pycairo: 1.26.0
- pydantic: 2.8.2
- pydantic-core: 2.20.1
- pydub: 0.25.1
- pygments: 2.18.0
- pyparsing: 3.1.2
- pyperclip: 1.8.2
- pytest: 8.2.1
- python-dateutil: 2.9.0
- python-dotenv: 1.0.1
- python-multipart: 0.0.9
- pytorch-lightning: 2.4.0
- pytz: 2024.1
- pyyaml: 6.0.1
- pyzmq: 26.0.3
- rdkit: 2024.3.2
- reportlab: 4.1.0
- requests: 2.32.2
- requests-cache: 1.2.0
- retrying: 1.3.4
- rich: 13.7.1
- rich-click: 1.8.2
- rlpycairo: 0.2.0
- rootutils: 1.0.7
- rotary-embedding-torch: 0.6.1
- ruff: 0.6.4
- scikit-learn: 1.5.0
- scipy: 1.13.1
- seaborn: 0.13.2
- semantic-version: 2.10.0
- sentry-sdk: 2.12.0
- setproctitle: 1.3.3
- setuptools: 70.0.0
- sh: 2.0.7
- shellingham: 1.5.4
- shortuuid: 1.0.13
- six: 1.16.0
- smmap: 5.0.1
- sniffio: 1.3.1
- soupsieve: 2.5
- sqlalchemy: 2.0.30
- stack-data: 0.6.3
- starlette: 0.38.4
- stevedore: 5.2.0
- suds-community: 1.1.2
- sympy: 1.12
- taylor-series-linear-attention: 0.1.12
- tenacity: 8.3.0
- threadpoolctl: 3.5.0
- timeout-decorator: 0.5.0
- tomli: 2.0.1
- tomlkit: 0.12.0
- torch: 2.3.0+cu118
- torch-geometric: 2.5.3
- torchaudio: 2.3.0+cu118
- torchmetrics: 1.4.1
- torchtyping: 0.1.4
- torchvision: 0.18.0+cu118
- tornado: 6.4
- tqdm: 4.66.4
- traitlets: 5.14.3
- triton: 2.3.0
- typeguard: 2.13.3
- typer: 0.12.5
- typing-extensions: 4.11.0
- tzdata: 2024.1
- unicodedata2: 15.1.0
- url-normalize: 1.4.3
- urllib3: 2.2.1
- userpath: 1.9.2
- uvicorn: 0.30.6
- virtualenv: 20.26.2
- wandb: 0.16.6
- wcwidth: 0.2.13
- websockets: 12.0
- wget: 3.2
- wheel: 0.43.0
- wrapt: 1.16.0
- xarray: 2024.3.0
- xmltodict: 0.13.0
- yarl: 1.9.4
- zope.event: 5.0
- zope.interface: 6.4.post2
* System:
- OS: Linux
- architecture:
- 64bit
- ELF
- processor: x86_64
- python: 3.10.14
- release: 4.18.0-553.16.1.el8_10.x86_64
- version: #1 SMP Thu Aug 8 07:11:46 EDT 2024

### More info

_No response_

cc @lantiga

Start at the DeepSpeedStrategy entry point and trace Stage 3 setup for nested nn.Modules and registered buffers. Create a minimal reproduction from the reported configuration, then verify that buffers and feature tensors share the training device while Stage 1–2 behavior remains unaffected.

Written by the indexing model from the issue text.