ray.train.lightning.RayFSDPStrategy
ray.train.lightning.RayFSDPStrategy#
- class ray.train.lightning.RayFSDPStrategy(*args: Any, **kwargs: Any)[source]#
Bases:
pytorch_lightning.strategies.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
PublicAPI (beta): This API is in beta and may change before becoming stable.
Methods
Gathers the full state dict to rank 0 on CPU.
Attributes