PyTorch Lightning Fabric中分布式保存的正确实践

在分布式训练场景下，模型和状态的保存是一个需要特别注意的操作。PyTorch Lightning的Fabric组件提供了便捷的分布式保存功能，但许多开发者在使用时容易陷入一些误区。

常见误区分析

许多开发者习惯性地认为在分布式环境中保存状态时，需要手动判断当前进程是否为rank 0进程，然后只在rank 0上执行保存操作。这种模式在直接使用PyTorch的分布式API时确实是常见的做法：

if rank == 0:
    torch.save(state_dict, "model.pth")

然而，当使用PyTorch Lightning的Fabric组件时，这种显式的条件判断是不必要的，甚至会导致错误。

Fabric.save()的设计原理

Fabric组件的save()方法已经内置了分布式逻辑，它会根据当前使用的策略自动决定保存行为：

1. 对于DDP策略，Fabric会自动在rank 0进程上执行保存操作
2. 对于FSDP策略，Fabric会协调所有rank进程保存各自的部分
3. 对于单机训练，行为与普通保存无异

这种设计封装了分布式细节，使开发者无需关心底层实现，只需统一调用fabric.save()即可。

正确使用示例

import lightning as L

def main(fabric):
    # 准备需要保存的状态
    state = {"model": model.state_dict(), "optimizer": optimizer.state_dict()}
    
    # 直接调用save，无需条件判断
    fabric.save("checkpoint.pth", state)

常见错误场景

如果在使用Fabric.save()时手动添加rank判断，可能会导致以下问题：

1. 非rank 0进程提前退出，导致分布式通信中断
2. 在需要多rank协同保存的策略(如FSDP)下，保存不完整
3. 屏障同步出现问题，引发连接错误

最佳实践建议

1. 始终使用Fabric提供的保存/加载API，而不是直接使用torch.save/load
2. 避免手动添加rank条件判断，相信框架的自动化处理
3. 对于复杂保存需求，可以通过Fabric策略配置来调整保存行为
4. 在保存前后通常不需要显式添加屏障(barrier)，除非有特殊同步需求

理解这些原理后，开发者可以更自信地在分布式环境中使用Fabric进行状态保存，避免常见的陷阱，提高开发效率和代码可靠性。