[RFC]: MULTI_STREAM_MEMORY_REUSE默认值修改为2-pytorch-AtomGit

已关闭

[RFC]: MULTI_STREAM_MEMORY_REUSE默认值修改为2 #4195

zzzzzzk创建于 8月19日关闭于 8月27日

8月19日创建

状态（Status）： Draft / Reviewing / Approved / Rejected / Superseded
作者（Authors）： @Your_community
创建日期（Created）： YYYY-MM-DD
更新日期（Updated）： YYYY-MM-DD
相关 Issue/PR： #123（关联 Issue/PR 以便追踪背景）

1. 概述

简介

本提案通过 MULTI_STREAM_MEMORY_REUSE 控制计算流与 HCCL 通信流之间的内存复用策略，解决集合通信相关内存释放不及时、内存峰值偏高的问题。核心价值是在保证通信正确和内存安全的前提下，提高设备内存利用率，支持更大模型或更大 Batch Size 的分布式训练。

动机

在大模型分布式训练中，计算和通信通常在不同流上并行执行。为保证通信期间 Tensor 不被提前释放，系统需要记录跨流依赖，但这也会导致部分已经使用完的内存不能及时返回内存池。当通信频繁或任务下发较快时，内存会持续滞留，增加峰值内存，甚至引发 OOM。
因此，需要在确认通信安全完成后及时解除不再需要的内存依赖，同时覆盖 Tensor 提前释放、兼容性回退等场景。不实施本提案，用户可能需要降低 Batch Size、缩小模型规模或增加设备资源，从而影响训练效率和资源利用率。

目标

在保证通信正确性和内存安全的前提下，提高多流内存复用率。
及时回收集合通信使用完毕的内存，降低训练峰值内存。
无需修改训练代码，通过环境变量即可控制和定位问题。
减少内存复用event的创建和使用

2. 用例分析

1.4 用例分析

本提案主要面向以下场景：

异步集合通信：用户发起 all_reduce 等异步通信，并能够在使用通信结果前明确调用 work.wait()。
Tensor 提前释放：用户在通信完成前释放 Tensor 引用。Mode 2 通过持有强引用，保证相关内存不会被提前回收或覆盖。
通信频繁的训练：在大模型分布式训练中，集合通信频繁，传统 recordStream 可能造成内存延迟回收。Mode 2 在通信同步后释放强引用，使内存返回内存池。

3. 方案设计

3.1 总体方案

Mode 2 面向 TorchNPU 分布式集合通信场景，采用“强引用托管 Tensor 生命周期”的设计：通信开始时不执行 recordStream，而是由通信任务保存输入、输出 Tensor 的强引用；达到安全同步点后释放引用，使内存可以返回缓存内存池。

整体设计涉及以下模块：

配置管理：读取并缓存 MULTI_STREAM_MEMORY_REUSE=2。
HCCL 通信层：识别 Mode 2，为参与通信的 Tensor 建立生命周期保护。
WorkHCCL：通过 TensorShelf 保存 Tensor 强引用。
缓存分配器：Tensor 引用释放后，按正常流程回收内存。
Watchdog：用户未调用 wait() 时，在检测到通信完成后触发延迟清理。

核心流程如下：

sequenceDiagram
    participant U as 用户程序
    participant PG as ProcessGroupHCCL
    participant W as WorkHCCL/TensorShelf
    participant H as HCCL通信流
    participant A as NPU缓存分配器

    U->>PG: 发起异步集合通信
    PG->>PG: 检查 Mode 2
    PG->>W: 保存通信Tensor强引用
    PG->>H: 下发通信任务
    PG-->>U: 返回Work对象

    alt 用户先调用 work.wait()
        U->>W: work.wait()
        W->>H: 当前计算流等待HCCL结束事件
        W->>W: 释放Tensor强引用
        U->>A: 释放Tensor

        alt Tensor无其他引用
            A-->>A: 内存返回缓存池
        else Tensor仍有其他引用
            A-->>U: 等待其他引用释放
        end

    else Tensor在wait前释放
        U->>U: 释放Tensor
        W-->>W: Work中的强引用继续保持Tensor存活
        Note over W,H: HCCL通信期间内存不会被提前回收

        U->>W: work.wait()
        W->>H: 当前计算流等待HCCL结束事件
        W->>W: 释放Tensor强引用
        A-->>A: Tensor无其他引用，内存返回缓存池
    end

