已关闭
【需求】Pipeline场景支持SWAP #174
DavidFFFan创建于  5月29日关闭于  13 天前
DavidFFFan
DavidFFFan成员
5月29日 创建

Pipeline 场景支持 Activation Swap

1. 基本信息

项目 内容
作者 DavidFFFan
相关模块 pipeline_parallel / activation_checkpoint / fully_shard
相关 issue / PR 本 issue:#174;实现 PR 待关联
适用后端 MindSpore/PyTorch

2. 背景

Pipeline Parallel(PP)训练会为多个 micro-batch 保留前向激活,直到对应反向到来。GPipe、Interleaved 1F1B、real overlap 等调度中,部分 (stage_index, micro_index) 的 FWD 与 BWD 之间存在较长空窗,激活持续驻留 device 会抬高峰值显存。

HyperParallel 已有 activation swap 基础能力,但基础 swap 本身不知道 PP 的 micro-batch、virtual stage、FSDP MetaStep、P2P container 和 overlap 线程边界,无法安全选择 D2H/H2D 时机。因此本特性不新增另一套 swap 接口,而是在用户通过 swap_wrapper 或 swap_tensor_wrapper 使能 swap 的基础上,由 PP scheduler 负责搬运调度。

本 RFC 要解决的问题:让 activation swap 支持 PP 调度,在 FWD/BWD 空窗期将激活卸载到 host,并在反向消费前恢复。
完成后的成功标准:训练精度与 no-swap 基线一致,无 hang;存在有效 swap window 的场景峰值显存下降;不改变原 PP/FSDP/P2P 语义。

3. 目标和非目标

3.1 目标

  1. 支持 GPipe、1F1B、Interleaved 1F1B、real overlap_b_f 和 dx/dw 拆分场景的 PP swap 调度。
  2. 为每个 (stage_index, micro_index) 建立独立 swap group,支持同一 rank 上的多个 virtual stage。
  3. 在有收益窗口时使用 copy stream 执行 D2H/H2D,并尽量与其它 chunk compute、FSDP all-gather 或 P2P 重叠。
  4. 不修改用户使能 swap 的方式:模块级使用 swap_wrapper,指定 tensor 使用 swap_tensor_wrapper;PP 只增加调度能力。
  5. 每次 schedule.run() 使用独立 generation,正常连续运行时 group 不串轮次。

3.2 非目标

  1. 本期不使用精确 cost model 保证每次 swap 都有净性能收益,当前使用固定门槛 MIN_SWAP_GAP = 4。
  2. 本期不自动解决任意跨 chunk 的底层 storage 共享。
  3. MPipe、PyTorch PP、recompute/mixed policy 不因静态 order 可生成就自动准入。
  4. 异常按训练终止处理,不承诺异常后复用同一进程或 schedule。
  5. 本期不修改 FSDP hook、post-backward、梯度同步或 P2P send/recv 配对语义。

4. 相关实现参考

来源 做法 限制 对本 RFC 的影响
HyperParallel activation swap SwapTensor、Storage、SwapGroup、SwapManager 管理 saved tensor、D2H/H2D、buffer 和 group 生命周期 不感知 PP chunk 和调度空窗 复用基础搬运与状态机,PP 不重复实现 tensor swap
HyperParallel PP scheduler 使用 rank-local finalized MetaStep order 描述 compute、P2P、FSDP 和 overlap container 原调度中没有 activation 搬运动作 在 finalized order 上注入 4 个 swap MetaStep
swap_wrapper 包裹模块或普通 callable,在 forward 内捕获需要保存的激活 只定义“交换哪些激活”,不决定 PP 搬运时机 作为模块级 swap 使能接口
swap_tensor_wrapper 在 forward/construct 中显式注册指定 tensor 必须处于有效 swap group context 作为 tensor 级 swap 使能接口;PP runtime 在 FWD leaf 建立 group context

5. 对外接口

5.1 接口定义

PP swap 由“激活选择”和“PP 调度”两部分组成,二者缺一不可:

from hyper_parallel.core.activation_checkpoint import swap_tensor_wrapper, swap_wrapper

# 方式一:模块级。包裹模块后,模块 forward 中需要保存的激活按 policy_fn 参与 swap。
wrapped_module = swap_wrapper(
    module,
    policy_fn=None,
    group_swap=False,
)

