Prhub

#32315 [AMD] Speed up DSV4 MoE weight loading from mmap views

原始 PR 作者 bingxche 合并时间 2026-08-02 14:47 文件变更 6 提交数 9 评论 5 代码增减 +91 / -0

执行摘要

AMD MoE 权重加载提速:H2D 前紧凑化 mmap view,加载快 5.6x

PR body 明确:DSV4-Pro TP8 每 rank 需要 140,544 次小 MoE H2D 拷贝,而 rank-local tensor view 仍引用 safetensors mmap 的大块存储,导致 TP0/TP7 在 H2D 上花 27-32 分钟、其他 rank 约 4 分钟。目标是在不改变默认行为的前提下,通过可选的 H2D 前紧凑化消除这些超大 storage 带来的额外拷贝开销。

值得精读。核心是理解 mmap-backed storage 与 torch view 对 H2D 拷贝性能的影响,以及如何用最小侵入的开关 + 针对性单测验证。关注 _copy_weight_view_before_h2d 的判断条件设计(storage_offset、untyped_storage().nbytes())和 CI 分层验证策略。

讨论亮点

amd-bot 在 CI 状态回复中强调 “PR CI is incomplete — do not read green as verified”:AMD 的 stage-b/stage-c 因 wait-for-stage-a-amd 超时全部跳过,新路径真实价值只在 nightly 验证;同时指出现有 4 个失败(XPU/NPU/NVIDIA/CPU)都不触及本 PR 路径。yctseng0211 补充 stage-b-test-1-gpu-large-amd 的失败将由 PR#32862 解决。

实现拆解

  1. python/sglang/srt/layers/moe/fused_moe_triton/layer.py 新增 _copy_weight_view_before_h2d_maybe_copy_weight_view_before_h2d。前者判断 CPU 张量是否满足「连续、storage offset 为 0、storage 恰好等于自身大小」,不满足则 clone(memory_format=torch.contiguous_format);后者受 SGLANG_MOE_COPY_WEIGHT_VIEWS_BEFORE_H2D 开关控制。
  2. _load_w13_load_w2_load_per_channel_weight_scaleexpert_data.copy_(loaded_weight) 前插入 _maybe_copy_weight_view_before_h2d 调用,覆盖路由专家与 per-channel scale 的加载路径。
  3. python/sglang/srt/environ.py 注册默认关闭的 EnvBool 环境变量,保持默认行为不变,降低回归风险。
  4. 在 AMD 相关 CI workflow(nightly-test-amd-rocm720.ymlpr-test-amd-rocm720.ymlnightly-test-amd.yml)中为 DSV4-flash/pro/pro-mtp、Qwen3.5-397B/235B、GPT-OSS-20B/120B 等大模型任务添加 -e SGLANG_MOE_COPY_WEIGHT_VIEWS_BEFORE_H2D=1,将新路径纳入 nightly 与部分 PR 验证。
  5. 新增单元测试 test_copy_weight_views_before_h2d.py,覆盖精确 storage(应跳过)、零 offset 大 storage view、非零 offset view、非连续张量四类场景,并注册到 CPU 基 CI。
文件 模块 状态 重要度
python/sglang/srt/layers/moe/fused_moe_triton/layer.py 权重加载 modified 7.13
test/registered/unit/layers/moe/test_copy_weight_views_before_h2d.py 单元测试 added 6.92
python/sglang/srt/environ.py 环境配置 modified 4.59
.github/workflows/nightly-test-amd-rocm720.yml CI 工作流 modified 3.2
.github/workflows/pr-test-amd-rocm720.yml CI 工作流 modified 2.38
.github/workflows/nightly-test-amd.yml CI 工作流 modified 2.24

关键符号

_copy_weight_view_before_h2d _maybe_copy_weight_view_before_h2d _load_w13 _load_w2 _load_per_channel_weight_scale

关键源码片段

python/sglang/srt/layers/moe/fused_moe_triton/layer.py core-logic

核心改动:新增 H2D 前紧凑化 helper,并接入 w13/w2/per-channel 三条 MoE 权重加载路径。

# 判定是否需要独立拷贝:仅当张量非连续、或 storage offset 非零、
# 或底层 storage 大于张量自身(典型是 mmap 大文件上的 view)时才 clone。
def _copy_weight_view_before_h2d(loaded_weight: torch.Tensor) -> torch.Tensor:
    """Copy a CPU tensor view into independent contiguous storage."""
    if loaded_weight.device.type != "cpu":
        return loaded_weight
    tensor_bytes = loaded_weight.numel() * loaded_weight.element_size()
    needs_copy = not (
        loaded_weight.is_contiguous()
        and loaded_weight.storage_offset() == 0
        and loaded_weight.untyped_storage().nbytes() == tensor_bytes
    )
    if not needs_copy:
        return loaded_weight
    return loaded_weight.clone(memory_format=torch.contiguous_format)
​
​
def _maybe_copy_weight_view_before_h2d(loaded_weight: torch.Tensor) -> torch.Tensor:
    # 默认关闭,仅当环境变量开启时启用紧凑化,保证默认路径零行为变化
    if not envs.SGLANG_MOE_COPY_WEIGHT_VIEWS_BEFORE_H2D.get():
        return loaded_weight
    return _copy_weight_view_before_h2d(loaded_weight)
