Prhub

#2840 [multi-lora v2 1/7] LoRA model extension and slot actor

原始 PR 作者 yueming-yuan 合并时间 2026-09-16 16:24 文件变更 12 提交数 1 评论 2 代码增减 +453 / -352

执行摘要

新增多 LoRA 槽位执行器与每租户独立优化器

PR body 给出的目标是 Expose command-level training on MultiLoRATrainRayActor in megatron_utils/lora/actor.py:让 Tinker 风格的网关可以像发命令一样驱动训练(累积梯度的 forward_backward、按槽位下发 Adam 超参的 optim_step、只做打分的 forward_only_logprobs,以及槽位级 load/unload),而不是沿用整批 rollout 驱动的 train 流程。它是堆叠在 #2839 之上的 1/7 连发,后续损失选择(#2841)、单槽位 checkpoint(#2842)、快照发布(#2843)、协议层(#2844)与 runtime(#2845)都依赖这里的命令面与优化器所有权模型。动机的第二层来自 review:nanjiangwill 指出池化优化器上的 reload_model_params() 会顺带把其他活跃槽位的 FP32 master 取整,破坏租户隔离,这直接催生了“每槽位一个独立 SlotOptimizer、随槽位销毁”的设计,也解释了为什么本 PR 连带删除 zero_optimizer_state_for_adapter 及其测试——状态不再跨租户复用,手工复位逻辑失去意义。

建议精读。这是整套 multi-LoRA v2 的地基(1/7),后续 6 个 PR 的损失选择、槽位 checkpoint、快照发布和网关 runtime 都建立在它的命令面与优化器所有权模型上,改动虽然只有 +453/-352,但触及所有 Megatron 训练路径。推荐阅读顺序:先看 megatron_utils/model.py 的抽函数是否严格等价(尤其被丢掉的 no_sync_func 断言),再看 lora/model.py 的 4 个槽位操作函数,最后细读 lora/optimizer.py 的 SlotOptimizer 与 step_slot_optimizers。需要重点确认三个设计决策:(1) 是否接受“pre-collective 失败 = cell 失败”的粗粒度恢复;(2) nonfinite 跳过时直接丢弃梯度是否符合网关的重放语义;(3) 每租户独立优化器的显存增长是否在容量预算内。同时留意 review 点名但未补齐的双槽位/两 rank 回归测试,以及 body 中“无新增测试文件”与实际 diff 的口径差异。

讨论亮点

本 PR 只有两条实质性行内意见,均来自 nanjiangwill、状态均为 COMMENTED,作者未在评论区回复,处理方式只能从最终代码反推:

关于(早期版本的)miles/backends/megatron_utils/lora/executor.py:“optimizer is the outer LayerWiseDistributedOptimizer, so reload_model_params() walks every child optimizer. Each BF16 child copies its model parameter back into its FP32 master. Loading or unloading slot A therefore rounds the FP32 masters for every live slot, changing their later Adam updates. Reload only for child in _slot_children(optimizer, slot) here and in unload_slot, and add a two-slot regression test where slot B’s FP32 master differs from its BF16 model value.”

关于 lora/optimizer.py 的 step_adapter_slots:“_step_one_slot() can fail on one rank during prepare_grads() or batch-size scaling before get_grad_norm_fp32() enters the world all-reduce. That rank catches the exception here and moves to the Gloo outcome gather while healthy ranks block in the norm all-reduce, turning per-slot recovery into a collective timeout. Add rank consensus after local preparation and before the norm collective, or treat any pre-collective failure as cell-fatal.”

结论与决策:第一条被采纳为架构级改动——池化优化器整体退场,SlotOptimizer 每槽位持有自己的 LayerWiseDistributedOptimizer,reload_masters 只作用于 self._inner,租户隔离从“约定”变成结构保证,并连带删除 zero_optimizer_state_for_adapter。第二条采纳了作者给出的第二个选项“pre-collective 失败即 cell-fatal”:step_slot_optimizers 不再捕获异常、注释写明 execution failures invalidate the cell,新测试断言异常向上逃逸且失败槽位不触发 allgather_params,代价是单个槽位故障会拖垮整个 cell。两条意见各自点名要求的回归测试(双槽位 FP32 master 差异、两 rank 在 get_grad_norm_fp32 之前失败)在本次改动中都没有落地,现有测试是单进程的 fake 版本。此外 yushengsu-thu 给了 APPROVED(checked),claude[bot] 只是仓库的 review 配置提示,nanjiangwill 的两条意见没有收到作者的文字回复。

