pytorch_lightning_deepspeed.py

from pytorch_lightning import Trainer
from pytorch_lightning.plugins import DeepSpeedPlugin

# Enable CPU Offloading
model = MyModel()
trainer = Trainer(gpus=4, plugins='deepspeed_stage_3_offload', precision=16)
trainer.fit(model)

# Enable CPU Offloading, and offload parameters to CPU
model = MyModel()
trainer = Trainer(
    gpus=4,
    plugins=DeepSpeedPlugin(
        stage=3,
        offload_optimizer=True,
        offload_parameters=True,
        remote_device='nvme',
        offload_params_device='nvme',
        offload_optimizer_device='nvme',
        nvme_path = '/local_nvme'
    ),
    precision=16
)
trainer.fit(model)