# 方式二:tensor 级。在 forward/construct 内显式注册指定 tensor。
target = swap_tensor_wrapper(
    target,
    tag=None,
    group_swap=False,
)

# PP 侧只负责调度。swap=True 不会替代上述 wrapper,也不会自动选择激活。
schedule = ScheduleGPipe(
    stage,
    micro_batch_num=8,
    swap=True,
)
入参 / 配置项 类型 默认值 是否必填 含义 合法范围 错误处理
swap_wrapper.module Cell / nn.Module / callable 无 是 需要模块级 swap 的对象 后端支持的模块或普通 callable 非法对象按 wrapper 原有校验报错
swap_wrapper.policy_fn callable / None None 否 决定保存或交换哪些激活 符合后端 policy 约定 policy 执行错误向上抛出
swap_wrapper.group_swap bool False 否 是否启用 group/bulk copy True / False 非法类型按后端执行报错
swap_tensor_wrapper.target Tensor 或嵌套 tensor 结构 无 是 显式注册参与 swap 的 tensor 后端可遍历的 tensor 结构 不在有效 group 中时告警并原样返回
swap_tensor_wrapper.tag str / None None 否 tensor 的可读标识 任意字符串 无
swap_tensor_wrapper.group_swap bool False 否 是否启用 group/bulk copy True / False 非法类型按后端执行报错
Schedule*.swap bool False 否 是否在 PP order 中注入 swap 调度 True / False 不支持的后端/组合在运行时拒绝

5.2 使用示例

5.2.1 模块级 swap

from hyper_parallel import PipelineStage, ScheduleInterleaved1F1B
from hyper_parallel.core.activation_checkpoint import swap_wrapper

# 先圈定需要 swap 的层;赋值回模型后 wrapper 才会生效。
for index, layer in enumerate(stage_model.layers):
    stage_model.layers[index] = swap_wrapper(layer, group_swap=True)

stage = PipelineStage(stage_model, stage_index, stage_num=pp_size)

# swap=True 仅开启 PP 搬运调度。
schedule = ScheduleInterleaved1F1B(
    stages=[stage],
    micro_batch_num=8,
    swap=True,
)
losses = schedule.run(*inputs)

5.2.2 指定 tensor swap

from hyper_parallel import PipelineStage, ScheduleGPipe
from hyper_parallel.core.activation_checkpoint import swap_tensor_wrapper

class TransformerBlock(nn.Cell):
    def construct(self, x):
        attn_out = self.attn(x)
        attn_out = swap_tensor_wrapper(attn_out, tag="attn_out", group_swap=True)
        x = self.norm1(x + attn_out)
        return x

stage = PipelineStage(stage_model, stage_index, stage_num=pp_size)
schedule = ScheduleGPipe(stage, micro_batch_num=8, swap=True)
losses = schedule.run(*inputs)

swap_tensor_wrapper 必须在 stage 的 forward/construct 路径内调用。PP runtime 会在 eligible FWD leaf 外建立 swap group context;如果该 chunk 没有有效 swap window,则不会创建物理 group,也不会执行搬运。

5.2.3 与 FSDP 组合

from hyper_parallel import PipelineStage, ScheduleGPipe, fully_shard
from hyper_parallel.core.activation_checkpoint import swap_wrapper

# 推荐先对 layer/子层应用 swap wrapper,再对整个 PP stage 应用 fully_shard。
for index, layer in enumerate(stage_model.layers):
    stage_model.layers[index] = swap_wrapper(layer)

fully_shard(stage_model, mesh=dp_mesh)
stage = PipelineStage(stage_model, stage_index, stage_num=pp_size)
schedule = ScheduleGPipe(stage, micro_batch_num=8, swap=True)

5.3 接口说明

为什么这样设计:复用统一 activation swap 接口,把“交换哪些激活”和“PP 中何时搬运”解耦。
和已有接口是否一致:一致。swap_wrapper / swap_tensor_wrapper 继续负责使能和选择激活。
PP 新增职责:Schedule 的 swap=True 只注入搬运/等待 MetaStep,不自动包裹模型,也不自动选择 tensor。