实现拆解

  1. 抽出共享执行原语(miles/backends/megatron_utils/model.py):把 train_one_step 里的 pipeline FWD/BWD 主体抽成新的 run_forward_backward_pass(args, dumper_phase_util, data_iterator, model, num_microbatches, num_rollouts),只返回 losses_reduced、不碰优化器;把 train() 内联的 DDP hook 配置抽成 setup_train_iteration_config(args, model, optimizer, disable_optimizer),统一设置 grad_scale_func、no_sync_func、grad_sync_func、param_sync_func 与 finalize_model_grads_func。train_one_step 与 train 改为调用这两个符号,行为等价,使多 LoRA 命令路径能复用同一套并行与梯度归约语义;注意抽取时去掉了原 overlap_grad_reduce 分支里的 assert config.no_sync_func is None 前置校验。

  2. 优化器所有权改为每槽位一个实例(lora/optimizer.py):删除 build_multi_lora_optimizer、step_adapter_slots、zero_adapter_slot_grads 这套“池化 + 按 slot 索引 chained children”的实现,新增 validate_multi_lora_optimizer_args,把 use_distributed_optimizer、fp16、非 Adam、optimizer_cpu_offload、stream_optimizer_state_to_disk、rematerialize_param_from_master_weight 等不支持项前移为启动期断言;setup_model_and_optimizer 的多 LoRA 分支现在只做校验并直接 return model, None, None,即没有池化优化器与 scheduler。新增 SlotOptimizer 类:构造时由 _build_slot_base_optimizers 在 _only_slot_trainable 上下文里为单槽位建 base optimizer(临时把 config.bf16 置 False 以拿到未包装的 torch optimizer),再包一层 LayerWiseDistributedOptimizer,并把每个 child 的 grad_stats_parallel_group 置 None;对外提供 apply_adam_params、prepare_grads、clip_and_step、allgather_params、reload_masters、zero_grads。新增 step_slot_optimizers 统一编排“按固定顺序 apply → prepare → step → zero → allgather”,注释写明执行期异常一律让 cell 失败。

  3. 新增槽位操作层(lora/model.py,新增 73 行):run_forward_backward 通过 get_data_iterator 取数据、setup_train_iteration_config(disable_optimizer=True) 配置本轮、reset_grad_metadata_keep_grads 只重置 DDP 记账而保留梯度累积窗口,并断言“一个 work unit 就是一次 FWD/BWD”,仅在 PP 末级聚合 loss;optim_step 只把 adam_params_by_slot 中出现的槽位交给 step_slot_optimizers(batch size 1,不做额外归一化,归一化交给客户端 loss 权重);load_slot 调 bridge 的 init_adapter_slot 后创建 SlotOptimizer 并清零梯度;unload_slot 调 clear_adapter_slot,由调用方丢弃实例、优化器状态随之消亡。

  4. Actor 与 worker 装配(lora/actor.py、actor.py、ray/specs/train.py、ray/train/group.py):新增 MultiLoRATrainRayActor,在 _init_training_state 里建 slot_optimizers 注册表,并实现 5 个带 @with_logs 与 _heartbeat.bump() 的命令方法;forward_only_logprobs 只在 PP 末级返回,并把结果 ray.put 成 CPU tensor 盒子。基类 MegatronTrainRayActor 把权重更新器初始化抽成 _init_training_state() 钩子,多 LoRA 子类覆写它从而不初始化普通 WeightUpdater。_compute_spec_trainer 仅在 train_backend == megatron、role == actor 且 is_multi_lora_enabled(args) 时替换 worker_class。miles/ray/train/group.py 新增 _execute_slots 与 6 个转发方法,要求 Tinker 训练组恰好一个存活 cell。

  5. 测试、清理与配套:删除 tests/fast/backends/megatron_utils/test_multi_lora_slot_cleanup.py(被测的 zero_optimizer_state_for_adapter 已随设计变更移除),新增 tests/fast/backends/megatron_utils/test_step_adapter_slots.py,用 _FakeSlotOptimizer 覆盖“未知 step 失败在报告成功前逃逸且不 all-gather”“nonfinite grad norm 跳过步进”“prepare 阶段失败逃逸”;扩展 tests/fast/ray/specs/test_train.py 的 worker class 选择用例并保留 critic 的后端选择;miles/utils/multi_lora.py 为 worker 选择提供 is_multi_lora_enabled 判定。body 声明无新增其它测试文件、pre-commit 与 git diff --check 通过、fresh CPU/GPU CI pending,验证与限制记录在 #2846 与 #3198。

