执行摘要
- 一句话:LoRA 按 attn-TP 切分并支持 dp attention 动态加载
- 推荐动作:值得精读。重点看三处设计:_effective_tp_size 如何把 moe_tp / attn_tp / tp 三种切分维度收敛到一个查询入口;RowParallel 归约组跟随 base layer 的 use_dp_attention_reduce 处理;_merge_lora_update_results 的失败优先与去重语义(以及全成功时返回首对象以保证 LRU 原地修改安全)。对要扩展 LoRA 到新并行组合的工程师,这是很好的参照。
功能与动机
PR body 明确说明:Attention layers shard on attn-TP (tp_size / dp_size), not the global rank, so adapter loads hit shape-mismatch asserts;动态 LoRA 端点此前 assert dp_size == 1,直接拒绝 dp attention 下的加载请求。此外 dp attention 会把 load/unload 扇出到每个 DP group 的一个 scheduler,原有 result[0] 读取会让非 0 号 rank 的失败被误报为成功,导致 tokenizer 侧注册表漂移。
实现拆解
- 新增 attn-TP 模块分类(python/sglang/srt/lora/utils.py):定义 ATTN_TP_LORA_MODULE_NAMES,收录 qkv_proj、o_proj、in_proj、in_proj_qkvz 等在 DP attention 下构建于 attn-TP 组的投影模块;注释说明 mamba.py 与 qwen3_5.py 在 dp attention 下用 attn_tp 参数构建这些层。
- 内存池感知 attn_tp_size(python/sglang/srt/lora/mem_pool.py + python/sglang/srt/lora/lora_manager.py):LoRAMemoryPool 构造函数新增必填 attn_tp_size;新增 _effective_tp_size(module_name) 统一按模块类型返回路由 MoE 的 moe_tp_size、注意力模块的 attn_tp_size 或默认外层 tp_size,并替换 get_lora_A_shape / get_lora_B_shape 中各自维护的内联分支。LoRAManager 在 init 时从 get_parallel() 提取 attn_tp_size(并行组在 init_torch_distributed 后冻结)并传入内存池;load_lora_weight_tensor 不再把 self.tp_rank 传给 slice 方法,切分 rank 由模块自身决定。
- 权重切片与归约组跟随 base layer(python/sglang/srt/lora/layers.py):所有 slice_lora_a_weights / slice_lora_b_weights 签名删除 tp_rank 参数,内部改取 self.base_layer.tp_rank(如 ColumnParallel / MergedColumnParallel / MLA / Inkling 系列);RowParallelLinearWithLoRA.forward 依据 base_layer.use_dp_attention_reduce 选择 attn_tp_group.all_reduce 或全局 tensor_model_parallel_all_reduce,避免跨 DP 组混和。
- 控制面合并与断言放宽(python/sglang/srt/managers/tokenizer_control_mixin.py):新增 _merge_lora_update_results,全部成功时返回首 rank 原对象(调用方 LRU 驱逐会原地修改 loaded_adapters),任一失败则失败优先、错误消息按出现顺序去重后以 | 连接;load/unload/load_from_tensors 三个端点统一走合并逻辑,并放宽动态 LoRA 断言为 dp_size == 1 or enable_dp_attention(纯 dp_size > 1 仍禁止)。
- 测试配套:新增 test/registered/unit/managers/test_lora_update_result_merge.py 覆盖合并语义(全成功返回首对象、部分失败优先、错误去重、无消息失败不崩溃);扩展 test/registered/unit/lora/test_mem_pool_ep_unit.py 增加 TestAttnModulesShardByAttnTp,覆盖 attn_tp=1 全宽、attn_tp=2 二分片、in_proj_qkvz 按 attn-TP 分类三组形状断言,并同步改造 _FakeDenseLayer 等桩签名。
关键文件:
python/sglang/srt/lora/layers.py(模块 LoRA 层;类别 source;类型 core-logic;符号 slice_lora_a_weights, slice_lora_b_weights, RowParallelLinearWithLoRA.forward): 核心切片与归约逻辑:所有 slice_lora_a_weights / slice_lora_b_weights 删除 tp_rank 参数,改由 base_layer.tp_rank 决定切分位置;RowParallelLinearWithLoRA.forward 按 use_dp_attention_reduce 选择 attn-TP 或全局 TP 归约组,是修复 shape-mismatch 与跨 DP 混和的关键。
python/sglang/srt/lora/mem_pool.py(模块 内存池;类别 source;类型 core-logic;符号 _effective_tp_size, get_lora_A_shape, get_lora_B_shape, load_lora_weight_tensor): LoRAMemoryPool 新增 attn_tp_size 构造参数与 _effective_tp_size 分类方法,统一决定 LoRA 缓冲区沿哪种 TP 维度切分;get_lora_A_shape / get_lora_B_shape 与 load_lora_weight_tensor 全部接入新逻辑。
python/sglang/srt/managers/tokenizer_control_mixin.py(模块 控制通道;类别 source;类型 core-logic;符号 _merge_lora_update_results, load_lora_adapter, unload_lora_adapter, load_lora_adapter_from_tensors): 新增 _merge_lora_update_results 合并 DP fan-out 的各 rank 回复,修正 result[0] 误读导致的注册表漂移;同时放宽 load/unload 的 dp_size 断言并接入三个端点。
python/sglang/srt/lora/utils.py(模块 工具库;类别 source;类型 configuration;符号 ATTN_TP_LORA_MODULE_NAMES): 新增 ATTN_TP_LORA_MODULE_NAMES 特例集合,标注哪些 LoRA 模块在 DP attention 下按 attn-TP 构建,是缓冲区分类的配置源头。
python/sglang/srt/lora/lora_manager.py(模块 LoRA 管理;类别 source;类型 dependency-wiring;符号 init_memory_pool, attn_tp_size): 从 get_parallel() 提取 attn_tp_size 并在 init_memory_pool 时传入 LoRAMemoryPool,完成构造链路的依赖接线。
test/registered/unit/managers/test_lora_update_result_merge.py(模块 单元测试;类别 test;类型 test-coverage;符号 TestMergeLoRAUpdateResults, test_all_success_returns_first_rank_result, test_any_rank_failure_wins, test_duplicate_error_messages_deduplicated): 新增测试覆盖 _merge_lora_update_results 的全部语义分支:全成功返回首对象、部分失败优先、错误去重、无消息失败不崩溃。
test/registered/unit/lora/test_mem_pool_ep_unit.py(模块 单元测试;类别 test;类型 test-coverage;符号 TestAttnModulesShardByAttnTp, test_attn_tp_1_keeps_attention_buffers_full_width, test_attn_tp_gt1_still_shards_attention_buffers, test_linear_attention_in_proj_shards_by_attn_tp): 新增 TestAttnModulesShardByAttnTp 回归测试,验证 attn_tp=1 全宽、attn_tp=2 二分片与 in_proj_qkvz 按 attn-TP 分类,并同步改造测试桩签名。
关键符号:slice_lora_a_weights, slice_lora_b_weights, _effective_tp_size, get_lora_A_shape, get_lora_B_shape, _merge_lora_update_results, load_lora_adapter, unload_lora_adapter, load_lora_adapter_from_tensors, init_memory_pool
关键源码片段
python/sglang/srt/lora/layers.py
核心切片与归约逻辑:所有 slice_lora_a_weights / slice_lora_b_weights 删除 tp_rank 参数,改由 base_layer.tp_rank 决定切分位置;RowParallelLinearWithLoRA.forward 按 use_dp_attention_reduce 选择 attn-TP 或全局 TP 归约组,是修复 shape-mismatch 与跨 DP 混和的关键。
# 片段一:RowParallelLinearWithLoRA.forward 中的归约组选择。
# use_dp_attention_reduce 表示 base layer 在 attn-TP 组内行并行归约,
# LoRA 增量必须归约到同一组,避免跨 DP 组混合中间结果。
if self.base_layer.use_dp_attention_reduce:
all_reduce = get_parallel().attn_tp_group.all_reduce
else:
all_reduce = tensor_model_parallel_all_reduce
# 片段二:MergedColumnParallel 的 LoRA-B 切片。
# 切分 rank 直接取 base_layer.tp_rank(attn-TP 局部 rank);外层全局
# TP rank 在 DP attention 下会越出本地权重宽度,导致 shape-mismatch。
def slice_lora_b_weights(self, B: torch.Tensor):
local_tp_rank = self.base_layer.tp_rank
partition_sizes = self.base_layer.output_partition_sizes
output_sizes = self.base_layer.output_sizes
slices = []
offset = 0
for full_size, part_size in zip(output_sizes, partition_sizes):
start_idx = local_tp_rank * part_size
end_idx = start_idx + part_size
slices.append(B[offset + start_idx : offset + end_idx, :])
offset += full_size
return torch.concat(slices, dim=0)
python/sglang/srt/lora/mem_pool.py
LoRAMemoryPool 新增 attn_tp_size 构造参数与 _effective_tp_size 分类方法,统一决定 LoRA 缓冲区沿哪种 TP 维度切分;get_lora_A_shape / get_lora_B_shape 与 load_lora_weight_tensor 全部接入新逻辑。
# 判断某个模块的 LoRA 缓冲区应沿哪个 TP 维度切分。
# 路由 MoE 专家按 moe_tp_size;注意力投影按 attn_tp_size(DP attention
# 下小于外层 tp_size);其余模块按外层 tp_size。
def _effective_tp_size(self, module_name: str) -> int:
if self.is_moe_module(module_name) and not self.is_shared_moe_module(
module_name
):
return self.moe_tp_size
if module_name in ATTN_TP_LORA_MODULE_NAMES:
return self.attn_tp_size
return self.tp_size
# get_lora_A_shape 中的用法:行并行模块的输入维按 effective_tp_size
# 均分,保证缓冲区宽度与 base layer 的实际权重切分一致。
effective_tp_size = self._effective_tp_size(module_name)
if (
effective_tp_size > 1
and module_name in ROW_PARALLELISM_LINEAR_LORA_NAMES
and module_name not in REPLICATED_LINEAR_LORA_NAMES
):
input_dim = divide(input_dim, effective_tp_size)
python/sglang/srt/managers/tokenizer_control_mixin.py
新增 _merge_lora_update_results 合并 DP fan-out 的各 rank 回复,修正 result[0] 误读导致的注册表漂移;同时放宽 load/unload 的 dp_size 断言并接入三个端点。
def _merge_lora_update_results(results: List[LoRAUpdateOutput]) -> LoRAUpdateOutput:
"""合并 DP fan-out 返回的各 rank LoRA 加载/卸载结果。
只有全部 rank 成功才算成功:把部分失败上报为成功会让 tokenizer 侧
LoRA 注册表与失败 rank 的实际状态脱节。失败优先:去重后的错误信息
以 | 连接,loaded_adapters 取第一个失败 rank 的现场。
"""
failed = [r for r in results if not r.success]
if not failed:
# 全成功时直接返回首 rank 原对象:调用方在 LRU 驱逐时会
# 原地修改 result.loaded_adapters,合成新对象会静默破坏该逻辑。
return results[0]
error_messages = list(
dict.fromkeys(r.error_message for r in failed if r.error_message)
)
return LoRAUpdateOutput(
success=False,
error_message=" | ".join(error_messages),
loaded_adapters=failed[0].loaded_adapters,
)
评论区精华
GitHub 界面上 Fridge003 直接 approve,copilot 因请求者配额耗尽未能审查;PR Test (Extra) 首轮失败后作者执行 /rerun-failed-ci。提交历史保留了 #32584 的两轮 review 反馈及其落地:
1) idle-forward 的 LoRA batch 重置应成为显式方法而非 prepare_lora_batch 的特殊分支(commit 70d8512 实现 reset_lora_batch);
2) 初版用 LORA_LINEAR_SHARD_SPECS 表统一分类并行度被建议收窄,最终改为 ATTN_TP_LORA_MODULE_NAMES 特例集合 + _effective_tp_size 查询,把 LoRA 编排重构留到单独计划(commit 1316bfc)。
- Idle-forward 的 LoRA batch 重置是否应为显式方法 (design): 已采纳并实现为显式 reset_lora_batch 方法,无行为变化。
- 用 shard-spec 表还是特例集合处理 attn-TP (design): 采用特例集合方案,恢复原有 ROW_PARALLELISM / REPLICATED 列表,无功能变化;shard-spec 化重构留待后续。
风险与影响
-
风险:
- 核心路径变更:LoRAMemoryPool 构造函数新增必填 attn_tp_size,所有直接构造点必须同步,仓库内仅 lora_manager 一处(已更新),外部自定义引擎会 TypeError。
- 破坏性签名变更:layers.py 的 slice_lora_a_weights / slice_lora_b_weights 删除 tp_rank 参数,仓库外 LoRA 子类若仍按旧签名实现会运行时报错;仓库内测试桩已同步。
- 属性访问无防护:RowParallelLinearWithLoRA.forward 直接读取 base_layer.use_dp_attention_reduce 而没有 hasattr 保护,若某 base layer 缺该属性将 AttributeError;当前 SGLang 内置层均有该属性,风险中等偏低。
- 新能力组合首次启用:dp attention + 动态 LoRA 的部分失败语义为"adapter 保持未注册但服务继续运行",用户侧需要感知该不一致窗口;纯 dp_size > 1 无 dp attention 仍被拒绝。
- 端到端 CI 覆盖偏薄:形状逻辑有 CPU 单元测试,真实多卡验证为手动 8x H200;PR Test (Extra) 首轮失败后 rerun,长期自动覆盖仍不足。
- 影响:用户影响:启用 --enable-dp-attention + --enable-lora 的部署可动态 load/unload LoRA,此前直接断言失败或在加载期崩溃;部分失败时 adapter 保持未注册、不可路由,错误信息去重后返回。
系统影响:LoRAMemoryPool 与各 LoRA 包装层对并行度的认知从单一外层 TP 演进为按模块分类(moe_tp / attn_tp / tp),缓冲区宽度更贴近实际权重切分,attn_tp 小于 tp 时注意力 LoRA 内存分配更紧凑。
团队影响:layers.py 的 slice 签名变更(删除 tp_rank)是破坏性 API,依赖旧签名的外部 LoRA 子类需同步;LoRAMemoryPool 构造新增必填参数,所有构造点需传入 attn_tp_size。该改动为后续 LoRA 编排重构(shard-spec 化)打下基础。
-
风险标记:核心路径变更, 破坏性签名变更, 端到端 CI 覆盖偏薄, 新能力组合首启, 属性访问无防护
关联脉络
- PR #32707 Split #32584 into 1/2: [LoRA] Guard DP-attention idle forwards against stale LoRA batch state: 本 PR 的前置、同源拆分,body 明确 stacked on #32707;其 idle-forward batch 状态守卫解决了 dp attention 下 LoRA 的第一处崩溃,本 PR 修复其余 shape-mismatch 与控制面问题。
- PR #32584 [title unavailable]: 被拆分的原始 PR,本 PR 为 2/2 的一半;提交历史中多处引用其 review feedback(显式 reset 方法、shard-spec 表收窄为特例集合)。
参与讨论