lightning_module_state_dict

RayFSDPStrategy.lightning_module_state_dict() → Dict[str, Any]

Gathers the full state dict to rank 0 on CPU.

FSDP checkpointing is broken in Lightning 2.0.x. This subclass patches the behavior to perform a full state dict checkpointing, gathering the checkpoint shards on rank 0 CPU. Upgrade to lightning>=2.1 to do sharded state dict checkpointing.

See the note in the class docstring for more details.

Returns:

The model’s state dict. With Lightning 2.0.x and PyTorch FSDP, it’s the full state dict gathered on rank 0 and offloaded to CPU, with the _forward_module. prefix removed from its keys. Otherwise, it’s the state dict that the parent FSDPStrategy returns.

Return type:

Dict[str, Any]