执行摘要

• 一句话：Inkling 原生 LoRA 支持：adapter-only GRPO 全链路打通
• 推荐动作：值得精读。重点关注三类设计决策：
  1) 插件式 LoRA 如何绕过 Megatron-Bridge 的 PEFT 集成，通过参数标签 + 显式梯度归约保证分布式正确性；
  2) _pp_assemble_full_adapter 与 _repack_onto_fresh_storage 解决的多维并行 + CUDA IPC 工程问题；
  3) SGLang 侧 lora_target_modules=["all"] 与 virtual-experts 的搭配。建议后续跟进 EP>TP 时 native shard 的加载验证、non-colocate 引擎的 LoRA 同步，以及 42/66 层真实配置的 CI 覆盖。

功能与动机

PR body 声明：Stacked on #1683（Inkling 模型 + 全参 RL），本 PR 在其上增加 LoRA 层，实现 adapter-only GRPO，复用同一套后端与并行堆栈。核心动机是训练成本：全参模式下每步都要把所有权重同步给 rollout 引擎（49.4s），而 LoRA 只需同步小型 adapter（2.5s），单步耗时降到全参的约 85%；同时遵循 Inkling 官方发布的 LoRA schema（r=32 all-linear），保证 adapter 可直接用于 serving（引擎侧原生 triton 后端 + virtual experts）。

实现拆解

实现按 5 个步骤拆解：

1. LoRA 插件：按官方 schema 逐模块注册 adapter

新增 miles_plugins/models/inkling/lora.py（759 行）。InklingLoRAAdapter 是纯参数容器，sharded_state_dict() 返回空字典，使 adapter 对 Megatron dist-checkpointing 完全不可见；每个 adapter 记录 load_meta（TP/EP rank、头维度、local intermediate 等导出元数据）。attention 因 MLA 结构拆 wq/wk/wv/wr 四组 A/B，并挂到 linear_qkv/linear_proj 的 forward 上叠加 delta，RMSNorm 已融合进 TELayerNormColumnParallelLinear，forward 里用 _rmsnorm 重算 norm 再喂 LoRA 输入；dense MLP 的 fc1/fc2 各一组 A/B，同样处理 fused layer-norm 与 sequence-parallel gather/reduce；routed experts 的 w1/w3/w2 三组参数中 w1_A/w3_A 共享输入投影，B 是三维参数 (num_local_experts, moe_intermediate, rank)，用 _grouped_linear（CUDA 走 F.grouped_mm，CPU 回退逐专家 linear）按 tokens_per_expert 做 grouped GEMM，并断言 expert_tensor_parallel_size == 1；shared experts 采用多专家共享 A、各专家独立 B 的 shared-outer 结构。参数约定：B 零初始化、A xavier，TP 复制参数标 _lora_grad_sum_group="tp"，EP 复制参数标 "ep" 且 allreduce=False。model.py 在 LoRA + Inkling 时用 wrap_model_provider_with_inkling_lora 包装 provider，initialize_model_and_optimizer 里支持 --lora-adapter-path 热启动（load_inkling_lora_adapter + optimizer.reload_model_params() 刷新 fp32 master，避免首步 optimizer 覆盖刚写入的 adapter 值）。

2. 复制参数梯度归约

lora_utils.py 新增 reduce_marked_lora_grads：在 finalize_model_grads 之前，按 _lora_grad_sum_group 标签把 "tp" 组、"ep" 组内各 rank 的 partial grad 做 all_reduce SUM（按 dtype 扁平化一次通信），随后才走常规 DDP 的 DP reduce-scatter。参数列表以 id(model[0]) 为 key 缓存，避免每步全模型扫描。该函数挂在 model.py::finalize_model_grads_with_empty_cache 中 finalize_model_grads 调用之前。

3. 导出与引擎同步