6. 方案设计

6.1 总体流程

flowchart TD
    A["用户应用 swap_wrapper 或 swap_tensor_wrapper"] --> B["构建 PP Schedule,swap=True"]
    B --> C["构建 compute / P2P order"]
    C --> D["注入本 rank FSDP actions"]
    D --> E["重写 P2P transport"]
    E --> F["分析 FWD/BWD leaf 与 gap"]
    F --> G{"gap >= 4?"}
    G -- "否" --> H["chunk 保持 device resident"]
    G -- "是" --> I["注入 4 个 swap MetaStep"]
    I --> J["run-scoped session 建立 chunk group"]
    J --> K["FWD 收集激活"]
    K --> L["D2H offload / H2D load"]
    L --> M["SWAP_WAIT_LOAD 后执行 BWD consumer"]
    M --> N["run close 回收 group"]

build_exec_order() 的顺序固定为:

construct compute/P2P order
  → inject local FSDP actions
  → rewrite P2P transport (plain / batch / boundary)
  → inject local PP swap actions

swap 最后注入,才能看到最终 FSDP lookahead 和 P2P container;swap MetaStep 不进入 BATCH_SEND_RECV.sub_steps。

6.2 架构设计

flowchart LR
    subgraph User["用户侧:选择激活"]
        SW["swap_wrapper"]
        STW["swap_tensor_wrapper"]
    end

    subgraph PP["PP 调度侧:选择时机"]
        Planner["inject_pipeline_swap_steps"]
        Session["PipelineSwapSession"]
        Executor["MetaStep executor"]
    end

    subgraph Base["基础 Swap Runtime"]
        Manager["SwapManager"]
        Group["SwapGroup / Storage / SwapTensor"]
        Copy["D2H / H2D copy stream"]
    end

    SW --> Group
    STW --> Group
    Planner --> Executor
    Executor --> Session
    Session --> Manager
    Manager --> Group
    Group --> Copy

职责边界:

层次 主要对象 职责
激活选择 swap_wrapper、swap_tensor_wrapper 使能 swap,定义需要管理的激活
基础 swap SwapTensor、Storage、SwapGroup、SwapManager saved tensor 注册、状态转换、D2H/H2D、buffer 和 group 生命周期
PP planner inject_pipeline_swap_steps() 从 finalized rank-local order 识别 leaf,为 eligible chunk 注入 4 个搬运动作
PP runtime PipelineSwapSession、MetaStep executor 建立 run/chunk group,执行 D2H/H2D launch/wait,并回收 group

6.3 单个 chunk 时序与 4 个 MetaStep

当前实现使用 4 个显式 MetaStep,不再是 3 个:

sequenceDiagram
    participant F as FWD leaf
    participant S as PP scheduler
    participant C as Copy stream
    participant B as BWD/BWD_INPUT

    F->>F: 在 chunk group context 中收集激活
    S->>C: SWAP_LAUNCH_OFFLOAD
    S->>S: SWAP_WAIT_OFFLOAD
    Note over S: 建立 event 依赖并释放 device storage
    S->>C: SWAP_LAUNCH_LOAD
    S->>S: SWAP_WAIT_LOAD
    Note over S: 在 backward consumer container 前建立 H2D 依赖
    S->>B: 执行 BWD 或 BWD_INPUT
FWD(collection)
  → SWAP_LAUNCH_OFFLOAD
  → SWAP_WAIT_OFFLOAD
  → SWAP_LAUNCH_LOAD
  → SWAP_WAIT_LOAD
  → BWD/BWD_INPUT
阶段 执行位置 动作 状态
FWD 收集 真实 FWD leaf 进入 chunk group context,wrapper 捕获或注册 tensor DEVICE
发起 D2H SWAP_LAUNCH_OFFLOAD copy stream 启动 offload DEVICE → D2H
完成 offload SWAP_WAIT_OFFLOAD 建立 event 依赖并释放 device storage D2H → HOST
发起 H2D SWAP_LAUNCH_LOAD 准备 device storage,copy stream 启动 load HOST → H2D
等待 H2D SWAP_WAIT_LOAD 在 BWD consumer container 前建立 load event 依赖 H2D → DEVICE
BWD 消费 真实 BWD / BWD_INPUT leaf 读取恢复后的 activation DEVICE
回收 wait_load() 与 run close 释放管理引用并删除 group group removed

