ray.train.lightning.RayFSDPStrategy#

class ray.train.lightning.RayFSDPStrategy(*args: Any, **kwargs: Any)[source]#

Bases: DDPFullyShardedStrategy

Subclass of FSDPStrategy to ensure compatibility with Ray orchestration.

For a full list of initialization arguments, please refer to: https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.strategies.FSDPStrategy.html

Note

It is recommended to upgrade lightning>=2.1 or above when using FSDP with Lightning, since Lightning starts to natively support state_dict_type, sharding_strategy, auto_wrap_policy and other FSDP configurations from 2.1.

PublicAPI (beta): This API is in beta and may change before becoming stable.

Methods

lightning_module_state_dict

Gathers the full state dict to rank 0 on CPU.

Attributes

`distributed_sampler_kwargs`
`root_device`