update_weight_from_tensor.py 是性能与正确性的关键：_pp_assemble_full_adapter 解决 exporter 的 bridge 只 gather TP/EP、不 gather PP 的问题，PP>1 时按 PP 组 all_gather_object 元数据 + 按 dtype 扁平 broadcast 完成全量组装；skip_base_sync 条件扩展 colocate_base_persistent（colocate 且未 offload_rollout 时 base 常驻引擎内存，可跳过 base 重传）；_repack_onto_fresh_storage 解决 torch_memory_saver.region() 的 cuMem MemPool 内存无法经 legacy CUDA IPC 导出的问题（offload_train 时把 adapter 张量拷到普通 caching-allocator 内存，按 (dtype, device) 合并成少量扁平 buffer）；同步结束后显式 ipc_collect() + empty_cache() 防 IPC 引用累积。sglang_engine.py 侧 load_lora_adapter_from_tensors 扩展 serialized_tensors 整包传输（与原 per-TP-rank 的 serialized_named_tensors 互斥二选一），Inkling checkpoint 时 lora_target_modules = ["all"] 交给 SGLang 自动识别模块名，dummy base load 时跳过启动 lora_paths 依赖 weight-sync。

4. 启动器与参数门控

scripts/run_inkling.py 新增 train_mode: "full" | "lora"，lora 模式拼装 --lora-rank 32 --lora-alpha 32 --target-modules all-linear --experts-shared-outer-loras --sglang-lora-backend triton --sglang-lora-use-virtual-experts --sglang-max-loras-per-batch 1 等参数，并设置 --sglang-ep-size rollout_num_gpus_per_engine；fully_async 断言仅支持 full。默认 lr 按模式区分（lora 5e-6，文档强调实际训练用 2e-4 才能让零初始化的 B 快速积累可见 delta-W）。utils/lora.py 新增 lora_rollout_enabled = is_lora_enabled and not lora_train_only，统一门控引擎侧 enable_lora、per-request lora_path 与 adapter 权重同步；actor.py 的 sleep/wake（torch_memory_saver tag）与更新器 is_lora 都改用它。arguments.py 新增 --lora-train-only，并修复 experts_shared_outer_loras 与 sglang_experts_shared_outer_loras 一致性检查在属性缺失时的崩溃。

5. 测试、CI 与文档

新增 tests/e2e/megatron/model_scripts/test_inkling_small_4layer_lora_ci.py：在 stage-c-4-gpu-h200 套件（约 30 分钟）跑 4 层 Inkling-Small LoRA 冒烟，开启 --check-lora-weight-equal（engine 侧 sha256 校验 adapter 一致性），skip list 处理 frozen towers 与 engine 派生 buffer，并注册 train/grad_norm、train/ppo_kl、rollout/raw_reward 等 CI gate；文档同步补充 train modes 与 lr 建议。

关键文件：