逻辑 chunk key 为 (stage_index, micro_index),物理 group name 为:

pp_swap_run{generation}_s{stage_index}_m{micro_index}

6.4 Planner 关键逻辑

6.4.1 Leaf 提取

  • 普通 FWD、BWD、BWD_INPUT、BWD_WEIGHT 是 compute leaf。
  • OVERLAP_B_F / OVERLAP_F_B 展开 sub_steps 用于配对,但原 composite container 保持不变。
  • 一个 composite 只计一个 compute slot,不能把并发 B/F 误算成两个隐藏窗口。
  • 普通路径按 FWD → BWD 配对;dx/dw 路径按 FWD → BWD_INPUT 配对。
  • BWD_WEIGHT 不是 activation first consumer,不参与 load 配对。

6.4.2 Eligibility

gap = backward.compute_index - forward.compute_index
eligible = gap >= MIN_SWAP_GAP  # 当前为 4

MIN_SWAP_GAP 是收益门槛,不是正确性同步条件:

gap 处理
1 BWD 紧邻 FWD,无法建立合法 load window
2 device storage 释放后立即重新申请,没有腾空收益
3 仅一个边缘窗口,通常不足以覆盖 copy 和 allocator 开销
≥ 4 注入完整 D2H/H2D 生命周期

6.4.3 四个静态 anchor

动作 anchor
SWAP_LAUNCH_OFFLOAD FWD 所属 top-level container 后;有 FSDP 时位于 FSDP_RESHARD 前
SWAP_WAIT_OFFLOAD FWD 后第一个 intervening compute 及连续通信块之后
SWAP_LAUNCH_LOAD 普通路径位于 BWD 前最后一个覆盖点;有 BWD lookahead unshard 时放到 FSDP_UNSHARD 前
SWAP_WAIT_LOAD 对应 BWD / BWD_INPUT top-level consumer container 前

6.5 Run-scoped session 与线程边界

每次 schedule.run() 根据当前 rank order 中的 SWAP_LAUNCH_OFFLOAD 建立 eligible key 集合。未入选 chunk 的 FWD 使用 nullcontext,不创建物理 group,也不执行 load wait。

SwapManager 使用 ContextVar 保存当前 group,支持嵌套恢复,并隔离不同 Python execution context:

with session.group_context(fwd_step):
    output = stage.forward_one_chunk(...)
    session.protect_aliases(fwd_step, output)

real overlap callback 必须调用统一的 execute_fwd_leaf() / execute_bwd_leaf(),不能绕过 leaf API 直接调用 stage。CommComputeOverlap.run() 在主线程执行 FWD、daemon worker 执行 BWD,并在返回前 join。双 Python 线程只表示 host 可以并发下发,device 是否重叠仍取决于 stream、依赖和硬件资源。

6.6 FSDP、P2P 与 dx/dw

FSDP 目标顺序:

位置 目标顺序
FWD 前 FSDP_UNSHARD → FWD(collection)
FWD 后 FWD → SWAP_LAUNCH_OFFLOAD → FSDP_RESHARD
BWD 前 SWAP_LAUNCH_LOAD → FSDP_UNSHARD → SWAP_WAIT_LOAD → BWD

如果目标 stage 参数已经保持 unsharded,没有可覆盖的 lookahead,则 H2D 延迟到 BWD container 前,避免 activation 过早回到 device。

P2P transport 关系:

模式 与 swap 的关系
plain 基础逐项 isend/irecv 路径
batch swap 不拆 fused container,wait anchor 放在完整通信块之后
boundary planner 在 boundary rewrite 后运行,并在 FWD leaf 持有 output 时显式保护 alias

dx/dw 拆分时,BWD_INPUT 是 activation first consumer,BWD_WEIGHT 只使用 dx 阶段保存的中间状态。因此只在 BWD_INPUT 前插入 SWAP_WAIT_LOAD,BWD_WEIGHT 不重复 load 或 release;stage 0 的 backward 保持统一 BWD。

6.7 Stream、内存与 alias 生命周期

D2H:

producer compute stream
  └─ compute_event
       └─ copy stream: D2H
            └─ offload_event
                 └─ SWAP_WAIT_OFFLOAD 建立依赖
                      └─ resize device storage to 0

H2D:

launch-side compute stream
  └─ compute_event
       └─ copy stream: H2D
            └─ load_event
                 └─ SWAP_WAIT_LOAD 建立依赖
                      └─ BWD/BWD_INPUT 读取 activation

正常 wait_offload() / wait_load() 使用 event 建立 device stream 依赖,不新增 host event.synchronize();异常 teardown 才会对 group 的 in-flight event 做定点同步。

当前 alias 保护包括:

  1. 同一 group 内按 tensor/storage 去重。
  2. D2H 前显式保护 stage 参数、first-stage 输入、recv buffer、当前 FWD output/loss root。
  3. group-managed tensor 在 H2D 后通过 set_() alias 到恢复后的连续 device buffer。

当前不做跨 group 自动 pointer 扫描;通用跨 chunk storage 共享需要后续引入明确的 ownership/generation 模型。

6.8 调度图

调度图使用 scheduler 生成的 finalized order,并基于跨 rank 数据依赖建立全局逻辑时间轴。同一列表示全局逻辑 compute slot,不代表实测 wall-clock;copy 条形表示静态 launch→wait 覆盖窗口。

  • C·F 表示 FWD leaf 在 collection context 中执行。
  • W·B / W·dx 表示对应 consumer container 前执行 SWAP_WAIT_LOAD。
  • offload 从 SWAP_LAUNCH_OFFLOAD 延伸到 SWAP_WAIT_OFFLOAD。
  • load 从 SWAP_LAUNCH_LOAD 延伸到 SWAP_WAIT_LOAD。
  • offload/load 为便于阅读分行绘制,当前实现共用同一条 copy stream,不会彼此并行。
  • 灰色斜纹表示该 rank 没有可执行 compute,即 pipeline bubble。
  • 图中不单独绘制跨 rank P2P、SEND/RECV 或 FSDP MetaStep。

6.8.1 GPipe + swap

GPipe + swap 按 rank 合并调度

GPipe 先执行全部 FWD,再执行全部 BWD,每个 micro-batch 通常有较长 activation 空窗。

6.8.2 VPP / Interleaved 1F1B + swap

VPP + swap 按 rank 合并调度

每个物理 rank 持有多个 virtual stage。group key 使用真实 stage_index,同一 micro id 在不同 virtual stage 之间不会混组。planner 对每个 (stage, micro) 独立判定,没有足够窗口的 chunk 保持 device resident。

6.8.3 real overlap + swap

real overlap + swap 按 rank 合并调度

OVERLAP_B_F(BWD_i, FWD_j) 是一个 top-level composite。主线程执行 FWD leaf,daemon worker 执行 BWD leaf,同一 slot 的 main/worker 属于同一个 composite,返回前 join。

6.8.4 real overlap + dx/dw + swap

real overlap + dx/dw + swap 按 rank 合并调度

SWAP_LAUNCH_LOAD(chunk i)
  → SWAP_WAIT_LOAD(chunk i)
  → OVERLAP_B_F(BWD_INPUT_i, FWD_j)
  → BWD_SEND_i / P2P gap
  → BWD_WEIGHT_i

6.9 代码改动点

模块 改动内容 是否影响已有行为
core/activation_checkpoint 复用 swap group、storage、event、alias 保护能力 否;非 PP 用法保持不变
core/pipeline_parallel/pipeline_swap.py planner、session、4 个 MetaStep handler 仅 swap=True 时生效
core/pipeline_parallel/scheduler.py 新增 swap MetaStep 类型、order 注入和运行时派发 默认 swap=False,不影响已有调度
platform/mindspore/activation_checkpoint swap_wrapper / swap_tensor_wrapper 后端实现 保持已有接口
tests/ut/core/pipeline_parallel order、anchor、session、异常和组合 UT 否
tests/mindspore/st/pipeline_parallel PP + swap 分布式精度/显存 E2E 否

6.10 方案取舍