文件 模块 状态 重要度
miles/backends/megatron_utils/lora/optimizer.py 优化器 modified 8.84
miles/backends/megatron_utils/lora/model.py 模型执行 added 8.78
miles/backends/megatron_utils/model.py 训练主循环 modified 8.4
miles/backends/megatron_utils/lora/actor.py 训练执行器 added 8.19
miles/backends/megatron_utils/actor.py 训练基类 modified 6.99
miles/ray/train/group.py 训练组 modified 7.22
miles/ray/specs/train.py 进程规格 modified 5.97
tests/fast/backends/megatron_utils/test_step_adapter_slots.py 优化器测试 added 6.83
miles/backends/megatron_utils/lora/slots.py 槽位管理 modified 6.49
tests/fast/backends/megatron_utils/test_multi_lora_slot_cleanup.py 槽位测试 removed 6.65
tests/fast/ray/specs/test_train.py 规格测试 modified 4.92
miles/utils/multi_lora.py 开关工具 modified 4.16

关键符号

MultiLoRATrainRayActor.forward_backward MultiLoRATrainRayActor.optim_step MultiLoRATrainRayActor.forward_only_logprobs MultiLoRATrainRayActor.load_slot MultiLoRATrainRayActor.unload_slot SlotOptimizer.__init__ SlotOptimizer.clip_and_step SlotOptimizer.zero_grads step_slot_optimizers validate_multi_lora_optimizer_args run_forward_backward run_forward_backward_pass setup_train_iteration_config MegatronTrainRayActor._init_training_state TrainRayGroup._execute_slots

关键源码片段

miles/backends/megatron_utils/lora/model.py core-logic

新增的槽位操作层,是 actor 命令与底层 optimizer/ 模型执行之间唯一的粘合点,定义了梯度累积、按槽位步进与槽位装载语义。

"""Multi-LoRA slot operations extending the shared Megatron model execution."""
​
​
def run_forward_backward(args, batch_id: int, model, rollout_data) -> dict:
    """一个 work unit 就是一次 forward/backward;梯度在此累积,不做 optimizer step。"""
    data_iterator, num_microbatches = get_data_iterator(args, model, rollout_data)
    assert len(num_microbatches) == 1, "a work unit is a single forward/backward pass"
​
    for iterator in data_iterator:
        iterator.reset()
    for model_chunk in model:
        model_chunk.train()
    # disable_optimizer:bf16 下 loss scaling 是恒等的,这次 pass 不需要 optimizer
    setup_train_iteration_config(args, model, None, disable_optimizer=True)
    # 只重置 DDP 的 per-iteration 记账,保留跨 batch 的梯度累积窗口
    reset_grad_metadata_keep_grads(model)
​
    dumper_phase_util = DumperMegatronUtil(args, model, DumperPhase.FWD_BWD, rollout_id=batch_id)
    losses_reduced = run_forward_backward_pass(
        args, dumper_phase_util, data_iterator, model, num_microbatches[0], num_rollouts=None
    )
    dumper_phase_util.finalize(model)
​
    if get_parallel_state().is_pp_last_stage:
        return aggregate_train_losses(losses_reduced, None)
    return {}
​
​
def optim_step(args, slot_optimizers: dict, adam_params_by_slot: dict) -> dict[int, dict]:
    # batch size 为 1:梯度按累积值步进,归一化交给客户端的 loss 权重
    stepped = {slot: slot_optimizers[slot] for slot in adam_params_by_slot}
    return step_slot_optimizers(stepped, adam_params_by_slot, clip_grad=args.clip_grad)
​
​
def load_slot(args, model, slot: int, rank: int, alpha: float) -> "SlotOptimizer":
    from megatron.bridge.peft.multi_lora_layers import init_adapter_slot
​
    init_adapter_slot(model, slot, rank=rank, alpha=alpha) # 在模型上挂载该槽位的 adapter
    slot_optimizer = SlotOptimizer(args, model, slot) # 优化器状态随槽位一起新建
    slot_optimizer.zero_grads()
    return slot_optimizer
​
​
def unload_slot(model, slot_optimizer) -> None:
    """调用方随后会丢弃该实例,其状态随之消亡。"""
    from megatron.bridge.peft.multi_lora_layers import clear_adapter_slot
​
    clear_adapter_slot(model, slot_optimizer.slot)
    slot_optimizer.zero_grads()
miles/backends/megatron_utils/lora/actor.py entrypoint

新增 MultiLoRATrainRayActor,是命令级训练接口(forward_backward/optim_step/forward_only_logprobs/load_slot/unload_slot)的唯一入口,也是后续 Tinker 网关调用的落点。

class MultiLoRATrainRayActor(MegatronTrainRayActor):
    """命令级训练 actor:每个命令一次远程调用,梯度跨调用累积。"""
