ray.train.lightning.RayFSDPStrategy

class ray.train.lightning.RayFSDPStrategy(*args: Any, **kwargs: Any)

基类：DDPFullyShardedStrategy

FSDPStrategy 的子类，以确保与 Ray 编排的兼容性。

如需查看完整的初始化参数列表，请参阅：https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.strategies.FSDPStrategy.html

备注

在使用 FSDP 与 Lightning 时，建议升级到 lightning>=2.1 或更高版本，因为从 2.1 版本开始，Lightning 开始原生支持 state_dict_type、sharding_strategy、auto_wrap_policy 和其他 FSDP 配置。

PublicAPI (测试版): 此API目前处于测试阶段，在成为稳定版本之前可能会发生变化。

方法

lightning_module_state_dict

在CPU上收集完整的状态字典到rank 0。

属性

distributed_sampler_kwargs
root_device