如果用户未调用 work.wait()，Watchdog 检测到通信完成后，会将待释放的 TensorShelf 交给主线程，并在后续调用 ProcessGroupHCCL 接口时完成清理。因此，推荐用户及时调用 wait()，以获得更及时、可预测的内存释放效果。

限制和约束：

仅适用于 TorchNPU HCCL 集合通信场景。
用户应保证异步通信任务最终调用 work.wait()。
Mode 2 对点对点通信不生效，将回退到常规 recordStream 机制。
强引用会延长 Tensor 生命周期；如果 wait() 调用较晚，可能增加而不是降低峰值内存。
不保证固定的性能或内存收益，实际效果取决于通信频率、Tensor 大小和同步时机。

mode0、mode1和mode2的区别
mode2从mode0和mode1演化而来，mode0不感知work.wait()中插入的流间依赖，总是会插入event，保证内存回收在集合通信完成之后，核心流程如下：

sequenceDiagram
    participant U as 用户程序
    participant PG as ProcessGroupHCCL
    participant W as WorkHCCL
    participant H as HCCL通信流
    participant A as NPU缓存分配器

    U->>PG: 发起异步集合通信
    PG->>PG: 检查 Mode 0
    PG->>A: recordStream记录HCCL通信流
    PG->>H: 下发通信任务
    PG-->>U: 返回Work对象

    opt 用户调用 work.wait()
        U->>W: work.wait()
        W->>H: 当前计算流等待HCCL结束事件
        Note over W,A: 仅建立流间同步<br/>不执行eraseStream
    end

    U->>A: Tensor释放
    A->>A: 检查block->stream_uses
    A->>H: 记录回收事件
    A->>A: 将Block加入待处理事件队列

    U->>A: 后续申请新的NPU内存
    A->>A: malloc()调用process_events()
    A->>H: 查询回收事件是否完成

    alt 回收事件已完成
        A->>A: event_count减一
        A->>A: 从待处理队列移除Event
        A-->>U: Block返回缓存池并可复用
    else 回收事件未完成
        A-->>A: Block继续保留，暂不可复用
    end

mode1 在mode0的基础上增加了work.wait()对流同步的感知，一旦用户调用了work.wait()，流间依赖得到保证，那么就会调用erase_stream()，这样在tensor释放的时候不回去创建event，核心流程如下：

sequenceDiagram
    participant U as 用户程序
    participant PG as ProcessGroupHCCL
    participant W as WorkHCCL
    participant H as HCCL通信流
    participant A as NPU缓存分配器

    U->>PG: 发起异步集合通信
    PG->>PG: 检查 Mode 1
    PG->>A: recordStream记录HCCL通信流
    PG->>W: 保存Tensor弱引用
    PG->>H: 下发通信任务
    PG-->>U: 返回Work对象

    alt 用户先调用 work.wait()
        U->>W: work.wait()
        W->>H: 当前计算流等待HCCL结束事件
        W->>A: eraseStream清除跨流标记
        U->>A: 释放Tensor
        A-->>A: 内存返回缓存池
    else Tensor在wait前释放
        U->>A: Tensor析构，触发内存释放
        A->>A: 检查block->stream_uses
        A->>H: 记录回收事件
        A->>A: 将Block加入待处理事件队列

        U->>A: 后续申请新的NPU内存
        A->>A: malloc()调用process_events()
        A->>H: 查询回收事件是否完成

        alt HCCL通信已完成
            A->>A: event_count减一
            A-->>U: Block返回缓存池并可复用
        else HCCL通信未完成
            A-->>A: Block继续保留，暂不可复用
        end
    end