​
    def _init_training_state(self) -> None:
        # 覆写基类钩子:多 LoRA 没有池化优化器,只有按槽位注册的 SlotOptimizer
        self.slot_optimizers: dict[int, SlotOptimizer] = {}
​
    @with_logs
    def forward_backward(self, batch_id: int, rollout_data_ref: Box) -> dict:
        self._heartbeat.bump()
        with ExitStack() as stack:
            rollout_data, store_get_result = get_rollout_data(self.args, rollout_data_ref)
            stack.enter_context(store_get_result)
            return lora_model.run_forward_backward(self.args, batch_id, self.model, rollout_data)
​
    @with_logs
    def optim_step(self, adam_params_by_slot: dict[int, dict]) -> dict[int, float]:
        # 只有请求里出现的槽位会被步进,其余槽位的梯度继续保留
        return lora_model.optim_step(self.args, self.slot_optimizers, adam_params_by_slot)
​
    @with_logs
    def forward_only_logprobs(self, batch_id: int, rollout_data_ref: Box) -> Box | None:
        self._heartbeat.bump()
        with ExitStack() as stack:
            rollout_data, store_get_result = get_rollout_data(self.args, rollout_data_ref)
            stack.enter_context(store_get_result)
            data_iterator, num_microbatches = get_data_iterator(self.args, self.model, rollout_data)
            outputs = self.compute_log_prob(data_iterator, num_microbatches, rollout_id=batch_id)
        if not get_parallel_state().is_pp_last_stage: # 非 PP 末级不返回结果
            return None
        # 只把 CPU tensor 放进 Ray 对象存储,避免序列化 GPU 张量
        return Box(ray.put({key: [t.cpu() for t in tensors] for key, tensors in outputs.items()}))
​
    @with_logs
    def load_slot(self, slot: int, rank: int, alpha: float) -> None:
        self.slot_optimizers[slot] = lora_model.load_slot(self.args, self.model, slot, rank, alpha)
​
    @with_logs
    def unload_slot(self, slot: int) -> None:
        # pop 出来后交给 unload_slot,实例被丢弃时它的优化器状态一起消亡
        lora_model.unload_slot(self.model, self.slot_optimizers.pop(slot))

评论区精华

槽位 load/unload 会改写其他租户的 FP32 master 正确性

nanjiangwill 指出当时的实现把外层 LayerWiseDistributedOptimizer 作为 reload 目标:reload_model_params() 会遍历所有 child,每个 BF16 child 把模型参数写回自己的 FP32 master,于是加载或卸载槽位 A 会把所有活跃槽位的 FP32 master 取整,改变它们后续的 Adam 更新;建议只对 _slot_children(optimizer, slot) 做 reload,并补一个“槽位 B 的 FP32 master 与其 BF16 模型值不一致”的双槽位回归测试。

结论:采纳并升级为架构级改动:池化优化器被删除,SlotOptimizer 每个槽位持有自己的 LayerWiseDistributedOptimizer,reload_masters 只作用于 self._inner,连带删除 zero_optimizer_state_for_adapter 及其测试;但评论要求的双槽位回归测试未加入。 · 已解决(测试缺口)

预集合通信失败导致 per-slot all-reduce 死锁 正确性

nanjiangwill 指出 step_adapter_slots 的 _step_one_slot() 可能在 prepare_grads() 或 batch-size 缩放阶段单 rank 失败,该 rank 走到 Gloo 结果汇聚,而健康 rank 阻塞在 get_grad_norm_fp32() 的 world all-reduce 上,把 per-slot 恢复变成集合通信超时;建议在本地准备之后、norm 集合通信之前加入 rank consensus,或把任何 pre-collective 失败视为 cell-fatal,并补两 rank 失败测试。

结论:选择了 cell-fatal 路线:step_slot_optimizers 不再捕获异常,docstring 写明 execution failures invalidate the cell,新测试断言异常在报告成功前逃逸且失败槽位不 allgather_params;代价是单个槽位故障会拖垮整个 Tinker cell,且两 rank 测试未落地。 · 已解决(选择 cell-fatal,测试缺口)

review 点名要求的回归测试缺失 测试

两条意见各自要求了具体测试(双槽位 FP32 master 不一致、两 rank 在 get_grad_norm_fp32 之前失败);本 PR 只提供了基于 _FakeSlotOptimizer 的单进程测试 test_step_adapter_slots.py,并删除了 test_multi_lora_slot_cleanup.py,同时 body 声明“未新增测试文件”、fresh CPU/GPU CI pending。