• miles_plugins/models/inkling/lora.py（模块 适配层；类别 source；类型 data-contract；符号 InklingLoRAAdapter, _rmsnorm, _new_param, _register_param）: 新增 759 行核心插件：定义 InklingLoRAAdapter 与 attention/dense MLP/routed experts/shared experts 的 LoRA 参数注册与 forward monkey-patch，全部 LoRA 参数对 Megatron dist-checkpoint 透明（sharded_state_dict 返回空），并定义 _lora_grad_sum_group 标签契约。
• miles/backends/megatron_utils/lora_utils.py（模块 适配工具；类别 source；类型 dependency-wiring；符号 sglang_lora_target_all_sentinel, reduce_marked_lora_grads, save_lora_checkpoint, load_lora_adapter）: 新增 sglang_lora_target_all_sentinel（Inkling 时让 SGLang 自动检测模块名）与 reduce_marked_lora_grads（TP/EP 复制参数梯度显式归约）；同时把 native checkpoint 分片从 tp/pp 命名改为 global_rank 命名并加 legacy fallback。
• miles/backends/megatron_utils/update_weight/update_weight_from_tensor.py（模块 权重同步；类别 source；类型 dependency-wiring；符号 _pp_assemble_full_adapter, _repack_onto_fresh_storage, _send_to_colocated_engine, _send_lora_params）: 权重同步核心：新增 _pp_assemble_full_adapter 补全 PP 维度、_repack_onto_fresh_storage 解决 cuMem IPC 导出失败、扩展 skip_base_sync 条件和同步后显存回收，是 weight update 49.4s→2.5s 的关键。
• miles/backends/sglang_utils/sglang_engine.py（模块 引擎适配；类别 source；类型 dependency-wiring；符号 load_lora_adapter_from_tensors, _compute_server_args）: 引擎侧适配：load_lora_adapter_from_tensors 支持整包 serialized_tensors 传输；Inkling checkpoint 时 lora_target_modules=['all'] 交给 SGLang 自动识别；dummy load 时跳过启动 lora_paths 依赖 weight-sync。
• miles/backends/megatron_utils/model.py（模块 模型构建；类别 source；类型 data-contract；符号 setup_model_and_optimizer, finalize_model_grads_with_empty_cache, initialize_model_and_optimizer）: 模型构建与训练钩子：LoRA+Inkling 时用 wrap_model_provider_with_inkling_lora 包装 provider；finalize_model_grads 前插入 reduce_marked_lora_grads；--lora-adapter-path 热启动后刷新 fp32 master。
• scripts/run_inkling.py（模块 启动脚本；类别 source；类型 core-logic；符号 ScriptArgs, _train）: launcher 核心：新增 --train-mode lora 及 LoRA 相关参数拼装，lora 模式使用 triton backend + virtual experts + ep-size；fully-async 断言仅支持 full，文档同步补 lr 建议。
• tests/e2e/megatron/model_scripts/test_inkling_small_4layer_lora_ci.py（模块 CI 测试；类别 test；类型 test-coverage；符号 _args, prepare, execute）: 4 层 LoRA 冒烟 CI：验证 checkpoint 转换、LoRA 训练、--check-lora-weight-equal 的 sha256 校验，注册 grad_norm/ppo_kl 等 CI gate；是功能而非精度的兜底。
• miles/utils/arguments.py（模块 参数解析；类别 source；类型 core-logic；符号 add_lora_arguments, miles_validate_args）: 新增 --lora-train-only 门控；修复 experts_shared_outer_loras 与 sglang_experts_shared_outer_loras 一致性检查在属性缺失时的崩溃。
• miles/utils/lora.py（模块 开关门控；类别 source；类型 core-logic；符号 lora_rollout_enabled）: 新增 lora_rollout_enabled 统一门控 rollout 侧 LoRA（enable_lora、lora_path、adapter weight sync），--lora-train-only 下关闭。
• miles/backends/megatron_utils/actor.py（模块 训练执行；类别 source；类型 dependency-wiring；符号 init, sleep, wake_up）: 训练 actor 的 sleep/wake 与权重同步改用 lora_rollout_enabled，决定 torch_memory_saver 的 tag 与是否走 adapter 同步。

关键符号：InklingLoRAAdapter, _rmsnorm, _new_param, _register_param, _dropout, _grouped_linear, _apply_attention_lora, _apply_dense_mlp_lora, _apply_expert_lora, _apply_shared_experts_lora, reduce_marked_lora_grads, sglang_lora_target_all_sentinel, _pp_assemble_full_adapter, _repack_onto_fresh_storage, lora_rollout_enabled, wrap_model_provider_with_inkling_lora, load_inkling_lora_adapter

关键源码片段

miles_plugins/models/inkling/lora.py

新增 759 行核心插件：定义 InklingLoRAAdapter 与 attention/dense MLP/routed experts/shared experts 的 LoRA 参数注册与 forward monkey-patch，全部 LoRA 参数对 Megatron dist-checkpoint 透明（sharded_state_dict 返回空），并定义 _lora_grad_sum_group 标签契约。

# miles_plugins/models/inkling/lora.py（节选）
# 核心参数工厂：所有 Inkling LoRA 参数都从这里创建。
# 关键约定：
# - tensor_model_parallel=False / partition_dim=-1 让 Megatron 把这个参数当作
# 非 TP 切分参数（对 replicated 参数非常重要）。
# - expert=True 时 allreduce=False，EP 下每个 rank 持有本地专家副本，由
# reduce_marked_lora_grads 在 DP reduce 之前显式做 ep 组求和。
# - _lora_grad_sum_group 标签是梯度归约契约：被标记的参数必须进入
# reduce_marked_lora_grads 的统计，否则复制的 LoRA 参数会梯度不一致。
​
def _new_param(ref_weight, shape, *, init, grad_sum_group=None, expert=False):
    tensor = torch.empty(*shape, dtype=ref_weight.dtype, device=ref_weight.device)
    if init == "zero":
        tensor.zero_() # B 矩阵零初始化，保证训练初期 delta 为 0
    elif tensor.ndim == 2:
        nn.init.xavier_uniform_(tensor) # A 矩阵 xavier
    else:
        for expert_tensor in tensor: # 3D 参数：逐专家 xavier
            nn.init.xavier_uniform_(expert_tensor)
    param = nn.Parameter(tensor)
    param.tensor_model_parallel = False
    param.partition_dim = -1
    param.partition_stride = 1
    if expert:
        param.allreduce = False
    if grad_sum_group is not None:
        # 被 reduce_marked_lora_grads 消费：复制的 adapter 梯度需要显式求和
        param._lora_grad_sum_group = grad_sum_group
    return param