3.2 技术选型

本提案选择“避免 recordStream，由 WorkHCCL 持有 Tensor 强引用”的方案。其他方案对比如下：

方案	优点	缺点	选择结论
mode 0	机制成熟、兼容性高	内存可能延迟回收	作为关闭优化和回退方案，不满足 Mode 2 的设计目标
mode 1	不延长 Tensor 生命周期	Tensor 提前析构后可能无法找到对应 Storage	作为 Mode 1 保留，不作为 Mode 2 方案
mode 2	不依赖 `recordStream`；实现简单；不需要访问分配器内部 Block	可能延长 Tensor 生命周期；依赖明确的完成或同步时机	选为 Mode 2 的核心方案
mode 3	Tensor 提前析构后仍可清除标记	与缓存分配器内部实现耦合，部分自定义分配器不支持	作为 Mode 3 提供，不纳入 Mode 2

Mode 2 的主要优势是不需要修改缓存分配器的 Block 管理逻辑，能够以较小的实现复杂度保证通信期间的内存安全。

3.3 功能与性能设计

功能实现包括：

进程首次读取环境变量时，将值 2 解析为 AVOID_RECORD_STREAM。
发起集合通信时，不对相关 Tensor 执行 recordStream。
将通信输入、输出及必要的中间 Tensor 保存到 TensorShelf。
用户调用 work.wait() 后，先建立计算流对通信流的事件等待关系，再清空 TensorShelf。
如果用户未调用 wait()，由 Watchdog 在通信完成后登记待清理对象，并由后续主线程调用完成释放。
点对点通信无法可靠判断用户是否会调用 wait()，因此回退到常规内存保护机制。

影响范围：

环境变量解析模块。
ProcessGroupHCCL 集合通信处理流程。
WorkHCCL 生命周期管理。
Watchdog 完成任务后的资源清理流程。

flowchart TD
    A[异步通信提交<br/>TensorShelf保存强引用] 
    --> B[Work加入Watchdog监控列表]

    B --> C{HCCL通信完成?}
    C -->|否| B
    C -->|是| D[Watchdog转移TensorShelf<br/>并删除已完成Work]

    D --> E[下一次主线程提交通信任务]
    E --> F[调用unstash<br/>清空TensorShelf]

    F --> G[释放Tensor强引用]
    G --> H[无其他引用时<br/>内存进入NPU缓存分配器]

性能设计：

预期收益：减少 recordStream 及相关事件管理开销，使通信同步后的内存释放路径更加直接。
潜在开销：保存、释放 Tensor 强引用以及 TensorShelf 加锁操作。
达成收益的前提：通信 Tensor 较大或通信频繁，并且用户能够及时调用 work.wait()。
负向场景：异步任务长期不等待时，Tensor 强引用持续存在，可能提高峰值内存。
本提案不设置固定的内存下降比例，应通过典型模型实测峰值内存和迭代耗时。

3.4 安全隐私与DFX设计

安全与可靠性：

通信完成前始终持有 Tensor 强引用，防止内存被提前释放或重新分配。
只有建立安全同步关系或确认通信完成后，才能释放强引用。
点对点通信自动回退，避免因用户未调用 wait() 引发越界访问或数据错误。

3.5 编程与调用设计

Mode 2 不新增 Python 或 C++ API。开发者通过环境变量启用功能，并继续使用现有 torch.distributed 编程模型。

3.5.1 编程模型基本设计

开发环境设计：明确好开发者使用的软/硬件环境、开发&调试工具链、编程框架、要提供的加速库或算子等。
开发约束：开发者使用过程中的约束和限制说明，如硬件平台、编程语言限制等。
可验收设计：提供相应功能、性能指标等的验收环境、标准或用例设计，保证最终的实现可达成既定目标。

