[BUG] `make_viewless_tensor()` causes GPU memory leak when `context_parallel_size > 1`
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 6h
- Merged PRs (30d)
- 271
Description
## Summary
When `context_parallel_size > 1`, `make_viewless_tensor()` in `megatron/core/utils.py` clones `hidden_states` (because it's a view from CP sequence
splitting). The cloned tensor is retained by PyTorch's C++ autograd engine after backward completes and never freed, causing ~0.5 GiB/step GPU memory leak
and eventual OOM.
This does **not** occur with `context_parallel_size = 1` because `hidden_states` is not a view, so `make_viewless_tensor()` returns it as-is without
cloning.
## Root Cause
`make_viewless_tensor()` ([utils.py#L720-L738](https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/utils.py)) with `keep_graph=True`:
```python
def make_viewless_tensor(inp, requires_grad, keep_graph):
if inp._base is None:
return inp # not a view → no-op (CP=1 hits this path)
if keep_graph:
return MakeViewlessTensor.apply(inp) # CP>1 hits this path → clone with independent storage
```
With `context_parallel_size > 1`:
1. `get_batch_on_this_cp_rank()` splits the sequence across CP ranks, making `hidden_states` a **view** (`._base is not None`)
2. `TransformerBlock.forward()`
([transformer_block.py#L752](https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/transformer/transformer_block.py#L752)) calls
`make_viewless_tensor(hidden_states, requires_grad=True, keep_graph=True)`
3. Since `hidden_states` is a view, `MakeViewlessTensor.apply()` clones it, creating a **new tensor with independent storage**
4. This cloned tensor flows through downstream TE modules (LayerNorm, Attention via `OperationFuser`), which save it via `save_for_backward()`
5. After backward completes, PyTorch's **C++ autograd engine retains a reference** to this tensor. The reference is invisible to Python's `gc`
(`gc.get_referrers()` returns 0 Python referrers, `gc.collect()` frees 0 bytes)
6. The tensor's GPU storage is **never freed**, leaking ~30-1200 MiB per micro-batch (varies with sequence length)
### Why C++ retains the reference
The cloned tensor enters TransformerEngine's `OperationFuser`, which sets `tensor._do_not_clear = True`
([fuser.py#L100](https://github.com/NVIDIA/TransformerEngine/blob/main/transformer_engine/pytorch/ops/fuser.py#L100)) to prevent premature clearing during
backward. This flag is **never removed** after backward. Combined with the C++ autograd engine holding onto the tensor, there is no path that releases
its GPU storage.
We have filed a related issue on TransformerEngine: https://github.com/NVIDIA/TransformerEngine/issues/2899
## Impact
- **Leak rate**: ~0.5 GiB per training step (16 micro-batches), accumulating linearly
- **OOM**: Training crashes within ~40 steps on 80GB GPUs
- **Leaked tensor characteristics**: shape `(T, 1, hidden_size)`, bf16, `requires_grad=True`, `grad_fn=MakeViewlessTensorBackward` (or `CloneBackward0`)
- **Affected callers**: All `make_viewless_tensor(..., keep_graph=True)` call sites where the input is a view — at least 15+ call sites across
`transformer_block.py`, `transformer_layer.py`, `fused_layer_norm.py`, `multi_token_prediction.py`, `mamba_block.py`, `pipeline_parallel/utils.py`, etc.
## Reproduction
### Minimal configuration
Any Megatron training with:
```
--context-parallel-size 2 # or any value > 1
--tensor-model-parallel-size 4
--pipeline-model-parallel-size 1
--micro-batch-size 1
--global-batch-size 16
```
### Observation
Monitor `torch.cuda.memory_allocated()` after each training step — it grows monotonically by ~0.5 GiB/step instead of staying stable.
### Diagnostic
Add the following after `backward_step()` in `schedules.py`:
```python
import gc
gc.collect()
count = 0
for obj in gc.get_objects():
if torch.is_tensor(obj) and obj.is_cuda and obj.dtype == torch.bfloat16:
if len(obj.shape) == 3 and obj.shape[2] == hidden_size and obj.requires_grad:
if obj.untyped_storage().size() > 0:
count += 1
print(f"Leaked hidden_states tensors: {count}") # grows by 1 per micro-batch
```
## Environment
| Component | Version |
|-----------|---------|
| Megatron Core | 0.18.0 |
| TransformerEngine | 2.13.0 |
| PyTorch | 2.10.0+cu128 |
| GPU | 8× A800-80GB |
## Workaround
Track cloned tensors in a global deque and release their storage after backward:
```python
# In megatron/core/utils.py
import collections
_leaked_viewless_tensors = collections.deque()
def _drain_leaked_viewless_tensors():
while _leaked_viewless_tensors:
t = _leaked_viewless_tensors.popleft()
if t.untyped_storage().size() > 0:
t.storage().resize_(0)
t.detach_()
del t
def make_viewless_tensor(inp, requires_grad, keep_graph):
if inp._base is None:
return inp
if keep_graph:
out = inp.clone().requires_grad_(requires_grad)
_leaked_viewless_tensors.append(out)
return out
else:
return _kernel_make_viewless_tensor(inp, requires_grad)
```
```python
# In megatron/core/pipeline_parallel/schedules.py, after each backward_step():
from megatron.core.utils import _drain_leaked_viewless_tensors
_drain_leaked_viewless_tensors()
```
This workaround has been validated for 40+ steps with `alloc(GiB)` completely stable (18.14 GiB, zero growth).
## Suggested Fix
Several options, in order of preference:
1. **Avoid cloning when possible**: When `hidden_states` is a view from CP splitting, `inp.contiguous()` may suffice to break the view relationship
without triggering C++ autograd retention. Needs testing.
2. **Integrate the drain mechanism**: Add `_drain_leaked_viewless_tensors()` as a first-class API in `megatron.core.utils` and call it from all scheduling
loops (`forward_backward_no_pipelining`, `forward_backward_pipelining_without_interleaving`, `hybrid_context_parallel_forward_backward`, etc.) after each
`backward_step()`.
3. **Coordinate with TE**: Request TransformerEngine to clean up `_do_not_clear` flags on input tensors after backward completes in `OperationFuser`,
which would allow the standard cleanup path to release these tensors.
---
Contributor guide
Assessment
This issue has not been assessed yet.