Feat/task25 mixture of lora - #255
Open
pophirasawa wants to merge 43 commits into
Open
Conversation
#### Summary 为 Mixture-of-LoRA 提供训练端和 rollout 端可复用的路由基础逻辑。本次提交只增加公共路由能力和 CPU 测试，不修改现有单 LoRA 路径。 #### Changes 新增 RoutingDecision 和 RoutingStatistics，统一保存 Top-K 前概率、expert 编号、归一化激活权重，以及可跨 rank 聚合的逐 expert 原始统计。实现 FP32 softmax、token-level Top-K 路由、response mask 过滤、归一化熵和 balance loss；expert 选择份额按 valid_tokens * K 归一化，离散选择统计不参与反向传播，router 梯度通过 Top-K 前概率计算。新增 17 个 CPU 测试，覆盖多个 K、temperature、mask、空 response、梯度、entropy 和非法参数。 #### Verification pytest：17 passed。Ruff、docformatter、通用文件检查和冲突标记检查均通过。gitleaks hook 因 H20 访问 proxy.golang.org 时 TLS 超时，未能完成环境安装；失败发生在依赖下载阶段，并非扫描发现问题。
#### Summary 完成 Mixture-of-LoRA 阶段 A 的公共配置、参数 schema 和无参数执行器接口，为 Megatron 与 SGLang 后续实现提供同一组结构和数学定义。 #### Changes 新增 MixtureLoraConfig 及 N、R、K、temperature、aux loss、alpha 和目标层校验；固定 expert A/B 与 router 的参数名称和全局 shape；新增 checkpoint/transport 共用的状态与 TP 分片描述；新增显式并行上下文、RoutedLoRAExecutor 协议和纯 PyTorch dense executor；增加逐 site balance loss 平均函数，并保证执行器不持有参数、不修改路由结果。 #### Verification Mixture-of-LoRA L0 测试 38 passed，覆盖 FP32、FP16、BF16、输出、梯度、mask、schema、可重复性和异常输入。现有单 LoRA 回归 43 passed、5 skipped。全仓 pre-commit 全部通过，包含 Ruff、docformatter 和 gitleaks。
#### Summary 接入 Mixture-of-LoRA 命令行参数和启动阶段校验，并保持单 expert 配置继续使用现有 LoRA rollout 路径。 #### Changes 新增 expert 数量、router Top-K、temperature 和 balance loss coefficient 参数；N 大于 1 时要求完整 router 配置、colocate、SGLang TP/DP 为 1，并拒绝 fully-async、merge mode、adapter mode 和不支持的目标层；新增 Mixture 启用判断和共享配置构造；将旧 LoRA rollout mode 校验集中到独立函数，N=1 仍自动选择现有 merge 路径。 #### Verification 参数、PEFT 和单 LoRA 回归共 101 passed、5 skipped。全仓 pre-commit 全部通过，包含 Ruff、docformatter 和 gitleaks。
Summary: - 接入 Megatron 训练端的 Mixture-of-LoRA 核心模块和 Bridge PEFT 注入路径。 - 保持单 expert 配置继续使用原有 LoRA 实现。 Changes: - 新增打包的 expert A/B 参数、FP32 token router、dense Top-K 执行和线性层包装器。 - 在模型 provider 中按 expert 数选择实现，并在 optimizer 建立前检查 base 冻结及 expert/router 可训练状态。 - 扩展参数分类与分项统计，保持 base state key 和 Megatron 线性层返回协议。 - 增加 CPU 数值、梯度、初始化、state key、Bridge 匹配、provider 路径和 CUDA profiler 测试。 Verification: - 相关 pytest：137 passed, 5 skipped。 - H20 CUDA forward/backward 与 profiler 测试通过。 - pre-commit run --all-files 全部通过。
Summary: - 为每个训练 microbatch 建立 Mixture-of-LoRA 路由上下文。 - 将逐 site balance loss 直接附着到模型激活，保持 policy loss 标量不变。 Changes: - 使用 full_loss_masks 过滤 prompt、padding 和 dummy token，并处理 batch-first mask 与 Megatron 激活布局。 - 普通 loss、per-token loss 和动态 batch 共用同一个 microbatch 缩放 helper。 - 按全模型 site 数归一化 aux loss，分离可导 aux tensor 与无梯度路由记录。 - 捕获 activation checkpoint 的路由上下文，并补齐冻结 base 时的 recompute input-grad 路径。 - 增加 mask、缩放、dummy、router 梯度、checkpoint 恢复、真实 Bridge 和真实 ColumnParallelLinear 测试。 Verification: - 相关 pytest：152 passed, 8 skipped。 - 真实 Megatron/Bridge 定向测试：2 passed。 - H20 CUDA profiler 测试通过。 - pre-commit run --all-files 全部通过。
Summary: - 在 optimizer step 结束后汇总每个 Mixture-of-LoRA site 的路由统计并接入现有训练日志。 - 保留实际训练目标中的 aux loss，同时输出用于判断 expert 塌缩的逐 site 与全局指标。 Changes: - 按固定 site_id 表打包 Top-K 前概率、Top-K 后权重、选择次数、Top-1 次数、熵和有效 token 数。 - 先在 DP/CP group 汇总原始统计，再在 pipeline group 合并各 stage 持有的 site。 - 输出 expert 平均概率、平均激活权重、选择份额、Top-1 比例、归一化熵、balance loss 和实际 aux loss。 - 使用训练目标权重汇总 balance loss，支持普通 loss、per-token loss 和 dummy microbatch。 - activation recompute 使用相同 key 覆盖记录，step 聚合完成后清理上下文。 - 增加手算指标、空记录、recompute 去重及两进程 DP/PP 聚合测试。 Verification: - 相关 pytest：108 passed, 2 skipped。 - 真实 Megatron/Bridge 定向测试：2 passed。 - 真实 Megatron/Bridge 环境导入训练模块成功。 - pre-commit run --all-files 全部通过。
Summary: - 为 Mixture-of-LoRA 增加 qkv 与 proj 的 Tensor Parallel、Sequence Parallel 执行路径。 - 保持 TP 分片后的 forward、expert/router 梯度和 aux loss 与单卡数学参考一致。 Changes: - qkv 的 A 沿 rank 维分片，B 沿输出维分片，router 在 TP rank 间复制并同步梯度。 - proj 的 A/router 沿输入维分片，B 沿输出维分片，显式汇总低秩中间结果和 router logits。 - 增加 sequence gather/scatter 及其反向通信，所有 collective 显式使用构造时保存的 TP group。 - qkv aux loss 按 TP size 缩放后再同步 replicated router 梯度；路由指标仅由 TP rank 0 记录。 - 保留 TP=1 参数形状与原有执行结果，并校验 rank、输入和输出维度的可分片性。 - 增加两进程 qkv/proj、SP 开关、policy gradient、aux gradient、mask 对齐和统计去重测试。 Verification: - 相关 pytest：109 passed, 2 skipped。 - 两进程 TP/SP gloo 对照测试通过。 - 真实 Megatron/Bridge 定向测试：2 passed。 - H20 CUDA 单卡 profiler 测试通过。 - pre-commit run --all-files 全部通过。 Limitations: - 当前 H20 节点仅暴露 1 张 GPU，尚未执行 NCCL TP=2 测试。
Summary: - 将 Mixture-of-LoRA expert、router 和配置接入 Megatron 原生 distributed checkpoint。 - 恢复时校验 Mixture 配置，并保持现有单 LoRA 的 HF PEFT 导出行为。 Changes: - qkv/proj 分别声明 A、B 和 router 的全局 shape 与 TP 分片轴。 - 以模块 extra state 保存 schema version、site_id、N/R/K/T/C、alpha、target、输入输出维度和 dtype。 - 加载配置不一致时列出具体字段并停止恢复。 - Mixture 模式跳过仅适用于标准单 LoRA 的 _save_lora_to_checkpoint 导出；N=1 路径不变。 - checkpoint 不记录 dense executor 类型，后续执行器可以共享同一参数结构。 - 增加配置冲突、真实 torch_dist 写盘/恢复、参数、路由、输出和 loss 一致性测试。 Verification: - 相关 pytest：110 passed, 2 skipped。 - 真实 Megatron distributed checkpoint 保存/恢复测试通过。 - 真实 Megatron/Bridge 定向测试通过。 - pre-commit run --all-files 全部通过。
新增 Qwen3 Mixture-of-LoRA SGLang 外部模型，在 qkv_proj 和 o_proj 上复用训练侧 Top-K 路由与 dense executor。 通过启动前 JSON 环境变量传递路由配置，并校验外部模型包冲突、参数名称、shape 和 dtype。补充单 LoRA 兼容、训练与 rollout 数值一致、BF16 router 及权重加载测试。
SGLang PP worker 加载 Mixture 参数时跳过不属于本 stage 的 layer，同时保留本地未知参数的严格报错。 新增 H20 BF16 CUDA Graph capture 测试和 PP layer 过滤测试，确认 routed adapter 可以沿用现有 graph 执行路径。
新增 Mixture-of-LoRA 权重收集器，按固定 schema 汇总 PP/TP 分片，并将 Megatron 的 GQA 分组 QKV 排布转换为 SGLang 的 Q/K/V 连续排布。 首轮同步发送 base、全部 expert 和 router，后续只发送 expert 和 router；只有最后一个 routed chunk 携带新 weight version。复用现有 pause、flush、IPC update 和 continue 流程，并从 raw/Bridge HF iterator 中排除 Mixture 参数。 补充 QKV 转换、HF 过滤、首轮与后续批次、CPU mirror、生命周期顺序和单 LoRA 回归测试。
SGLang 按 Hugging Face checkpoint 的 architecture 名称注册外部模型。本次将入口类名改为 Qwen3ForCausalLM，确保 Mixture-of-LoRA 实现覆盖内置 Qwen3，而不是启动后继续使用原模型。 Mixture-of-LoRA 是纯文本模型，启动子进程前会清理多模态外部处理器环境变量，并补充入口注册和环境隔离测试。H20 实机已验证预填充、CUDA Graph 解码及连续三次 router/expert 在线更新。
Mixture-of-LoRA 权重更新按暂停、发送和恢复三个阶段同步各训练 rank 的失败状态。任一阶段失败时恢复 rollout generation，并保持原 weight version 与 base 同步状态，避免服务停在暂停状态或发布半完成版本。 IPC 流水等待上一 chunk 失败时仍先完成既有 chunk barrier，防止其他 rank 永久等待。新增 engine 更新失败和 chunk barrier 失败注入测试；Task 25 回归 173 passed，全量 pre-commit 通过。
静态 CP 下按 site 在 context-parallel group 汇总 Top-K 入选次数、路由概率和有效 response token 数。每个 rank 只保留本地概率和的可导贡献，随后由现有 DP/CP 梯度归约合成完整序列的 balance loss。 普通 loss 与 per-token loss 分别使用正确的样本或全局 token 缩放，路由日志中的 aux loss 不随 CP 数放大。新增 CP=2、空 token rank、梯度和指标参考测试，并在启动阶段明确拒绝首版不支持的动态 CP。Task 25 完整回归通过，全量 pre-commit 通过。
补充 Mixture-of-LoRA checkpoint 恢复后的下一步一致性测试。第一个 optimizer step 后保存 expert/router、AdamW、LR scheduler、iteration 和 Torch RNG，再比较连续训练与恢复训练的第二步结果。 测试启用 dropout 并确认 router 在第二步发生更新；恢复后的 loss、全部参数、优化器动量、scheduler 和 iteration 与不中断路径一致。Task 25 回归 176 passed，全量 pre-commit 通过。
新增 Qwen3-4B、DAPO math、GRPO、8 卡 colocate 的 Mixture-of-LoRA recipe，显式配置 N=4、R=16、Top-K、temperature 和 balance coefficient。训练侧使用 TP=2/SP，rollout 启动八个独立的 TP=1/DP=1 SGLang engine，并保留命令行覆盖参数。 新增中英文使用文档和 VitePress 导航，说明启动条件、支持范围、路由公式、逐 site 指标、权重同步和 checkpoint 恢复。Bash 语法检查及全量 pre-commit 通过。
Summary 修复 Mixture-of-LoRA 在 PP/VPP 训练中的 site 标识、权重转换和 rollout 同步问题，使各 pipeline stage 的 expert 与 router 能稳定聚合并更新到 SGLang。 Changes - 使用全局层偏移生成稳定 site_id，并传递 virtual pipeline stage。 - base 权重仅由参数所属 stage 执行 Bridge 转换，转换后在 PP group 广播；修正转换耗时统计。 - 汇总各 PP stage 的 Mixture 参数元数据，补充 selective CPU offload 和 SGLang base CPU backup。 - 增加 rank/TP 整除、VLM/MoE 拒绝和 Bridge 可选依赖保护，并增强同步失败日志。 - 补充 PP site、权重同步、offload、参数校验和 rollout 回归测试。 Verification - Task 25、两卡分布式及单 LoRA 兼容回归：192 passed。 - Bridge 与权重同步定向回归：33 passed。 - pre-commit：除独立执行的 gitleaks hook 外，其余 hook 全部通过。 - Gitleaks v8.24.2 staged scan：no leaks found。 - A800 8 卡 PP=2、TP=2、SP、colocate 两步训练完成，包含两次训练后权重同步和 iteration 0/1 checkpoint。
Summary 对齐 Megatron 原生分布式 checkpoint 设计，使 Mixture-of-LoRA 在保持 TP/PP 不变时支持调整 DP size 后恢复 optimizer、expert、router 和训练进度。 Changes - recipe 启用 fully reshardable distributed optimizer，并在中英文文档说明跨 DP 恢复条件。 - 补齐 TP 路由指标同步，使日志 rank 在任意 TP rank 上都能获得完整指标。 - 增加 DP 梯度与指标、PP stage 梯度与指标、PP checkpoint 归属及 DP 缩容恢复测试。 - 更新 VPP Mixture factory 回归测试，并让纯 CPU checkpoint 测试不依赖 GPU 空闲显存。 Verification - Task 25、单 LoRA 及相关 VPP 定向回归：232 passed。 - ruff、冲突标记和 whitespace 检查通过。 - Qwen3-4B、DAPO math、colocate 从 4 进程 DP=2 checkpoint 恢复到 2 进程 DP=1，完成 step 2 并保存 iteration 2。 - 恢复后四个 expert 的全局平均激活权重约为 24.5%、26.2%、22.6%、26.7%，未出现单 expert 塌缩。
将 Qwen3-4B Mixture-of-LoRA recipe 从 FP16 切换为 BF16，并移除仅适用于 FP16 动态缩放的 loss scale 参数。同步更新中英文使用文档，使公开 recipe 与已完成的 200-step 验收实验配置一致。
为 SGLang router registration fixture 补充 Mixture 配置与 external model stub，并为 SFT loss 测试提供有效的 DP world size。修复全量 pytest 中由新增 Mixture import 和严格 loss 缩放校验触发的测试错误。
默认 CPU CI 未安装 Megatron 或 SGLang 时，按仓库现有测试约定跳过依赖真实后端的新增用例，避免在测试收集阶段失败。 安装完整后端的环境仍会执行参数校验、权重同步、SGLang 配置和模型一致性测试。
默认 CPU CI 未安装 Megatron 时，张量并行路由指标用例仍会失败：mp.spawn 拉起的子进程是全新解释器，不加载 pytest 与 conftest，worker 内的 from megatron.core import parallel_state 直接抛 ModuleNotFoundError，并以 ProcessRaisedException 冒泡，使三个 Python 版本的测试 job 全部中断。 在 mp.spawn 之前补上 megatron.core 与 megatron.training 的 importorskip，并为该文件补充缺失的 pytest 导入。守卫模块按 worker 的实际依赖链选取：parallel_state 来自 megatron.core，_reduce_mixture_lora_routing_metrics 所在的 relax.backends.megatron.model 在模块级还需要 megatron.training。 这是 f9ce3fb 守卫补齐的延续，该 commit 覆盖了模块级可见后端导入的四个文件，遗漏了本文件这种藏在子进程里的导入。写法与 test_mixture_lora_checkpoint_distributed.py 现有约定一致，安装完整后端的官方镜像仍会完整执行该用例。 Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
pophirasawa
marked this pull request as ready for review
August 11, 2026 17:05
pophirasawa
marked this pull request as draft
August 11, 2026 17:05
pophirasawa
marked this pull request as ready for review
August 11, 2026 17:05
同步最新 main 的 S3 模型来源、HF/FP8 导出和权重分桶更新。 冲突处理说明： - 模型保存路径保留上游 patch_megatron_model、FP8 streaming writer 和 post hook，同时让 Mixture 模式跳过仅适用于单 adapter 的 LoRA 导出。 - Tensor 权重同步保留上游不等长 bucket padding，并允许 Mixture 分块传输在最终分块前使用空 weight version。 - SGLang 启动保留上游 ModelSource/S3 逻辑，并在子进程启动前注入 Mixture external model 配置及 base-once 同步所需的 CPU 权重备份。 本提交只解决以上三个代码冲突；S3 loader 测试 fixture 的兼容修复留给后续独立提交。
最新 main 新增的 S3 loader 测试会构造精简版 megatron_peft_utils 模块，再单独导入 sglang_engine。Task 25 为 sglang_engine 增加 Mixture 配置与模式判断导入后，原 fixture 缺少对应符号，导致测试在收集阶段失败。 为 fixture 补充 build_mixture_lora_config 和 is_mixture_lora_enabled 两个 stub，并固定返回未启用 Mixture，使测试继续覆盖原有 S3 模型加载路径。该修改只影响测试替身，不改变 S3、SGLang 或 Mixture 的运行时行为。 验证：tests/test_s3_model_loader.py 共 51 个用例全部通过。
SGLang 的 PP 模型使用全局 layer 表，并以占位模块表示其他 stage 的层。Mixture 安装逻辑改为只遍历当前 stage 的 [start_layer, end_layer) 区间，直接使用全局 layer id，避免访问 PPMissingLayer 或重复偏移层编号。 新增 PP stage 回归测试，使用无 self_attn 的占位层验证只为本 stage 真实层安装 qkv/proj adapter，并检查同步 schema 使用正确的全局 site id。 验证：tests/models/qwen3_mixture_lora/test_sglang_model.py，8 passed。
量化 SGLang linear 的首个参数可能是 int32 或 int8 的 packed 权重，不能作为 LoRA expert 和 router 的参数类型。安装 adapter 时继续从 base parameter 获取设备，但优先使用 linear.params_dtype 作为浮点计算类型。 当 linear 无法提供浮点 params_dtype，且首个参数也是整数类型时，在模型构造阶段给出明确错误，避免创建需要梯度的整数 Parameter。 新增量化 linear 测试，覆盖 int32 packed 权重配合 BF16 params_dtype，以及缺少浮点 dtype 时的失败路径。验证：SGLang 模型测试 10 passed。
带版本号的最终分块会在所有 rank 确认前序请求成功后再发送，避免失败更新被发布为完整策略。 同步失败后保持 rollout 暂停并保留原版本，补充故障注入测试及中英文恢复说明。
EP metadata 合并时为 replicated non-expert 参数选择最小全局 rank，保证每个 converted slot 只有一个 owner。 补充 metadata 单测和真实 EP=2 Gloo collective 回归，同时验证不同 expert shard 仍完整保留。
Mixture expert 复用 MCore 的 CPU master-weight 切分和 CUDA model-parallel RNG tracker，避免 qkv rank-axis shard 重复。 新增 CPU Gloo TP2 与 CUDA NCCL TP2 初始化回归，并保留单卡初始化及 LoRA B 零初始化行为。
SGLangEngine 默认关闭 Mixture 自动注入，仅常规 policy rollout actor 显式开启，避免 GenRM 和 teacher 继承 policy adapter。 无 Mixture 配置时清理残留的 router 环境变量，并补充 policy 与辅助 engine 角色隔离测试。
Mixture rollout mapper 尚未定义 MTP adapter site，启动校验在 mtp_num_layers 大于零时直接报告不支持。 补充启用 MTP 的失败测试和 mtp_num_layers=None 的兼容测试。
当前 routing context 只包装 MCore tensor-parallel checkpoint，TE FP8/FP4 full recompute 会绕过 aux-loss 重计算路径。 仅拒绝 FP8/FP4 与 full recompute 的组合，并覆盖 selective 与 BF16 full recompute 的兼容边界。
训练注入和 SGLang rollout mapper 当前只实现 Qwen3 布局，HF 参数校验要求文本 config 的 model_type 为 qwen3。 补充 dense Llama 拒绝与 dense Qwen3 接受测试，并明确既有正向 fixture 的模型类型。
Routing context 使用 qkv_format 区分 bshd 与 thd；bshd 始终将 batch-first mask 转为 sequence-first，避免 B 等于 S 时跳过转置。 补充方阵非对称 mask、packed THD 及既有 DP/PP/CP/TP 路径回归。
SGLang 的 tensor 更新接口按 TP rank 选择序列化 bucket。rollout 使用 TP=1、PP=2 时，两个 pipeline stage 都会读取索引 0，原有 CUDA IPC 句柄只能由首个 stage 所在 GPU 正确打开。 仅对 Mixture-of-LoRA 且 SGLang PP 大于 1 的 colocate 同步改用 host flattened bucket，并临时切换 file_system 共享策略后恢复原值。全参数、单 LoRA 和 PP=1 继续使用原有 CUDA IPC 路径。 增加 host/device 传输与共享策略恢复单测；完成 Qwen3-4B BF16、训练 TP=2、SGLang TP=1/PP=2 的完整 GRPO step，覆盖初始同步、rollout、反向、checkpoint 和训练后同步。
Author
8 月 13 日代码自查及后续修复8 月 13 日对 Task 25 当前实现b363c18进行了一轮代码自查，发现以下问题。 自查发现的问题
对应修改
以上问题分别修复并拆成独立提交。修改完成后，我们重新执行了定向测试、全量 pytest，以及 Mixture、单 LoRA、全参数和 PP 场景的短训练。Task 25 的设计与验收范围见 Issue #220。 SGLang rollout PP=2复测时发现 SGLang 的 tensor 更新接口按
新增单测覆盖 device/host 两条传输分支以及共享策略恢复。Qwen3-4B BF16 colocate 短训练使用训练 step 0 的全局 expert selection share 为 短训练回归
固定 smoke 数据用于控制生成长度和测试耗时，训练、反向和权重同步均为真实执行。Mixture 基线与训练侧 PP=2 另外使用真实 DAPO 数据完成验证。 单 LoRA 短训所用容器的 Megatron-Bridge 缺少 Relax 维护的 GitHub CPU CI 修复上述修改推送后，GitHub Actions 的 Python 3.10、3.11 和 3.12 测试均在同一个 TP 分布式用例失败。
默认无 Megatron 环境下，两个 CPU TP 用例为 最终检查
本轮共增加 12 个独立修复提交，当前最终提交为 |
TP 初始化回归依赖 Megatron Core 的并行线性层实现，但默认 CPU CI 不安装该可选后端。 在三个 TP 初始化测试入口增加函数级 importorskip：无 Megatron 环境只跳过相关用例，完整训练环境仍执行真实前向、反向和初始化检查。
Contributor
|
这一系列 commit 完成了 Mixture-of-LoRA 从训练、权重同步到 SGLang rollout的端到端接入，整体工程质量很高。但下面几处问题需要进一步完善，以及几处架构改进建议需要进一步跟进。 代码问题：
架构改进：
辛苦进一步修复代码问题并适当完善代码架构。 |
hf_weight_iterator_bridge 在模块导入时执行 from megatron.core import mpu，绑定的是导入那一刻 sys.modules 里的对象。同目录的 test_dtype_codes.py 与 test_broadcast_converted.py 会在各自导入阶段装入 megatron.core 桩模块，所以按目录运行 tests/backends/megatron/weight_update/ 时，bridge 绑定到的是桩 mpu，而用例改写的是 megatron.core.mpu，两者并非同一个对象，两个 bridge 用例因此报 AttributeError；单独运行该文件反而正常。 改为用 patch.object 直接替换 bridge 模块上的 mpu 属性，无论它绑定的是真实模块还是桩模块都能命中。打桩返回值与原先完全一致，只是换了作用目标，未改动任何被测代码。 按目录运行 tests/backends/megatron/weight_update/ 由 73 passed、2 failed 变为 75 passed；单文件运行 23 passed。
relax/utils/mixture_lora.py 与 relax/backends/megatron/mixture_lora.py 同名不同包，导入语句和调用栈里难以一眼分辨，而 SGLang 侧还要再加一个同主题模块，混淆只会更严重。 前者改名为 relax/utils/mixture_lora_common.py，只放后端无关的配置、路由计算、参数命名与传输描述，不允许导入任何训练或推理后端；后者改名为 relax/backends/megatron/mixture_lora_modules.py，只放依赖 Megatron 并行状态的模块与 Bridge 注入。两个文件的模块说明同步写清这条边界，其余改动全部是导入路径同步，没有行为变化。 pytest tests/utils/test_mixture_lora_routing.py tests/backends/megatron/test_mixture_lora.py tests/backends/sglang/test_router_registration.py 为 97 passed。
张量并行开启时，专家权重通过按专家切出的视图逐个初始化。视图与 Parameter 共享存储，数值能写回去，但 Megatron 在视图上设置的 tensor_model_parallel 属性会随视图一起消失。结果是 lora_A、lora_B 以及行并行站点上的 router.weight 都保持默认的未标记状态，梯度裁剪把它们当成各 rank 的重复参数，全局梯度范数只统计其中一份，裁剪系数偏大；行并行与列并行的切分轴不同，这个偏差还会随站点类型变化。 初始化完成后显式给 lora_A、lora_B 打标记，router 在构造时按站点类型打标记：列并行站点的 router 是复制参数，标记为重复；行并行站点的 router 沿输入轴切分，每一份都要计入全局范数。已有标记的参数不重复覆盖。 切分轴此前在适配器初始化、sharded checkpoint 的 axis_map、rollout 权重同步三处各写了一份，容易改一处漏两处。统一收敛到 relax/utils/mixture_lora_common.py 的 mixture_lora_tp_partition_dims，三处都从这张表读取，None 表示复制参数。 新增两个单测覆盖行并行与列并行下三类参数的标记结果；分布式用例补充 param_is_not_tensor_parallel_duplicate 断言，确认切分参数确实参与全局梯度范数。8 卡环境下 tests/backends/megatron 三个 Mixture 分布式用例集 24 passed；CPU 侧 tests/backends/megatron/test_mixture_lora.py 与 tests/utils/test_mixture_lora_routing.py 88 passed，权重同步与 SGLang 用例 29 passed。
权重同步有三处结构相同的流水发送循环：先把上一块的引用交给引擎确认，再发下一块，最后统一等待。等待这一步并不是各 rank 对称的，只有 IPC gather 源那个 rank 持有 object ref，引擎侧出错也只会在它上面抛出。base 与 adapter 两条路径的循环直接裸调 ray.get，完全没有把这个失败告诉别的 rank；Mixture 那条只在发布版本号的最后一块之前做了一次失败标志 all-reduce，中间块的异常仍然是就地抛出。结果是某一块传输失败时，其余 rank 察觉不到，会继续走进下一块的集合通信，一直阻塞到分布式超时，日志里先看到的是超时而不是真正出错的那一块。 新增 relax/backends/megatron/weight_update/synchronized_send.py，提供三个原语：run_synchronized_phase 负责阶段级操作的失败同步，raise_on_any_rank_failure 把任意 rank 的本地异常经 gloo all-reduce 变成全体 rank 一起抛出，send_chunks_pipelined 负责流水发送与确认。update_weight_from_tensor 的三处循环统一改为调用 send_chunks_pipelined，中间块、发布版本号前的确认和收尾的最后一块都经过失败同步；阶段级操作改用 run_synchronized_phase，删除原有的 _run_synchronized_weight_update_phase 与随之不再使用的导入。 成功路径的行为没有变化：发送顺序、张量存活范围和每块的 barrier 都与原来一致。变化的是失败路径，从只有 gather 源 rank 抛出、其余 rank 等待超时，变成任意 rank 出错时全体一起抛出同一个异常。 新增 tests/backends/megatron/weight_update/test_synchronized_send.py 覆盖正常发送顺序、单 rank 失败时全体抛出、以及失败后不再发布后续分块；同步用例中的打桩目标同步指向新模块。
路由权重的形状、dtype 和缺键检查原本写在生成器循环里，与 PP 广播、TP all-gather 交替执行。某个 rank 的分片有问题时，它在广播之前就地抛出异常直接退出，而其余 rank 已经进入同一轮集合通信，只能等到分布式超时才失败，日志里最先看到的是超时而不是那个真正出错的张量，排查成本很高。 改为在任何集合通信之前先完成本 rank 全部分片的校验，并借助上一提交的 raise_on_any_rank_failure 把校验结果广播给所有 rank：只要有一个 rank 校验失败，全体 rank 立即一起抛错，不会有人卡在集合通信上。校验通过后再进入原有的广播与 all-gather 流程，检查逻辑与报错信息保持不变，只是拆成 _collect_local_tensors 与 _iter_weight_chunks 两步。 新增两个单测：一个确认缺键和形状不符会在任何广播、all-gather 发出之前抛错，另一个确认本 rank 分片正常但其他 rank 失败时同样会抛错退出。
relax/models/qwen3_mixture_lora/sglang/model.py 里同时放着两类代码：路由注入、前向改写、权重加载这些对任何架构都一样的逻辑，以及 Qwen3 特有的注入点声明。再接入一个模型架构就得整体复制一遍，两份实现随后各自演进，很快就会出现只修了一边的情况。 架构无关的部分移到 relax/models/mixture_lora_sglang.py：运行时配置读取、路由线性层前向、适配器挂载、路由权重加载，以及供各架构复用的 MixtureLoraSGLangModelMixin（包含 TP 必须为 1 的检查，以及把 mixture_lora 权重从常规权重里分流后再交给父类加载）。qwen3_mixture_lora/sglang/model.py 只剩注入点声明与 EntryClass，代码从 213 行降到 22 行。行为和对外类名保持不变。 测试同步拆分：通用逻辑的用例移到 tests/models/test_mixture_lora_sglang.py，tests/models/qwen3_mixture_lora/test_sglang_model.py 只保留 Qwen3 注入点与注册相关断言。
Author
|
感谢建议🫡，目前针对 ad3e145 之后 review 提出的 3 个代码问题和 3 条架构改进建议，逐条做了修改，另外自查时发现两个测试问题，一并处理。本轮共 6 个独立提交。@GUOGUOPOT 代码问题
结构改进
对应修改
关于问题 2 与建议 5 的改动范围问题 2 和建议 5 都与权重发送流程有关，因此 新的发送函数会在每个分块结束后同步一次执行状态，因此各训练进程必须产生相同数量的分块。现有的两个 HF 权重迭代器均满足这一要求：
单元测试和三种模式的短训练均验证了这套流程可以正常完成。
关于问题 1 的 TP 切分规则处理问题 1 时我们发现，适配器初始化、sharded checkpoint 的 该函数按注入位置返回每个参数的切分轴，轴的选择与被注入线性层的并行方式一致：
列并行层的输入完整、输出被切分，因此 有三处需要按同一个轴处理这些参数，任何一处与其他两处不一致都不会报错，只会得到错误结果：
统一之后，这三处都从同一个函数读取切分轴，新增参数或调整注入位置时只需修改一处。 自查补充的修改
测试
短训练测试由于本轮修改会影响参数的 TP 属性，以及全参数、单 LoRA 和 Mixture 共用的权重发送流程，我们在最终提交
三组训练均正常结束， Mixture 在 step 0 的四个 expert 平均选中占比分别为 检查
本轮共 6 个独立提交，当前最终提交为 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
feat(lora): 支持 token 级 Mixture-of-LoRA 训练
Note
完整实验数据、统计脚本和图片保存在公开分支
task25-experiment-artifacts。收尾实验总览见followup/README.md，三组性能统计见three-way-window-2-11.json，actor 阶段显存统计见actor-memory-window-2-11.json。Important
Mixture-of-LoRA 核心功能、CPU/CUDA 回归、200-step 主实验、checkpoint 真实续训，以及全参/单 LoRA/Mixture 性能对照均已完成。
Note
200-step 主实验、checkpoint 续训和三组性能对比基于
feat/task25-mixture-lora@a82eb62，当时的上游基线为main@5b23011。后续通过b363c18将main@050ab04合入开发分支，并使用5d86614补齐 S3 loader 测试 fixture。最新合并版本没有重复运行 200-step 验收，已完成全量 pytest、Task 25 Megatron CUDA 定向测试和三种训练模式的两步端到端回归，结果见下文。验收状态
1312 passed, 15 skipped, 2 failed，两个失败均在当时的 upstream 基线复现。main@050ab04后运行全量 pytest，并分别完成全参、单 LoRA、Mixture-LoRA 的 8 卡 BF16 两步训练。已经确认的 upstream 基线问题：
test_registry_reinforce_plus_plus.py与test_registry_sft.py连续运行时存在 2 个 enum identity failure，已在未包含本 PR 改动的upstream/main@5b23011复现，不属于本 PR 待修内容。改动内容
本 PR 为 Relax 增加 token 级 Mixture-of-LoRA 训练和 colocate SGLang rollout 支持。
--lora-num-experts未设置或等于 1 时继续使用现有单 LoRA 路径。关联设计 Issue：#220。
参与成员：@pophirasawa、@zTonyZhao。
改动清单
relax/utils/arguments.py、relax/utils/megatron_peft_utils.pyrelax/utils/mixture_lora.pyMixtureLoraConfig、state/transport schema、route_topk()、路由统计和 balance loss；提供训练端与 rollout 端共用的RoutedLoRAExecutor接口。relax/backends/megatron/mixture_lora.pyMixtureLoRAExperts、MixtureLoRARouter、MixtureLoRAAdapter和MixtureParallelLinearAdapter；处理 qkv/proj 的 TP/SP 布局、response mask、aux loss 和 checkpoint metadata。relax/backends/megatron/model_provider.py、model.py、loss.pyrelax/backends/megatron/mixture_lora.py、model.pyrelax/models/qwen3_mixture_lora/sglang/model.pySGLangMixtureLoRA；实现权重加载、prefill、decode 和 CUDA graph 下的 routed forward。relax/backends/megatron/weight_update/mixture_lora_sync.py、update_weight_from_tensor.pyrelax/backends/sglang/sglang_engine.py、relax/utils/env.pyscripts/training/text/run-qwen3-4B-mixture-lora-8xgpu.sh、docs/*/guide/mixture-lora.mdtests/utils/*mixture_lora*、tests/backends/*mixture_lora*、tests/models/qwen3_mixture_lora/*本分支相对
main@050ab04新增或修改 37 个文件。第 37 个文件是独立提交5d86614补充的 S3 loader 测试 fixture，用于给伪造的 PEFT 模块提供新增 Mixture 符号。代码、测试和文档均包含在本 PR 中，提交历史按配置、训练端、分布式、checkpoint、rollout、同步、recipe 和修复逐步拆分，便于分别审查。官方验收标准对应关系
--lora-num-experts 4 --lora-rank 16可启用requires_grad，optimizer 只接收可训练参数N=1继续使用原有ParallelLinearAdapter、checkpoint 和 rollout 路径[0.2502, 0.2457, 0.2501, 0.2540]，72 个 site 无单 expert 塌缩查看 36 个变更文件的逐文件说明
docs/.vitepress/config.mtsdocs/en/guide/mixture-lora.mddocs/zh/guide/mixture-lora.mdrelax/utils/arguments.pyrelax/backends/megatron/arguments.pyrelax/utils/megatron_peft_utils.pyrelax/utils/mixture_lora.pyrelax/backends/megatron/mixture_lora.pyrelax/backends/megatron/model_provider.pyrelax/backends/megatron/model.pyrelax/backends/megatron/loss.pyrelax/utils/megatron_bridge_utils.pyrelax/backends/megatron/weight_update/common.pyrelax/backends/megatron/weight_update/hf_weight_iterator_bridge.pyrelax/backends/megatron/weight_update/hf_weight_iterator_direct.pyrelax/backends/megatron/weight_update/mixture_lora_sync.pyrelax/backends/megatron/weight_update/update_weight_from_tensor.pypause/flush/update/continue权重更新流程。relax/utils/env.pyrelax/backends/sglang/sglang_engine.pyrelax/models/qwen3_mixture_lora/__init__.pyrelax/models/qwen3_mixture_lora/sglang/__init__.pyrelax/models/qwen3_mixture_lora/sglang/model.pyscripts/training/text/run-qwen3-4B-mixture-lora-8xgpu.shtests/utils/test_arguments_mixture_lora.pytests/utils/test_megatron_peft_utils.pytests/utils/test_mixture_lora_routing.pytests/backends/megatron/test_mixture_lora.pytests/backends/megatron/test_mixture_lora_arguments.pytests/backends/megatron/test_mixture_lora_distributed.pytests/backends/megatron/test_mixture_lora_checkpoint_distributed.pytests/backends/megatron/test_model_provider_vpp.pytests/backends/megatron/test_sft_chunked_ce.pytests/backends/megatron/weight_update/test_mixture_lora_weight_sync.pytests/backends/sglang/test_mixture_lora.pytests/backends/sglang/test_router_registration.pytests/models/qwen3_mixture_lora/test_sglang_model.py背景
单个 LoRA adapter 对复杂 RL 数据的容量有限。Task 25 希望在 base model 保持冻结的前提下，让多个 LoRA expert 共同参与训练，并由 router 根据 token 动态组合 expert。同时还需要保留单 LoRA 兼容性、支持 colocate GRPO、输出可判断路由是否塌缩的指标，并保证 checkpoint 可以正确保存和恢复。
使用的主要技术
nn.Module与 autogradlinear_qkv、linear_proj外包一层MixtureParallelLinearAdapter，保持 base 模型构建和参数名称稳定。torch.distributedcollectivepause -> flush -> update -> continue流程，不新增 manager、service 或 barrier。数据流
flowchart LR X[Token hidden state] --> Base[冻结的 base linear] X --> Router[可训练线性 router] Router --> Prob[FP32 softmax / temperature] Prob --> TopK[Top-K 选择与重新归一化] X --> Experts[LoRA A/B experts] TopK --> Combine[按选中权重组合 expert 输出] Experts --> Combine Base --> Add[Base 输出 + routed LoRA delta] Combine --> Add TopK --> Stats[response token 路由统计] Stats --> Aux[balance loss] Stats --> Metrics[逐 site 与全局指标]一次 colocate optimizer step 的权重流如下：
flowchart LR Train[Megatron actor] --> Update[更新 expert + router] Update --> Gather[合并 TP shard / 收集 PP stage] Gather --> Transport[Mixture transport schema] Transport --> Pause[暂停 SGLang generation] Pause --> Load[SGLang 加载 expert + router] Load --> Resume[恢复 rollout] Update --> Ckpt[Megatron distributed checkpoint]端到端调用链
relax/utils/arguments.py解析 LoRA 与 router 参数，build_mixture_lora_config()生成MixtureLoraConfig，无效的 expert、Top-K、rank 和并行组合会在模型创建前报错。model_provider.py调用build_mixture_lora_peft()，通过 Megatron Bridge PEFT 在每个目标linear_qkv和linear_proj上注入MixtureParallelLinearAdapter，然后检查 base 冻结状态和 expert/router 可训练参数数量。model.py为当前 microbatch 创建MixtureLoRARoutingContext，其中包含 response mask、objective scale 和本次 forward 的路由记录容器。MixtureLoRARouter生成 logits，route_topk()完成 FP32 softmax、Top-K 选择和权重重新归一化，MegatronDenseRoutedLoRAExecutor计算并组合各 expert 的 LoRA delta。compute_routing_statistics()按 response mask 计算每个 site 的入选份额、概率均值、激活权重和熵，mean_routing_balance_loss()得到当前 microbatch 的 aux loss。_AttachAuxLoss把 aux loss 的梯度接到当前 PP stage 的 adapter 输出。get_microbatch_objective_scale()复用主 loss 的 token/microbatch 缩放信息，保证 dynamic batch、dummy batch 和 recompute 下不重复或遗漏 aux 梯度。pack_mixture_lora_routing_records()将本地记录整理为固定形状 tensor，_reduce_mixture_lora_routing_metrics()按 TP/DP/CP/PP 归约，再将逐 site 和全局指标交给 Metrics Service。MixtureLoraSync收集各 PP stage 的 expert/router，merge_mixture_lora_tp_shards()还原全局逻辑 tensor，然后按TransportTensorSpec分块发送给 rollout。sglang_engine.py在 SGLang 子进程启动前注入 external model package 和 Mixture 配置。SGLangMixtureLoRA安装到 Qwen3 的 qkv/proj，load_sglang_mixture_lora_weights()验证 schema、site、shape 和 dtype 后更新参数。实现方式
启用参数
--lora-num-experts--lora-rank--lora-alpha--lora-router-top-k--lora-router-temperature--lora-router-aux-loss-coef--lora-target-moduleslinear_qkv linear_proj这些值是当前 Qwen3-4B 验收 recipe 的配置，不是全部参数的通用默认值。Mixture 模式要求显式提供 router 参数，
Top-K必须满足1 <= K <= N；column-parallel 路径还要求 rank 可以按 TP size 正确切分。运行 200-step 训练
在具有 8 张可用 GPU 的训练节点上，进入 Relax 仓库根目录，替换模型、数据和输出路径后执行：
cd /path/to/Relax MODEL_PATH=/path/to/Qwen3-4B \ PROMPT_DATA=/path/to/dapo-math-17k.jsonl \ OUTPUT_DIR=/path/to/mixture-lora-output \ NUM_ROLLOUT=200 \ SAVE_INTERVAL=50 \ bash scripts/training/text/run-qwen3-4B-mixture-lora-8xgpu.sh该 recipe 直接设置 4 个 expert、rank 16、Top-K 2、BF16、TP=2 和 SP。actor 与 rollout 都申请 8 张 GPU，并通过
--colocate共用同一组设备，因此实际需要 8 张 GPU。本次实验使用 8×A800 80GB，完成 200/200 step。完整 recipe 位于
scripts/training/text/run-qwen3-4B-mixture-lora-8xgpu.sh：参数与兼容性
当
--lora-num-experts大于 1 时启用 Mixture 模式。--lora-rank、--lora-alpha、--lora-target-modules和--lora-dropout保持原有含义，并新增 expert 数量、router Top-K、router temperature 和 aux loss coefficient。参数解析阶段会拒绝Top-K大于 expert 数量等无效配置。expert 数量未设置或等于 1 时，Relax 继续使用现有
ParallelLinearAdapter注入、checkpoint 导出和 SGLang LoRA 路径，不启用 Mixture external model 和 Mixture 权重传输格式。参数、checkpoint 与传输布局
每个 routed site 都保存三组可训练参数：
mixture_lora.experts.lora_A[num_experts, rank, input_size]mixture_lora.experts.lora_B[num_experts, output_size, rank]mixture_lora.router.weight[num_experts, input_size]linear_qkv遵循 column-parallel 布局：lora_A的 rank 维按 TP 切分，lora_B的 output 维按 TP 切分，router 在 TP rank 上复制并 all-reduce 梯度。linear_proj遵循 row-parallel 布局：lora_A和 router 的 input 维按 TP 切分，lora_B的 output 维按 TP 切分。executor 使用与 Megatron linear 层一致的 gather/reduce 约定恢复完整 routed delta。checkpoint 中的参数 key 与上表名称一致，每个 adapter 还保存
schema_version、site_id、expert 数量、rank、Top-K、temperature、aux coefficient、alpha、target modules、输入/输出尺寸和 dtype。权重同步通过MixtureLoraStateSpec和TransportTensorSpec携带相同的 site、参数类型、全局 shape 和 dtype，SGLang 在写入前逐项校验。训练端
每个 routed projection 包含多组 LoRA A/B expert 参数和一个可训练线性 router。router 为每个 token 计算 expert logits，选取 Top-K expert，对入选权重重新归一化，再组合对应 LoRA 输出。
首版执行器使用 dense 计算，并通过独立 executor 接口与路由语义、参数 schema、checkpoint key 和权重传输名称隔离。以后替换为 grouped sparse kernel 时，可以继续使用相同的模型参数与 checkpoint。
每个 token 的路由过程为：
MixtureLoRARouter使用输入 hidden state 计算 FP32 logits。route_topk()在 FP32 中执行 temperature softmax 和 Top-K。MegatronDenseRoutedLoRAExecutor计算各 expert 的B(A(x))，再按路由权重求和。alpha / rank，与冻结 base linear 的输出相加。首版 dense executor 会计算全部 expert，再将未入选 expert 的权重置零。这保证公式、梯度和分布式布局先稳定下来，训练端 grouped sparse kernel 属于后续性能优化，不影响当前 checkpoint 和 rollout 参数格式。
Mixture 模式下 base model 参数保持冻结。expert 和 router 使用独立参数分类，并作为模型中仅有的可训练参数加入 optimizer。routed module 通过现有 Megatron Bridge PEFT 流程注入，避免改变 base 参数名称和模型构建方式。
router balance loss 使用 expert 入选频率和平均路由概率计算，只统计有效 response token。每个 router 的 aux loss 在其所属 PP stage 本地接入反向传播，复制参数的梯度继续由现有分布式同步处理。
TP/SP 下 routed projection 遵循周围 Megatron 线性层的数据布局。静态 CP 会在 CP group 内汇总路由计数，同时保留各 rank 本地可导的概率贡献。PP 下 routed module、expert 和 router 留在目标 transformer layer 所在 stage，不跨 stage 传输 router tensor。
路由指标
每个 optimizer step 都会记录 balance loss、逐 site 指标和全局指标，包括：
是否发生 expert 塌缩主要根据跨 token 聚合后的激活权重和入选份额判断。逐 token 熵用于观察 router 的选择置信度，不单独作为塌缩门槛。
balance loss 按 routed site 独立计算后取平均。对 expert
e：F_e = selection_count_e / (valid_response_tokens * K)，表示 Top-K 入选份额，作为离散统计量不参与反向传播；P_e是有效 response token 上 Top-K 前 router 概率的均值，保留梯度；N * sum(F_e * P_e)。F_e的分母包含K，因此均匀路由时的 loss 基线不会随 Top-K 改变。SGLang rollout 与权重同步
SGLang external model 与训练端共用路由函数和参数 schema。首次同步发送 base、全部 expert 和 router；后续 optimizer step 只发送 expert 和 router，并复用现有
pause -> update -> resume流程。Mixture 配置会在 SGLang 子进程启动前注入。权重加载会检查 schema version、site、参数类型、全局 shape、dtype 和 Mixture 配置。prefill、decode、CUDA graph capture 和在线权重更新均使用与训练端一致的 token Top-K 语义。
Checkpoint
expert 和 router 进入现有 Megatron distributed checkpoint。恢复时会校验 Mixture 配置和 site schema；TP、PP 不变时，fully reshardable distributed optimizer 支持改变 DP size 后继续恢复。
checkpoint 覆盖模型参数、optimizer、scheduler、iteration、RNG 和数据进度。Mixture 模式不复用标准单 adapter PEFT 导出，因为每个 routed site 包含多个 expert 和 router；现有单 LoRA 导出行为保持不变。
测试
已经完成：
[0.2502, 0.2457, 0.2501, 0.2540]。72 个 routed site 均未塌缩到单一 expert，单个 site 的最大 expert 平均权重为 0.4110。1312 passed, 15 skipped, 2 failed，两个失败均在 upstream5b23011复现。87dfe2f的干净 worktree 上完成全仓库 pre-commit：除 Gitleaks 外的全部 hook 通过；Gitleaks 因离线环境无法初始化远端 hook，改为单独运行仓库自带的 tracked-files 包装脚本，扫描约 9.26 MB，结果为no leaks found。两次检查后 worktree 保持干净。1307 passed, 17 skipped, 5 deselected，耗时 6 分 47 秒。main@050ab04后重新运行完整pytest tests/。测试主体完成1568 passed, 12 skipped；一个分布式测试首次运行遇到随机端口占用，相关文件独立重跑3 passed，因此合并版本的全量测试通过。61 passed，无 skip、无失败，覆盖真实 Bridge PEFT、CUDA forward/backward profiler、分布式 checkpoint、DP/TP/SP/CP/PP 和权重同步。tests/test_s3_model_loader.py完成51 passed。pytest tests/，并与 upstream 对照失败项最新 main 合并验证
b363c18将main@050ab04合入开发分支，三个冲突分别位于模型导出、Tensor 权重同步和 SGLang 启动。合并时同时保留最新 main 的 S3 模型来源、HF/FP8 导出、post hook、不等长 bucket padding 和 CPU 权重备份逻辑，以及 Task 25 的 Mixture 导出分支、分块 weight version 和 external model 配置。5d86614单独修复tests/test_s3_model_loader.py的测试 fixture。该 fixture 使用精简版megatron_peft_utils导入sglang_engine.py，因此需要补充build_mixture_lora_config和is_mixture_lora_enabled两个 stub。两个 stub 固定返回未启用 Mixture，只影响测试替身，不改变 S3、SGLang 或 Mixture 的运行时行为。合并版本的验证结果如下：
pytest tests/1568 passed, 12 skipped；一个用例首次遇到随机端口占用，所属文件重跑3 passed61 passed51 passed两步训练用于确认最新 main 合并后不同训练模式均可启动并完成完整 step，不替代下方基于
a82eb62的 200-step 稳定性和 step 2–11 性能实验。200-step 实验数据
succeeded6d1d751；实际使用的 BF16 recipe 与文档随后提交为317cc44Mixture 200-step 作业的全程单卡峰值显存为 64,084 MiB（8 卡范围 63,554–64,084 MiB），全程 GPU 平均利用率为 83.92%。三组性能对比使用下面单独定义的统一窗口。
全参、单 LoRA 与 Mixture 性能对比
三组使用相同的模型、数据、batch、response length、TP/SP、BF16 和
sglang-mem-fraction-static=0.7。全参和单 LoRA 各运行 12 step，Mixture 使用200-step 主实验的相同下标，统一统计 step 2–11。三组平均 response 长度相差不到
0.3%，窗口内没有 checkpoint 保存。
单 LoRA 与全参在本次十步窗口中的 step time 相差 0.06%，未观察到明显的系统吞吐差异。
当前 dense Mixture 路径比单 LoRA 慢 28.56%，其中 rollout 增加 41.82 秒、actor train
增加 12.99 秒；rollout 占额外 step time 的 66.8%。Mixture 的 Top-K 选择依赖 token，
无法像固定单 LoRA 一样静态合并进 base，但当前训练端和 SGLang rollout 端还会计算全部
四个 expert，因此该开销包含可由 grouped sparse 或融合 kernel 优化的部分。
Actor 阶段根据 TensorBoard 的
actor_train_time、update_weights_time和指标写入时间重建区间，再过滤 5 秒间隔的 NVML 样本。单 LoRA 和 Mixture-LoRA 的 actor 阶段系统
显存均值比全参分别低 30.3% 和 23.8%。三组都设置了
--sglang-mem-fraction-static 0.7，SGLang rollout 会预留相近大小的静态显存池，生成阶段的 KV cache 和 batch 又会把显存推到该容量附近，因此端到端峰值都接近 60–62 GiB。
这个峰值主要反映统一的 rollout 显存配额，用于判断完整 workload 是否会 OOM；actor
阶段统计才用于比较不同训练方式的显存差异。完整方法、边界和逐项数据见
followup/03_吞吐实验.md与
followup/04_三组对比结论.md。Checkpoint 真实续训
在 8×A800 80GB 环境中，从 200-step 主实验保存的 iteration 49 独立快照恢复训练，并继续完成 step 50–52。恢复过程加载了模型参数、fully reshardable optimizer 状态、scheduler、iteration、RNG 和数据状态，没有设置跳过 optimizer/RNG 或重置 optimizer 的参数。
sample_offset=816、sample_index=6528恢复任务在完成 step 52 后由 watcher 主动停止，因此 Ray 状态为
STOPPED；这是预定的短续训结束方式，不是训练错误。测试命令
Mixture 核心功能、并行布局、checkpoint、SGLang 和单 LoRA 回归可以使用以下命令集中运行：
离线容器中的全仓库静态检查使用：
Gitleaks 使用服务器已安装的 v8.24.2，第二条命令扫描全部 tracked files。完整 CPU 可运行回归由
CUDA_VISIBLE_DEVICES='' pytest tests/开始，定位并修复本 PR 引起的 14 个 fixture 问题后，对 3 个需要可见 CUDA 的 checkpoint 执行阶段和 2 个 upstream registry 顺序用例显式 deselect，得到1307 passed, 17 skipped, 5 deselected。全量回归发现与修复
第一次在
317cc44clean worktree 中运行pytest tests/，结果为1293 passed, 17 skipped, 8 failed, 11 errors。逐项单独重跑并与upstream/main@5b23011对比后，结论如下：sglang_engine.py新增了build_mixture_lora_config、is_mixture_lora_enabled和configure_mixture_lora_external_modelimport，旧测试使用的伪模块没有提供这些符号，导致 fixture 在导入被测模块时失败。87dfe2f已补齐对应 stub。87dfe2f在该 fixture 中固定 DP world size 为 1，使测试输入满足真实训练前提。CUDA_VISIBLE_DEVICES=''下无法执行，因为当前 Megatron distributed checkpoint 实现即使保存 CPU tensor，也会在 finalize 阶段调用torch.cuda.current_device()。它们随后已在 GPU 可见环境全部通过。test_registry_reinforce_plus_plus.py会删除并重新导入relax.core.registry，导致后续已收集测试持有旧 enum 引用。相同命令已在未包含 Task 25 改动的upstream/main@5b23011复现，本 PR 不修改该无关问题。最终在
87dfe2fclean worktree 中重跑完整 CPU 可运行测试集，显式 deselect 上述 3 个需要可见 CUDA 的用例和 2 个 upstream 基线用例，结果为：随后在 8×A800 可见 GPU 环境运行完整测试集，结果为：
两个失败为 registry enum identity/order 问题，已使用相同命令在
upstream/main@5b23011复现；Task 25 新增的 3 个 CUDA checkpoint 用例全部通过。兼容性与边界
N=1或不设置--lora-num-experts时不创建 router，不启动 external model，继续走现有单 LoRA 注入、导出和 rollout 路径。linear_qkv与linear_proj上启用；其他模型和 projection 不会静默套用。建议重点审查
site_id、state spec、transport spec 和 checkpoint metadata 是否足够稳定且能在不一致时立即报错。N=1路径是否完全绕开 Mixture model、checkpoint 和 SGLang external model。1307 passed, 17 skipped, 5 deselected）pytest tests/已运行（1312 passed, 15 skipped, 2 failed；2 个失败均在 upstream 复现）61 passed、S3 loader51 passed）变更类型
截图与日志
200-step 训练稳定性
图注：Qwen3-4B DAPO GRPO 使用 BF16 连续完成 200 step，全部 scalar 保持有限值，最后 50 step 的 raw reward 均值为 0.3475。
Expert 路由
图注：最后 50 step 的全局 expert 激活权重为
[0.2502, 0.2457, 0.2501, 0.2540]，72 个 routed site 中单个 expert 的最大平均权重为 0.4110。显存与吞吐
图注：Mixture 实验采样到的单卡峰值显存为 64,084 MiB，最后 50 step 平均耗时为 234.19 秒/step。
Checkpoint 恢复
iteration-50 快照已完成真实续训：完整 distributed checkpoint 和数据状态从 iteration 49 恢复，随后完成 step 50–52。命令、验收点和日志摘要见
followup/02_断点续训.md。全参、单 LoRA 与 Mixture 对比
图注：三组 step 2–11 的逐步耗时、rollout 吞吐和 actor train 吞吐；同一步下三组平均 response 长度相差不到 0.3%。
图注：左图展示 step 2–11 的 actor 阶段单卡系统显存均值变化，中图汇总均值、P95 和峰值，右图对比 actor 阶段峰值与端到端峰值。三组统一设置
--sglang-mem-fraction-static 0.7，端到端峰值主要由 SGLang rollout 静态显存池、KV cache 和生成 batch 决定，因此数值接近；单 LoRA 和 Mixture-LoRA 的 actor 阶段显存均值比全参分别低 30.3% 和 23.8%。单 LoRA 对照所用镜像未应用 Relax 的 Megatron-Bridge PEFT 补丁，因此缺少
create_peft；本次实验临时使用镜像中已有的LoRA类完成构建。该兼容修改未纳入本 PR，使用仓库要求的docker/patch/megatron/20260506-85bced0ae.patch后无需回退。