方案 优点 缺点 是否选择 原因
仅依赖 layer prefetch hook 复用普通 swap 路径,改动少 不感知 micro-batch、virtual stage、FSDP/P2P 和 overlap container 否 无法保证 PP 生命周期和同步位置正确
在 finalized PP order 注入 4 个 MetaStep 调度可观察、可测试;能与 FSDP/P2P anchor 协同 scheduler 需要维护 planner 与 session 是 明确分离激活选择和搬运调度,适配所有 PP schedule
为每个 FWD/BWD 机械插入同步点 实现直接 无窗口 chunk 也付出开销,破坏 overlap 否 不能满足显存/性能收益目标

主要代价是 scheduler order 增加 4 类控制 MetaStep,并需要维护 group generation、alias 保护和异常清理。

7. 组件依赖

依赖组件 强依赖 / 弱依赖 当前状态 未 ready 时本期能力
Activation swap(SwapManager、wrapper) 强依赖 已 ready 无法选择和搬运激活
PP scheduler 强依赖 已 ready 无 PP swap 调度能力
FSDP 弱依赖 GPipe / Interleaved 已验证 无 FSDP 时仍支持 PP swap
P2P transport 强依赖 plain / batch 支持,boundary 设计支持 可退回已验证 transport
checkpoint / recompute 弱依赖 mixed policy + real overlap 待准入 纯 swap 场景可交付
optimizer 不涉及 不涉及 不影响
MindSpore 后端 强依赖 已支持 本期无可交付运行后端
PyTorch 后端 弱依赖 PP + swap 暂未准入 MindSpore 能力不受影响
完整能力需要:activation wrapper + PP finalized order + SwapManager event/stream 生命周期。
本期最小可交付能力:MindSpore 上 GPipe / 1F1B / Interleaved / real overlap 的 PP swap 调度。

8. 约束与兼容性

8.1 支持矩阵

组合 当前状态 看护
GPipe + swap MindSpore 支持 order UT;无独立无 FSDP E2E
1F1B + swap MindSpore 支持 order UT;可能没有有效 swap window
Interleaved 1F1B + swap MindSpore 支持 order UT;E2E 由 FSDP + Interleaved 覆盖
real overlap + swap MindSpore 支持 8 卡 E2E
real overlap + dx/dw + swap MindSpore 支持 8 卡 E2E
FSDP + GPipe / Interleaved + swap MindSpore 支持 4 卡 E2E
FSDP + real overlap + swap 设计支持,不禁入 专项 E2E 待补
boundary + swap 设计支持,不禁入 专项 E2E 待补
PyTorch PP + swap 运行时拒绝 独立准入前不开放
MPipe + swap 不支持 preprocess/body group ownership 尚未定义
recompute/mixed policy + real overlap 待准入 collection 边界待验证
forward-only 不启用 没有 backward consumer

8.2 兼容性与收益

类型 内容
不支持项 PyTorch PP、MPipe、forward-only;recompute/mixed policy + real overlap 待准入
性能收益 D2H/H2D 尽量覆盖其它 chunk compute、FSDP all-gather 或 P2P;固定 gap 门槛不保证所有模型净收益
显存收益 有有效 window 的 chunk 在 FWD/BWD 空窗期释放 device activation storage
性能劣化 copy、allocator 和 MetaStep 调度有额外开销;gap < 4 时不注入以规避明显无收益场景
PT / MS 差异 本期仅 MindSpore PP 准入;PyTorch PP 运行时拒绝
已有行为 swap=False 为默认值;不应用 wrapper 时 swap=True 只有调度、没有可搬运激活;非 PP swap 接口不变

8.3 正确性不变量

  1. 每次正常运行使用独立 generation,上一轮 active group 回到基线。
  2. 只有 eligible FWD leaf 能收集到自己的 chunk group。
  3. D2H 必须发生在 FWD 后,H2D 必须在对应 BWD / BWD_INPUT 消费前建立 stream 依赖。
  4. 每个 eligible chunk 必须按 4 个 MetaStep 完整执行 launch/wait 生命周期。
  5. pipeline-owned alias 不允许被 resize_(0)。
  6. BWD_WEIGHT 不重复执行 activation load/release。
  7. swap 不拆 P2P batch,不改变跨 rank send/recv 配对顺序。
  8. 异常按训练终止处理并执行 best-effort 清理,不保证异常后的同进程重试。

9. 验证设计

