perf: chunk Megatron TP-side grad all-reduce to bound peak memory

[aoshen02]

What

Backports slime #1899.

Patches _allreduce_non_tensor_model_parallel_grads (and its legacy alias _allreduce_layernorm_grads) in megatron.core.distributed.finalize_model_grads to coalesce/all-reduce TP-side gradients in size-bounded chunks (SLIME_GRAD_COALESCE_CHUNK_BYTES, default 1 GiB) instead of a single large _flatten_dense_tensors(grads).

Why

The default path flattens all TP-side grads into one contiguous buffer before the all-reduce. For large models that buffer can be huge, and under allocator fragmentation the contiguous allocation OOMs even when total free memory is sufficient. Chunking bounds the peak contiguous allocation.

SUM/AVG reductions are element-wise, so chunking is mathematically equivalent — no numerical change, purely a memory-pressure mitigation.

Shape / safety

• Self-contained: new slime/backends/megatron_utils/megatron_patch/ subpackage, wired via a single from . import megatron_patch in megatron_utils/__init__.py.
• The patch module guards its megatron imports in try/except; if megatron is unavailable or its symbol layout changed, it logs a warning and is a no-op (never crashes import).
• Cross-compatible across the core_v0.13.0 and post-core_v0.15.0rc7 Megatron lines — API differences (use_custom_fsdp vs use_megatron_fsdp, _get_main_grad_attr arity, target-fn signature) are resolved at runtime, no version-conditional imports.

Patch file is byte-identical to slime upstream.

Scope note

This is not a precision change — it's OOM/numerical-safety infra (the P2 item from the upstream survey). Opened separately per request.

Decoupling

Independent of #92 / #93 / #94 / #95 — new subpackage + one import line; only shared file is __init__.py which none of the others touch.

🤖 Generated with Claude Code

[gemini-code-assist]

Code Review

This pull request introduces a patch to coalesce and all-reduce tensor parallel (TP) gradients in size-bounded chunks instead of one large flattened tensor, helping to lower peak contiguous-memory allocation and avoid out-of-memory (OOM) errors. The review feedback suggests using strict=True in zip calls to ensure parameter and gradient alignment, and explicitly deleting loop variables (synced, param, buf) to prevent PyTorch from retaining memory across chunk iterations.

[gemini-code-assist]

PyTorch Memory Retention Issue

In Python, loop variables (such as synced, param, and buf) remain in the local scope even after the loop finishes.

Because _unflatten_dense_tensors returns views of the coalesced tensor, the loop variable synced (which holds the last view in the loop) keeps a reference to the underlying Storage of coalesced alive. As a result, even though del coalesced is called at the end of the chunk iteration, the GPU memory allocated for coalesced (up to 1 GiB by default) is not freed until synced is overwritten or deleted.

In the next iteration of the outer loop, when the next chunk's coalesced tensor is allocated, the previous chunk's memory is still active in GPU memory. This effectively doubles the peak memory overhead of the coalescing buffer, partially defeating the purpose of chunking.

Solution

Explicitly delete synced, param, and buf along with coalesced at the end of each chunk iteration to ensure the memory is immediately reclaimed by PyTorch's caching allocator. Additionally, we should use strict=True in zip to ensure that the unflattened tensors match the chunk size exactly.

Suggested change

	for param, buf, synced in zip(
	p_chunk, g_chunk, _unflatten_dense_tensors(coalesced, g_chunk), strict=False
	):
	buf.copy_(synced)
	grad_attr = _grad_attr(param, fsdp_on)
	orig_grad = getattr(param, grad_attr)
	if fsdp_on and hasattr(orig_grad, "_local_tensor"):
	# buf already aliases orig_grad._local_tensor.data;
	# restore original DTensor wrapper (post-rc7 dev semantics).
	setattr(param, grad_attr, orig_grad)
	else:
	setattr(param, grad_attr, _reshard_if_dtensor(buf, orig_grad))
	del coalesced
	for param, buf, synced in zip(
	p_chunk, g_chunk, _unflatten_dense_tensors(coalesced, g_chunk), strict=True
	):
	buf.copy_(synced)
	grad_attr = _grad_attr(param, fsdp_on)
	orig_grad = getattr(param, grad_attr)
	if fsdp_on and hasattr(orig_grad, "_local_tensor"):
	# buf already aliases orig_grad._local_tensor.data;
	# restore original DTensor wrapper (post-rc7 dev semantics).
	setattr(param, grad_attr, orig_grad)
	else:
	setattr(param, grad_attr, _reshard_if_dtensor(buf, orig_grad))
	del coalesced, synced, param, buf

[gemini-code-assist]

Using strict=True is safer here to ensure that params and grads are perfectly aligned and of equal length. Since Megatron-LM core v0.13+ requires Python 3.10+, strict=True is fully supported and prevents potential silent bugs if the lists ever diverge in length.

Suggested change

	for p, g in zip(params, grads, strict=False):
	for p, g in zip(params, grads, strict=True):

[CalvinXKY]

LGTM