Prhub

#2843 [multi-lora v2 4/7] Publish immutable adapter snapshots to disk

原始 PR 作者 yueming-yuan 合并时间 2026-09-16 16:25 文件变更 14 提交数 1 评论 5 代码增减 +141 / -44

执行摘要

新增 LoRA 适配器不可变磁盘快照发布路径

PR body 给出的目标是「Publish Tinker sampler adapters through WeightPublisher」「Every trainer rank joins the shared HF iterator's gathers; rank 0 writes the selected slot's rank-trimmed tensors, adapter configuration, and metadata into an immutable directory」,并强调「an existing version cannot be overwritten」。旧路径要求 trainer 主动把 adapter push 给 SGLang,依赖 update window 与 engine session,检查点 IO 失败还会被折成结果字典静默继续。Tinker 网关需要推理引擎加载与训练一致的冻结 base,且在没有 trainer push 的情况下也能被 server-cell 生命周期标记为 ready,因此本 PR 把「发布」重写成一次磁盘事务,并为引擎侧的按需加载提供不可变目录。

值得精读。三个最值得关注的设计点:一是 joyliu-q 发现的生成器惰性求值 bug 及其修复方式(把循环变量改为函数形参绑定,而非依赖 itertools.chain 的求值时机),这是 Python 里极易复现的陷阱;二是 materialize_adapter 的集合通信契约——所有 rank 必须走同一路径、只有 writer 产出数据;三是 staged-push 被整体删除后,失败与重试语义转移到了哪里(#2845 的 runtime 与 SGLang on-demand loading),以及删除 slot_lora_name 后全仓引用是否已清理。如果所在团队还在用旧的多 LoRA push 行为,需要先确认配置不再显式设置 update_weights/model_path。

讨论亮点

1) joyliu-q 在 hf_weight_iterator/__init__.py 上给出最小复现,指出生成器惰性求值会让多个 adapter 共享最后一个 lora_name:消费发生在循环之后,输出 ['B:w1', 'B:w2', 'B:w1', 'B:w2'] 而不是 ['A:w1', 'A:w2', 'B:w1', 'B:w2'],并附 Go loopvar-preview 作为更详细解释。最终版本把拼接抽成局部函数 prefixed_units(lora_name, adapter)(形参绑定),并补 test_each_adapter_keeps_its_own_prefix 回归。
2) nanjiangwill 在 updater.py 上指出 _registered_adapters 在 staged stream 与 commit 完成前就写入缓存:失败重试会命中缓存跳过注册、只发 begin/end;复用已发布名字还会跳过 defer_publish 的新鲜度校验,建议不要缓存、让 SGLang 自己保证唯一性并补失败重试/重名测试。最终版本选择整体删除 Tinker staged-push 路径。
3) nanjiangwill 在 actor.py 的 push_slot 上指出 wait_futures() 只捕获传输层异常,SGLang 可以返回 HTTP 200 但 body 为 success: false,这些结果在广播成功前被丢弃,建议走 check_weight_sync_results(..., is_lora=True)。该 push 路径同样被删除。
4) nanjiangwill 在 group.py 的非阻塞生命周期问题:失败热推送的恢复假设 abort 让引擎回到干净状态,但 EngineWeightUpdateSession.__exit__() 会吞掉 abort 失败且 _updatable_engines() 仍标记 ready,可能让后续同名磁盘回填失败;建议补 abort 失败测试或明确该场景需要替换引擎。最终版本移除了 _updatable_engines/update window,恢复语义转移到 #2845 的 runtime,本 PR 未给出引擎替换策略。
5) yushengsu-thu 建议对照 sglang 的 #38165 与 miles 的 #3127(immutable adapter versions alongside generation)的实现方式;随后给出 APPROVED。

实现拆解

1. 新增发布器抽象

  • 新增 miles/backends/training_utils/weight_update/snapshot_publisher.py:WeightPublisher.__init__ 断言 iterator.placement.is_full_gather(rank 0 必须持有完整 adapter),publish_adapter(AdapterSpec, path, metadata) 内定义 write_shards(tmp_dir),由 rank 0 把 materialize_adapter 的结果搬到 CPU、经 safetensors.torch.save 序列化,连同 adapter_config.json(self._adapter_config | {"r": ..., "lora_alpha": ...})写入临时目录,最终交给 write_checkpoint_dir(path, write_shards, metadata=..., overwrite=False) 原子发布。
  • 关键点:所有 rank 都必须进入 materialize_adapter(内部是集合通信),只有 rank 0 落盘;overwrite=False 保证已存在版本不可覆盖。