​
​
def _grouped_linear(inputs, weights, tokens_per_expert):
    """对 permuted 后的 token 缓冲做逐本地专家 matmul，一次 grouped GEMM 完成。"""
    if inputs.is_cuda:
        offsets = torch.as_tensor(list(tokens_per_expert), device=inputs.device, dtype=torch.int32).cumsum(
            0, dtype=torch.int32
        )
        return F.grouped_mm(inputs, weights.transpose(1, 2), offs=offsets)
    # CPU 回退：逐 segment 做普通 linear 再拼接
    segments = torch.split(inputs, list(tokens_per_expert), dim=0)
    return torch.cat([F.linear(segment, weights[idx]) for idx, segment in enumerate(segments)], dim=0)
​
​
def _apply_expert_lora(moe, args, hf_prefix, *, scale, dropout, a_init):
    # Inkling 的专家 LoRA 按官方 schema 拆成 w1 / w3（gate+up）与 w2（down），
    # 且与外层 routed experts 共用一份 A 矩阵。
    config = moe.config
    assert (getattr(config, "expert_tensor_parallel_size", 1) or 1) == 1, "Inkling LoRA assumes ETP=1"
    rank = int(args.lora_rank)
    hidden_size = config.hidden_size
    moe_intermediate = config.moe_ffn_hidden_size
    experts = moe.experts
    num_local_experts = experts.num_local_experts
    is_ep = parallel_state.get_expert_model_parallel_world_size() > 1
    ep_group = "ep" if is_ep else None
​
    adapter = InklingLoRAAdapter("experts", hf_prefix + "mlp.experts.")
    fc1_ref, fc2_ref = experts.linear_fc1.weight0, experts.linear_fc2.weight0
    _register_param(adapter, "w1_A", fc1_ref, (rank, hidden_size), init=a_init, grad_sum_group=ep_group, expert=is_ep)
    _register_param(adapter, "w3_A", fc1_ref, (rank, hidden_size), init=a_init, grad_sum_group=ep_group, expert=is_ep)
    _register_param(adapter, "w1_B", fc1_ref, (num_local_experts, moe_intermediate, rank), expert=is_ep)
    _register_param(adapter, "w3_B", fc1_ref, (num_local_experts, moe_intermediate, rank), expert=is_ep)
    _register_param(adapter, "w2_A", fc2_ref, (num_local_experts, rank, moe_intermediate), init=a_init, expert=is_ep)
    _register_param(adapter, "w2_B", fc2_ref, (hidden_size, rank), grad_sum_group=ep_group, expert=is_ep)
    experts.lora_adapter = adapter
​
    fc1 = experts.linear_fc1
    original_fc1 = fc1.forward
​
    def expert_fc1_forward(inputs, tokens_per_expert, *forward_args, **forward_kwargs):
        output, bias = original_fc1(inputs, tokens_per_expert, *forward_args, **forward_kwargs)
        dropped = _dropout(inputs, dropout, fc1.training)
        # 共享的 A 矩阵拼起来走一次线性，再由 grouped GEMM 落到各自专家
        joint = F.linear(dropped, torch.cat([adapter.w1_A, adapter.w3_A], dim=0))
        gate = _grouped_linear(joint[..., :rank].contiguous(), adapter.w1_B, tokens_per_expert)
        up = _grouped_linear(joint[..., rank:].contiguous(), adapter.w3_B, tokens_per_expert)
        delta = torch.cat([gate, up], dim=-1)
        return torch.add(output, delta, alpha=scale), bias
