Patch Megatron TP grad coalesce to chunked all-reduce

[jingshenghang]

Summary

Patch megatron.core.distributed.finalize_model_grads._allreduce_non_tensor_model_parallel_grads (and its legacy alias _allreduce_layernorm_grads) to coalesce TP-side grads in size-bounded chunks instead of one large _flatten_dense_tensors(grads) buffer.

Why

The original implementation builds a single contiguous coalesce buffer via _flatten_dense_tensors(grads), then runs one all_reduce. On large models with many TP-side grads (LayerNorm sequence_parallel, q_layernorm/k_layernorm, average_gradients_across_tp_domain) this single buffer can be many GiB. Under allocator fragmentation (e.g. PYTORCH_ALLOC_CONF=expandable_segments:True) the contiguous allocation can fail even when total free memory is large; empty_cache() does not help because no segment can be fully released.

Concretely on GLM-4.6 750B with --accumulate-allreduce-grads-in-fp32, this batch is ~12.43 GiB per GPU, which OOMs reliably under expandable_segments.

What changes

• New slime/backends/megatron_utils/megatron_patch/ subpackage containing the patch
• slime/backends/megatron_utils/__init__.py imports the subpackage on slime startup — same monkey-patch pattern as the existing deep_ep.Buffer.__init__ patch already in that file
• Patched function greedy-splits the grad list into chunks of at most SLIME_GRAD_COALESCE_CHUNK_BYTES (default 1 GiB) and runs flatten -> all_reduce -> unflatten -> copy_back per chunk

Mathematical equivalence

SUM and AVG are element-wise:

all_reduce(concat(g_1, ..., g_N))[i] == all_reduce(g_i)[i]

Every grad receives bit-identical values vs. the original.

Cross-version compatibility

The patch resolves API differences at runtime via inspect.signature and hasattr — no version-conditional imports. Verified to apply on:

Megatron line	_get_main_grad_attr	target sig	DDP FSDP field
core_v0.13.0 line	(param, use_custom_fsdp)	(model, config)	use_custom_fsdp
post-core_v0.15.0rc7 dev	(param)	(model, config, tp_group)	use_megatron_fsdp

Performance impact

Each call issues ceil(total_grad_bytes / chunk_bytes) all_reduce calls instead of one — for the GLM-4.6 750B case roughly 12 launches instead of 1. The original was bandwidth-bound; extra launch overhead is sub-millisecond, well below per-step variance. Total bytes on the wire are unchanged.

Configuration

Item
chunk size	SLIME_GRAD_COALESCE_CHUNK_BYTES, default 1 << 30 (1 GiB)
disable	remove the from . import megatron_patch line in slime/backends/megatron_utils/__init__.py, or set SLIME_GRAD_COALESCE_CHUNK_BYTES larger than the total grad bytes

Safety

The try/except ImportError block calls warnings.warn (not silent pass), so a future Megatron upgrade that renames a helper surfaces the patch failure in logs rather than silently bringing back the OOM. On successful application an INFO log line is emitted, easily greppable in training logs:

slime grad coalesce patch applied to megatron.core.distributed.finalize_model_grads._allreduce_non_tensor_model_parallel_grads (chunk=1024 MiB)

Test plan

• Patch applies cleanly across both Megatron lines (core_v0.13.0 line and post-core_v0.15.0rc7 dev line); INFO log line emitted in each case
• End-to-end: a previously-OOM training run on GLM-4.6 750B passes the grad-sync step under the patched function