3.5.2 接口定义与设计

本提案不新增编程 API，仅提供环境变量配置接口。

3.5.2.1 MULTI_STREAM_MEMORY_REUSE

接口描述：配置 TorchNPU 多流内存复用策略；取值 2 时启用强引用托管方案。

接口原型：

export MULTI_STREAM_MEMORY_REUSE=2

输入参数：

参数名称	输入/输出	类型	描述	取值范围
`MULTI_STREAM_MEMORY_REUSE`	输入	整数字符串	多流内存复用模式	`0、1、2、3`；本提案使用 `2`

返回参数：无。

异常处理：设置为有效范围之外的数值时，初始化或首次读取配置时抛出参数错误。

约束说明：

应在 HCCL 进程组初始化前设置。
首次读取后配置被缓存，运行过程中修改不会生效。
Mode 2 对点对点通信不生效。

变更说明：不新增或修改 torch.distributed API，仅增加对应的内部内存管理策略。

调用参考：

export MULTI_STREAM_MEMORY_REUSE=2
torchrun --nproc_per_node=8 train.py

3.5.3 编程手册设计

为了帮助开发者能快速上手开发，要设计好本提案相关特性/功能的《编程手册》，要包含哪些内容和章节，单独输出还是共用，在已有的手册中更新还是输出等。确保最后输出的《编程手册》中有相关变更内容。

4. 测试设计

4.1 端到端测试

端到端测试应选择具有代表性的分布式训练任务，例如 DDP 大模型训练。
选择的模型：deepseek_16b
开启：DP =2,FSDP=2,TP=2,EP=2,确保all_reduce, all_gather, reduce_scatter, all_to_all都能运行到

测试步骤：

分别使用 Mode 1 和 Mode 2 运行相同的训练任务1000步。
保持模型、数据集、Batch Size、设备数量和随机种子一致。
完成预热后运行多个稳定迭代。
记录训练结果、峰值内存、单步耗时和吞吐率。
进行多轮或长时间稳定性测试。

验收内容：

Loss、模型参数及训练结果与基线保持一致。
无 NaN、通信超时、内存越界或进程异常。
调用 wait() 后，通信 Tensor 不再被无期限持有。

5. 缺点和风险

风险类型	风险说明	应对措施
内存占用增加	Mode 2 持有 Tensor 强引用。如果用户延迟或未调用 `work.wait()`，相关内存将长期无法回收，峰值内存可能上升。
通信兼容性风险	点对点异步通信无法保证用户一定调用 `wait()`，直接使用 Mode 2 存在提前释放或长期持有风险。
内存泄漏风险	如果通信异常、Watchdog 未执行，且用户未调用 `wait()`，强引用可能长时间保留。此外该需求依赖开启watchdog

6. 现有技术

本提案主要参考 PyTorch ProcessGroupNCCL 的多流内存管理机制。

在 CUDA 多流场景中，PyTorch 通常使用 record_stream 防止 Tensor 内存在异步任务完成前被复用。该机制安全性较高，但会增加事件查询开销，并使内存复用时机不够可预测。PyTorch 官方也指出：当框架能够管理 Tensor 完整生命周期和流间同步时，可以通过显式同步替代 record_stream。PyTorch record_stream 文档

ProcessGroupNCCL 采用了与 Mode 2 类似的方案：不执行 recordStream，而是通过 TensorShelf 保存通信 Tensor 的强引用；在用户计算流等待 NCCL 通信结束事件后，再释放这些引用，从而兼顾跨流安全和内存复用。当前 PyTorch 主线已将该行为设为默认，并将 TORCH_NCCL_AVOID_RECORD_STREAMS 标记为废弃配置。ProcessGroupNCCL 设计、ProcessGroupNCCL 实现