FEATURED · 精选文章

LlamaFactory FSDPTurbo EP/EFSDP 设计指南:MoE 专家并行、专家参数分片与外层 FSDP2/CP 的集成实现

发布时间 / 2026/9/5 18:00:21
来源 / 创域科博编辑部
栏目 / 资讯中心
LlamaFactory FSDPTurbo EP/EFSDP 设计指南:MoE 专家并行、专家参数分片与外层 FSDP2/CP 的集成实现 LlamaFactory FSDPTurbo EP/EFSDP 设计指南MoE 专家并行、专家参数分片与外层 FSDP2/CP 的集成实现【免费下载链接】LlamaFactoryUnified Efficient Fine-Tuning of 100 LLMs VLMs (ACL 2024)项目地址: https://gitcode.com/GitHub_Trending/ll/LlamaFactory本文基于 LlamaFactory 仓库中的设计文档 FSDPTurbo EP/EFSDP 设计说明英文中文版展开结合 fsdpturbo 分布式插件源码、参数解析接口 与 回归测试系统讲解fsdpturbo插件如何在 Qwen3/Qwen3.5 MoE 模型上组合专家并行EP、专家参数分片EFSDP、Ulysses CP 与外层 FSDP2。读完本文你将掌握该插件的配置边界、Mesh 拓扑计算、模型切分顺序、跨 Mesh 梯度范数、权重加载机制以及 FLA kernel 插件的搭配方式与官方验证矩阵。1. 总体设计双侧职责分离fsdpturbo插件的核心原则是保持 FSDPTurbo 与 LlamaFactory 两侧职责清晰FSDPTurbo 负责专家并行EP、专家参数分片EFSDP和设备算子注册。LlamaFactory 负责进程初始化、基础 DeviceMesh、外层 FSDP2、CP、模型初始化与权重加载。集成层负责把两套参数布局组合起来并处理跨 Mesh 的梯度范数。从源码结构看这一职责划分落实为继承关系fsdpturbo.py 中的FSDPTurboFSDP2Engine直接继承自同目录下的FSDP2Engine仅在shard_model、prepare_model、_copy_weights等少数方法上做覆写。这样外层 FSDP2 的 checkpoint、保存流程可以原样复用MoE 专家侧的拓扑则被封装在插件私有的FSDPTurboParallelState中不污染 LlamaFactory 的全局DistributedInterface接口。该依赖的上游版本被固定在 requirements/fsdpturbo.txt 中通过 git 地址锁定到可复现的 main 分支提交属于可选安装组件——这也解释了第 4 节中为何所有 FSDPTurbo 导入都采用延迟导入。2. 配置边界公共拓扑与 FSDPTurbo 私有参数的严格隔离公共并行拓扑放在训练参数顶层FSDPTurbo 私有参数保留在dist_config中。最小可用配置如下cp_size: 1 dist_config: name: fsdpturbo ep_size: 16 ep_dispatcher: eager最小示例中各字段的职责ep_size专家并行组大小。ep_dispatcherFSDPTurbo EP dispatcher默认为eager。一个更完整的真实示例见 examples/v1/train_full/train_full_qwen3_moe_fsdpturbo_ep_fsdp.yaml它使用Qwen/Qwen3.5-35B-A3B、init_config: init_on_meta并同时配置了 FLA kernel 插件model: Qwen/Qwen3.5-35B-A3B model_class: llm kernel_config: name: auto, flash-linear-attention include_kernels: chunk_gated_delta_rule, fused_recurrent_gated_delta_rule chunk_size: 64 dist_config: name: fsdpturbo ep_size: 16 ep_dispatcher: eager cp_size: 1 init_config: name: init_on_meta2.1 参数解析FSDPTurboParams与严格校验在 interface.py 中FSDPTurboParams是一个 dataclass完整字段与默认值如下字段类型默认值说明nameLiteral[fsdpturbo]fsdpturbo插件标识reshard_after_forwardboolTrue继承自 FSDP2 参数offload_paramsboolFalse继承自 FSDP2 参数pin_memoryboolTrue继承自 FSDP2 参数dcp_pathstr \| NoneNone继承自 FSDP2 参数ep_sizeint1专家并行组大小__post_init__中校验必须为正数ep_dispatcherstreagerEP dispatcher 选择fsdp_ignored_moduleslist[str][]额外排除在外层 FSDP2 之外的模块hook_moduleslist[str][]EFSDP hook 的可选模块模式fsdp_implementationstrnativeEFSDP 实现可选native或custom三个要点值得注意公共拓扑字段留在顶层。dp_size、cp_size、cp_mode、mp_replicate_size、mp_shard_size和dist_timeout属于公共拓扑字段继续放在训练参数顶层。dist_config被严格解析为FSDPTurboParams如果把公共拓扑字段误放进dist_config会直接报Unknown params错误而不是静默忽略。这一行为在 test_fsdpturbo_ep.py 的test_fsdpturbo_uses_class_plugin_and_strict_backend_params中有显式断言。ep_modules/ep_fsdp_modules不是 YAML 接口。模型相关的模块路径统一由FSDPTurboEPModelSpec注册表管理严格参数解析会拒绝用户直接填写这两个字段避免配置与模型实际结构失配。测试同样覆盖了这条拒绝路径。bf16控制参数存储与计算 dtype。顶层训练参数bf16由 backend 在 FSDP materialization 之前完成模型 dtype 转换FSDPTurboFSDP2Engine.shard_model中param_dtype torch.bfloat16 if self.mixed_precision bf16 else torch.float32因此ModelEngine不需要感知 distributed backend 的私有配置。单测test_fsdpturbo_sets_storage_dtype_inside_backend验证了 fp32 输入模型在mixed_precisionbf16时被转换为 bfloat16。2.2 EFSDP 目标由模型规格决定而非用户配置EFSDP 的目标模块由模型规格决定Attention、Embedding、LM Head 等非专家参数不进入 FSDPTurbo EFSDP plan继续由 LlamaFactory 外层 FSDP2 管理。模型相关的模块路径和准备逻辑统一由FSDPTurboEPModelSpec注册表管理。当前内置两个规格qwen3_moeep_modules [model.layers.{*}.mlp.experts]ep_fsdp_modules [model.layers.{*}.mlp]准备逻辑会把 sparse expert 模块的hidden_dim补一份hidden_size别名因为 FSDPTurbo 的 eager EP dispatcher 期望专家块暴露hidden_size。qwen3_5_moeep_modules [model.language_model.layers.{*}.mlp.experts]ep_fsdp_modules [model.language_model.layers.{*}.mlp]。未注册的model_type会抛出明确的ValueErrorNo FSDPTurbo EP spec is registered for model_type...。测试test_qwen35_ep_model_spec对注册表内容和查询路径做了断言。3. Mesh 初始化插件私有的四维专家父 MeshLlamaFactory 的DistributedInterface只初始化自身原有的 model/data mesh。它不感知 EP、EFSDP也不为 distributed plugin 提供额外的 mesh 注册接口。FSDPTurbo 的专家拓扑由插件文件内的FSDPTurboParallelState独立创建并持有。调用链为run_sft / run_dpo / run_rm - DistributedInterface(training_args) - 初始化 LlamaFactory model/data mesh - DistributedPlugin(fsdpturbo).shard_model(...) - FSDPTurboFSDP2Engine.__init__() - FSDPTurboParallelState.initialize() - 初始化并保存 expert parent mesh 及其子 meshFSDPTurboParallelState创建的专家侧父 Mesh 为四维(edp, efsdp, ep, expert_cp)当前尺寸计算方式dp_size world_size / cp_size ep_fsdp_size dp_size / ep_size edp_size dp_size / (ep_size * ep_fsdp_size) mesh_shape (edp_size, ep_fsdp_size, ep_size, cp_size)各维度含义edpexpert data parallel是专家侧的额外数据并行副本维度efsdp是专家参数分片维度ep是专家并行维度expert_cp与外层 CP 对齐保证专家侧张量在 CP 维上的通信与注意力侧一致。从源码看fsdpturbo.py 的FSDPTurboParallelState.initialize初始化包含以下约束ep_size必须为正数且必须整除dp_size否则抛出ValueErrorep_size 1时要求 DP mesh 已初始化随后用init_device_mesh创建父 mesh并通过self.expert_mesh[dim]取出edp_mesh、efsdp_mesh、ep_mesh、expert_cp_mesh四个子 mesh重复初始化时会校验拓扑(dp_size, cp_size, ep_size)与首次一致拓扑发生变化会抛RuntimeError该状态是模块级单例_FSDPTURBO_PARALLEL_STATE通过get_fsdpturbo_parallel_state()访问。测试test_fsdpturbo_owns_expert_mesh_topology用一个 16 卡、cp_size1、ep_size8的假接口验证了计算结果最终init_device_mesh收到的参数恰为mesh_shape(1, 2, 8, 1)维度名为(edp, efsdp, ep, expert_cp)即edp1、efsdp16/82、ep8、expert_cp1。4. 模型切分顺序专家侧优先外层 FSDP2 随后模型包装顺序必须保持专家侧优先外层 FSDP2 随后DistributedPlugin(fsdpturbo) - FSDPTurboFSDP2Engine.shard_model(model) - prepare_model_ep(model) - expert_parallelize_modules(model, ep_mesh, ep_plan) - expert_fully_shard_modules(model, efsdp_mesh, ep_plan, fsdp_plan) - 收集专家参数作为 ignored_params - FSDP2Engine.prepare_model(model, ignored_params...) - 对剩余 Transformer Layer 和根模块执行 outer fully_shard源码中prepare_model_ep的具体流程是通过FSDPTurboEPModelSpec.get(model)取得模型规格并执行spec.prepare(model)构造EPPlanConfig(apply_modulesep_modules, dispatcher..., apply_efsdp_modules...)其中gradient_divide_factor被显式设置为ep_size * efsdp_size用于在 EP/EFSDP 混合切分下对专家梯度做正确的缩放测试test_fsdpturbo_sets_public_efsdp_gradient_divide_factor验证了ep_size4, efsdp_size2时该值为8.0构造FSDPPlanConfig——注意 FSDPTurbo 只用它来放置 EFSDP hook 和选择实现EFSDP 目标实际来自ep_plan.apply_efsdp_modules调用expert_parallelize_modules(model, ep_mesh, ep_plan)完成专家并行化当ep_fsdp_size 1时再调用expert_fully_shard_modules(model, efsdp_mesh, ep_plan, fsdp_plan)完成专家参数分片收集ignored_params用户fsdp_ignored_modules之外ep_size 1时自动把ep_modules加入忽略集合然后用module_name_match逐一匹配named_modules()并收集其全部参数。rank 0 会打印被忽略的 EP 参数数量。这样做的意义在于避免同一专家参数同时被 EFSDP 和外层 FSDP2 管理两套切分方案重叠会导致参数布局冲突和通信重复。外层 FSDP2 仍复用 LlamaFactory 原有的初始化、checkpoint 和保存流程。4.1 四种 EP dispatcher实现边界与验证状态LlamaFactory 集成层接受eager、fused、mc2和domino四种ep_dispatcher并将选项原样传给 FSDPTurbo。它们的实现边界与验证状态如下Dispatcher主要路径额外要求验证状态eager使用 PyTorch 实现 permute、unpermute 和 grouped matmul张量仍在当前加速设备上通过标准 AllToAll 完成 token dispatch/combine依赖最少用作参考实现已在 Ascend A3 上完成端到端数值与性能验证fused保持相同的 AllToAll 拓扑将 permute、unpermute 和 grouped matmul 切换为设备融合算子需要匹配的设备算子、dtype 和 layout存在空专家时可回退 eager 局部算子已在 Ascend A3 上完成端到端数值与性能验证mc2使用专用算子融合 AllToAllV 与 grouped matmul减少通信与计算之间的中间开销依赖 MC2 NPU 算子、HCCL communicator 及 shape/dtype 约束FSDPTurbo 已实现设计文档范围未做端到端验证domino将专家模块输入的第一维分成两片用独立通信流和 event 重叠 AllToAll 与专家计算需要异步 stream/event 支持且两个分片都要有足够 token 工作量覆盖调度开销FSDPTurbo 已实现设计文档范围未做端到端验证只验证eager与fused的原因在于它们分别覆盖参考实现和 A3 常用设备融合路径足以隔离并验证 LlamaFactory 与 FSDPTurbo 之间 EP/EFSDP 集成的正确性。mc2与domino引入了额外的算子、通信调度和输入形状约束需要独立的数值比较、长步稳定性和 profiler 分析因此不应从已有实验结果推断它们已达到相同的稳定性、精度或性能水平。5. FSDPTurbo 依赖入口延迟导入与直接模块导入LlamaFactory 从各功能的定义模块直接导入所需对象而不是走fsdp_turbo.distributed.__init__的聚合导出以避免 package 初始化期间的额外依赖和潜在循环导入from fsdp_turbo.distributed.expert_parallel.expert_fully_shard_parallel import ( expert_fully_shard_modules, ) from fsdp_turbo.distributed.expert_parallel.expert_parallel import expert_parallelize_modules from fsdp_turbo.fsdp_turbo_config import EPPlanConfig, FSDPPlanConfig from fsdp_turbo.utils.str_match import module_name_match这些导入写在prepare_model_ep()函数体内。因此没有安装 FSDPTurbo 时其他 distributed backend如 interface.py 中的fsdp2、deepspeed仍可正常导入和运行。FSDPTurboDistributed插件注册时也遵循同样的模式shard_model、clip_grad_norm等静态方法内部再from .fsdpturbo import ...。6. 跨 Mesh 梯度范数按所属 Mesh 分组归约外层参数和专家参数可能属于不同的 DTensor mesh不能直接放入一次标准clip_grad_norm_()调用混合 DTensor mesh 可能触发 DTensor stack 传播失败。fsdpturbo的 clip_grad_norm_ 采用分组局部求和 全归约方案分组遍历model.parameters()检查每个有梯度参数的grad.device_mesh.mesh_dim_names若与{ep, efsdp}相交则归入ep_params否则归入non_ep_params。局部 p 次方和_local_pth_sum先把 DTensor 梯度to_local()并转为 fp32再累加各梯度范数的norm_type次幂默认norm_type2.0。分路归约非专家参数的局部和沿DP group 和 CP group做all_reduce(SUM)专家参数的局部和沿FSDPTurboParallelState保存的EFSDP、EP、expert-CP group做all_reduce(SUM)CP 只在cp_size 1时提供对应 group否则传None跳过。统一裁剪两个局部和相加后取1/norm_type次方得到全局范数clip_coef min(max_norm / (norm 1e-6), 1.0)对全部本地梯度乘以同一个系数返回全局范数值。启动阶段还会执行一次零梯度 warmup_warmup_grad_norm给所有requires_grad参数写入零梯度并调用一次clip_grad_norm_(model, 1.0)使涉及的 collective 通信在正式训练开始前完成初始化随后立即清空梯度。当前这是fsdpturbobackend 的专用实现注册在FSDPTurboDistributed.clip_grad_norm上其他 backend 继续保留原有梯度范数路径等待上游 distributed plugin 接口解耦后再统一。7. 权重加载混合 Shard 布局 DTensor 的分片拷贝LlamaFactory 保留init_on_meta与 safetensors 加载流程。父类FSDP2Engine的加载器通过self._copy_weights(...)动态调用 FSDPTurbo engine 的覆写实现因此该方法不是死代码。覆写的_copy_weights支持包含多个Shardplacement 的 DTensor逻辑为若加载张量与目标参数 dtype 不一致先转换 dtype对 DTensor 参数枚举所有Shardplacement取当前 rank 的 mesh 坐标按各 mesh 维度依次narrow出当前 rank 应持有的本地切片chunk 大小向上取整处理无法整除的情形若某维度上start full_size尾部空分片则直接返回不做拷贝最后把切片拷贝进param.to_local()对应位置非 DTensor 参数则退化为普通的param.data.copy_(loaded_tensor)。模型保存和 checkpoint 接口继续复用 LlamaFactory 的 FSDP2 实现FSDPTurboDistributed的save_model、save_checkpoint、load_checkpoint均直接转发到fsdp2模块的同名函数。8. Kernel 插件FLA 算子与分布式配置的解耦FLAflash-linear-attention算子不属于 distributed config算子选择通过独立的kernel_config完成kernel_config: name: auto, flash-linear-attention include_kernels: chunk_gated_delta_rule, fused_recurrent_gated_delta_rule chunk_size: 32调用链ModelEngine - apply_kernels(auto, flash-linear-attention) - LlamaFactory 当前加速器对应的 auto kernels - KernelPlugin(flash-linear-attention).apply(...) - fsdp_turbo.ops.get_op() - FSDPTurbo device operator registry - fsdp_turbo.utils.patch.patch_model_members() - FLA backend implementation关键规则chunk_size当前支持16、32、64默认值为64kernel plugin 与 distributed plugin 彼此独立。name: flash-linear-attention只安装所选 FLA 算子逗号分隔的name: auto, flash-linear-attention会在分布式切分前组合 LlamaFactory 当前加速器的 auto kernels 与 FLA pluginLlamaFactory 负责算子名到模型属性的映射FLA 适配逻辑见 fla.py和chunk_size参数绑定FSDPTurbo 负责设备算子注册、选择和通用 callable patchFLA 依赖可选的外部三方件因此保持显式选择不属于内置auto集合FSDPTurbo 随后会替换目标专家模块的forward所以专家计算的最终路径由ep_dispatcher决定auto 阶段应用的 MoE kernel 不会作为独立的第二条专家执行路径保留。9. CP 运行约束与端到端验证矩阵9.1 CPUlysses相关约束使用 CP 时有三条必须满足的运行时约束init_on_meta构造模型时必须与from_pretrained路径一样传递attn_implementation否则模型会回退到非 FlashAttention 实现Ulysses CP 无法启动Ulysses 在调用 Hugging Face FlashAttention 前重建全局 attention mask只有二维position IDs 才参与 packed-sequence 检测Qwen3.5 mRoPE 等多轴 position IDs 已在 rotary embedding 中消费不应传入 FlashAttention 的 packed-sequence 检测逻辑。9.2 官方验证结果Qwen3.5-35B-A3BBF16 AdamW full SFT当前实现已在 Atlas 900 A3 SuperPoD 和 Atlas 950 SuperPoD 上完成以下验证使用 FSDPTurbo0e96fbc。A3 环境为 CANN 9.0.0、PyTorch 2.7.1、torch-npu 2.7.1.post4A5 环境为 CANN 9.1.0-beta.3、PyTorch 2.10.0、torch-npu 2.10.0.post2。表中性能按第 1 步至第 100 步的日志时间戳计算不包含首步前的初始化、编译和训练后的模型保存时间机器型号CPEPEFSDPCheckpointKernel / Dispatcher步数Loss首步 → 末步性能结果Atlas 900 A3 SuperPoD1161关闭FLAchunk size 16/ eager1001.3361 → 0.07932.51 s/it通过并完成保存Atlas 900 A3 SuperPoD1161关闭FLAchunk size 16/ fused1001.3354 → 0.11792.17 s/it通过并完成保存Atlas 900 A3 SuperPoD242关闭auto FLAchunk size 64/ fused1001.8114 → 0.52607.65 s/it通过并完成保存Atlas 900 A3 SuperPoD242关闭auto FLAchunk size 64/ eager1001.8095 → 0.55965.88 s/it通过并完成保存Atlas 950 SuperPoD181关闭未配置 kernel plugin / eager1001.3575 → 0.44392.68 s/it通过并完成保存结论与适用前提五组训练的 loss 和 grad norm 均保持有限值全部完成 100 步训练和模型保存同一切分下EP16 eager/fused 的逐步 loss 相关系数为 0.997CP2/EP4/EFSDP2 为 0.977说明两种 dispatcher 的优化轨迹一致性能收益与切分相关EP16 纯专家并行下 fused 比 eager 快约 13%而加入 CP 和 EFSDP 后 fused 反而慢约 30%因此不能把fused视为所有 mesh 的默认最优选择应按自身拓扑实测训练超参EP16 两组使用 global batch 16、cutoff length 256CP2 两组使用 global batch 8、cutoff length 128A5 组使用 global batch 8、cutoff length 256。首末 loss 仅用于验证各组自身收敛趋势不同切分组之间的绝对 loss 不应直接作为精度等价结论。10. 延伸阅读原始设计文档中英英文、中文插件实现fsdpturbo.py、distributed 插件接口可运行的完整示例Qwen3.5 MoE FSDPTurbo EP/EFSDP 全参 SFT回归测试test_fsdpturbo_ep.py覆盖模型规格注册、严格参数解析、存储 dtype 切换、gradient_divide_factor与专家 mesh 拓扑计算依赖锁定fsdpturbo 上游版本。【免费下载链接】LlamaFactoryUnified Efficient Fine-Tuning of 100 LLMs VLMs (ACL 2024)项目地址: https://gitcode.com/GitHub_Trending/ll/LlamaFactory创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
RELATED — 相关阅读

相关资讯

LATEST — 最新资讯

最新发布

TODAY — 本日精选

新闻

WEEKLY — 本周精选

新闻

MONTHLY — 本月精选

新闻