开源项目
TRL修正FSDP2异步训练的权重同步阻塞
Cast FSDP2 weight-transfer payloads to the vLLM manifest dtype
概述
TRL的AsyncGRPO和AsyncDistillation使用vLLM生成样本。FSDP2在指定bfloat16或从BF16检查点自动推断类型时,vLLM初始化得到低精度权重清单,Accelerate随后却可能把训练参数升为FP32,导致实际广播字节数与清单不同,权重同步阻塞。默认float32配置不触发同一问题。 主线在WeightTransferClient.send_weights中,按清单为每个张量转换类型,再交给打包NCCL广播。最终合入版保留类型转换,未保留提案中的名称、形状校验和有界等待,因此不能宣称提供了完整传输验证协议。作者用Qwen3-0.6B、FSDP2 BF16和vLLM 0.30完成单训练进程验证,三次同步均完成;分布式较大规模结果来自较早版本,不作为最终版本同等验证。