结论:槽位隔离与失败传播目前只被 fake 单测覆盖,真实 LayerWise/collective 行为没有回归保护,需要在后续 PR 或专门测试中补齐;body 的测试声明与 diff 口径也建议向作者确认。 · 未解决

风险与影响

  1. 共享路径回归面大:miles/backends/megatron_utils/model.py 是所有 Megatron 训练(整模、单 LoRA、critic)都会经过的文件,本次把 train_one_step 与 train 的代码搬家;抽取后的 setup_train_iteration_config 丢掉了原实现里的 assert config.no_sync_func is None(overlap_grad_reduce 场景),等于移除了“自定义 no_sync_func 不该被覆盖”的守卫,属于顺带的行为收窄,需要回归确认不会掩盖配置错误。
  2. 集合通信强耦合:step_slot_optimizers 依赖所有 rank 使用相同的 sorted(adam_params_by_slot) 顺序与相同槽位集合,否则会在 get_grad_norm_fp32 的 world all-reduce 上错位或卡死;当前没有 rank consensus,任何 pre-collective 失败都让 cell 失败——这是 review 两个选项中的一个,但恢复粒度确实变粗,单个坏槽位会拖垮整个 Tinker cell。
  3. 梯度处理语义:clip_and_step 在 norm 非有限时返回 skipped_nonfinite,随后 step_slot_optimizers 仍会对该槽位 zero_grads(),即异常 batch 的累积梯度被直接丢弃;上游必须把该 outcome 当作需要重放或显式失败来消费,否则数据会被静默吞掉,这一点要在 #2844/#2845 的网关实现里重点核对。
  4. 显存与启动约束:每个活跃租户各持一套 LayerWise 优化器与 FP32 master,显存随槽位数线性增长;validate_multi_lora_optimizer_args 又把 use_distributed_optimizer、optimizer_cpu_offload、stream_optimizer_state_to_disk、rematerialize_param_from_master_weight 等省显存/省内存开关全部禁掉,多租户并发下的显存上限很可能先于算力成为瓶颈。
  5. 单 cell 假设:miles/ray/train/group.py 的 _execute_slots 用 (cell,) = self._cells,多 LoRA 训练组必须恰好一个存活 cell,否则直接抛 ValueError/RuntimeError,与常规 rollout 组多 cell 并发模型不同,部署时需要显式约束。
  6. 测试与 CI 缺口:新增测试全部基于 _FakeSlotOptimizer,不覆盖真实 LayerWise/collective 行为;review 点名要求的双槽位与两 rank 回归测试缺失;body 声明 fresh CPU/GPU CI pending,同时“未新增测试文件”的表述与 diff 中新增 test_step_adapter_slots.py 不完全一致,建议向作者确认测试落盘状态。
  7. 小问题:MultiLoRATrainRayActor.optim_step 的返回类型标注是 dict[int, float],而 lora_model.optim_step 实际返回 dict[int, dict](grad_norm 或 skipped_nonfinite),标注与真实契约不一致,容易误导后续网关调用方。

对用户与上层服务:多 LoRA 服务首次具备“命令式训练”接口,Tinker 网关(#2844–#2849)可以按请求驱动 forward_backward/optim_step,按租户下发 Adam 超参;普通整模与单 LoRA 训练的启动方式、参数与行为不变,但仍会经过被重构的 model.py 代码路径,因此回归验证需要覆盖非多 LoRA 场景。

对系统:多 LoRA 模式的启动校验明显收紧(必须 bf16、纯 Adam、关闭 dist opt 与 offload 系列开关),不满足会在初始化阶段直接断言失败,属于有意为之的显式失败;每个槽位的优化器与 FP32 master 使显存占用随活跃租户数线性增长,容量规划需要按最大并发槽位数预留。

对团队与运维:优化器状态的所有权从“池化 + 手工复位”变成“随槽位生命周期”,后续按槽位 checkpoint(#2842)、不可变快照发布(#2843)可以直接复用 slot_optimizers 注册表;但单 cell 约束、nonfinite 跳过的重放策略、以及 cell 级失败恢复粒度都需要网关侧与运维一起对齐。

影响程度:触及 Megatron 后端核心训练路径(优化器构造、FWD/BWD 抽函数、actor 初始化钩子)并新增一个模块,属于跨模块重构叠加新特性地基,影响面偏大;对非 megatron 后端与不启用 multi-LoRA 的部署无功能影响。

核心训练路径重构 集合通信顺序强耦合 槽位故障拖垮整个 cell 显存随活跃租户线性增长 缺少双槽位与多 rank 回归测试 合并前 fresh CI 未跑完

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接,后续同步到相关引用后会出现在这里。

完整报告

参与讨论