2. 改造 HF 权重迭代器契约

  • miles/backends/training_utils/weight_update/hf_weight_iterator/__init__.py:WeightUpdatePlacement 新增 is_full_gather property;iter_hf_weights 里的前缀拼接抽成局部函数 prefixed_units(lora_name, adapter);新增 materialize_adapter(adapter, materialize=...),返回裸 HF 名字的 {hf_key: tensor},绕开 pack_units_by_size 的传输分桶——这正是 safetensors 需要的键形态。
  • 抽象钩子 _iter_hf_adapter_units 去掉 lora_name 形参、改为产出裸 hf_name,前缀职责上移到基类;Megatron 实现 miles/backends/megatron_utils/update_weight/hf_weight_iterator.py 同步改签名。这是对后续 rollout 侧(#2844/#2845)的接口收口。

3. 训练侧接入与导出边界

  • miles/backends/megatron_utils/lora/actor.py:_init_training_state() 里用 get_hf_weight_iterator(args, self.model, required_placement=WeightUpdatePlacement(gather_pp=True), model_name=..., quantization_config=...) 加 build_lora_sync_config(args) 构造 self.weight_publisher,与传输层完全解耦(无 update window、无 engine session)。
  • 新增 export_slot(slot, rank, alpha, path, metadata),带 @with_logs 与 self._heartbeat.bump(),直接把异常抛给 TrainerCell,不再转成结果字典;miles/ray/train/group.py 增加同名 async RPC,通过 _execute_slots 广播到所有 rank。

4. 引擎侧与配置约束

  • miles/backends/sglang_utils/sglang_config.py:ModelConfig.resolve 在 is_multi_lora_enabled(args) 时断言 update_weights is not True 与 effective_model_path == args.hf_checkpoint,随后强制 update_weights = False;原来的 update_weights is None 分支保留给普通路径。
  • miles/backends/sglang_utils/sglang_engine.py:多 LoRA 时拒绝 load_format == "dummy";max_loaded_loras 未显式设置时默认 2 * max_loras_per_batch(含 TODO:动态分配),显式 CLI/group 覆盖仍保留。

5. 清理与配套

  • 删除 Tinker staged-push 路径:weight_update/updater.py 恢复为纯 base 同步(连带单 LoRA adapter);相关 engine-session 错误类、update-window 改动、staged-publication 专用测试一并移除,miles/utils/multi_lora.py 的 __all__ 删掉 slot_lora_name 并重写若干断言文案。
  • 新增公共写入助手 save_adapter_to_disk(out_dir, config, tensors)(miles/utils/lora.py,safetensors 延迟 import,避免不碰权重的路径背依赖)。
  • 测试:tests/fast/backends/sglang_utils/test_sglang_config.py 对 multi_lora=True/False 参数化并断言 update_weights is not multi_lora;tests/fast/backends/training_utils/weight_update/test_hf_weight_iterator.py 新增 test_each_adapter_keeps_its_own_prefix 锁住前缀回归。文档 docs/advanced/lora.md 同步更新冻结 base 与不可变目录说明。
  • PR body 提示:发布器的参数化回归(进程在链接替换前后退出、随后重试)落在 #2849 的 checkpoint 测试文件里,不在本 PR 内。
文件 模块 状态 重要度
miles/backends/training_utils/weight_update/snapshot_publisher.py 权重发布 added 8.41
miles/backends/training_utils/weight_update/hf_weight_iterator/__init__.py 权重迭代 modified 8.05
miles/backends/megatron_utils/lora/actor.py 槽位训练 modified 6.97
miles/backends/sglang_utils/sglang_config.py 引擎配置 modified 6.11
miles/backends/sglang_utils/sglang_engine.py 推理引擎 modified 5.84
miles/utils/multi_lora.py 适配器配置 modified 6.08
miles/backends/megatron_utils/update_weight/hf_weight_iterator.py 权重导出 modified 5.66
miles/ray/train/group.py 训练编排 modified 5.73
miles/backends/training_utils/weight_update/updater.py 权重同步 modified 5.23
miles/utils/lora.py 适配器落盘 modified 5.97
tests/fast/backends/training_utils/weight_update/test_hf_weight_iterator.py 迭代器测试 modified 5.38
tests/fast/backends/sglang_utils/test_sglang_config.py 配置测试 modified 5.18
docs/advanced/lora.md 适配器文档 modified 1.96

关键符号

WeightPublisher.__init__ WeightPublisher.publish_adapter HfWeightIteratorBase.materialize_adapter HfWeightIteratorBase.iter_hf_weights HfWeightIteratorBase._iter_hf_adapter_units MultiLoRATrainRayActor._init_training_state MultiLoRATrainRayActor.export_slot TrainGroup.export_slot save_adapter_to_disk ModelConfig.resolve

关键源码片段

miles/backends/training_utils/weight_update/snapshot_publisher.py core-logic

本 PR 的核心新增文件:定义 WeightPublisher,把「一个 slot 的 adapter 落成引擎可加载的不可变目录」这一件事封装成单次原子发布事务,是整条磁盘快照路线的基础。

# 新增文件:miles/backends/training_utils/weight_update/snapshot_publisher.py
import json
from pathlib import Path
​
import safetensors.torch
import torch
import torch.distributed as dist
​
from miles.backends.training_utils.checkpoint_io import write_checkpoint_dir
from miles.backends.training_utils.weight_update.hf_weight_iterator import HfWeightIteratorBase
from miles.utils.multi_lora import AdapterSpec
​
​
class WeightPublisher:
    def __init__(self, iterator: HfWeightIteratorBase, adapter_config: dict) -> None:
        # 发布前提:rank 0 必须持有完整 adapter。PP/TP/EP 任一维没有 gather 时,
        # rank 0 手里只有分片,写出的 safetensors 会缺权重,所以在构造期就 fail fast。
        assert iterator.placement.is_full_gather, "publishing requires the full adapter on rank 0"
        self._iterator = iterator
        self._adapter_config = adapter_config
​
    @torch.no_grad()
    def publish_adapter(self, adapter: AdapterSpec, path: str, metadata: dict | None = None) -> None:
        # 所有 rank 都必须进入 materialize_adapter,因为里面的 gather 是集合通信;
        # 只有 rank 0(is_writer)把 tensor 搬到 CPU 并落盘,其余 rank 走 materialize=False
        # 分支:仍然参与通信,但返回空字典。
        is_writer = dist.get_rank() == 0
​
        def write_shards(tmp_dir: Path):
            tensors = {
                name: tensor.detach().contiguous().cpu()
                for name, tensor in self._iterator.materialize_adapter(adapter, materialize=is_writer).items()
            }
            data = safetensors.torch.save(tensors) if is_writer else None
​
            if is_writer:
                # r / lora_alpha 直接取自 AdapterSpec,保证落盘的缩放系数与训练侧 slot 配置一致,
                # 而不是依赖训练启动参数里可能被改写过的默认值。
                config = self._adapter_config | {"r": adapter.rank, "lora_alpha": adapter.alpha}
                (tmp_dir / "adapter_config.json").write_text(json.dumps(config))
                (tmp_dir / "adapter_model.safetensors").write_bytes(data)
​
        # overwrite=False 配合临时目录 + 原子改名:目录要么是完整的新版本,要么保持旧版本,
        # 引擎侧永远读不到写了一半的 adapter,且已存在的版本号无法被覆盖。
        write_checkpoint_dir(path, write_shards, metadata=metadata, overwrite=False)
miles/backends/training_utils/weight_update/hf_weight_iterator/__init__.py core-logic

迭代器契约改造的落点:新增 is_full_gather 与 materialize_adapter,前缀拼接上移到 prefixed_units,_iter_hf_adapter_units 去掉 lora_name;同时是 joyliu-q 报告的多 adapter 前缀串扰 bug 所在文件与修复位置。

# miles/backends/training_utils/weight_update/hf_weight_iterator/__init__.py(改造后)
@dataclasses.dataclass(frozen=True)
class WeightUpdatePlacement:
    gather_pp: bool
    gather_tp: bool = True
    gather_ep: bool = True
​
    @property
    def is_full_gather(self) -> bool:
        # WeightPublisher 用它判断「rank 0 手上的 adapter 是否完整」;
        # 未来若某后端只做部分 gather,构造期就会断言失败而不是写出残缺权重。
        return self.gather_pp and self.gather_tp and self.gather_ep
​
​
class HfWeightIteratorBase(ABC):
    def iter_hf_weights(self, weights, *, include_base=True, adapters=(), materialize=True):
        def prefixed_units(lora_name, adapter):
            # 关键修复:把 lora_name 作为函数形参绑定,而不是在 for 循环里用生成器表达式
            # 捕获循环变量。生成器是惰性求值的,等到下游消费时 lora_name 已经是最后一个值,
            # 多个 adapter 会全部顶着同一个前缀(joyliu-q 在 review 中给出的复现就是这个)。
            for unit in self._iter_hf_adapter_units(adapter, materialize=materialize):
                yield [(f"{lora_name}:{name}", tensor) for name, tensor in unit]
​
        hf_param_units = self._iter_hf_param_units(weights, materialize=materialize) if include_base else iter(())
        for lora_name, adapter in adapters:
            hf_param_units = itertools.chain(hf_param_units, prefixed_units(lora_name, adapter))
        atomic_update_groups = self._hf_atomic_update_groups() if include_base and materialize else []
        hf_param_units = assemble_atomic_update_groups(hf_param_units, atomic_update_groups)
        yield from pack_units_by_size(hf_param_units, self.args.update_weight_buffer_size)
​
    def materialize_adapter(self, adapter, *, materialize: bool = True) -> dict[str, torch.Tensor]:
        """One adapter as ``{hf_key: tensor}``, skipping the transport bucketing.
        Collective: ``materialize=False`` joins the gathers, returns {}."""
        # 绕过 pack_units_by_size 的分桶,直接得到裸 HF 名字的字典,
        # 正好对应 adapter_model.safetensors 的键名;
        # non-writer rank 传 materialize=False:参与集合通信但返回空字典。
        return {
            name: tensor
            for unit in self._iter_hf_adapter_units(adapter, materialize=materialize)
            for name, tensor in unit
        }
​
    @abstractmethod
    def _iter_hf_adapter_units(self, adapter, *, materialize: bool):
        # 契约变化:不再接收 lora_name、不再负责加前缀,只产出 rank-trimmed 的裸 hf_key。
        # 前缀拼接连同 adapter 归属判断一起收归基类,rollout 侧的命名由调用方决定。
        raise NotImplementedError
miles/backends/megatron_utils/lora/actor.py core-logic

多 LoRA 训练 actor 的接线点:_init_training_state 里独立于传输层构造 WeightPublisher,新增 export_slot 作为唯一的导出边界并让异常穿透 TrainerCell。

# miles/backends/megatron_utils/lora/actor.py(多 LoRA 训练 actor)
class MultiLoRATrainRayActor(MegatronTrainRayActor):
    def _init_training_state(self) -> None:
        args = self.args
        self.slot_optimizers: dict[int, SlotOptimizer] = {}
        # 发布器与传输层解耦:只依赖 HF 权重迭代器 + LoRA 配置,
        # 不再有 update window / engine session,因此也不需要在推送热路径上缓存注册名。
        iterator = get_hf_weight_iterator(
            args,
            self.model,
            required_placement=WeightUpdatePlacement(gather_pp=True),
            model_name=type(self.hf_config).__name__.lower() if args.model_name is None else args.model_name,
            quantization_config=getattr(self.hf_config, "quantization_config", None),
        )
        self.weight_publisher = WeightPublisher(iterator, build_lora_sync_config(args))
​
    @with_logs
    def export_slot(self, slot: int, rank: int, alpha: float, path: str, metadata: dict | None = None) -> None:
        """Write the slot's adapter as an engine-loadable dir."""
        self._heartbeat.bump()
        # export_slot 自己拥有导出边界:检查点 IO 失败直接抛异常穿过 TrainerCell,
        # 不再像旧的 staged-push 路径那样被折叠成结果字典后静默继续。
        self.weight_publisher.publish_adapter(AdapterSpec(slot=slot, rank=rank, alpha=alpha), path, metadata=metadata)

评论区精华

生成器惰性求值导致多 adapter 前缀串扰 正确性

joyliu-q 在 hf_weight_iterator/__init__.py 指出 for 循环里用生成器表达式捕获 lora_name,消费发生在循环之后会让所有 adapter 共用最后一个前缀,并给出最小复现(期望 ['A:w1','A:w2','B:w1','B:w2'],实际 ['B:w1','B:w2','B:w1','B:w2']),另附 Go loopvar-preview 作为更详细解释。

结论:修复方式是把拼接抽成局部函数 prefixed_units(lora_name, adapter)(形参绑定),并在 TestIterHfWeightsTemplate 增加 test_each_adapter_keeps_its_own_prefix 回归。 · 已解决

staged 发布名缓存导致失败重试跳过注册 正确性

nanjiangwill 在 updater.py 指出 _registered_adapters 在 staged stream 与 commit 完成前写入:失败后 SGLang 移除 pending 注册,但重试命中缓存只发 begin/end;复用已发布名字还会跳过 defer_publish 的新鲜度校验。建议不要缓存注册名并补失败重试、重名测试。

结论:最终版本整体删除 Tinker staged-push 路径,updater 回归普通 base 同步,该问题随代码删除而消失。 · 已解决

HTTP 200 + success=false 的推送结果被丢弃 正确性

nanjiangwill 在 actor.py 指出 wait_futures() 只报传输 / 任务异常,SGLang 可能返回 HTTP 200 但 body 为 success: false,代码在广播成功前丢弃这些结果,建议走 check_weight_sync_results(..., is_lora=True)。

结论:对应的 push_slot 代码在最终版本被删除;该结果校验要求需要由后续 runtime(#2845)的模型加载路径承接,本 PR 未直接验证。 · deferred

abort 失败后引擎仍被标记 ready,磁盘回填可能失败 设计

nanjiangwill 在 group.py 提出:失败热推送的恢复假设 abort 让引擎回到干净状态,但 EngineWeightUpdateSession.__exit__() 会 suppress abort 异常,而 _updatable_engines() 仍标记 ready;若引擎残留 open update session 或 pending adapter 名,同名磁盘回填可能失败。建议补 abort 失败测试或明确该场景需要替换引擎。

结论:最终版本移除了 _updatable_engines/update window 相关改动,恢复语义转移给 #2845 的 runtime;本 PR 未定义引擎替换策略,属于残留未决项。 · 待处理

对照 sglang / miles 的既有实现路线 question

yushengsu-thu 建议参考 sglang 的 #38165(async LoRA rollout and weight sync)与 miles 的 #3127(publish immutable adapter versions alongside generation)的做法。

结论:本 PR 采纳了「不可变版本目录 + 磁盘回填、引擎按需加载」的路线,并在此后给出 APPROVED。 · 已解决

风险与影响

  1. 发布链路失败即终止训练:export_slot 不再把检查点 IO 失败转成结果字典(PR body:「export_slot no longer converts checkpoint IO failures into result dictionaries」),write_checkpoint_dir、safetensors 序列化、文件写入、目录发布任一环节出错都会作为异常穿过 TrainerCell 终止服务,文件系统抖动会被放大为训练中断。
  2. 多 LoRA 配置语义收紧:ModelConfig.resolve 在 is_multi_lora_enabled(args) 时强制 update_weights=False 并断言 effective_model_path == args.hf_checkpoint;原先自定义 model_path 只打 warning 并默认 False,现在会直接 assert 失败,属于对既有 multi-LoRA 配置的兼容性破坏。
  3. 公共符号 slot_lora_name 从 miles/utils/multi_lora.py 的 __all__ 移除,本 PR 未在同仓给出替代定义;若还有 rollout、session 或脚本侧未迁移的 import,会直接 ImportError。
  4. materialize=is_writer 的隐式集合通信契约:non-writer rank 走 materialize=False 必须仍执行同样的 gather,任一侧提前返回或抛错都会让 NCCL 挂死;当前实现正确,但后续修改 materialize_adapter 时容易踩坑。
  5. max_loaded_loras 默认 2 * max_loras_per_batch 是带 TODO 的启发式,slot 数超过该值时 SGLang 会触发淘汰/重载,表现为 on-demand load 抖动,且该值直接影响显存占用。
  6. 测试覆盖缺口:本 PR 仅更新两个 fast 测试,WeightPublisher.publish_adapter 与 export_slot 的回归(进程在发布前后退出、重复保存)按 PR body 位于 #2849 的 checkpoint 测试文件,单独合并本 PR 时磁盘发布路径缺少直接单测。
  7. 对已有脚本的迁移:miles/utils/lora.py 新增 save_adapter_to_disk 与 WeightPublisher 内部写盘逻辑存在两套写 adapter 目录的实现,若长期并存容易出现配置字段(如 r/lora_alpha 是否覆盖)不一致。

用户侧:Tinker 用户可以在训练进行中让引擎加载某个固定版本 adapter,不再受 push window 与 update session 影响;但多 LoRA 用户必须满足「引擎加载与训练一致的冻结 base」这一新约束,显式 update_weights: true 或自定义 model_path 的配置会启动失败。系统侧:训练与推理通过文件系统解耦,engine 从被动 push 变为按需拉取,减少了训练侧同步阻塞,代价是 rank 0 上多一次完整 adapter 的 D2H 拷贝与 safetensors 写盘(由增量目录+原子改名保证读侧一致性)。团队侧:发布器、迭代器、runtime 三层职责被切开,后续 5/7、6/7 只消费 export_slot 与快照路径;同时失败语义从「返回结果字典」变成「抛异常终止 cell」,运维需要据此调整告警与重启策略。

权重发布路径重写 公共符号 slot_lora_name 删除 多 LoRA 强制 update_weights=False IO 失败直接终止 TrainerCell 本 PR 缺少发布器直接测试 max_loaded_loras 启发式默认值

关联 Issue

#2845 [multi-lora v2 6/7] Tinker gateway runtime and entry point
#2849 [multi-lora v2 7/7] Tinker gateway example and tests

完整报告

参与讨论