​
    fc1.forward = expert_fc1_forward
    # fc2 同理：w2_A 按专家 grouped，w2_B 直接线性，略

miles/backends/megatron_utils/lora_utils.py

新增 sglang_lora_target_all_sentinel（Inkling 时让 SGLang 自动检测模块名）与 reduce_marked_lora_grads（TP/EP 复制参数梯度显式归约）；同时把 native checkpoint 分片从 tp/pp 命名改为 global_rank 命名并加 legacy fallback。

# miles/backends/megatron_utils/lora_utils.py（节选）
# Inkling LoRA 与标准 LoRA 的最大差异：标准路径靠 Megatron-Bridge 的 PEFT
# 集成自动处理梯度归约，而插件直接注册的 replicated 参数（TP 复制、EP 复制）
# 不会被 Megatron 自动 reduce。这里在 finalize_model_grads 之前按参数上的
# _lora_grad_sum_group 标签显式做 all_reduce SUM，随后再走 DDP 的 DP reduce-scatter。
​
_marked_lora_grad_params_cache: dict[int, list] = {}
​
​
def reduce_marked_lora_grads(model) -> None:
    from megatron.core import parallel_state as ps
​
    key = id(model[0]) if model else 0
    marked = _marked_lora_grad_params_cache.get(key)
    if marked is None:
        marked = []
        for chunk in model:
            for param in chunk.parameters():
                group_name = getattr(param, "_lora_grad_sum_group", None)
                if group_name is not None and param.requires_grad:
                    marked.append((param, group_name))
        _marked_lora_grad_params_cache[key] = marked
    if not marked:
        return
    groups = {
        "tp": (ps.get_tensor_model_parallel_group(), ps.get_tensor_model_parallel_world_size()),
        "ep": (ps.get_expert_model_parallel_group(), ps.get_expert_model_parallel_world_size()),
    }
    for group_name in ("tp", "ep"):
        group, size = groups[group_name]
        if size <= 1:
            continue # 单 rank 组无需通信
        grads = []
        for param, g_name in marked:
            if g_name != group_name:
                continue
            grad = getattr(param, "main_grad", None) # Megatron 优化器路径优先
            if grad is None:
                grad = param.grad
            if grad is not None:
                grads.append(grad)
        # 按 dtype 分组后扁平化，一次 all_reduce 处理整组梯度，减少通信次数
        for dt in {g.dtype for g in grads}:
            gs = [g for g in grads if g.dtype == dt]
            if len(gs) == 1:
                dist.all_reduce(gs[0], op=dist.ReduceOp.SUM, group=group)
                continue
            flat = torch._utils._flatten_dense_tensors(gs)
            dist.all_reduce(flat, op=dist.ReduceOp.SUM, group=group)
            for g, red in zip(gs, torch._utils._unflatten_dense_tensors(flat, gs), strict=False):
                g.copy_(red)

miles/backends/megatron_utils/update_weight/update_weight_from_tensor.py

权重同步核心：新增 _pp_assemble_full_adapter 补全 PP 维度、_repack_onto_fresh_storage 解决 cuMem IPC 导出失败、扩展 skip_base_sync 条件和同步后显存回收，是 weight update 49.4s→2.5s 的关键。

# miles/backends/megatron_utils/update_weight/update_weight_from_tensor.py（节选）
# LoRA adapter 参数在 torch_memory_saver.region()（torch.cuda.use_mem_pool）里创建，
# 存储来自 cuMem MemPool。cuMem 分配无法通过 legacy CUDA IPC 导出，
# 直接交给 SGLang 引擎会在第一次同步时报 "CUDA error: invalid argument"。
# 这里在 region 之外重新分配普通 caching-allocator 内存，把张量拷贝过去，
# 让 MultiprocessingSerializer 能拿到可导出的 IPC handle。
# 按 (dtype, device) 合并成一份平铺 buffer，pickler 会 memoize storage，
# 引擎收到的是一把 IPC handle 而不是每个 adapter 张量一个 handle。
​
def _repack_onto_fresh_storage(named_tensors):
    groups: dict[tuple[torch.dtype, torch.device], list[tuple[str, torch.Tensor]]] = {}
    for name, tensor in named_tensors:
        if tensor.is_cuda:
            groups.setdefault((tensor.dtype, tensor.device), []).append((name, tensor))
