ray.train.lightning.RayFSDPStrategy.lightning_module_state_dict#

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

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.