在大型模型 TP 训练中，原版将梯度合并为一个连续 buffer 的 all_reduce 分配了巨量连续内存（GLM-4.6 750B 约 12.43 GiB），在 PYTORCH_ALLOC_CONF=expandable_segments:True 下因内存碎片稳定 OOM。通过分块避免一次大连续分配，解决可靠 OOM 问题。
建议精读此 PR，尤其是 _split_into_chunks 贪心分块策略、_grad_attr 和 _fsdp_flag 的跨版本兼容写法，以及 monkey-patch 嵌入方式。这些设计模式可复用于其他 Megatron 内部函数修补场景。
此 PR 无公开 review 讨论，直接由维护者审核并合并。PR body 详细说明了设计权衡和数学等价性。
slime/backends/megatron_utils/megatron_patch/megatron_chunked_grad_coalesce_patch.py 定义了核心函数 _split_into_chunks（贪心分块）、_grad_attr（兼容两个版本的 grad 属性获取）、_fsdp_flag（兼容两个 FSDP 配置名），以及补丁版本的 _allreduce_non_tensor_model_parallel_grads。average_gradients_across_tp_domain 和 sequence_parallel / qk_layernorm 分为 sum 和 avg 两类梯度列表；对每类调用 _split_into_chunks 按 SLIME_GRAD_COALESCE_CHUNK_BYTES（默认 1 GiB）分块，每块单独执行 _flatten_dense_tensors -> all_reduce -> _unflatten_dense_tensors 并写回原参数 grad。inspect.signature 判断 _get_main_grad_attr 是否接受 use_custom_fsdp 参数；通过函数签名判断 tp_group 是否传入；无需版本条件导入。megatron_patch/__init__.py 导入补丁模块，slime/backends/megatron_utils/__init__.py 在初始化末尾添加 from . import megatron_patch，沿用已有的 monkey-patch 模式。SLIME_GRAD_COALESCE_CHUNK_BYTES 可调（默认 1073741824 字节）；try/except ImportError 捕获导入失败并发出警告，不会静默失败。| 文件 | 模块 | 状态 | 重要度 |
|---|---|---|---|
slime/backends/megatron_utils/megatron_patch/megatron_chunked_grad_coalesce_patch.py |
梯度归约 | added | 8.9 |
slime/backends/megatron_utils/megatron_patch/__init__.py |
补丁入口 | added | 4.59 |
slime/backends/megatron_utils/__init__.py |
初始化 | modified | 4.89 |
slime/backends/megatron_utils/megatron_patch/megatron_chunked_grad_coalesce_patch.py
core-logic
新核心补丁文件，实现可配置分块的 TP 梯度归约，包含跨版本兼容适配和贪心分块算法。
# 该文件是 slime/backends/megatron_utils/megatron_patch/megatron_chunked_grad_coalesce_patch.py
# 运行时动态检测 Megatron 版本，适配 API 差异。
import inspect
import os
import torch
from megatron.core.distributed.finalize_model_grads import (
_flatten_dense_tensors, _get_main_grad_attr, _unflatten_dense_tensors,
)
# 检测 _get_main_grad_attr 是否接受 fsdp 参数 (post-core_v0.15.0rc7 版本 )
_gma_takes_fsdp_arg = len(inspect.signature(_get_main_grad_attr).parameters) >= 2
def _grad_attr(param, fsdp_on):
"""兼容两个版本：若支持 fsdp 标志则传入，否则仅传 param。"""
if _gma_takes_fsdp_arg:
return _get_main_grad_attr(param, fsdp_on)
return _get_main_grad_attr(param)
def _fsdp_flag(ddp_config):
"""兼容两个版本的 FSDP 属性名 (use_megatron_fsdp / use_custom_fsdp)。"""
return bool(getattr(ddp_config, "use_megatron_fsdp", False) or
getattr(ddp_config, "use_custom_fsdp", False))
# 环境变量控制块大小，默认 1 GiB
_chunk_bytes = int(os.environ.get("SLIME_GRAD_COALESCE_CHUNK_BYTES") or (1 << 30))
def _split_into_chunks(params, grads, target_bytes):
"""贪婪分块：确保每个块总大小不超过 target_bytes，
单个超大梯度单独成块。返回 (param_chunk, grad_chunk) 列表。"""
chunks, cur_p, cur_g, cur_b = [], [], [], 0
for p, g in zip(params, grads, strict=False):
gb = g.numel() * g.element_size()
if cur_g and cur_b + gb > target_bytes:
chunks.append((cur_p, cur_g))
cur_p, cur_g, cur_b = [], [], 0
cur_p.append(p)
cur_g.append(g)
cur_b += gb
if cur_g:
chunks.append((cur_p, cur_g))
return chunks
当前评论区没有形成足够清晰的争议点或结论，后续有更多讨论时会体现在这里。
1）性能风险：分块增加了 all_reduce 调用次数（约 12 vs 1），但 body 评估额外开销 <1ms，可忽略。2）兼容性风险：补丁依赖 Megatron 内部 _get_main_grad_attr、_flatten_dense_tensors 等私有函数，未来版本升级可能导致失效，但 try/except 会输出警告日志，避免静默失败。3）配置风险：若 SLIME_GRAD_COALESCE_CHUNK_BYTES 设置过小（如 <1 MB）会导致大量小 all_reduce，降低性能；默认值合理。4）未添加测试覆盖此补丁的集成行为，回归依赖运行时日志。
对用户：使用 TP 并行且梯度总量大的大模型训练（如 GLM-4.6 750B）可直接避免 OOM，训练稳定性显著提升。对系统：无外部依赖变更，仅增强现有补丁机制。对团队：需维护跨版本兼容的补丁代码，但通过 inspect 动态适配减轻了负担。影响范围中等，但影响程度较高。
当前没有检测到明确关联的 Issue 链接，后续同步到相关引用后会出现在这里。
参与讨论