​
    views: dict[str, torch.Tensor] = {}
    for (dtype, device), items in groups.items():
        flat = torch.empty(sum(t.numel() for _, t in items), dtype=dtype, device=device)
        offset = 0
        for name, tensor in items:
            view = flat[offset : offset + tensor.numel()].view(tensor.shape)
            view.copy_(tensor)
            views[name] = view
            offset += tensor.numel()
​
    return {name: views.get(name, tensor) for name, tensor in named_tensors}

评论区精华

本 PR 无实质 review 讨论：仅 yueming-yuan 的 APPROVED（空正文），Issue 中只有 gemini-code-assist 的自动下线公告。设计决策主要沉淀在代码注释里，尤其是 _repack_onto_fresh_storage 对 torch_memory_saver MemPool 与 CUDA IPC 冲突的说明，以及 reduce_marked_lora_grads 对复制参数梯度归约契约的注释。

• 暂无高价值评论线程

风险与影响

• 风险：
  1. 同步链路影响面：update_weight_from_tensor.py 的 skip_base_sync 条件与 _send_to_colocated_engine 的 LoRA 分支是全仓库 LoRA 共用路径，本 PR 改动影响既有 LoRA 场景（如 #2075 session 侧 adapter 同步），需要回归。
  2. checkpoint 格式变化：native 分片从 adapter_megatron_tp{tp}_pp{pp}.pt 改为 adapter_megatron_rank{rank}.pt，旧格式靠 legacy fallback 读取且注释警告 only valid when EP<=TP；新格式在 EP>TP 等复杂并行下的加载一致性需额外验证。
  3. 梯度归约正确性：reduce_marked_lora_grads 依赖每个 replicated 参数都带 _lora_grad_sum_group 标签，漏标会静默产生错误梯度；缓存以 id(model[0]) 为键，模型重建且 id 复用时会命中陈旧缓存。
  4. 覆盖风险：CI 仅覆盖 4 层切片且定位是功能冒烟（非精度），42/66 层真实配置（TP4 SP PP8 EP4 等）未进 CI；non-colocate（distributed engine）的 LoRA 同步仍是 NotImplementedError。
  5. 性能：每步同步后新增 ipc_collect/empty_cache，对大模型步频的影响需结合 wandb 数据持续观察。
     - 影响：对用户：Inkling / Inkling-Small 用户可按官方 LoRA schema 低成本训练，显著降低权重同步带宽与显存压力（weight update 49.4s→2.5s，step 约全参 85%）。对系统：改动横跨模型插件、Megatron 后端、SGLang 引擎、启动脚本与 CI，LoRA 共用链路的 save/load 命名与同步逻辑发生变化，其他 LoRA 后端需回归。对团队：确立了「每模块 plugin adapter + 显式 TP/EP 梯度归约 + PP 级全量组装 + CUDA IPC 直传引擎」模式，可为后续模型的原生 LoRA 提供标准范式。
     - 风险标记：核心权重同步链路改动, 新增 TP/EP 梯度归约协议, CI 仅覆盖 4 层切片, distributed engine LoRA 未支持, 旧格式 checkpoint 兼容依赖 legacy fallback

关联脉络

• PR #1683 Inkling model + full-parameter RL: PR body 声明 stacked on 该 PR，是本次 LoRA 支持的基座；全参 GRPO 提供了后端与并行堆栈
• PR #2013 [tito] Add the Inkling TITO family (Inkling / Inkling-Small): 同一 Inkling 模型线的会话与 tokenizer 支持，后续 LoRA 会话训练可复用
• PR #2009 [docs] Add Inkling-Small model page: Inkling-Small 文档页面，本 PR 同步更新 inkling-small.md 的 LoRA 训练说明
• PR #1794 feat(multi-lora): enable and validate MoE expert adapters: MoE 专家层级 LoRA 的先例，涵盖 (tp,pp,ep) 分片与校验，本 PR 的专家 LoRA 是其补充
• PR #2075 session: apply the trained LoRA adapter to session-server rollouts: LoRA 训练与 rollout 链路共用 lora_path/adapter 同步逻辑，本 PR 修改同一链路需回归验证