​
​
def _load_w13(
    self,
    expert_data: torch.Tensor,
    shard_dim: int,
    shard_id: str,
    loaded_weight: torch.Tensor,
    tp_rank: int,
    is_bias: bool = False,
):
    # ... 经过 tp sharding 与 transpose/narrow 后,loaded_weight 往往是
    # 大 mmap storage 上的小 view;在 expert_data.copy_ 前先紧凑化,
    # 避免 H2D 阶段按底层大 storage 整体搬运。
    loaded_weight = _maybe_copy_weight_view_before_h2d(loaded_weight)
    expert_data.copy_(loaded_weight)
test/registered/unit/layers/moe/test_copy_weight_views_before_h2d.py test-coverage

新增单测覆盖四种 storage/view 场景,验证 helper 行为。

import torchfrom sglang.srt.layers.moe.fused_moe_triton.layer import (
    _copy_weight_view_before_h2d,
)
from sglang.test.ci.ci_register import register_cpu_ciregister_cpu_ci(est_time=2, suite="base-a-test-cpu")
​
​
def _assert_independent_copy(source: torch.Tensor, result: torch.Tensor) -> None:
    # 校验拷贝结果:值相等、新对象、连续、offset 为 0、storage 恰好等于自身大小
    assert torch.equal(result, source)
    assert result is not source
    assert result.is_contiguous()
    assert result.storage_offset() == 0
    assert result.untyped_storage().nbytes() == result.numel() * result.element_size()
​
​
def test_skips_copy_for_exact_storage():
    # 已是精确独立 storage 的张量应原样返回,不产生额外拷贝
    source = torch.arange(16).reshape(4, 4)
    assert _copy_weight_view_before_h2d(source) is source
​
​
def test_copies_zero_offset_storage_view():
    # 零 offset 但底层 storage 更大的 view 需要拷贝
    backing = torch.arange(32).reshape(8, 4)
    source = backing.narrow(0, 0, 2)
    assert source.storage_offset() == 0
    assert source.untyped_storage().nbytes() > source.numel() * source.element_size()
    _assert_independent_copy(source, _copy_weight_view_before_h2d(source))
​
​
def test_copies_nonzero_offset_storage_view():
    # 非零 offset 的 view 需要拷贝
    backing = torch.arange(32).reshape(8, 4)
    source = backing.narrow(0, 6, 2)
    assert source.storage_offset() != 0
    _assert_independent_copy(source, _copy_weight_view_before_h2d(source))
​
​
def test_copies_noncontiguous_tensor():
    # 非连续张量(如 transpose)需要拷贝
    source = torch.arange(16).reshape(4, 4).transpose(0, 1)
    assert not source.is_contiguous()
    _assert_independent_copy(source, _copy_weight_view_before_h2d(source))

评论区精华

PR CI 未覆盖 AMD 路径,AMD stage 全部跳过 测试

amd-bot 指出:AMD 的 stage-b/stage-c 因 `wait-for-stage-a-amd` 超时全部跳过,真正验证本 PR 价值的 nightly 套件不在 PR CI 运行;已执行的 4 个失败(XPU/NPU/NVIDIA/CPU)均不触及本 PR 代码路径。

结论:本 PR 的新路径只通过了 CPU 单测与 nightly 验证,PR CI 绿灯不能视为 AMD 已验证;该 CI 基础设施问题由后续 PR(#32862)修复。 · acknowledged

stage-b-test-1-gpu-large-amd 失败将由 PR#32862 解决 infra

yctseng0211 在评论中指出 stage-b-test-1-gpu-large-amd (linux-mi300-1gpu-sglang) 的失败将由 PR#32862 解决。

结论:该失败与本 PR 改动无关,属于 CI runner 基础设施问题,由独立 PR 修复。 · 已解决

风险与影响

启用后每个 MoE 权重 view 会先 CPU clone,新增约 10 秒 wall time 与临时内存峰值;对于本就精确 storage 的张量会跳过。新逻辑只接入了 fused_moe_triton/layer.py_load_w13/_load_w2/_load_per_channel_weight_scale,其他 MoE 后端或非 triton 路径不受影响,效果覆盖不全。untyped_storage().nbytes() 比较对 mmap 大 storage 判断正确,但零维张量等边界未在单测中覆盖。PR CI 的 AMD 阶段全部跳过,本路径在 CI 绿灯下仍未得到 GPU 级验证,依赖 nightly(后续由 PR#32862 保障)。

默认关闭,对现有用户无行为变化。对 AMD 平台大模型部署(DSV4-Pro、Qwen3.5-397B/235B、GPT-OSS)启用后模型加载时间大幅缩短(5.6x / 1.9-7.8x),显著减少夜间验证超时;同时为后续所有依赖 mmap 权重加载的大型 MoE 模型提供了可复用的优化开关。团队收益:nightly CI 稳定性提升、runner 占用时间减少。

默认关闭降低回归面 PR CI 未覆盖 AMD 路径 CPU clone 增加内存峰值 仅覆盖 fused_moe_triton 路径

关联 Issue

未识别关联 Issue

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

完整报告

参与讨论