NVIDIA / NVIDIA/Megatron-LM

parallel_state.destroy_model_parallel leaks NCCL process groups

Open
#6,897 0 comments 0 reactions 1 assignee Claimed by @Connor-XY View on GitHub
Dominant language
Python
Stars
17.9k
Forks
4.5k
Avg merge
4d 3h
Merged PRs (30d)
272

Description

`parallel_state.destroy_model_parallel()` releases only the three Gloo groups. Every NCCL
group is set to `None` without `torch.distributed.destroy_process_group()`, so the
communicator stays alive in c10d's global `_world.pg_map` — dropping Megatron's reference
does not free it.

Each `initialize_model_parallel()` / `destroy_model_parallel()` cycle therefore leaks
~20 process groups.

## Repro

```python
# torchrun --nproc-per-node=8 leak.py
import torch.distributed as dist
from torch.distributed.distributed_c10d import _world
from tests.unit_tests.test_utilities import Utils

Utils.initialize_model_parallel(1, 1)
for cycle in range(1, 7):
if dist.get_rank() == 0:
print(f"cycle {cycle}: pg_map={len(_world.pg_map)}")
Utils.destroy_model_parallel()
Utils.initialize_model_parallel(1, 1)
```

```
cycle 1: pg_map=24
cycle 2: pg_map=44
cycle 3: pg_map=70
cycle 4: pg_map=89
cycle 5: pg_map=115
cycle 6: pg_map=134
```

A final `destroy_model_parallel()` drops it to 131 — the three Gloo groups.

## Impact

Only processes that re-initialize are affected, i.e. test suites; production initializes
once. It surfaces in the `tests/unit_tests/distributed/mfsdp_v2/**` CI bucket, where
`nccl_allocator.init()` sets `NCCL_NVLS_ENABLE=1`
([nccl_allocator.py#L154](../blob/main/megatron/core/nccl_allocator.py#L154)) for the rest
of the process. Every leaked communicator then reserves NVLink SHARP multicast memory,
and later tests fail to bind it:

```
Failed to bind NVLink SHARP (NVLS) Multicast memory of size 2097152 : CUDA error 2 'out of memory'
```

Reproduced on 8×H100 running that directory: 15 failed / 102 passed, the failures
clustered in `test_overlap.py`, `test_memory.py` and `test_symmetric_memory.py` —
i.e. whatever runs after the pool is exhausted, not tests that are themselves wrong.

Code: [`parallel_state.destroy_model_parallel`](../blob/main/megatron/core/parallel_state.py)
destroys `_DATA_PARALLEL_GROUP_GLOO`, `_DATA_PARALLEL_GROUP_WITH_CP_GLOO` and
`_EXPERT_DATA_PARALLEL_GROUP_GLOO` only.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.