文章目录
- [1. ZeRO-1 流程。，特别是关于通信量和为什么 `Reduce-Scatter` 就足够了的思考。](#1. ZeRO-1 流程。，特别是关于通信量和为什么
Reduce-Scatter就足够了的思考。) -
-
- 图一：前向与后向传播阶段 (Forward & Backward)
- 图二：梯度同步与参数更新阶段 (Gradient Sync & Weight Update)
-
- [**步骤 (2): 梯度同步 (Gradient Synchronization)**](#步骤 (2): 梯度同步 (Gradient Synchronization))
- [**对通信量的思考：为什么 `Reduce-Scatter` 就够了？**](#对通信量的思考：为什么
Reduce-Scatter就够了？) - [**步骤 (3): 参数更新 (Weight Update)**](#步骤 (3): 参数更新 (Weight Update))
- [**缺失的最后一步：参数同步 (Weight Synchronization)**](#缺失的最后一步：参数同步 (Weight Synchronization))
- [完整 ZeRO-1 流程总结](#完整 ZeRO-1 流程总结)
-
- [2. 在混合精度训练的前提下，这边我们对【只能将相应的W（蓝色部分）进行更新】这句话做更严谨的说明， 什么意思？](#2. 在混合精度训练的前提下，这边我们对【只能将相应的W（蓝色部分）进行更新】这句话做更严谨的说明， 什么意思？)
-
-
- [1. 混合精度训练的标准流程](#1. 混合精度训练的标准流程)
- [2. 将混合精度训练与 ZeRO-1 结合](#2. 将混合精度训练与 ZeRO-1 结合)
- [3. 对"只能将相应的W（蓝色部分）进行更新"的严谨说明](#3. 对“只能将相应的W（蓝色部分）进行更新”的严谨说明)
-
- 代码实战
- 3.火焰图分析
-
-
- 图表解读 (时间线分析)
-
- [**阶段一：时间 0ms - ~280ms (静态内存布局)**](#阶段一：时间 0ms - ~280ms (静态内存布局))
- [**阶段二：时间 ~280ms - ~420ms (反向传播 -> 梯度分配)**](#阶段二：时间 ~280ms - ~420ms (反向传播 -> 梯度分配))
- [**阶段三：时间 ~420ms - ~480ms (优化器步骤)**](#阶段三：时间 ~420ms - ~480ms (优化器步骤))
- [**阶段四：时间 ~480ms 之后 (后续的训练步骤)**](#阶段四：时间 ~480ms 之后 (后续的训练步骤))
- 总结：为什么图是这个样子的？
- [1. `model_engine.step()` 是在什么时候执行的？](#1.
model_engine.step()是在什么时候执行的？) - [2. 蓝色块和"10条"：是什么？为什么这么大？](#2. 蓝色块和“10条”：是什么？为什么这么大？)
-
- [A. 巨大的蓝色块 (GRADIENT)](#A. 巨大的蓝色块 (GRADIENT))
- [B. "10条"细长的尖峰 (多个训练 Step)](#B. “10条”细长的尖峰 (多个训练 Step))
- [第一步：计算我们代码中 `SimpleModel` 的总参数量](#第一步：计算我们代码中
SimpleModel的总参数量) - [第二步：根据参数量计算 `SimpleModel` 的梯度大小](#第二步：根据参数量计算
SimpleModel的梯度大小) - 第三步：对比理论计算与图中观测结果
- 最终结论
-
- 问题定位
- 稳定后内存分析
- 启动前内存
- 稳定后
https://zhuanlan.zhihu.com/p/618865052
学习这篇文章遇到的问题
https://docs.pytorch.org/memory_viz 分析pickle 内存占用
torch memory 工具使用https://pytorch.org/blog/understanding-gpu-memory-1/
https://pytorch.cadn.net.cn/docs_en/2.5/torch_cuda_memory.html
chrome://tracing/ json 分析
1. ZeRO-1 流程。，特别是关于通信量和为什么 Reduce-Scatter 就足够了的思考。
图一：前向与后向传播阶段 (Forward & Backward)
这张图清晰地展示了 ZeRO-1 在进行 loss.backward() 之后，但在 optimizer.step() 之前的状态。
- GPU 数量: N=3 (GPU1, GPU2, GPU3)
- 输入数据 (X) : 一个 batch 的数据被分成了三份 (X1, X2, X3)，每个 GPU 处理一份。这是数据并行 (Data Parallelism) 的核心。
- 模型参数 (W) : 每个 GPU 都拥有一份完整的、相同的模型参数副本 (W)。这是标准数据并行和 ZeRO-1/2 的共同点。在图中用 FP16 表示，是混合精度训练的体现。
- 优化器状态 (O) : 这是 ZeRO-1 的关键创新点！ 优化器状态被分区 (Partitioned) 。
- GPU1 只存储优化器状态的前 1/3 (O1)。
- GPU2 只存储中间的 1/3 (O2)。
- GPU3 只存储最后的 1/3 (O3)。
- 由于优化器状态 (O) 是 FP32，并且通常是参数量的两倍（Adam），这个分区操作极大地节省了每个 GPU 的显存。
- 梯度 (G) :
- 经过前向传播 (
forward) 和后向传播 (backward)，每个 GPU 都基于自己的数据分片 (X1, X2, X3) 计算出了一份完整的梯度。 - 我们称之为本地梯度
G1,G2,G3。这三份梯度内容不同，但形状和大小都与完整模型参数 W 相同。
- 经过前向传播 (
小结: 图一的状态是"万事俱备，只欠同步"。每个 GPU 都有了本地梯度，但需要将这些本地梯度聚合成一个全局一致的梯度，才能进行统一的参数更新。
图二：梯度同步与参数更新阶段 (Gradient Sync & Weight Update)
这张图展示了 optimizer.step() 内部发生的核心操作。
步骤 (2): 梯度同步 (Gradient Synchronization)
- 操作 : 对本地梯度
G1, G2, G3执行一次All-Reduce操作。 All-Reduce的过程 :- Reduce : 所有 GPU 将它们的本地梯度相加，得到一个全局的总梯度
G_sum = G1 + G2 + G3。 - Broadcast : 将这个
G_sum(或者除以N后的平均梯度G_avg) 分发回每一个 GPU。
- Reduce : 所有 GPU 将它们的本地梯度相加，得到一个全局的总梯度
- 结果 :
All-Reduce完成后，GPU1, GPU2, GPU3 上都拥有了同一份完整的、全局同步好的梯度 G。 - 通信量 :
- 如您所述，一次高效的 Ring-AllReduce 的通信量大约是
2 * (N-1)/N * M，其中M是梯度张量的总大小。为了简化，我们常说它与模型参数量成正比，约为2 * M。
- 如您所述，一次高效的 Ring-AllReduce 的通信量大约是
对通信量的思考：为什么 Reduce-Scatter 就够了？
这是您提到的一个非常深刻的洞见！
All-Reduce的问题 : 它给了每个 GPU 一份完整的梯度 G，但实际上每个 GPU 真的需要完整的 G 吗？- 回顾 ZeRO-1 的本质 : GPU1 只负责更新 W 的前 1/3，因为它只持有 O1。因此，理论上 GPU1 只需要 G 的前 1/3 就可以完成它的更新任务。
Reduce-Scatter的优势 : 这个操作正是为此而生。它会计算出全局总梯度G_sum，然后立即将其切分 ，只把每个 GPU 需要的那一部分发回给它。- GPU1 收到
G_sum的前 1/3。 - GPU2 收到
G_sum的中间 1/3。 - GPU3 收到
G_sum的后 1/3。
- GPU1 收到
Reduce-Scatter的通信量 : 大约是(N-1)/N * M，几乎是All-Reduce的一半。- 结论 : 没错，对于 ZeRO-1 来说，使用
Reduce-Scatter来同步梯度是完全足够且更高效 的。这实际上已经预示了 ZeRO-2 的核心思想！可以说，ZeRO-2 就是将 ZeRO-1 中这个"理论上可以但早期没实现"的优化给正式化和系统化了。
步骤 (3): 参数更新 (Weight Update)
- 操作 : 每个 GPU 使用它拥有的完整的全局梯度 G 和它自己保管的部分的优化器状态 O，来更新它所负责的那部分参数 W。
- GPU1 :
- 输入: 完整的 G, 部分的状态 O1。
- 计算: 使用 O1 和 G 的前 1/3，计算出 W 的前 1/3 的更新量。
- 更新: 将 W 的前 1/3（图中蓝色部分）进行更新。
- GPU2 和 GPU3: 同理，分别更新 W 的中间 1/3 和最后 1/3。
- 问题出现 : 更新后，每个 GPU 上的 W 变得不一致 了！
- GPU1 上的 W:
[W1_new, W2_old, W3_old] - GPU2 上的 W:
[W1_old, W2_new, W3_old] - GPU3 上的 W:
[W1_old, W2_new, W3_new]
- GPU1 上的 W:
缺失的最后一步：参数同步 (Weight Synchronization)
为了开始下一次的前向传播（它要求所有 GPU 上的模型参数 W 必须是相同的），我们还需要最后一步，这在图二中没有画出来：
- 操作 : 对更新后的、不一致的 W 执行一次
All-Gather操作。 - 过程: 每个 GPU 将自己更新的那一小部分 W 发送给所有其他 GPU。
- 结果 :
All-Gather结束后，所有 GPU 都重新拥有了一份完整的、最新的、完全一致的模型参数 W。 - 通信量 :
All-Gather的通信量也大约是(N-1)/N * M。
完整 ZeRO-1 流程总结
- Fwd/Bwd : 各自计算本地完整梯度
Gi。 - 梯度同步 :
All-Reduce梯度 (或者更优的Reduce-Scatter)。 - 本地更新: 各自用完整的 G 和部分的状态 O，更新部分参数 W。
- 参数同步 :
All-Gather更新后的参数 W，使所有 GPU 上的 W 再次保持一致。
这套流程清晰地展示了 ZeRO-1 是如何通过分区优化器状态来节省显存，以及为了维护数据一致性所引入的通信开销。您对于通信量和 Reduce-Scatter 优化的思考非常到位，这正是从 ZeRO-1 迈向 ZeRO-2 的关键一步。
2. 在混合精度训练的前提下，这边我们对【只能将相应的W（蓝色部分）进行更新】这句话做更严谨的说明， 什么意思？
好的，我们来对"只能将相应的W（蓝色部分）进行更新 "这句话，在混合精度训练 (Mixed-Precision Training) 的背景下，做一个更严谨和深入的说明。
这句话的背后，其实隐藏了一个混合精度训练的关键实现细节：FP32主参数副本 (FP32 Master Parameters)。
1. 混合精度训练的标准流程
在标准的混合精度训练中（不考虑 ZeRO），一个 optimizer.step() 的流程是这样的：
-
FP16 前向/后向传播:
- 模型参数
W是 FP16 格式，以节省显存和加速计算。 loss.backward()计算出的梯度G也是 FP16 格式。
- 模型参数
-
梯度解缩放 (Unscaling):
- 为了防止数值下溢，梯度在计算时被一个大的缩放因子
S乘以了。在更新前，需要将梯度除以S恢复原值。
- 为了防止数值下溢，梯度在计算时被一个大的缩放因子
-
参数更新 (在 FP32 上进行):
- 关键 : 为了避免在多次迭代中微小的梯度更新量因为 FP16 精度不足而被"舍入"为零，参数的更新操作是在一个高精度 (FP32) 的主副本上进行的。
- 优化器维护着一份与模型参数
W完全相同的、但数据类型为 FP32 的主参数副本 (我们称之为W_master_fp32)。 - 更新公式 :
W_master_fp32_new = optimizer_update(W_master_fp32_old, G_fp16, O_fp32)- 优化器读取旧的
W_master_fp32。 - 将 FP16 的梯度
G_fp16转换为 FP32。 - 结合 FP32 的优化器状态
O_fp32(momentum, variance)。 - 计算出新的 FP32 主参数
W_master_fp32_new。
- 优化器读取旧的
-
同步 FP16 参数:
- 将更新后的
W_master_fp32_new拷贝并转换 回 FP16 格式，覆盖模型原本的 FP16 参数W。 W_fp16 = cast_to_fp16(W_master_fp32_new)- 这样，下一次的前向传播就可以使用最新的 FP16 参数了。
- 将更新后的
2. 将混合精度训练与 ZeRO-1 结合
现在，我们把这个流程应用到 ZeRO-1 的场景中。
- ZeRO-1 的核心 : 分区 (Partition) 了优化器状态
O_fp32。 - 混合精度的要求 : 仍然需要一份
W_master_fp32来进行精确更新。
那么，这份 W_master_fp32 应该如何存放呢？ZeRO-1 的设计者做出了一个合乎逻辑的选择：
既然优化器状态 O_fp32 已经被分区了，那么与之紧密相关的 W_master_fp32 也应该以同样的方式进行分区。
所以，在 ZeRO-1 中：
GPU1持有: 完整的W_fp16, 完整的G_fp16, 分区的O1_fp32, 以及分区的W1_master_fp32。GPU2持有: 完整的W_fp16, 完整的G_fp16, 分区的O2_fp32, 以及分区的W2_master_fp32。GPU3同理。
3. 对"只能将相应的W（蓝色部分）进行更新"的严谨说明
结合以上背景，我们可以对这句话进行精确的拆解和说明了。
这句话实际上描述了两个层面的更新：
- 对 FP32 主参数副本 的更新。
- 对 FP16 模型参数 的更新。
以 GPU1 为例，它的更新流程是：
-
准备更新主参数:
- GPU1 拥有完整的、全局同步好的梯度
G_fp16。 - 它从
G_fp16中取出前 1/3 的部分，即G1_fp16。 - 它拥有自己的优化器状态分片
O1_fp32。 - 它拥有自己的主参数分片
W1_master_fp32。
- GPU1 拥有完整的、全局同步好的梯度
-
执行主参数更新 (在 FP32 上):
- GPU1 只能 使用它拥有的这三样东西 (
G1_fp16,O1_fp32,W1_master_fp32) 来计算新的主参数分片。 W1_master_fp32_new = optimizer_update(W1_master_fp32_old, G1_fp16, O1_fp32)- 由于 GPU1 缺少
O2_fp32,O3_fp32以及W2_master_fp32,W3_master_fp32，它无法 计算出模型参数中间 1/3 和最后 1/3 的精确更新。这就是"只能"的根本原因。
- GPU1 只能 使用它拥有的这三样东西 (
-
更新 FP16 模型参数
W:- 此时，GPU1 计算出了
W1_master_fp32_new。 - 它将这个新的 FP32 分片转换成 FP16 格式，即
W1_fp16_new。 - 然后，它用
W1_fp16_new去覆盖 它本地存储的那份完整W_fp16的前 1/3 部分。 - 这就是图中"将相应的W（蓝色部分）进行更新"的直观体现。
- 此时，GPU1 计算出了
所以，更严谨的说法是:
"在混合精度训练的背景下，由于每个 GPU 只持有一部分（1/N）的 FP32 优化器状态 和FP32 主参数副本 ，因此在 optimizer.step() 阶段，它只能计算出并更新自己所负责的那一部分主参数副本。随后，这个更新后的 FP32 参数分片被转换回 FP16，并用于覆盖本地存储的、完整的 FP16 模型参数 W 中相应的位置（即图中的蓝色部分）。"
这个过程确保了参数更新的数值稳定性（在 FP32 上进行），同时通过分区 O_fp32 和 W_master_fp32 实现了显存的节省。而代价就是，更新后每个 GPU 上的 W_fp16 变得不一致，需要后续的 All-Gather 来同步。
好的，我们来详细解读这张 PyTorch Memory Snapshot 的火焰图。这张图非常直观地展示了 ZeRO-1 在混合精度训练下的内存布局和动态变化，完全印证了我们之前的理论分析。
作者的核心观点:
实际上，zero1并不是直接更新蓝色权重（fp16），而是直接更新红色优化器中维护的权重(fp32)，蓝色权重是由红色权重cast而来。
这个观点是完全正确 的，它精准地描述了混合精度训练中"FP32主参数副本"的工作机制。图中的"红色权重"就是我们之前讨论的 W_master_fp32。
解读火焰图的各个部分
这张图是一个时间线上的内存占用图，Y轴代表显存占用大小，X轴代表时间。不同的色块代表不同的张量（Tensor）占用的内存。
底部的、持久存在的内存块
-
灰色块:
fp16 完整的weight，参与fwd计算- 内容 : 这就是模型本身的参数，
W_fp16。 - 特点 :
- 它是 FP16 格式，所以相对较小。
- 它是完整的，每个 GPU 都有一份。
- 它在整个训练过程中持久存在于显存中，因为每一次前向和后向传播都需要用到它。
- 这就是图中的"蓝色权重"所指的实体。
- 内容 : 这就是模型本身的参数，
-
橙色块:
fp32 切块的weight，由optimizer维护，只负责做更新 (detach了)- 内容 : 这就是我们讨论的FP32主参数副本
W_master_fp32的分片。 - 特点 :
- 它是 FP32 格式，用于保证更新的精度。
- 它是切块的 (partitioned) 。这张图是
rank1的快照，所以这个橙色块只代表了W_master_fp32的一部分（例如，在一共2个GPU的情况下，是后半部分）。这就是 ZeRO-1 节省显存的关键之一。 - 它由优化器维护 ，并且是持久存在的。
detach()表示这个张量不参与梯度计算（它不是计算图的一部分），它的唯一作用就是作为一个高精度的"账本"，在optimizer.step()时被读取和更新。
- 这就是图中的"红色权重"所指的实体。
- 内容 : 这就是我们讨论的FP32主参数副本
-
红色块:
切块的adam动量1- 内容 : 这是 Adam 优化器的一阶矩 (momentum) 状态的分片。
- 特点 :
- 它是 FP32 格式。
- 它是切块的 ，与橙色块的
W_master_fp32分片一一对应。rank1只持有自己负责的那部分动量。 - 这也是 ZeRO-1 节省显存的核心。
-
绿色块:
切块的adam动量2- 内容 : 这是 Adam 优化器的二阶矩 (variance) 状态的分片。
- 特点: 与红色块完全相同，都是分区的、FP32 的优化器状态。
小结: 火焰图的底部稳定区域完美地展示了 ZeRO-1 在一个 GPU 上的静态内存布局：
- 一份完整的、低精度的模型副本 (灰色
W_fp16) 用于计算。 - 一份分区的、高精度的模型主副本 (橙色
W_master_fp32_partition) 用于精确更新。 - 一份分区的、高精度的优化器状态 (红色/绿色
O_fp32_partition) 用于历史信息记录。
追踪整个过程（结合火焰图的动态变化）
现在，让我们沿着时间轴（X轴）来看那些动态出现和消失的、更高层的色块。这些通常是梯度、激活值和临时缓冲区。
-
前向/后向传播阶段 (Forward/Backward):
- 当训练开始一个迭代时，你会看到内存占用开始出现许多尖峰（spikes）。
- 这些尖峰主要是激活值 (activations)。前向传播时，每一层的输出（激活值）被创建并保留在显存中，导致内存占用逐步上升。
- 后向传播时，这些激活值被用来计算梯度。一旦某个激活值用完，它占用的内存就可以被释放，所以你会看到内存占用有升有降。
- 在后向传播结束时，会产生一份完整的 FP16 梯度，它也会短暂地占用一块较大的显存。
-
梯度同步阶段 (
All-Reduce):- 在
optimizer.step()之前，会有一个All-Reduce操作。 - 这个操作会创建通信缓冲区，你可能会在火焰图上看到一些与通信相关的、短暂的大块内存分配。
All-Reduce结束后，每个 GPU 都有了一份完整的、同步好的 FP16 梯度。
- 在
-
参数更新阶段 (
optimizer.step()):- 这是最关键的阶段，它解释了作者的观点。
- 读取 :
optimizer读取：- 本地的橙色块 (
W_master_fp32_partition)。 - 本地的红色/绿色块 (
O_fp32_partition)。 - 完整的、同步好的梯度 (FP16)，并将其转换为 FP32。
- 本地的橙色块 (
- 计算 :
optimizer只在 GPU 内部 ，使用上述数据，计算出新的W_master_fp32_partition_new。 - 更新 :
- 这个新的
W_master_fp32_partition_new直接覆盖 了旧的橙色块 。这是对"红色权重"的直接更新。 - 同时，新的优化器状态也覆盖了旧的红色/绿色块。
- 这个新的
-
参数同步阶段 (FP16 Weight Synchronization):
- 在
optimizer.step()之后，需要将 FP32 主参数的更新反映到 FP16 模型参数上。 cast操作 :optimizer会将更新后的橙色块 (W_master_fp32_partition_new) 转换 (cast) 为 FP16 格式，得到W_fp16_partition_new。- 这个
W_fp16_partition_new是一个临时的张量。
All-Gather操作 :- 所有 GPU 对它们各自的
W_fp16_partition_new执行All-Gather。 - 操作结束后，每个 GPU 都拥有了一份完整的、最新的 FP16 参数
W_fp16_new_full。
- 所有 GPU 对它们各自的
- 覆盖 :
- 最后，用这份
W_fp16_new_full完全覆盖 掉底部的灰色块 (W_fp16)。 - 这样，用于下一次前向传播的模型参数就更新了。
- 最后，用这份
- 在
结论
这张火焰图通过可视化的方式，完美地展示了：
- ZeRO-1 的内存节省原理 :
W_master_fp32和O_fp32都是分区的（橙色、红色、绿色块都很小）。 - 混合精度训练的流程 : 更新发生在 FP32 的主参数（橙色块）上，然后通过
cast和All-Gather同步回用于计算的 FP16 参数（灰色块）。 - 作者观点的正确性 : 更新的核心目标是橙色块 ，而不是直接修改灰色块。灰色块的更新只是整个流程最后一步的同步结果。
通过将理论与这张实际的内存快照相结合，我们可以对 ZeRO-1 的工作机制有一个非常具体和深刻的理解。
代码实战
cpp
import torch
import torch.nn as nn
import deepspeed
import argparse
import os
# 导入 Profiler 相关的类
from torch.profiler import profile, record_function, ProfilerActivity
# 1. 定义模型 (和之前一样)
class SimpleModel(nn.Module):
def __init__(self):
super().__init__()
# 增加一层以创建更多事件
self.layer1 = nn.Linear(1024, 2048)
self.relu1 = nn.ReLU()
self.layer2 = nn.Linear(2048, 4096)
self.relu2 = nn.ReLU()
self.layer3 = nn.Linear(4096, 1024)
def forward(self, x):
x = self.relu1(self.layer1(x))
x = self.relu2(self.layer2(x))
x = self.layer3(x)
return x
def get_args():
parser = argparse.ArgumentParser(description='Profiled DeepSpeed ZeRO-1 Example')
parser.add_argument('--local_rank', type=int, default=-1,
help='local rank passed from distributed launcher')
parser = deepspeed.add_config_arguments(parser)
return parser.parse_args()
def main():
args = get_args()
# 2. 初始化 DeepSpeed
deepspeed.init_distributed()
torch.cuda.set_device(args.local_rank)
rank = torch.distributed.get_rank()
# 3. 创建模型
model = SimpleModel()
# 4. 使用 deepspeed.initialize
model_engine, _, _, _ = deepspeed.initialize(
args=args,
model=model,
model_parameters=model.parameters()
)
# 5. 准备 Profiler
# 我们只在 rank 0 上进行 profiling，以生成一个清晰的 trace 文件
if rank == 0:
print("Profiler is configured on rank 0. Starting training...")
prof = profile(
activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
record_shapes=True,
profile_memory=True, # 开启内存分析
with_stack=True,
with_flops=True,
with_modules=True
)
else:
# 其他 rank 创建一个空的上下文管理器，什么都不做
prof = torch.profiler.ExecutionTrace()
# 6. 开始 profiling 上下文
with prof:
# 运行几个 step 来捕获有意义的数据
for step in range(5):
if rank == 0:
print(f"--- Step {step} ---")
# 模拟输入数据
inputs = torch.randn(
model_engine.train_micro_batch_size_per_gpu(),
1024,
device=model_engine.device
)
# 前向、后向、更新
loss = model_engine(inputs).mean()
model_engine.backward(loss)
model_engine.step()
# 7. Profiling 结束后，导出 JSON 文件
if rank == 0:
print("Profiling finished. Exporting trace to JSON file...")
# 导出的 JSON 文件可以被 chrome://tracing 或 Netron 打开
prof.export_chrome_trace("ds_zero1_trace.json")
print("Trace saved to ds_zero1_trace.json")
# 也可以打印一个总结
print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))
print(f"Rank {rank} finished.")
if __name__ == "__main__":
main()好的，这个 RuntimeError: mat1 and mat2 must have the same dtype, but got Float and Half 是一个在进行混合精度训练时非常典型的错误。
问题根源 : 这个错误发生在 F.linear (矩阵乘法) 函数中。它告诉你，参与矩阵乘法的两个张量（mat1 和 mat2）的数据类型不匹配：一个是 Float (FP32) ，另一个是 Half (FP16)。在 CUDA kernel 中，矩阵乘法要求两个输入的数据类型必须完全相同。
错误分析
让我们来追踪一下数据流，看看哪里出了问题：
-
配置文件 : 你的
ds_config_zero1.json中设置了"fp16": {"enabled": true}。 -
deepspeed.initialize: 当你调用这个函数时，DeepSpeed 会自动将你的模型（SimpleModel）的参数（self.weight,self.bias）转换为 FP16 (Half)。这是混合精度训练的第一步，目的是节省模型参数占用的显存。 -
输入数据 (
inputs) : 在你的训练循环中，你创建了输入数据：pythoninputs = torch.randn( model_engine.train_micro_batch_size_per_gpu(), 1024, device=model_engine.device )这里的
torch.randn默认创建的是torch.float32(Float) 类型的张量。 -
前向传播 :
- 当模型执行
self.layer1(inputs)时，F.linear函数被调用。 mat1就是你的输入数据inputs，它的类型是 FP32 (Float)。mat2是layer1的权重self.weight，它已经被 DeepSpeed 转换成了 FP16 (Half)。
- 当模型执行
-
冲突 :
F.linear接收到一个 FP32 的输入和一个 FP16 的权重，于是抛出了数据类型不匹配的RuntimeError。
解决方案
解决方案非常直接：确保传递给模型的输入数据也和模型参数的数据类型一致，即 FP16 (Half)。
你需要对创建 inputs 的那一行代码做一个小小的修改。
修改前的代码:
python
inputs = torch.randn(
model_engine.train_micro_batch_size_per_gpu(),
1024,
device=model_engine.device
)修改后的代码:
python
inputs = torch.randn(
model_engine.train_micro_batch_size_per_gpu(),
1024,
device=model_engine.device
).half() # <--- 在这里添加 .half()或者，使用 to() 方法，这样更通用一些：
python
inputs = torch.randn(
model_engine.train_micro_batch_size_per_gpu(),
1024,
device=model_engine.device
).to(model_engine.dtype) # <--- 使用 model_engine 的 dtype，这是最佳实践为什么 .to(model_engine.dtype) 是最佳实践？
model_engine.dtype会自动反映 DeepSpeed 配置的数据类型。- 如果你在
ds_config.json中将fp16改为bf16，或者禁用了混合精度，这段代码不需要任何修改，它会自动创建正确类型的输入数据。 - 而硬编码
.half()或.bfloat16()会使得代码在配置改变时需要手动修改，不够灵活。
修正后的完整 main 函数
cpp
import torch
import torch.nn as nn
import deepspeed
import argparse
import os
import contextlib
# 导入 Profiler 相关的类
from torch.profiler import profile, record_function, ProfilerActivity
# 1. 定义模型 (和之前一样)
class SimpleModel(nn.Module):
def __init__(self):
super().__init__()
# 增加一层以创建更多事件
self.layer1 = nn.Linear(1024, 2048)
self.relu1 = nn.ReLU()
self.layer2 = nn.Linear(2048, 4096)
self.relu2 = nn.ReLU()
self.layer3 = nn.Linear(4096, 1024)
def forward(self, x):
x = self.relu1(self.layer1(x))
x = self.relu2(self.layer2(x))
x = self.layer3(x)
return x
def get_args():
parser = argparse.ArgumentParser(description='Profiled DeepSpeed ZeRO-1 Example')
parser.add_argument('--local_rank', type=int, default=-1,
help='local rank passed from distributed launcher')
parser = deepspeed.add_config_arguments(parser)
return parser.parse_args()
def main():
args = get_args()
# 2. 初始化 DeepSpeed
deepspeed.init_distributed()
torch.cuda.set_device(args.local_rank)
rank = torch.distributed.get_rank()
# 3. 创建模型
model = SimpleModel()
# 4. 使用 deepspeed.initialize
model_engine, _, _, _ = deepspeed.initialize(
args=args,
model=model,
model_parameters=model.parameters()
)
# 5. 准备 Profiler
# 我们只在 rank 0 上进行 profiling，以生成一个清晰的 trace 文件
if rank == 0:
print("Profiler is configured on rank 0. Starting training...")
prof = profile(
activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
record_shapes=True,
profile_memory=True, # 开启内存分析
with_stack=True,
with_flops=True,
with_modules=True
)
else:
# 其他 rank 创建一个空的上下文管理器，什么都不做
# 用标准的空上下文管理器，兼容旧版 PyTorch
prof = contextlib.nullcontext()
# 6. 开始 profiling 上下文
with prof:
# 运行几个 step 来捕获有意义的数据
for step in range(5):
if rank == 0:
print(f"--- Step {step} ---")
model_dtype = next(model_engine.module.parameters()).dtype
# 模拟输入数据
inputs = torch.randn(
model_engine.train_micro_batch_size_per_gpu(),
1024,
device=model_engine.device
).to(model_dtype)
# 前向、后向、更新
loss = model_engine(inputs).mean()
model_engine.backward(loss)
model_engine.step()
# 7. Profiling 结束后，导出 JSON 文件
if rank == 0:
print("Profiling finished. Exporting trace to JSON file...")
# 导出的 JSON 文件可以被 chrome://tracing 或 Netron 打开
prof.export_chrome_trace("ds_zero1_trace.json")
print("Trace saved to ds_zero1_trace.json")
# 也可以打印一个总结
print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))
print(f"Rank {rank} finished.")
if __name__ == "__main__":
main()
这里是修正了输入数据类型的完整 `main` 函数部分，你可以直接替换。
这个修改非常简单，但它点出了混合精度训练中一个非常核心的原则：**不仅模型参数需要转换，输入数据流也必须匹配。**
修改后，再次运行你的 `deepspeed` 启动命令，程序应该就能正常执行了。
rofiler is configured on rank 0. Starting training...
--- Step 0 ---
[2025-11-02 12:11:31,049] [INFO] [loss_scaler.py:190:update_scale] [deepspeed] OVERFLOW! Rank 0 Skipping step. Attempted loss scale: 65536, but hysteresis is 2. Reducing hysteresis to 1
--- Step 1 ---
[2025-11-02 12:11:31,053] [INFO] [loss_scaler.py:183:update_scale] [deepspeed] OVERFLOW! Rank 0 Skipping step. Attempted loss scale: 65536, reducing to 32768
--- Step 2 ---
--- Step 3 ---
--- Step 4 ---
Rank 1 finished.
Profiling finished. Exporting trace to JSON file...
Trace saved to ds_zero1_trace.json
------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------
Name Self CPU % Self CPU CPU total % CPU total CPU time avg Self CUDA Self CUDA % CUDA total CUDA time avg CPU Mem Self CPU Mem CUDA Mem Self CUDA Mem # of Calls Total FLOPs
------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------
record_param_comms 0.27% 964.681us 0.36% 1.280ms 39.987us 1.865ms 41.29% 1.865ms 58.277us 0 b 0 b 0 b 0 b 32 --
c10d::allreduce_ 0.08% 294.774us 0.36% 1.294ms 99.551us 0.000us 0.00% 1.472ms 113.269us 0 b 0 b 0 b 0 b 13 --
nccl:all_reduce 0.00% 0.000us 0.00% 0.000us 0.000us 1.472ms 32.60% 1.472ms 113.269us 0 b 0 b 0 b 0 b 13 --
ncclDevKernel_AllReduce_Sum_f16_RING_LL(ncclDevKerne... 0.00% 0.000us 0.00% 0.000us 0.000us 1.258ms 27.85% 1.258ms 251.550us 0 b 0 b 0 b 0 b 5 --
aten::copy_ 0.17% 621.840us 0.45% 1.596ms 12.668us 848.852us 18.79% 866.195us 6.875us 0 b 0 b 0 b 0 b 126 --
Optimizer.step#FusedAdam.step 0.00% 0.000us 0.00% 0.000us 0.000us 635.036us 14.06% 635.036us 211.679us 0 b 0 b 0 b 0 b 3 --
aten::to 0.03% 122.049us 0.52% 1.849ms 19.886us 0.000us 0.00% 508.376us 5.466us 0 b 0 b 533.52 Mb 0 b 93 --
aten::_to_copy 0.12% 427.348us 0.48% 1.727ms 23.662us 0.000us 0.00% 508.376us 6.964us 0 b 0 b 533.52 Mb 0 b 73 --
c10d::_allgather_base_ 0.02% 60.331us 0.08% 291.778us 97.259us 0.000us 0.00% 392.379us 130.793us 0 b 0 b 0 b 0 b 3 --
ncclDevKernel_AllGather_RING_LL(ncclDevKernelArgsSto... 0.00% 0.000us 0.00% 0.000us 0.000us 392.379us 8.69% 392.379us 130.793us 0 b 0 b 0 b 0 b 3 --
------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------
Self CPU time total: 357.859ms
Self CUDA time total: 4.517ms
Rank 0 finished.
[rank0]:[W1102 12:11:32.188913801 ProcessGroupNCCL.cpp:1496] Warning: WARNING: destroy_process_group() was not called before program exit, which can leak resources. For more info, please see https://pytorch.org/docs/stable/distributed.html#shutdown (function operator())
notebook-76ec3aa0aa0d-worker-0:90410:91095 [0] NCCL INFO [Service thread] Connection closed by localRank 1
[2025-11-02 12:11:33,088] [INFO] [launch.py:351:main] Process 90410 exits successfully.
[2025-11-02 12:11:33,089] [INFO] [launch.py:351:main] Process 90411 exits successfully.
(python3.10) bash-4.4$
代码`import torch
import torch.nn as nn
import deepspeed
import argparse
import os
import contextlib
# 导入 Profiler 相关的类
from torch.profiler import profile, record_function, ProfilerActivity
# 1. 定义模型 (和之前一样)
class SimpleModel(nn.Module):
def __init__(self):
super().__init__()
# 增加一层以创建更多事件
self.layer1 = nn.Linear(1024, 2048)
self.relu1 = nn.ReLU()
self.layer2 = nn.Linear(2048, 4096)
self.relu2 = nn.ReLU()
self.layer3 = nn.Linear(4096, 1024)
def forward(self, x):
x = self.relu1(self.layer1(x))
x = self.relu2(self.layer2(x))
x = self.layer3(x)
return x
def get_args():
parser = argparse.ArgumentParser(description='Profiled DeepSpeed ZeRO-1 Example')
parser.add_argument('--local_rank', type=int, default=-1,
help='local rank passed from distributed launcher')
parser = deepspeed.add_config_arguments(parser)
return parser.parse_args()
def main():
args = get_args()
# 2. 初始化 DeepSpeed
deepspeed.init_distributed()
torch.cuda.set_device(args.local_rank)
rank = torch.distributed.get_rank()
# 3. 创建模型
model = SimpleModel()
# 4. 使用 deepspeed.initialize
model_engine, _, _, _ = deepspeed.initialize(
args=args,
model=model,
model_parameters=model.parameters()
)
# 5. 准备 Profiler
# 我们只在 rank 0 上进行 profiling，以生成一个清晰的 trace 文件
if rank == 0:
print("Profiler is configured on rank 0. Starting training...")
prof = profile(
activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
record_shapes=True,
profile_memory=True, # 开启内存分析
with_stack=True,
with_flops=True,
with_modules=True
)
else:
# 其他 rank 创建一个空的上下文管理器，什么都不做
# 用标准的空上下文管理器，兼容旧版 PyTorch
prof = contextlib.nullcontext()
torch.cuda.memory._record_memory_history()
# 6. 开始 profiling 上下文
with prof:
# 运行几个 step 来捕获有意义的数据
for step in range(5):
if rank == 0:
print(f"--- Step {step} ---")
model_dtype = next(model_engine.module.parameters()).dtype
# 模拟输入数据
inputs = torch.randn(
model_engine.train_micro_batch_size_per_gpu(),
1024,
device=model_engine.device
).to(model_dtype)
# 前向、后向、更新
loss = model_engine(inputs).mean()
model_engine.backward(loss)
model_engine.step()
if rank == 0:
torch.cuda.memory._dump_snapshot("my_snapshottest2.pickle")
print(f"Rank {rank} finished.")
if __name__ == "__main__":
main()在这里插入图片描述
cpp
import torch
import torch.nn as nn
import deepspeed
import argparse
import os
import socket
from datetime import datetime
# --- Model 和 get_args 定义 (与之前相同) ---
class SimpleModel(nn.Module):
def __init__(self):
super().__init__()
self.layer1 = nn.Linear(1024, 2048)
self.relu1 = nn.ReLU()
self.layer2 = nn.Linear(2048, 4096)
self.relu2 = nn.ReLU()
self.layer3 = nn.Linear(4096, 1024)
def forward(self, x):
return self.relu2(self.layer2(self.relu1(self.layer1(x))))
def get_args():
parser = argparse.ArgumentParser(description='Profile All Ranks with DeepSpeed ZeRO-1')
parser.add_argument('--local_rank', type=int, default=-1, help='local rank passed from distributed launcher')
parser = deepspeed.add_config_arguments(parser)
args = parser.parse_args()
return args
# --- 【核心修改】trace_handler 现在需要知道当前的 rank ---
def get_trace_handler(rank):
"""
返回一个为特定 rank 定制的 on_trace_ready handler.
"""
# 获取主机名和时间戳，用于生成唯一的文件前缀
host_name = socket.gethostname()
timestamp = datetime.now().strftime("%b_%d_%H_%M_%S")
# 每个 rank 都有自己独立的文件前缀
file_prefix = f"{host_name}_{timestamp}_rank{rank}"
def trace_handler(prof: torch.profiler.profile):
print(f"Rank {rank}: Profiler is ready. Exporting traces to files with prefix: {file_prefix}")
# 1. 导出 Chrome trace JSON 文件 (可选，但推荐)
try:
prof.export_chrome_trace(f"{file_prefix}.json.gz")
except Exception as e:
print(f"Rank {rank}: Failed to export chrome trace: {e}")
# 2. 【关键】导出该 rank 对应的 GPU 的内存时间线 HTML 文件
try:
# device 参数在这里不需要，Profiler 会自动使用它正在监控的设备
prof.export_memory_timeline(f"{file_prefix}.html")
except Exception as e:
# 在某些旧版本或特定配置下，这个函数可能不可用
print(f"Rank {rank}: Failed to export memory timeline: {e}")
return trace_handler
def main():
args = get_args()
deepspeed.init_distributed()
rank = torch.distributed.get_rank()
torch.cuda.set_device(args.local_rank)
model = SimpleModel()
model_engine, _, _, _ = deepspeed.initialize(
args=args, model=model, model_parameters=model.parameters()
)
# --- Profiler 设置 (所有 rank 都开启) ---
# 每个 rank 使用自己定制的 trace handler
profiler_handler = get_trace_handler(rank)
# 我们让 profiler 只在第2和第3步激活，以跳过初始化的内存波动
profiler_schedule = torch.profiler.schedule(wait=0, warmup=0, active=6, repeat=1)
print(f"Rank {rank}: Initializing profiler...")
# 使用 with 语句来管理 profiler 的生命周期
with torch.profiler.profile(
activities=[
torch.profiler.ProfilerActivity.CPU,
torch.profiler.ProfilerActivity.CUDA,
],
# schedule=profiler_schedule,
on_trace_ready=profiler_handler,
record_shapes=True,
profile_memory=True,
with_stack=True,
) as prof:
# 运行足够的 step 来触发 profiler
for step in range(10):
print(f"Rank {rank}, starting Step {step}...")
model_dtype = next(model_engine.module.parameters()).dtype
inputs = torch.randn(
model_engine.train_micro_batch_size_per_gpu(),
1024,
device=model_engine.device
).to(model_dtype)
loss = model_engine(inputs).mean()
model_engine.backward(loss)
model_engine.step()
# 手动通知 profiler 完成了一个 step
prof.step()
print(f"Rank {rank} finished.")
if __name__ == "__main__":
main()3.火焰图分析
工具https://pytorch.cadn.net.cn/docs_en/2.5/torch_cuda_memory.html
https://pytorch.org/blog/understanding-gpu-memory-1
太棒了！这张由 prof.export_memory_timeline() 生成的图，是理解 DeepSpeed ZeRO-1 内存工作原理的绝佳教材。它看起来是这个样子，是因为它以一种非常直观的方式，完美地捕捉了一个训练 step 中各个阶段的内存动态变化。
让我们像看电影慢放一样，从左到右（按时间顺序）来解剖这张图。
图表解读 (时间线分析)
这张图展示了 GPU 显存的使用情况。Y 轴是显存占用量（GB），X 轴是时间（ms）。不同的颜色代表不同类型的内存分配，图例在右上角。
阶段一：时间 0ms - ~280ms (静态内存布局)
在这个阶段，我们看到两条持久存在的、平坦的内存块：
- 绿色 (PARAMETER) : 这是完整的 FP16 模型参数 。根据 ZeRO-1 的原理，每个 GPU 都需要一份完整的、低精度的模型参数来执行前向和后向计算。这块内存在
deepspeed.initialize时被分配，并一直存在于整个训练过程中。 - 灰色 (Unknown) : 这是 ZeRO-1 的核心优化所在 。这块内存代表了被分区（Partitioned）的 FP32 主参数 和被分区的 FP32 优化器状态 （Adam 的 momentum 和 variance）。
- 为什么是灰色的 "Unknown"？ 因为这些内存是由 DeepSpeed 的底层 C++ 或 CUDA 代码直接管理的，PyTorch Profiler 的自动分类器有时无法精确地识别它们的用途，就将它们归为"未知"。
- 关键点 ：这块灰色内存的大小，仅仅是完整 FP32 参数和优化器状态的 1/N（N是你的 GPU 数量，这里是 1/2）。这就是 ZeRO-1 节省大量静态内存的地方。
阶段二：时间 ~280ms - ~420ms (反向传播 -> 梯度分配)
我们看到了一个巨大的蓝色方块拔地而起，然后又消失。
- 深蓝色 (GRADIENT) : 这是在执行
model_engine.backward(loss)时，为模型参数分配的完整的 FP16 梯度 。- 为什么是完整的？ 因为在计算梯度时，每个 GPU 仍然需要计算出对应于完整模型参数的梯度。
- 为什么会消失？ 在
model_engine.step()执行完毕后，优化器已经使用完这些梯度，会通过optimizer.zero_grad()将其释放。所以它是一个动态的、临时的大块内存。 - 内存峰值 : 这个蓝色方块的出现，导致了总内存占用达到峰值。图顶部的
Max memory allocated: 1.14 GiB就是在这个时刻达到的。
阶段三：时间 ~420ms - ~480ms (优化器步骤)
在巨大的蓝色梯度块消失后，出现了一些小而杂乱的内存块。
- 黄色 (OPTIMIZER_STATE) , 灰色 (Unknown) , 浅蓝色 (AUTOGRAD_DETAIL) : 这个阶段对应
model_engine.step()。- 在这个过程中，DeepSpeed 会执行梯度的 AllReduce、更新分区后的 FP32 主参数、更新分区后的优化器状态、再将更新后的参数 AllGather 回完整的 FP16 模型参数。
- 这个过程涉及到大量的临时缓冲区（用于通信、类型转换等），因此你会看到很多短暂的、不同类型的内存分配。
阶段四：时间 ~480ms 之后 (后续的训练步骤)
我们看到了一系列重复的、更窄的"尖峰"图案。
- 重复的模式 : 每一个"尖峰"都代表一个新的训练 step (
step=1,step=2, ...)。 - 模式内容: 每个尖峰内部，都重复着"阶段二"和"阶段三"的模式：分配巨大的蓝色梯度块 -> 释放梯度块 -> 杂乱的优化器步骤。
- 为什么第一个 step 看起来更宽？ 第一个训练 step 通常会包含一些一次性的操作，比如 CUDA kernel 的 JIT 编译、缓存预热等，所以它的耗时会比后续稳定状态下的 step 更长。
总结：为什么图是这个样子的？
这张图完美地回答了"ZeRO-1 是如何工作的？"这个问题：
-
持久内存 (图的底部):
- 绿色 : 每个 GPU 都有一份完整的 FP16 参数 (用于计算)。
- 灰色 : 每个 GPU 只持有1/N 的 FP32 参数和优化器状态 (用于更新，这是省内存的关键)。
-
动态内存 (图中的尖峰):
- 蓝色 : 在
backward()期间，会临时分配一块完整的 FP16 梯度，这是训练过程中的内存峰值点。 step()之后，这块巨大的梯度内存会被释放，为下一个 step 腾出空间。
- 蓝色 : 在
你的图清晰地展示了： ZeRO-1 的内存优化不是 在梯度计算阶段（因为梯度仍然是完整的），而是在静态的、持久的优化器状态和主参数上。通过将这部分巨大的内存分区存储，它极大地降低了模型训练的基线内存占用。
好的，这个问题提得非常深入，直指这张图的核心！你观察到的"蓝色块"和"10条"现象，是解密 backward() 和 step() 内部机制的关键。
让我们来详细拆解。
1. model_engine.step() 是在什么时候执行的？
在你的代码中，每个循环的顺序是：
loss = model_engine(inputs).mean()(前向传播)model_engine.backward(loss)(反向传播)model_engine.step()(优化器更新)prof.step()(通知 Profiler 这一步结束了)
现在，我们把这个顺序和图上的事件对应起来：
backward()阶段 : 对应图中巨大的蓝色方块 (GRADIENT) 出现和存在的时期。在这个阶段，PyTorch 的自动求导引擎计算出每一层参数的梯度，并为它们分配显存。step()阶段 : 对应蓝色方块之后 ，那一小段混乱、包含很多细小尖峰的区域（包括你说的"10条"可能就位于此）。在这个阶段，DeepSpeed ZeRO 优化器会执行以下复杂操作：- 规约梯度 (Reduce Gradients) : 将每个 GPU 上完整的梯度进行
AllReduce操作，使得每个 GPU 最终只得到它负责的那部分参数的、聚合后的梯度。 - 更新参数 (Update Parameters): 使用聚合后的梯度，去更新它自己持有的那一小块 FP32 主参数和优化器状态。
- 同步模型 (Update Model Weights) : 将更新后的一小块 FP32 参数转换回 FP16，然后通过
AllGather操作广播给所有 GPU，让每个 GPU都能重建一份完整且最新的 FP16 模型参数，为下一次前向传播做准备。
- 规约梯度 (Reduce Gradients) : 将每个 GPU 上完整的梯度进行
optimizer.zero_grad(): 这个操作通常在step()之后隐式或显式调用，它会释放巨大的蓝色梯度块，所以我们能看到蓝色块在step()阶段结束后就消失了。
结论 : model_engine.step() 发生在巨大蓝色梯度块出现之后，消失之前的那段"混乱"时期。
2. 蓝色块和"10条"：是什么？为什么这么大？
你观察到的现象可以分为两部分：一个主要的、巨大的蓝色块，以及一系列重复出现的窄"尖峰"（你可能数了大约10个，但实际上是 active=6 步，只是第一个 step 比较宽）。
A. 巨大的蓝色块 (GRADIENT)
-
它是什么？
这个蓝色块代表为模型所有参数的梯度 (
.grad) 所分配的内存。当backward()被调用时，PyTorch 会为每一个需要梯度的nn.Parameter创建一个和它形状、大小完全相同的张量来存储梯度值。 -
为什么这么大？
这是理解 ZeRO-1 的一个最关键、最核心的知识点！
- 梯度大小 = 参数大小: 一个模型的梯度所占用的内存，和它的参数本身占用的内存是一样大的。如果你的模型有 10 亿个 FP16 参数（占用 2GB），那么它的梯度也会占用 2GB。
- ZeRO-1 中梯度是完整的 : ZeRO-1 的优化点在于分区存储优化器状态和 FP32 主参数 ，但是，在
backward()这一步，每个 GPU 仍然需要计算和存储一份完整的、未分区的梯度。 - 内存峰值来源 : 这就导致了训练过程中的内存峰值。总内存 = (持久内存：FP16完整参数 + 分区FP32参数 + 分区优化器状态) + (动态内存：完整的梯度)。
- 你的模型有多大？ : 我们来估算一下你代码中模型的梯度大小：
layer1 (1024, 2048): ~2.1M paramslayer2 (2048, 4096): ~8.4M paramslayer3 (4096, 1024): ~4.2M params- 总参数量: ~14.7M
- FP16 梯度大小 : 14.7M * 2 bytes/param ≈ 29.4 MB
- 你的图显示了什么？ : 图中的蓝色块高度大约是 0.9 GiB (约 920 MB) ！这说明你生成这张图时所用的模型，比我们代码里的
SimpleModel大得多（大约是 30 多倍）。这完全正常，对于像 GPT 这样的大模型，梯度本身占用几个 GB 是很常见的。
一句话总结蓝色块 : 它是完整的模型梯度，在 ZeRO-1 中它并未被分区，因此非常巨大，构成了训练时的内存峰值。
B. "10条"细长的尖峰 (多个训练 Step)
你看到的重复的、像柱子一样的"尖峰"，其实是多个训练 step 的内存使用模式 。由于你的 schedule 设置为 active=6，Profiler 记录了 6 个 step（或者你的循环只跑了5步，就记录了5个）。
-
每一"条"/"尖峰"是什么？
每一条尖峰就是一个完整的backward()+step()过程。它内部的模式都是一样的：
- 内存突然飙升（巨大的蓝色梯度块被分配）。
- 内存维持在高位一小段时间（梯度计算和
step准备）。 - 内存突然下降（梯度被释放）。
-
为什么第一个尖峰比后面的宽？
这是典型的"热身"现象。第一个被分析的 step (在你的例子中是
step=0) 通常会包含一些一次性的 CUDA 操作（如 kernel JIT 编译、缓存初始化等），所以它的执行时间会比后续已经"热身"完毕的 step 更长。
结论 : 你看到的"10条"或"6条"蓝色相间的尖峰，是 Profiler 记录的多个连续训练步骤的内存快照。每一个尖峰都代表一次"分配完整梯度 -> 使用梯度 -> 释放梯度"的循环，直观地展示了模型训练的动态内存行为。
好的，您提出的这个验证非常关键，它揭示了一个核心事实：我们代码中的 SimpleModel 和生成这张图所用的模型，根本不是同一个模型。
计算结果与图中完全不吻合。这恰好证明了我们之前的推断是正确的------这张图分析的是一个参数量大得多的模型。
下面是详细的计算过程，它将清晰地展示这个巨大的差异。
第一步：计算我们代码中 SimpleModel 的总参数量
我们来逐层计算 Simple-Model 的参数数量。公式是：参数量 = 输入维度 × 输出维度 + 输出维度 (偏置项)。
-
self.layer1 = nn.Linear(1024, 2048)- 权重 (Weights):
1024 × 2048 = 2,097,152 - 偏置 (Bias):
2048 - 该层参数 :
2,097,152 + 2048 = 2,099,200
- 权重 (Weights):
-
self.layer2 = nn.Linear(2048, 4096)- 权重 (Weights):
2048 × 4096 = 8,388,608 - 偏置 (Bias):
4096 - 该层参数 :
8,388,608 + 4096 = 8,392,704
- 权重 (Weights):
-
self.layer3 = nn.Linear(4096, 1024)- 权重 (Weights):
4096 × 1024 = 4,194,304 - 偏置 (Bias):
1024 - 该层参数 :
4,194,304 + 1024 = 4,195,328
- 权重 (Weights):
SimpleModel 的总参数量 :
2,099,200 + 8,392,704 + 4,195,328 = 14,687,232
约等于 14.7M (1470万) 个参数。
第二步：根据参数量计算 SimpleModel 的梯度大小
现在我们用这个参数量来计算它在 FP16 模式下应该产生的梯度大小。
- 公式 :
梯度大小 (Bytes) = 总参数量 × 2 Bytes/参数 - 计算 :
14,687,232 × 2 = 29,374,464 Bytes
接下来，我们将这个结果转换为更易读的单位 (MB)：
1 MB = 1024 × 1024 = 1,048,576 Bytes- 梯度大小 (MB) :
29,374,464 / 1,048,576 ≈ 28.01 MB
理论计算结论 : 如果我们运行代码中的 SimpleModel，它产生的梯度大小应该大约是 28 MB。
第三步：对比理论计算与图中观测结果
现在，我们把两个数字放在一起对比：
- 理论计算出的梯度大小 (
SimpleModel) : 28 MB - 图中观测到的梯度大小 (蓝色块) : ~1.08 GiB (约 1106 MB)
对比结果 ：
28 MB (理论) vs 1106 MB (观测)
两者相差了近 40 倍！
最终结论
这个巨大的差异无可辩驳地证明了：
您提供的内存分析图，并非由我们示例代码中的 SimpleModel 生成的。
这张图是分析一个参数量约为 5.8 亿 的大模型时得到的。我们代码中的 SimpleModel（约 1500 万参数）只是一个用于演示如何运行 Profiler 的"玩具"模型。
这个验证过程非常有价值，它告诉我们：
- 实践出真知: 不能想当然地认为示例代码和示例结果是匹配的。
- 量化分析的力量: 通过精确计算，我们可以从 Profiler 的结果中反推出模型的关键属性（如参数量），并验证我们的假设。
- 梯度是内存大户: 这个对比也再次强调了，梯度是深度学习训练中动态内存消耗的主要来源，其大小与模型参数量严格挂钩。
#以上代码不小心注释掉了
schedule=profiler_schedule,
on_trace_ready=profiler_handler,
问题定位
根据你提供的内存分配信息，我看到了问题的关键所在。让我分析一下为什么14.7M参数的模型会分配1GB的梯度内存。
问题分析
从堆栈跟踪可以看到，这个1GB的内存分配发生在：
/home/wangshangquan.wsq/.local/lib/python3.10/site-packages/deepspeed/runtime/zero/stage_1_and_2.py:2050:backward这行代码在DeepSpeed的ZeRO阶段1和2实现中。让我们计算一下：
参数计算
- 你的模型参数：14.7M (14,700,000个参数)
- 如果使用FP16训练：每个参数2字节
- 预期梯度大小：14.7M × 2字节 ≈ 29.4MB
但实际分配了1GB (953.7MiB ≈ 1,000,000,000字节)
可能的原因
-
ZeRO-2的梯度分区：在ZeRO-2中，梯度是分区的，每个rank可能分配了完整的梯度缓冲区
-
内存对齐和预分配：DeepSpeed可能为了性能进行了内存预分配
-
优化器状态：可能包含了优化器状态的内存分配
诊断和验证代码
让我们修改代码来验证实际的内存使用情况：
python
import torch
import torch.nn as nn
import deepspeed
import argparse
import os
import socket
from datetime import datetime
class SimpleModel(nn.Module):
def __init__(self):
super().__init__()
self.layer1 = nn.Linear(1024, 2048)
self.relu1 = nn.ReLU()
self.layer2 = nn.Linear(2048, 4096)
self.relu2 = nn.ReLU()
self.layer3 = nn.Linear(4096, 1024)
def forward(self, x):
return self.relu2(self.layer2(self.relu1(self.layer1(x))))
def get_args():
parser = argparse.ArgumentParser(description='Profile All Ranks with DeepSpeed ZeRO-1')
parser.add_argument('--local_rank', type=int, default=-1, help='local rank passed from distributed launcher')
parser = deepspeed.add_config_arguments(parser)
args = parser.parse_args()
return args
def print_memory_usage(rank, message):
if torch.cuda.is_available():
allocated = torch.cuda.memory_allocated(rank) / 1024**3 # GB
reserved = torch.cuda.memory_reserved(rank) / 1024**3 # GB
print(f"Rank {rank}: {message} - Allocated: {allocated:.2f} GB, Reserved: {reserved:.2f} GB")
def main():
args = get_args()
deepspeed.init_distributed()
rank = torch.distributed.get_rank()
torch.cuda.set_device(args.local_rank)
# 清空GPU内存
torch.cuda.empty_cache()
print_memory_usage(rank, "After empty cache")
model = SimpleModel()
# 计算理论参数和梯度大小
total_params = sum(p.numel() for p in model.parameters())
print(f"Rank {rank}: Model parameters: {total_params:,} = {total_params/1e6:.2f}M")
print(f"Rank {rank}: Expected gradient size (FP16): {total_params * 2 / 1024**2:.2f} MB")
print(f"Rank {rank}: Expected gradient size (FP32): {total_params * 4 / 1024**2:.2f} MB")
print_memory_usage(rank, "After model creation")
# 初始化DeepSpeed
model_engine, optimizer, _, _ = deepspeed.initialize(
args=args, model=model, model_parameters=model.parameters()
)
print_memory_usage(rank, "After DeepSpeed initialize")
# 检查DeepSpeed配置
zero_stage = model_engine.zero_optimization_stage()
print(f"Rank {rank}: ZeRO stage: {zero_stage}")
# 打印优化器信息
if hasattr(optimizer, 'param_groups'):
print(f"Rank {rank}: Optimizer has {len(optimizer.param_groups)} parameter groups")
# 检查实际管理的参数数量
total_params_in_engine = sum(p.numel() for p in model_engine.module.parameters())
print(f"Rank {rank}: Parameters in engine: {total_params_in_engine:,}")
# 运行一个训练步骤来观察内存变化
model_dtype = next(model_engine.module.parameters()).dtype
print(f"Rank {rank}: Model dtype: {model_dtype}")
inputs = torch.randn(32, 1024, device=model_engine.device).to(model_dtype)
print_memory_usage(rank, "Before forward")
outputs = model_engine(inputs)
loss = outputs.mean()
print_memory_usage(rank, "After forward")
model_engine.backward(loss)
print_memory_usage(rank, "After backward")
# 检查梯度状态
for name, param in model_engine.module.named_parameters():
if param.grad is not None:
grad_size = param.grad.numel() * param.grad.element_size()
print(f"Rank {rank}: Gradient {name}: {grad_size / 1024**2:.2f} MB")
break # 只打印第一个以节省空间
model_engine.step()
print_memory_usage(rank, "After step")
if __name__ == "__main__":
main()可能解决方案
- 检查DeepSpeed配置文件：
json
{
"zero_optimization": {
"stage": 1,
"reduce_bucket_size": 5e8,
"allgather_bucket_size": 5e8
}
}-
调整bucket大小 ：减少
reduce_bucket_size和allgather_bucket_size的值 -
使用ZeRO-1而不是ZeRO-2：ZeRO-1的梯度内存使用更接近理论值
-
检查是否意外使用了ZeRO-Offload：这可能会导致额外的内存分配
运行上面的诊断代码，它会告诉你：
- 每个阶段的实际内存使用
- DeepSpeed的ZeRO阶段
- 梯度的实际大小
- 模型的数据类型
这样我们就能确定1GB分配的具体原因。
Rank 0: After empty cache - Allocated: 0.00 GB, Reserved: 0.00 GB
Rank 0: Model parameters: 14,687,232 = 14.69M
Rank 0: Expected gradient size (FP16): 28.01 MB
Rank 0: Expected gradient size (FP32): 56.03 MB
Rank 0: After model creation - Allocated: 0.00 GB, Reserved: 0.00 G
Loading extension module fused_adam...
Time to load fused_adam op: 27.383007049560547 seconds
Rank 1: After DeepSpeed initialize - Allocated: 0.05 GB, Reserved: 0.11 GB
Rank 1: ZeRO stage: 1
Rank 1: Optimizer has 1 parameter groups
Rank 1: Parameters in engine: 14,687,232
Rank 1: Model dtype: torch.float16
Rank 1: Before forward - Allocated: 0.05 GB, Reserved: 0.11 GB
Loading extension module fused_adam...
Time to load fused_adam op: 27.499696254730225 seconds
2025-11-02 14:21:41,985 INFO logging.py:128:log_dist Rank 0 Using DeepSpeed Optimizer param name adam as basic optimizer
2025-11-02 14:21:41,985 INFO logging.py:128:log_dist Rank 0 Removing param_group that has no 'params' in the basic Optimizer
2025-11-02 14:21:41,985 INFO logging.py:128:log_dist Rank 0 DeepSpeed Basic Optimizer = FusedAdam
2025-11-02 14:21:41,986 INFO utils.py:59:is_zero_supported_optimizer Checking ZeRO support for optimizer=FusedAdam type=<class 'deepspeed.ops.adam.fused_adam.FusedAdam'>
2025-11-02 14:21:41,986 INFO logging.py:128:log_dist Rank 0 Creating torch.float16 ZeRO stage 1 optimizer
2025-11-02 14:21:41,986 INFO stage_1_and_2.py:149:**init** Reduce bucket size 500000000
2025-11-02 14:21:41,986 INFO stage_1_and_2.py:150:**init** Allgather bucket size 500000000
2025-11-02 14:21:41,986 INFO stage_1_and_2.py:151:**init** CPU Offload: False
2025-11-02 14:21:41,986 INFO stage_1_and_2.py:152:**init** Round robin gradient partitioning: False
Rank 1: After forward - Allocated: 0.09 GB, Reserved: 0.14 GB
Rank 1: After backward - Allocated: 0.15 GB, Reserved: 1.12 GB
我们注意到在 backward 之后，Reserved 内存从 0.14 GB 增加到了 1.12 GB，而 Allocated 内存只增加了 0.05 GB。这表明确实有大约 1 GB 的内存被预留（reserved）了，但实际分配（allocated）的并没有那么多。
日志中我们也可以看到，在 backward 之后，分配的内存（allocated）只有0.14GB，而预留（reserved）是1.12GB。这意味着大部分内存并没有被实际使用，而是被缓存（预留）以供后续使用。这通常是深度学习框架的内存管理策略，为了避免频繁分配和释放内存
我们尝试修改配置文件，将 reduce_bucket_size 设置为一个较小的值，比如50MB，然后观察内存变化。
在，我们更关心的是为什么在 backward 之后预留了 1.12 GB 的内存。这很可能是由于 DeepSpeed 为通信桶预留的内存。在 ZeRO-1 中，通信桶用于梯度的 All-Reduce。默认的桶大小是 500MB，但是这里我们看到了两倍于这个值的预留内存，可能的原因如下：
可能同时有两个桶（例如，一个用于当前步骤，一个用于下一步骤）？或者为了重叠通信和计算而预留了多个缓冲区。
可能还包括了优化器状态的内存。在 ZeRO-1 中，优化器状态是分片的，所以每个进程只存储一部分优化器状态。但是，我们的模型只有14.7M参数，即使使用FP16，优化器状态（包括参数、梯度、动量、方差）也不会这么大。
让我们计算一下优化器状态的内存：
参数：14.7M * 2字节（FP16） = 29.4 MB
梯度：14.7M * 2字节（FP16） = 29.4 MB
优化器状态（Adam）：对于每个参数，有2个状态（动量、方差），每个状态是FP32，所以是4字节2，所以每个参数需要8字节的优化器状态。
14.7M * 8字节 = 117.6 MB
所以，即使加上优化器状态，总共也就大约 29.42 + 117.6 = 176.4 MB。所以1GB的预留内存显然不是由模型参数、梯度和优化器状态直接引起的。
因此，最可能的原因还是 DeepSpeed 的通信缓冲区。我们可以在配置中显式设置 reduce_bucket_size 和 allgather_bucket_size 来减小这个值。
另外，我们注意到在配置中使用了 ZeRO 阶段1，并且没有指定桶大小。DeepSpeed 默认的 reduce_bucket_size 是 500MB，而 allgather_bucket_size 在阶段1中并不使用（因为阶段1不需要allgather参数）。所以，这里可能只是 reduce_bucket_size 的影响。
修改后的配置：
json
{
"zero_optimization": {
"stage": 1,
"reduce_bucket_size": 5e7
}
}
所以，如果我们减少 reduce_bucket_size，我们应该会看到预留内存的减少。
另外，我们还可以尝试在代码中强制清理缓存，看看预留内存是否会释放：
修改配置
cpp
{
"train_batch_size": 256,
"train_micro_batch_size_per_gpu": 128,
"steps_per_print": 10,
"optimizer": {
"type": "Adam",
"params": {
"lr": 0.001,
"betas": [0.9, 0.999],
"eps": 1e-08,
"weight_decay": 0
}
},
"fp16": {
"enabled": true,
"loss_scale": 0,
"loss_scale_window": 1000,
"hysteresis": 2,
"min_loss_scale": 1
},
"zero_optimization": {
"stage": 1,
"reduce_bucket_size": 5e7
}
}
python3.10) bash-4.4$ deepspeed --num_gpus=2 testzerodsTest.py --deepspeed_config ds_config_zero1.json
/opt/conda/envs/python3.10.13/lib/python3.10/site-packages/torch/cuda/__init__.py:61: FutureWarning: The pynvml package is deprecated. Please install nvidia-ml-py instead. If you did not install pynvml directly, please report this to the maintainers of the package that installed pynvml for you.
import pynvml # type: ignore[import]
[2025-11-02 14:41:28,045] [INFO] [real_accelerator.py:222:get_accelerator] Setting ds_accelerator to cuda (auto detect)
[2025-11-02 14:41:37,757] [WARNING] [runner.py:215:fetch_hostfile] Unable to find hostfile, will proceed with training with local resources only.
[2025-11-02 14:41:37,758] [INFO] [runner.py:607:main] cmd = /opt/conda/envs/python3.10.13/bin/python3.10 -u -m deepspeed.launcher.launch --world_info=eyJsb2NhbGhvc3QiOiBbMCwgMV19 --master_addr=127.0.0.1 --master_port=29500 --enable_each_rank_log=None testzerodsTest.py --deepspeed_config ds_config_zero1.json
/opt/conda/envs/python3.10.13/lib/python3.10/site-packages/torch/cuda/__init__.py:61: FutureWarning: The pynvml package is deprecated. Please install nvidia-ml-py instead. If you did not install pynvml directly, please report this to the maintainers of the package that installed pynvml for you.
import pynvml # type: ignore[import]
[2025-11-02 14:41:40,005] [INFO] [real_accelerator.py:222:get_accelerator] Setting ds_accelerator to cuda (auto detect)
[2025-11-02 14:41:49,809] [INFO] [launch.py:139:main] 0 NCCL_IB_GID_INDEX=3
[2025-11-02 14:41:49,809] [INFO] [launch.py:139:main] 0 NCCL_NET_PLUGIN=none
[2025-11-02 14:41:49,809] [INFO] [launch.py:139:main] 0 NCCL_DEBUG=INFO
[2025-11-02 14:41:49,809] [INFO] [launch.py:146:main] WORLD INFO DICT: {'localhost': [0, 1]}
[2025-11-02 14:41:49,809] [INFO] [launch.py:152:main] nnodes=1, num_local_procs=2, node_rank=0
[2025-11-02 14:41:49,810] [INFO] [launch.py:163:main] global_rank_mapping=defaultdict(<class 'list'>, {'localhost': [0, 1]})
[2025-11-02 14:41:49,810] [INFO] [launch.py:164:main] dist_world_size=2
[2025-11-02 14:41:49,810] [INFO] [launch.py:168:main] Setting CUDA_VISIBLE_DEVICES=0,1
[2025-11-02 14:41:49,810] [INFO] [launch.py:256:main] process 155969 spawned with command: ['/opt/conda/envs/python3.10.13/bin/python3.10', '-u', 'testzerodsTest.py', '--local_rank=0', '--deepspeed_config', 'ds_config_zero1.json']
[2025-11-02 14:41:49,811] [INFO] [launch.py:256:main] process 155970 spawned with command: ['/opt/conda/envs/python3.10.13/bin/python3.10', '-u', 'testzerodsTest.py', '--local_rank=1', '--deepspeed_config', 'ds_config_zero1.json']
/opt/conda/envs/python3.10.13/lib/python3.10/site-packages/torch/cuda/__init__.py:61: FutureWarning: The pynvml package is deprecated. Please install nvidia-ml-py instead. If you did not install pynvml directly, please report this to the maintainers of the package that installed pynvml for you.
import pynvml # type: ignore[import]
/opt/conda/envs/python3.10.13/lib/python3.10/site-packages/torch/cuda/__init__.py:61: FutureWarning: The pynvml package is deprecated. Please install nvidia-ml-py instead. If you did not install pynvml directly, please report this to the maintainers of the package that installed pynvml for you.
import pynvml # type: ignore[import]
[2025-11-02 14:41:54,099] [INFO] [real_accelerator.py:222:get_accelerator] Setting ds_accelerator to cuda (auto detect)
[2025-11-02 14:41:54,099] [INFO] [real_accelerator.py:222:get_accelerator] Setting ds_accelerator to cuda (auto detect)
[2025-11-02 14:42:11,131] [INFO] [comm.py:658:init_distributed] cdb=None
[2025-11-02 14:42:11,131] [INFO] [comm.py:689:init_distributed] Initializing TorchBackend in DeepSpeed with backend nccl
[2025-11-02 14:42:11,131] [INFO] [comm.py:658:init_distributed] cdb=None
Rank 0: After empty cache - Allocated: 0.00 GB, Reserved: 0.00 GB
Rank 0: Model parameters: 14,687,232 = 14.69M
Rank 0: Expected gradient size (FP16): 28.01 MB
Rank 0: Expected gradient size (FP32): 56.03 MB
Rank 0: After model creation - Allocated: 0.00 GB, Reserved: 0.00 GB
[2025-11-02 14:42:11,214] [INFO] [logging.py:128:log_dist] [Rank 0] DeepSpeed info: version=0.16.4, git-hash=unknown, git-branch=unknown
[2025-11-02 14:42:11,217] [INFO] [config.py:734:__init__] Config mesh_device None world_size = 2
Rank 1: After empty cache - Allocated: 0.00 GB, Reserved: 0.00 GB
Rank 1: Model parameters: 14,687,232 = 14.69M
Rank 1: Expected gradient size (FP16): 28.01 MB
Rank 1: Expected gradient size (FP32): 56.03 MB
Rank 1: After model creation - Allocated: 0.00 GB, Reserved: 0.00 GB
[2025-11-02 14:42:12,476] [INFO] [config.py:734:__init__] Config mesh_device None world_size = 2
notebook-76ec3aa0aa0d-worker-0:155969:155969 [0] NCCL INFO Bootstrap: Using eth0:33.203.57.184<0>
notebook-76ec3aa0aa0d-worker-0:155969:155969 [0] NCCL INFO ACCL_TUNING_LEVEL=2
notebook-76ec3aa0aa0d-worker-0:155969:155969 [0] NCCL INFO cudaDriverVersion 12040
notebook-76ec3aa0aa0d-worker-0:155969:155969 [0] NCCL INFO NCCL version 2.26.5.12-accl-n+cuda12.8, COMMIT_ID 2e6879b700b6cd1510bf9dc29a9db0132609d4cc, BUILD_TIME 2025-05-21 11:34:07
notebook-76ec3aa0aa0d-worker-0:155969:155969 [0] NCCL INFO Comm config Blocking set to 1
notebook-76ec3aa0aa0d-worker-0:155969:155969 [0] NCCL INFO Comm config Traffic class set to 1818586738
notebook-76ec3aa0aa0d-worker-0:155969:156631 [0] NCCL INFO C4 stats mode none, reduce 1, send/recv 0.
notebook-76ec3aa0aa0d-worker-0:155970:155970 [1] NCCL INFO cudaDriverVersion 12040
notebook-76ec3aa0aa0d-worker-0:155970:155970 [1] NCCL INFO Bootstrap: Using eth0:33.203.57.184<0>
notebook-76ec3aa0aa0d-worker-0:155970:155970 [1] NCCL INFO ACCL_TUNING_LEVEL=2
notebook-76ec3aa0aa0d-worker-0:155970:155970 [1] NCCL INFO NCCL version 2.26.5.12-accl-n+cuda12.8, COMMIT_ID 2e6879b700b6cd1510bf9dc29a9db0132609d4cc, BUILD_TIME 2025-05-21 11:34:07
notebook-76ec3aa0aa0d-worker-0:155970:155970 [1] NCCL INFO Comm config Blocking set to 1
notebook-76ec3aa0aa0d-worker-0:155970:155970 [1] NCCL INFO Comm config Traffic class set to 1030188649
notebook-76ec3aa0aa0d-worker-0:155970:156634 [0] NCCL INFO C4 stats mode none, reduce 1, send/recv 0.
notebook-76ec3aa0aa0d-worker-0:155969:156631 [0] NCCL INFO NET/Plugin: Could not find: libnccl-net-none.so. Using internal net plugin.
notebook-76ec3aa0aa0d-worker-0:155969:156631 [0] NCCL INFO NET/IB : No device found.
notebook-76ec3aa0aa0d-worker-0:155969:156631 [0] NCCL INFO NET/IB : Using [RO]; OOB eth0:33.203.57.184<0>
notebook-76ec3aa0aa0d-worker-0:155969:156631 [0] NCCL INFO NET/Socket : Using [0]eth0:33.203.57.184<0>
notebook-76ec3aa0aa0d-worker-0:155969:156631 [0] NCCL INFO PROFILER/Plugin: Could not find: libnccl-profiler.so.
notebook-76ec3aa0aa0d-worker-0:155969:156631 [0] NCCL INFO Using network Socket
notebook-76ec3aa0aa0d-worker-0:155970:156634 [1] NCCL INFO NET/Plugin: Could not find: libnccl-net-none.so. Using internal net plugin.
notebook-76ec3aa0aa0d-worker-0:155970:156634 [1] NCCL INFO NET/IB : No device found.
notebook-76ec3aa0aa0d-worker-0:155970:156634 [1] NCCL INFO NET/IB : Using [RO]; OOB eth0:33.203.57.184<0>
notebook-76ec3aa0aa0d-worker-0:155970:156634 [1] NCCL INFO NET/Socket : Using [0]eth0:33.203.57.184<0>
notebook-76ec3aa0aa0d-worker-0:155970:156634 [1] NCCL INFO PROFILER/Plugin: Could not find: libnccl-profiler.so.
notebook-76ec3aa0aa0d-worker-0:155970:156634 [1] NCCL INFO Using network Socket
notebook-76ec3aa0aa0d-worker-0:155969:156631 [0] NCCL INFO ncclCommInitRankConfig comm 0x9fac070 rank 0 nranks 2 cudaDev 0 nvmlDev 0 busId 109000 commId 0x62380d103bc2ab46 commHash 7077421177785985862 - Init START
notebook-76ec3aa0aa0d-worker-0:155970:156634 [1] NCCL INFO ncclCommInitRankConfig comm 0xb34f1a0 rank 1 nranks 2 cudaDev 1 nvmlDev 1 busId 17f000 commId 0x62380d103bc2ab46 commHash 7077421177785985862 - Init START
notebook-76ec3aa0aa0d-worker-0:155970:156634 [1] NCCL INFO RAS client listening socket at ::1<28028>
notebook-76ec3aa0aa0d-worker-0:155969:156631 [0] NCCL INFO RAS client listening socket at ::1<28028>
notebook-76ec3aa0aa0d-worker-0:155970:156634 [1] NCCL INFO Bootstrap timings total 0.001037 (create 0.000026, send 0.000118, recv 0.000212, ring 0.000344, delay 0.000001)
notebook-76ec3aa0aa0d-worker-0:155969:156631 [0] NCCL INFO Bootstrap timings total 0.208179 (create 0.000025, send 0.000126, recv 0.207310, ring 0.000357, delay 0.000001)
notebook-76ec3aa0aa0d-worker-0:155970:156634 [1] NCCL INFO gTaskid = b839cb210000b7ec.
notebook-76ec3aa0aa0d-worker-0:155969:156631 [0] NCCL INFO gTaskid = b839cb210000b7ec.
notebook-76ec3aa0aa0d-worker-0:155969:156631 [0] NCCL INFO hpn_license_register error: Error occurred during zmq_msg_recv(): Resource temporarily unavailable.
notebook-76ec3aa0aa0d-worker-0:155970:156634 [1] NCCL INFO hpn_license_register error: Error occurred during zmq_msg_recv(): Resource temporarily unavailable.
notebook-76ec3aa0aa0d-worker-0:155969:156631 [0] NCCL INFO MNNVL busId 0x109000 fabric UUID 0.0 cliqueId 0x0 state 3 healthMask 0x0
notebook-76ec3aa0aa0d-worker-0:155970:156634 [1] NCCL INFO MNNVL busId 0x17f000 fabric UUID 0.0 cliqueId 0x0 state 3 healthMask 0x0
notebook-76ec3aa0aa0d-worker-0:155969:156631 [0] NCCL INFO Setting affinity for GPU 0 to ffffffff,ffff0000,00000000,ffffffff,ffff0000,00000000
notebook-76ec3aa0aa0d-worker-0:155970:156634 [1] NCCL INFO Setting affinity for GPU 1 to ffffffff,ffff0000,00000000,ffffffff,ffff0000,00000000
notebook-76ec3aa0aa0d-worker-0:155969:156631 [0] NCCL INFO comm 0x9fac070 rank 0 nRanks 2 nNodes 1 localRanks 2 localRank 0 MNNVL 0
notebook-76ec3aa0aa0d-worker-0:155970:156634 [1] NCCL INFO comm 0xb34f1a0 rank 1 nRanks 2 nNodes 1 localRanks 2 localRank 1 MNNVL 0
notebook-76ec3aa0aa0d-worker-0:155969:156631 [0] NCCL INFO Channel 00/16 : 0 1
notebook-76ec3aa0aa0d-worker-0:155969:156631 [0] NCCL INFO Channel 01/16 : 0 1
notebook-76ec3aa0aa0d-worker-0:155969:156631 [0] NCCL INFO Channel 02/16 : 0 1
notebook-76ec3aa0aa0d-worker-0:155969:156631 [0] NCCL INFO Channel 03/16 : 0 1
notebook-76ec3aa0aa0d-worker-0:155969:156631 [0] NCCL INFO Channel 04/16 : 0 1
notebook-76ec3aa0aa0d-worker-0:155969:156631 [0] NCCL INFO Channel 05/16 : 0 1
notebook-76ec3aa0aa0d-worker-0:155970:156634 [1] NCCL INFO Trees [0] -1/-1/-1->1->0 [1] -1/-1/-1->1->0 [2] -1/-1/-1->1->0 [3] -1/-1/-1->1->0 [4] 0/-1/-1->1->-1 [5] 0/-1/-1->1->-1 [6] 0/-1/-1->1->-1 [7] 0/-1/-1->1->-1 [8] -1/-1/-1->1->0 [9] -1/-1/-1->1->0 [10] -1/-1/-1->1->0 [11] -1/-1/-1->1->0 [12] 0/-1/-1->1->-1 [13] 0/-1/-1->1->-1 [14] 0/-1/-1->1->-1 [15] 0/-1/-1->1->-1
notebook-76ec3aa0aa0d-worker-0:155969:156631 [0] NCCL INFO Channel 06/16 : 0 1
notebook-76ec3aa0aa0d-worker-0:155969:156631 [0] NCCL INFO Channel 07/16 : 0 1
notebook-76ec3aa0aa0d-worker-0:155970:156634 [1] NCCL INFO P2P Chunksize set to 524288
notebook-76ec3aa0aa0d-worker-0:155969:156631 [0] NCCL INFO Channel 08/16 : 0 1
notebook-76ec3aa0aa0d-worker-0:155969:156631 [0] NCCL INFO Channel 09/16 : 0 1
notebook-76ec3aa0aa0d-worker-0:155969:156631 [0] NCCL INFO Channel 10/16 : 0 1
notebook-76ec3aa0aa0d-worker-0:155969:156631 [0] NCCL INFO Channel 11/16 : 0 1
notebook-76ec3aa0aa0d-worker-0:155969:156631 [0] NCCL INFO Channel 12/16 : 0 1
notebook-76ec3aa0aa0d-worker-0:155969:156631 [0] NCCL INFO Channel 13/16 : 0 1
notebook-76ec3aa0aa0d-worker-0:155969:156631 [0] NCCL INFO Channel 14/16 : 0 1
notebook-76ec3aa0aa0d-worker-0:155969:156631 [0] NCCL INFO Channel 15/16 : 0 1
notebook-76ec3aa0aa0d-worker-0:155969:156631 [0] NCCL INFO Trees [0] 1/-1/-1->0->-1 [1] 1/-1/-1->0->-1 [2] 1/-1/-1->0->-1 [3] 1/-1/-1->0->-1 [4] -1/-1/-1->0->1 [5] -1/-1/-1->0->1 [6] -1/-1/-1->0->1 [7] -1/-1/-1->0->1 [8] 1/-1/-1->0->-1 [9] 1/-1/-1->0->-1 [10] 1/-1/-1->0->-1 [11] 1/-1/-1->0->-1 [12] -1/-1/-1->0->1 [13] -1/-1/-1->0->1 [14] -1/-1/-1->0->1 [15] -1/-1/-1->0->1
notebook-76ec3aa0aa0d-worker-0:155969:156631 [0] NCCL INFO P2P Chunksize set to 524288
notebook-76ec3aa0aa0d-worker-0:155969:156631 [0] NCCL INFO Check P2P Type intraNodeP2pSupport 1 directMode 0
notebook-76ec3aa0aa0d-worker-0:155970:156687 [1] NCCL INFO [Proxy Service] Device 1 CPU core 82
notebook-76ec3aa0aa0d-worker-0:155970:156689 [1] NCCL INFO [Proxy Service UDS] Device 1 CPU core 83
notebook-76ec3aa0aa0d-worker-0:155969:156688 [0] NCCL INFO [Proxy Service] Device 0 CPU core 164
notebook-76ec3aa0aa0d-worker-0:155969:156690 [0] NCCL INFO [Proxy Service UDS] Device 0 CPU core 165
notebook-76ec3aa0aa0d-worker-0:155970:156634 [1] NCCL INFO threadThresholds 8/8/64 | 16/8/64 | 512 | 512
notebook-76ec3aa0aa0d-worker-0:155970:156634 [1] NCCL INFO 16 coll channels, 16 collnet channels, 0 nvls channels, 16 p2p channels, 16 p2p channels per peer
notebook-76ec3aa0aa0d-worker-0:155969:156631 [0] NCCL INFO threadThresholds 8/8/64 | 16/8/64 | 512 | 512
notebook-76ec3aa0aa0d-worker-0:155969:156631 [0] NCCL INFO 16 coll channels, 16 collnet channels, 0 nvls channels, 16 p2p channels, 16 p2p channels per peer
notebook-76ec3aa0aa0d-worker-0:155969:156631 [0] NCCL INFO CC Off, workFifoBytes 1048576
notebook-76ec3aa0aa0d-worker-0:155969:156631 [0] NCCL INFO TUNER/Plugin: Could not find: libnccl-tuner.so. Using internal tuner plugin.
notebook-76ec3aa0aa0d-worker-0:155970:156634 [1] NCCL INFO TUNER/Plugin: Could not find: libnccl-tuner.so. Using internal tuner plugin.
notebook-76ec3aa0aa0d-worker-0:155969:156631 [0] NCCL INFO ncclCommInitRankConfig comm 0x9fac070 rank 0 nranks 2 cudaDev 0 nvmlDev 0 busId 109000 commId 0x62380d103bc2ab46 commHash 7077421177785985862 - Init COMPLETE
notebook-76ec3aa0aa0d-worker-0:155970:156634 [1] NCCL INFO ncclCommInitRankConfig comm 0xb34f1a0 rank 1 nranks 2 cudaDev 1 nvmlDev 1 busId 17f000 commId 0x62380d103bc2ab46 commHash 7077421177785985862 - Init COMPLETE
notebook-76ec3aa0aa0d-worker-0:155969:156631 [0] NCCL INFO Init timings - ncclCommInitRankConfig: rank 0 nranks 2 total 7.23 (kernels 1.31, alloc 0.22, bootstrap 0.21, allgathers 0.01, topo 0.37, graphs 0.00, connections 0.11, rest 5.01)
notebook-76ec3aa0aa0d-worker-0:155970:156634 [1] NCCL INFO Init timings - ncclCommInitRankConfig: rank 1 nranks 2 total 6.91 (kernels 1.18, alloc 0.23, bootstrap 0.00, allgathers 0.01, topo 0.37, graphs 0.00, connections 0.09, rest 5.02)
notebook-76ec3aa0aa0d-worker-0:155970:156691 [1] NCCL INFO Channel 00/0 : 1[1] -> 0[0] via P2P/CUMEM
notebook-76ec3aa0aa0d-worker-0:155970:156691 [1] NCCL INFO Channel 01/0 : 1[1] -> 0[0] via P2P/CUMEM
notebook-76ec3aa0aa0d-worker-0:155970:156691 [1] NCCL INFO Channel 02/0 : 1[1] -> 0[0] via P2P/CUMEM
notebook-76ec3aa0aa0d-worker-0:155969:156692 [0] NCCL INFO Channel 00/0 : 0[0] -> 1[1] via P2P/CUMEM
notebook-76ec3aa0aa0d-worker-0:155970:156691 [1] NCCL INFO Channel 03/0 : 1[1] -> 0[0] via P2P/CUMEM
notebook-76ec3aa0aa0d-worker-0:155969:156692 [0] NCCL INFO Channel 01/0 : 0[0] -> 1[1] via P2P/CUMEM
notebook-76ec3aa0aa0d-worker-0:155970:156691 [1] NCCL INFO Channel 04/0 : 1[1] -> 0[0] via P2P/CUMEM
notebook-76ec3aa0aa0d-worker-0:155969:156692 [0] NCCL INFO Channel 02/0 : 0[0] -> 1[1] via P2P/CUMEM
notebook-76ec3aa0aa0d-worker-0:155970:156691 [1] NCCL INFO Channel 05/0 : 1[1] -> 0[0] via P2P/CUMEM
notebook-76ec3aa0aa0d-worker-0:155969:156692 [0] NCCL INFO Channel 03/0 : 0[0] -> 1[1] via P2P/CUMEM
notebook-76ec3aa0aa0d-worker-0:155970:156691 [1] NCCL INFO Channel 06/0 : 1[1] -> 0[0] via P2P/CUMEM
notebook-76ec3aa0aa0d-worker-0:155969:156692 [0] NCCL INFO Channel 04/0 : 0[0] -> 1[1] via P2P/CUMEM
notebook-76ec3aa0aa0d-worker-0:155970:156691 [1] NCCL INFO Channel 07/0 : 1[1] -> 0[0] via P2P/CUMEM
notebook-76ec3aa0aa0d-worker-0:155969:156692 [0] NCCL INFO Channel 05/0 : 0[0] -> 1[1] via P2P/CUMEM
notebook-76ec3aa0aa0d-worker-0:155970:156691 [1] NCCL INFO Channel 08/0 : 1[1] -> 0[0] via P2P/CUMEM
notebook-76ec3aa0aa0d-worker-0:155969:156692 [0] NCCL INFO Channel 06/0 : 0[0] -> 1[1] via P2P/CUMEM
notebook-76ec3aa0aa0d-worker-0:155970:156691 [1] NCCL INFO Channel 09/0 : 1[1] -> 0[0] via P2P/CUMEM
notebook-76ec3aa0aa0d-worker-0:155969:156692 [0] NCCL INFO Channel 07/0 : 0[0] -> 1[1] via P2P/CUMEM
notebook-76ec3aa0aa0d-worker-0:155970:156691 [1] NCCL INFO Channel 10/0 : 1[1] -> 0[0] via P2P/CUMEM
notebook-76ec3aa0aa0d-worker-0:155969:156692 [0] NCCL INFO Channel 08/0 : 0[0] -> 1[1] via P2P/CUMEM
notebook-76ec3aa0aa0d-worker-0:155970:156691 [1] NCCL INFO Channel 11/0 : 1[1] -> 0[0] via P2P/CUMEM
notebook-76ec3aa0aa0d-worker-0:155969:156692 [0] NCCL INFO Channel 09/0 : 0[0] -> 1[1] via P2P/CUMEM
notebook-76ec3aa0aa0d-worker-0:155970:156691 [1] NCCL INFO Channel 12/0 : 1[1] -> 0[0] via P2P/CUMEM
notebook-76ec3aa0aa0d-worker-0:155969:156692 [0] NCCL INFO Channel 10/0 : 0[0] -> 1[1] via P2P/CUMEM
notebook-76ec3aa0aa0d-worker-0:155970:156691 [1] NCCL INFO Channel 13/0 : 1[1] -> 0[0] via P2P/CUMEM
notebook-76ec3aa0aa0d-worker-0:155969:156692 [0] NCCL INFO Channel 11/0 : 0[0] -> 1[1] via P2P/CUMEM
notebook-76ec3aa0aa0d-worker-0:155970:156691 [1] NCCL INFO Channel 14/0 : 1[1] -> 0[0] via P2P/CUMEM
notebook-76ec3aa0aa0d-worker-0:155969:156692 [0] NCCL INFO Channel 12/0 : 0[0] -> 1[1] via P2P/CUMEM
notebook-76ec3aa0aa0d-worker-0:155970:156691 [1] NCCL INFO Channel 15/0 : 1[1] -> 0[0] via P2P/CUMEM
notebook-76ec3aa0aa0d-worker-0:155969:156692 [0] NCCL INFO Channel 13/0 : 0[0] -> 1[1] via P2P/CUMEM
notebook-76ec3aa0aa0d-worker-0:155969:156692 [0] NCCL INFO Channel 14/0 : 0[0] -> 1[1] via P2P/CUMEM
notebook-76ec3aa0aa0d-worker-0:155969:156692 [0] NCCL INFO Channel 15/0 : 0[0] -> 1[1] via P2P/CUMEM
notebook-76ec3aa0aa0d-worker-0:155969:156692 [0] NCCL INFO Connected all rings, use ring PXN 0 GDR 1
notebook-76ec3aa0aa0d-worker-0:155970:156691 [1] NCCL INFO Connected all rings, use ring PXN 0 GDR 1
[2025-11-02 14:42:21,114] [INFO] [logging.py:128:log_dist] [Rank 0] DeepSpeed Flops Profiler Enabled: False
Using :/usr/local/ninja as PyTorch extensions root...
Using :/usr/local/ninja as PyTorch extensions root...
Detected CUDA files, patching ldflags
Emitting ninja build file :/usr/local/ninja/fused_adam/build.ninja...
/opt/conda/envs/python3.10.13/lib/python3.10/site-packages/torch/utils/cpp_extension.py:2059: UserWarning: TORCH_CUDA_ARCH_LIST is not set, all archs for visible cards are included for compilation.
If this is not desired, please set os.environ['TORCH_CUDA_ARCH_LIST'].
warnings.warn(
Building extension module fused_adam...
Allowing ninja to set a default number of workers... (overridable by setting the environment variable MAX_JOBS=N)
[1/3] /usr/local/cuda/bin/nvcc --generate-dependencies-with-compile --dependency-output multi_tensor_adam.cuda.o.d -DTORCH_EXTENSION_NAME=fused_adam -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\"_gcc\" -DPYBIND11_STDLIB=\"_libstdcpp\" -DPYBIND11_BUILD_ABI=\"_cxxabi1011\" -I/home/wangshangquan.wsq/.local/lib/python3.10/site-packages/deepspeed/ops/csrc/includes -I/home/wangshangquan.wsq/.local/lib/python3.10/site-packages/deepspeed/ops/csrc/adam -isystem /opt/conda/envs/python3.10.13/lib/python3.10/site-packages/torch/include -isystem /opt/conda/envs/python3.10.13/lib/python3.10/site-packages/torch/include/torch/csrc/api/include -isystem /opt/conda/envs/python3.10.13/lib/python3.10/site-packages/torch/include/TH -isystem /opt/conda/envs/python3.10.13/lib/python3.10/site-packages/torch/include/THC -isystem /usr/local/cuda/include -isystem /opt/conda/envs/python3.10.13/include/python3.10 -D_GLIBCXX_USE_CXX11_ABI=0 -D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ -D__CUDA_NO_BFLOAT16_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ --expt-relaxed-constexpr -gencode=arch=compute_90,code=compute_90 -gencode=arch=compute_90,code=sm_90 --compiler-options '-fPIC' -O3 -DVERSION_GE_1_1 -DVERSION_GE_1_3 -DVERSION_GE_1_5 -lineinfo --use_fast_math -gencode=arch=compute_90,code=sm_90 -gencode=arch=compute_90,code=compute_90 -DBF16_AVAILABLE -U__CUDA_NO_BFLOAT16_OPERATORS__ -U__CUDA_NO_BFLOAT162_OPERATORS__ -U__CUDA_NO_BFLOAT16_CONVERSIONS__ -std=c++17 -c /home/wangshangquan.wsq/.local/lib/python3.10/site-packages/deepspeed/ops/csrc/adam/multi_tensor_adam.cu -o multi_tensor_adam.cuda.o
[2/3] c++ -MMD -MF fused_adam_frontend.o.d -DTORCH_EXTENSION_NAME=fused_adam -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\"_gcc\" -DPYBIND11_STDLIB=\"_libstdcpp\" -DPYBIND11_BUILD_ABI=\"_cxxabi1011\" -I/home/wangshangquan.wsq/.local/lib/python3.10/site-packages/deepspeed/ops/csrc/includes -I/home/wangshangquan.wsq/.local/lib/python3.10/site-packages/deepspeed/ops/csrc/adam -isystem /opt/conda/envs/python3.10.13/lib/python3.10/site-packages/torch/include -isystem /opt/conda/envs/python3.10.13/lib/python3.10/site-packages/torch/include/torch/csrc/api/include -isystem /opt/conda/envs/python3.10.13/lib/python3.10/site-packages/torch/include/TH -isystem /opt/conda/envs/python3.10.13/lib/python3.10/site-packages/torch/include/THC -isystem /usr/local/cuda/include -isystem /opt/conda/envs/python3.10.13/include/python3.10 -D_GLIBCXX_USE_CXX11_ABI=0 -fPIC -std=c++17 -O3 -std=c++17 -g -Wno-reorder -DVERSION_GE_1_1 -DVERSION_GE_1_3 -DVERSION_GE_1_5 -DBF16_AVAILABLE -c /home/wangshangquan.wsq/.local/lib/python3.10/site-packages/deepspeed/ops/csrc/adam/fused_adam_frontend.cpp -o fused_adam_frontend.o
[3/3] c++ fused_adam_frontend.o multi_tensor_adam.cuda.o -shared -L/opt/conda/envs/python3.10.13/lib/python3.10/site-packages/torch/lib -lc10 -lc10_cuda -ltorch_cpu -ltorch_cuda -ltorch -ltorch_python -L/usr/local/cuda/lib64 -lcudart -o fused_adam.so
Loading extension module fused_adam...
Time to load fused_adam op: 27.844947814941406 seconds
Rank 1: After DeepSpeed initialize - Allocated: 0.05 GB, Reserved: 0.11 GB
Rank 1: ZeRO stage: 1
Rank 1: Optimizer has 1 parameter groups
Rank 1: Parameters in engine: 14,687,232
Rank 1: Model dtype: torch.float16
Rank 1: Before forward - Allocated: 0.05 GB, Reserved: 0.11 GB
Loading extension module fused_adam...
Time to load fused_adam op: 27.943018198013306 seconds
[2025-11-02 14:42:49,058] [INFO] [logging.py:128:log_dist] [Rank 0] Using DeepSpeed Optimizer param name adam as basic optimizer
[2025-11-02 14:42:49,059] [INFO] [logging.py:128:log_dist] [Rank 0] Removing param_group that has no 'params' in the basic Optimizer
[2025-11-02 14:42:49,059] [INFO] [logging.py:128:log_dist] [Rank 0] DeepSpeed Basic Optimizer = FusedAdam
[2025-11-02 14:42:49,059] [INFO] [utils.py:59:is_zero_supported_optimizer] Checking ZeRO support for optimizer=FusedAdam type=<class 'deepspeed.ops.adam.fused_adam.FusedAdam'>
[2025-11-02 14:42:49,059] [INFO] [logging.py:128:log_dist] [Rank 0] Creating torch.float16 ZeRO stage 1 optimizer
[2025-11-02 14:42:49,059] [INFO] [stage_1_and_2.py:149:__init__] Reduce bucket size 50000000
[2025-11-02 14:42:49,059] [INFO] [stage_1_and_2.py:150:__init__] Allgather bucket size 500000000
[2025-11-02 14:42:49,059] [INFO] [stage_1_and_2.py:151:__init__] CPU Offload: False
[2025-11-02 14:42:49,060] [INFO] [stage_1_and_2.py:152:__init__] Round robin gradient partitioning: False
Rank 1: After forward - Allocated: 0.09 GB, Reserved: 0.14 GB
Rank 1: After backward - Allocated: 0.15 GB, Reserved: 0.29 GB
[2025-11-02 14:42:49,408] [INFO] [utils.py:781:see_memory_usage] Before initializing optimizer states
[2025-11-02 14:42:49,408] [INFO] [utils.py:782:see_memory_usage] MA 0.05 GB Max_MA 0.07 GB CA 0.08 GB Max_CA 0 GB
[2025-11-02 14:42:49,408] [INFO] [utils.py:789:see_memory_usage] CPU Virtual Memory: used = 23.44 GB, percent = 2.7%
[2025-11-02 14:42:49,533] [INFO] [utils.py:781:see_memory_usage] After initializing optimizer states
[2025-11-02 14:42:49,534] [INFO] [utils.py:782:see_memory_usage] MA 0.05 GB Max_MA 0.08 GB CA 0.11 GB Max_CA 0 GB
[2025-11-02 14:42:49,534] [INFO] [utils.py:789:see_memory_usage] CPU Virtual Memory: used = 23.44 GB, percent = 2.7%
[2025-11-02 14:42:49,534] [INFO] [stage_1_and_2.py:550:__init__] optimizer state initialized
[2025-11-02 14:42:49,651] [INFO] [utils.py:781:see_memory_usage] After initializing ZeRO optimizer
[2025-11-02 14:42:49,652] [INFO] [utils.py:782:see_memory_usage] MA 0.05 GB Max_MA 0.05 GB CA 0.11 GB Max_CA 0 GB
[2025-11-02 14:42:49,652] [INFO] [utils.py:789:see_memory_usage] CPU Virtual Memory: used = 23.44 GB, percent = 2.7%
[2025-11-02 14:42:49,652] [INFO] [logging.py:128:log_dist] [Rank 0] DeepSpeed Final Optimizer = DeepSpeedZeroOptimizer
[2025-11-02 14:42:49,652] [INFO] [logging.py:128:log_dist] [Rank 0] DeepSpeed using configured LR scheduler = None
[2025-11-02 14:42:49,652] [INFO] [logging.py:128:log_dist] [Rank 0] DeepSpeed LR Scheduler = None
[2025-11-02 14:42:49,652] [INFO] [logging.py:128:log_dist] [Rank 0] step=0, skipped=0, lr=[0.001], mom=[[0.9, 0.999]]
[2025-11-02 14:42:49,652] [INFO] [config.py:1001:print] DeepSpeedEngine configuration:
[2025-11-02 14:42:49,653] [INFO] [config.py:1005:print] activation_checkpointing_config {
"partition_activations": false,
"contiguous_memory_optimization": false,
"cpu_checkpointing": false,
"number_checkpoints": null,
"synchronize_checkpoint_boundary": false,
"profile": false
}
[2025-11-02 14:42:49,653] [INFO] [config.py:1005:print] aio_config ................... {'block_size': 1048576, 'queue_depth': 8, 'intra_op_parallelism': 1, 'single_submit': False, 'overlap_events': True, 'use_gds': False}
[2025-11-02 14:42:49,653] [INFO] [config.py:1005:print] amp_enabled .................. False
[2025-11-02 14:42:49,653] [INFO] [config.py:1005:print] amp_params ................... False
[2025-11-02 14:42:49,653] [INFO] [config.py:1005:print] autotuning_config ............ {
"enabled": false,
"start_step": null,
"end_step": null,
"metric_path": null,
"arg_mappings": null,
"metric": "throughput",
"model_info": null,
"results_dir": "autotuning_results",
"exps_dir": "autotuning_exps",
"overwrite": true,
"fast": true,
"start_profile_step": 3,
"end_profile_step": 5,
"tuner_type": "gridsearch",
"tuner_early_stopping": 5,
"tuner_num_trials": 50,
"model_info_path": null,
"mp_size": 1,
"max_train_batch_size": null,
"min_train_batch_size": 1,
"max_train_micro_batch_size_per_gpu": 1.024000e+03,
"min_train_micro_batch_size_per_gpu": 1,
"num_tuning_micro_batch_sizes": 3
}
[2025-11-02 14:42:49,653] [INFO] [config.py:1005:print] bfloat16_enabled ............. False
[2025-11-02 14:42:49,653] [INFO] [config.py:1005:print] bfloat16_immediate_grad_update False
[2025-11-02 14:42:49,653] [INFO] [config.py:1005:print] checkpoint_parallel_write_pipeline False
[2025-11-02 14:42:49,653] [INFO] [config.py:1005:print] checkpoint_tag_validation_enabled True
[2025-11-02 14:42:49,653] [INFO] [config.py:1005:print] checkpoint_tag_validation_fail False
[2025-11-02 14:42:49,653] [INFO] [config.py:1005:print] comms_config ................. <deepspeed.comm.config.DeepSpeedCommsConfig object at 0x7facd29c0bb0>
[2025-11-02 14:42:49,653] [INFO] [config.py:1005:print] communication_data_type ...... None
[2025-11-02 14:42:49,653] [INFO] [config.py:1005:print] compression_config ........... {'weight_quantization': {'shared_parameters': {'enabled': False, 'quantizer_kernel': False, 'schedule_offset': 0, 'quantize_groups': 1, 'quantize_verbose': False, 'quantization_type': 'symmetric', 'quantize_weight_in_forward': False, 'rounding': 'nearest', 'fp16_mixed_quantize': False, 'quantize_change_ratio': 0.001}, 'different_groups': {}}, 'activation_quantization': {'shared_parameters': {'enabled': False, 'quantization_type': 'symmetric', 'range_calibration': 'dynamic', 'schedule_offset': 1000}, 'different_groups': {}}, 'sparse_pruning': {'shared_parameters': {'enabled': False, 'method': 'l1', 'schedule_offset': 1000}, 'different_groups': {}}, 'row_pruning': {'shared_parameters': {'enabled': False, 'method': 'l1', 'schedule_offset': 1000}, 'different_groups': {}}, 'head_pruning': {'shared_parameters': {'enabled': False, 'method': 'topk', 'schedule_offset': 1000}, 'different_groups': {}}, 'channel_pruning': {'shared_parameters': {'enabled': False, 'method': 'l1', 'schedule_offset': 1000}, 'different_groups': {}}, 'layer_reduction': {'enabled': False}}
[2025-11-02 14:42:49,653] [INFO] [config.py:1005:print] curriculum_enabled_legacy .... False
[2025-11-02 14:42:49,653] [INFO] [config.py:1005:print] curriculum_params_legacy ..... False
[2025-11-02 14:42:49,653] [INFO] [config.py:1005:print] data_efficiency_config ....... {'enabled': False, 'seed': 1234, 'data_sampling': {'enabled': False, 'num_epochs': 1000, 'num_workers': 0, 'curriculum_learning': {'enabled': False}}, 'data_routing': {'enabled': False, 'random_ltd': {'enabled': False, 'layer_token_lr_schedule': {'enabled': False}}}}
[2025-11-02 14:42:49,654] [INFO] [config.py:1005:print] data_efficiency_enabled ...... False
[2025-11-02 14:42:49,654] [INFO] [config.py:1005:print] dataloader_drop_last ......... False
[2025-11-02 14:42:49,654] [INFO] [config.py:1005:print] disable_allgather ............ False
[2025-11-02 14:42:49,654] [INFO] [config.py:1005:print] dump_state ................... False
[2025-11-02 14:42:49,654] [INFO] [config.py:1005:print] dynamic_loss_scale_args ...... {'init_scale': 65536, 'scale_window': 1000, 'delayed_shift': 2, 'consecutive_hysteresis': False, 'min_scale': 1}
[2025-11-02 14:42:49,654] [INFO] [config.py:1005:print] eigenvalue_enabled ........... False
[2025-11-02 14:42:49,654] [INFO] [config.py:1005:print] eigenvalue_gas_boundary_resolution 1
[2025-11-02 14:42:49,654] [INFO] [config.py:1005:print] eigenvalue_layer_name ........ bert.encoder.layer
[2025-11-02 14:42:49,654] [INFO] [config.py:1005:print] eigenvalue_layer_num ......... 0
[2025-11-02 14:42:49,654] [INFO] [config.py:1005:print] eigenvalue_max_iter .......... 100
[2025-11-02 14:42:49,654] [INFO] [config.py:1005:print] eigenvalue_stability ......... 1e-06
[2025-11-02 14:42:49,654] [INFO] [config.py:1005:print] eigenvalue_tol ............... 0.01
[2025-11-02 14:42:49,654] [INFO] [config.py:1005:print] eigenvalue_verbose ........... False
[2025-11-02 14:42:49,654] [INFO] [config.py:1005:print] elasticity_enabled ........... False
[2025-11-02 14:42:49,654] [INFO] [config.py:1005:print] flops_profiler_config ........ {
"enabled": false,
"recompute_fwd_factor": 0.0,
"profile_step": 1,
"module_depth": -1,
"top_modules": 1,
"detailed": true,
"output_file": null
}
[2025-11-02 14:42:49,654] [INFO] [config.py:1005:print] fp16_auto_cast ............... False
[2025-11-02 14:42:49,654] [INFO] [config.py:1005:print] fp16_enabled ................. True
[2025-11-02 14:42:49,654] [INFO] [config.py:1005:print] fp16_master_weights_and_gradients False
[2025-11-02 14:42:49,654] [INFO] [config.py:1005:print] global_rank .................. 0
[2025-11-02 14:42:49,654] [INFO] [config.py:1005:print] grad_accum_dtype ............. None
[2025-11-02 14:42:49,654] [INFO] [config.py:1005:print] gradient_accumulation_steps .. 1
[2025-11-02 14:42:49,654] [INFO] [config.py:1005:print] gradient_clipping ............ 0.0
[2025-11-02 14:42:49,654] [INFO] [config.py:1005:print] gradient_predivide_factor .... 1.0
[2025-11-02 14:42:49,654] [INFO] [config.py:1005:print] graph_harvesting ............. False
[2025-11-02 14:42:49,654] [INFO] [config.py:1005:print] hybrid_engine ................ enabled=False max_out_tokens=512 inference_tp_size=1 release_inference_cache=False pin_parameters=True tp_gather_partition_size=8
[2025-11-02 14:42:49,654] [INFO] [config.py:1005:print] initial_dynamic_scale ........ 65536
[2025-11-02 14:42:49,654] [INFO] [config.py:1005:print] load_universal_checkpoint .... False
[2025-11-02 14:42:49,654] [INFO] [config.py:1005:print] loss_scale ................... 0
[2025-11-02 14:42:49,655] [INFO] [config.py:1005:print] memory_breakdown ............. False
[2025-11-02 14:42:49,655] [INFO] [config.py:1005:print] mics_hierarchial_params_gather False
[2025-11-02 14:42:49,655] [INFO] [config.py:1005:print] mics_shard_size .............. -1
[2025-11-02 14:42:49,655] [INFO] [config.py:1005:print] monitor_config ............... tensorboard=TensorBoardConfig(enabled=False, output_path='', job_name='DeepSpeedJobName') comet=CometConfig(enabled=False, samples_log_interval=100, project=None, workspace=None, api_key=None, experiment_name=None, experiment_key=None, online=None, mode=None) wandb=WandbConfig(enabled=False, group=None, team=None, project='deepspeed') csv_monitor=CSVConfig(enabled=False, output_path='', job_name='DeepSpeedJobName')
[2025-11-02 14:42:49,655] [INFO] [config.py:1005:print] nebula_config ................ {
"enabled": false,
"persistent_storage_path": null,
"persistent_time_interval": 100,
"num_of_version_in_retention": 2,
"enable_nebula_load": true,
"load_path": null
}
[2025-11-02 14:42:49,655] [INFO] [config.py:1005:print] optimizer_legacy_fusion ...... False
[2025-11-02 14:42:49,655] [INFO] [config.py:1005:print] optimizer_name ............... adam
[2025-11-02 14:42:49,655] [INFO] [config.py:1005:print] optimizer_params ............. {'lr': 0.001, 'betas': [0.9, 0.999], 'eps': 1e-08, 'weight_decay': 0}
[2025-11-02 14:42:49,655] [INFO] [config.py:1005:print] pipeline ..................... {'stages': 'auto', 'partition': 'best', 'seed_layers': False, 'activation_checkpoint_interval': 0, 'pipe_partitioned': True, 'grad_partitioned': True}
[2025-11-02 14:42:49,655] [INFO] [config.py:1005:print] pld_enabled .................. False
[2025-11-02 14:42:49,655] [INFO] [config.py:1005:print] pld_params ................... False
[2025-11-02 14:42:49,655] [INFO] [config.py:1005:print] prescale_gradients ........... False
[2025-11-02 14:42:49,655] [INFO] [config.py:1005:print] scheduler_name ............... None
[2025-11-02 14:42:49,655] [INFO] [config.py:1005:print] scheduler_params ............. None
[2025-11-02 14:42:49,655] [INFO] [config.py:1005:print] seq_parallel_communication_data_type torch.float32
[2025-11-02 14:42:49,655] [INFO] [config.py:1005:print] sparse_attention ............. None
[2025-11-02 14:42:49,655] [INFO] [config.py:1005:print] sparse_gradients_enabled ..... False
[2025-11-02 14:42:49,655] [INFO] [config.py:1005:print] steps_per_print .............. 10
[2025-11-02 14:42:49,655] [INFO] [config.py:1005:print] tensor_parallel_config ....... dtype=torch.float16 autotp_size=0 tensor_parallel=TPConfig(tp_size=1, tp_grain_size=1, mpu=None, tp_group=None) injection_policy_tuple=None keep_module_on_host=False replace_with_kernel_inject=False
[2025-11-02 14:42:49,655] [INFO] [config.py:1005:print] timers_config ................ enabled=True synchronized=True
[2025-11-02 14:42:49,655] [INFO] [config.py:1005:print] train_batch_size ............. 256
[2025-11-02 14:42:49,655] [INFO] [config.py:1005:print] train_micro_batch_size_per_gpu 128
[2025-11-02 14:42:49,655] [INFO] [config.py:1005:print] use_data_before_expert_parallel_ False
[2025-11-02 14:42:49,655] [INFO] [config.py:1005:print] use_node_local_storage ....... False
[2025-11-02 14:42:49,655] [INFO] [config.py:1005:print] wall_clock_breakdown ......... False
[2025-11-02 14:42:49,656] [INFO] [config.py:1005:print] weight_quantization_config ... None
[2025-11-02 14:42:49,656] [INFO] [config.py:1005:print] world_size ................... 2
[2025-11-02 14:42:49,656] [INFO] [config.py:1005:print] zero_allow_untested_optimizer False
[2025-11-02 14:42:49,656] [INFO] [config.py:1005:print] zero_config .................. stage=1 contiguous_gradients=True reduce_scatter=True reduce_bucket_size=50000000 use_multi_rank_bucket_allreduce=True allgather_partitions=True allgather_bucket_size=500000000 overlap_comm=False load_from_fp32_weights=True elastic_checkpoint=False offload_param=None offload_optimizer=None sub_group_size=1000000000 cpu_offload_param=None cpu_offload_use_pin_memory=None cpu_offload=None prefetch_bucket_size=50000000 param_persistence_threshold=100000 model_persistence_threshold=9223372036854775807 max_live_parameters=1000000000 max_reuse_distance=1000000000 gather_16bit_weights_on_model_save=False module_granularity_threshold=0 use_all_reduce_for_fetch_params=False stage3_gather_fp16_weights_on_model_save=False ignore_unused_parameters=True legacy_stage1=False round_robin_gradients=False zero_hpz_partition_size=1 zero_quantized_weights=False zero_quantized_nontrainable_weights=False zero_quantized_gradients=False zeropp_loco_param=None mics_shard_size=-1 mics_hierarchical_params_gather=False memory_efficient_linear=True pipeline_loading_checkpoint=False override_module_apply=True log_trace_cache_warnings=False
[2025-11-02 14:42:49,656] [INFO] [config.py:1005:print] zero_enabled ................. True
[2025-11-02 14:42:49,656] [INFO] [config.py:1005:print] zero_force_ds_cpu_optimizer .. True
[2025-11-02 14:42:49,656] [INFO] [config.py:1005:print] zero_optimization_stage ...... 1
[2025-11-02 14:42:49,656] [INFO] [config.py:991:print_user_config] json = {
"train_batch_size": 256,
"train_micro_batch_size_per_gpu": 128,
"steps_per_print": 10,
"optimizer": {
"type": "Adam",
"params": {
"lr": 0.001,
"betas": [0.9, 0.999],
"eps": 1e-08,
"weight_decay": 0
}
},
"fp16": {
"enabled": true,
"loss_scale": 0,
"loss_scale_window": 1000,
"hysteresis": 2,
"min_loss_scale": 1
},
"zero_optimization": {
"stage": 1,
"reduce_bucket_size": 5.000000e+07
}
}
Rank 0: After DeepSpeed initialize - Allocated: 0.05 GB, Reserved: 0.11 GB
Rank 0: ZeRO stage: 1
Rank 0: Optimizer has 1 parameter groups
Rank 0: Parameters in engine: 14,687,232
Rank 0: Model dtype: torch.float16
Rank 0: Before forward - Allocated: 0.05 GB, Reserved: 0.11 GB
Rank 0: After forward - Allocated: 0.09 GB, Reserved: 0.14 GB
Rank 0: After backward - Allocated: 0.14 GB, Reserved: 0.28 GB
[2025-11-02 14:42:49,890] [INFO] [loss_scaler.py:190:update_scale] [deepspeed] OVERFLOW! Rank 0 Skipping step. Attempted loss scale: 65536, but hysteresis is 2. Reducing hysteresis to 1
Rank 1: After step - Allocated: 0.15 GB, Reserved: 0.29 GB
Rank 0: After step - Allocated: 0.14 GB, Reserved: 0.28 GB
[rank0]:[W1102 14:42:50.589162769 ProcessGroupNCCL.cpp:1496] Warning: WARNING: destroy_process_group() was not called before program exit, which can leak resources. For more info, please see https://pytorch.org/docs/stable/distributed.html#shutdown (function operator())
[2025-11-02 14:42:51,870] [INFO] [launch.py:351:main] Process 155969 exits successfully.
[2025-11-02 14:42:51,871] [INFO] [launch.py:351:main] Process 155970 ex稳定后内存分析
在这里插入图片描述
由rofiler_schedule = torch.profiler.schedule(wait=0, warmup=0, active=6, repeat=1)
改为rofiler_schedule = torch.profiler.schedule(wait=2, warmup=2, active=6, repeat=1)
启动前内存
cpp
import torch
import torch.nn as nn
import deepspeed
import argparse
import os
import socket
from datetime import datetime
# --- Model 和 get_args 定义 (与之前相同) ---
class SimpleModel(nn.Module):
def __init__(self):
super().__init__()
self.layer1 = nn.Linear(1024, 2048)
self.relu1 = nn.ReLU()
self.layer2 = nn.Linear(2048, 4096)
self.relu2 = nn.ReLU()
self.layer3 = nn.Linear(4096, 1024)
def forward(self, x):
return self.relu2(self.layer2(self.relu1(self.layer1(x))))
def get_args():
parser = argparse.ArgumentParser(description='Record Memory Snapshot with DeepSpeed')
parser.add_argument('--local_rank', type=int, default=-1, help='local rank passed from distributed launcher')
parser = deepspeed.add_config_arguments(parser)
args = parser.parse_args()
return args
# --- 从 Meta 示例中适配的辅助函数 ---
# 每个快照中记录的最大内存事件数
MAX_NUM_OF_MEM_EVENTS_PER_SNAPSHOT: int = 100000
def start_record_memory_history():
if not torch.cuda.is_available():
print("CUDA unavailable. Not recording memory history")
return
print("Starting snapshot record_memory_history")
torch.cuda.memory._record_memory_history(
max_entries=MAX_NUM_OF_MEM_EVENTS_PER_SNAPSHOT
)
def stop_record_memory_history():
if not torch.cuda.is_available():
return
print("Stopping snapshot record_memory_history")
torch.cuda.memory._record_memory_history(enabled=None)
def export_memory_snapshot(rank):
"""
为特定的 rank 导出内存快照。
"""
if not torch.cuda.is_available():
return
# 为每个 rank 生成唯一的文件名
host_name = socket.gethostname()
timestamp = datetime.now().strftime("%b_%d_%H_%M_%S")
file_prefix = f"{host_name}_{timestamp}_rank{rank}"
filename = f"{file_prefix}.pickle"
try:
print(f"Rank {rank}: Saving snapshot to local file: {filename}")
torch.cuda.memory._dump_snapshot(filename)
except Exception as e:
print(f"Rank {rank}: Failed to capture memory snapshot {e}")
# --- 主函数 ---
def main():
args = get_args()
deepspeed.init_distributed()
rank = torch.distributed.get_rank()
torch.cuda.set_device(args.local_rank)
model = SimpleModel()
model_engine, _, _, _ = deepspeed.initialize(
args=args, model=model, model_parameters=model.parameters()
)
# 在所有 rank 上启动内存历史记录
start_record_memory_history()
# 运行几个 step
for step in range(5):
print(f"Rank {rank}, starting Step {step}...")
model_dtype = next(model_engine.module.parameters()).dtype
inputs = torch.randn(
model_engine.train_micro_batch_size_per_gpu(),
1024,
device=model_engine.device
).to(model_dtype)
loss = model_engine(inputs).mean()
model_engine.backward(loss)
model_engine.step()
# 在所有 rank 上导出各自的内存快照
export_memory_snapshot(rank)
# 停止记录
stop_record_memory_history()
print(f"Rank {rank} finished.")
if __name__ == "__main__":
main()稳定后
cpp
{
"train_batch_size": 256,
"train_micro_batch_size_per_gpu": 128,
"steps_per_print": 10,
"optimizer": {
"type": "Adam",
"params": {
"lr": 1e-3,
"betas": [0.9, 0.999],
"eps": 1e-8,
"weight_decay": 0
}
},
"fp16": {
"enabled": true,
"loss_scale": 0,
"loss_scale_window": 1000,
"hysteresis": 2,
"min_loss_scale": 1
},
"zero_optimization": {
"stage": 1,
"reduce_bucket_size": 5e7
}
}
cpp
import torch
import torch.nn as nn
import deepspeed
import argparse
import os
import socket
from datetime import datetime
# --- Model 和 get_args 定义 (与之前相同) ---
class SimpleModel(nn.Module):
def __init__(self):
super().__init__()
self.layer1 = nn.Linear(1024, 2048)
self.relu1 = nn.ReLU()
self.layer2 = nn.Linear(2048, 4096)
self.relu2 = nn.ReLU()
self.layer3 = nn.Linear(4096, 1024)
def forward(self, x):
return self.relu2(self.layer2(self.relu1(self.layer1(x))))
def get_args():
parser = argparse.ArgumentParser(description='Record Memory Snapshot with DeepSpeed')
parser.add_argument('--local_rank', type=int, default=-1, help='local rank passed from distributed launcher')
parser = deepspeed.add_config_arguments(parser)
args = parser.parse_args()
return args
# --- 从 Meta 示例中适配的辅助函数 ---
# 每个快照中记录的最大内存事件数
MAX_NUM_OF_MEM_EVENTS_PER_SNAPSHOT: int = 100000
def start_record_memory_history():
if not torch.cuda.is_available():
print("CUDA unavailable. Not recording memory history")
return
print("Starting snapshot record_memory_history")
torch.cuda.memory._record_memory_history(
max_entries=MAX_NUM_OF_MEM_EVENTS_PER_SNAPSHOT
)
def stop_record_memory_history():
if not torch.cuda.is_available():
return
print("Stopping snapshot record_memory_history")
torch.cuda.memory._record_memory_history(enabled=None)
def export_memory_snapshot(rank):
"""
为特定的 rank 导出内存快照。
"""
if not torch.cuda.is_available():
return
# 为每个 rank 生成唯一的文件名
host_name = socket.gethostname()
timestamp = datetime.now().strftime("%b_%d_%H_%M_%S")
file_prefix = f"{host_name}_{timestamp}_rank{rank}"
filename = f"{file_prefix}.pickle"
try:
print(f"Rank {rank}: Saving snapshot to local file: {filename}")
torch.cuda.memory._dump_snapshot(filename)
except Exception as e:
print(f"Rank {rank}: Failed to capture memory snapshot {e}")
# --- 主函数 ---
def main():
args = get_args()
deepspeed.init_distributed()
rank = torch.distributed.get_rank()
torch.cuda.set_device(args.local_rank)
model = SimpleModel()
model_engine, _, _, _ = deepspeed.initialize(
args=args, model=model, model_parameters=model.parameters()
)
# --- 【核心修改】将训练循环分为预热和记录两个阶段 ---
# 1. 预热阶段 (前 10 个 step)
warmup_steps = 10
print(f"Rank {rank}: Starting {warmup_steps} warmup steps...")
for step in range(warmup_steps):
print(f"Rank {rank}, starting Warmup Step {step}...")
model_dtype = next(model_engine.module.parameters()).dtype
inputs = torch.randn(
model_engine.train_micro_batch_size_per_gpu(),
1024,
device=model_engine.device
).to(model_dtype)
loss = model_engine(inputs).mean()
model_engine.backward(loss)
model_engine.step()
print(f"Rank {rank}: Warmup finished.")
# 2. 记录阶段 (在预热后开始)
# 先启动内存记录器
start_record_memory_history()
# 再运行几个 step 以捕获数据
record_steps = 5
print(f"Rank {rank}: Starting {record_steps} recording steps...")
for step in range(record_steps):
print(f"Rank {rank}, starting Recording Step {step}...")
# 训练逻辑与预热阶段完全相同
model_dtype = next(model_engine.module.parameters()).dtype
inputs = torch.randn(
model_engine.train_micro_batch_size_per_gpu(),
1024,
device=model_engine.device
).to(model_dtype)
loss = model_engine(inputs).mean()
model_engine.backward(loss)
model_engine.step()
print(f"Rank {rank}: Recording finished.")
# 记录阶段结束后，导出快照
export_memory_snapshot(rank)
# 停止记录
stop_record_memory_history()
print(f"Rank {rank} finished training.")
if __name__ == "__main__":
main()