9.1 用例分层

用例级别 数量 覆盖内容 通过标准
UT / order 多组 4 个 MetaStep 顺序、gap、composite、FSDP anchor、dx/dw、batch/boundary、session generation、alias/event order 与不变量一致,测试全部通过
Level0 2 组 FSDP + GPipe/Interleaved、real overlap + dx/dw loss/梯度一致,无 hang,峰值显存下降
Level1 1 组 real overlap + swap,PP=4、EP=2 loss/全部本地梯度一致,worker 正常,峰值显存下降

9.2 交互验证

组合 是否验证 通过标准
GPipe / 1F1B / Interleaved + swap 是 4 个 MetaStep order 合法;无 window 时不注入
swap + FSDP 是 精度一致;LAUNCH_LOAD → FSDP_UNSHARD → WAIT_LOAD → BWD
swap + real overlap 是 loss/梯度一致,无 hang,线程完成
swap + dx/dw 是 BWD_INPUT 前 wait;BWD_WEIGHT 不重复 load
swap + P2P plain/batch 是 send/recv 配对不变,无 hang
swap + boundary 待补 精度、无 hang、异常 drain、多轮运行
PT / MS 对齐 否 本期 PT 未准入,运行时明确拒绝

9.3 性能 / 显存验证

场景 基线 开启本特性 指标 已验证结果 / 通过标准
FSDP + GPipe no-swap 5 step swap 5 step loss、梯度、峰值显存 精度一致;峰值显存下降约 17.0%~19.8%
FSDP + Interleaved no-swap 5 step swap 5 step loss、梯度、峰值显存 精度一致;峰值显存下降约 3.4%~11.0%
real overlap no-swap 3 step swap 3 step loss、全部本地梯度、峰值显存 精度一致;峰值显存下降约 19.8%~49.4%
real overlap + dx/dw no-swap 3 step swap 3 step loss、全部本地梯度、峰值显存 精度一致;峰值显存下降约 26.1%~52.1%

real-overlap 两个 E2E 使用默认 auto → batch transport,不覆盖 FSDP + real overlap 或 boundary + swap。

10. 实现计划

PR / 阶段 内容 依赖 验证 状态
PR1 PP swap planner、run-scoped session、4 个 MetaStep 基础 activation swap UT / order 已完成
PR2 GPipe / 1F1B / Interleaved、FSDP、P2P 组合 PR1 UT + Level0 已完成
PR3 real overlap 与 dx/dw PR2 Level0 + Level1 已完成
后续 1 FSDP + real overlap、dx/dw 变体、boundary 专项 E2E PR3 Level1 待补
后续 2 PyTorch PP + swap 独立准入 后端实现 UT + Level0 + Level1 待补
后续 3 recompute/mixed policy + real overlap collection 边界验证 Level1 待补
后续 4 activation bytes、带宽和 compute time 驱动的 cost model 性能数据 性能回归 待补
后续 5 composite 内 FWD 完成后提前 D2H、多 copy stream、buffer 池优化 cost model 性能/显存 待评估
likedislike
DavidFFFanDavidFFFan成员
5月29日 修改了issue 的描述
此处折叠了12条事件消息 查看更多
DavidFFFanDavidFFFan成员
8月10日 修改了issue 的描述
hui-zhang940
8月12日 评论:

如何确认某个 tensor 被正确 swap?是否有验证方式?

likedislike
hui-zhang940
8月12日 评论:

swap_tensor_wrapper没转测过,是新增的接口么

likedislike
hui-zhang940
8月12日 评论:

swap_tensor_wrapper没转测过,是新增的接口么

@hui-zhang940

应客户要求的临时特性,不是一开始排的需求,组件这边对于该接口未验收过,可以泛化验证一下这个接口

likedislike
hui-zhang940
8月12日 评论:

如何确认某个 tensor 被正确 swap?是否有验证方式?

@hui-zhang940

@hui-zhang940 自能通过显存变化判断 tensor 被swap

likedislike
DavidFFFanDavidFFFan成员
8月23日 修改了issue 的描述
fangwenyifangwenyi成员
13 天前 issue状态由 ACCEPTED 改变为 DONE
fangwenyifangwenyi成员
13 天前 关闭了 issue