From 8a0e1d3401c7734fcb94f9933da5f6060a0926c1 Mon Sep 17 00:00:00 2001 From: bsarwari Date: Wed, 28 Oct 2026 17:21:53 -0400 Subject: [PATCH 1/2] Destroy the model-parallel process groups between pipeline fwd/bwd test iterations Six NcclPipelineParallelForwardBackwardTest subtests fail on ROCm 10.3.0-20261027 on a 2x 16 GB gfx1200 host with "HIP failure: 'out of memory'" / "NCCL split failed ... Failed to CUDA calloc 536870912 bytes". _forward_backward_test_impl() runs initialize_model_parallel() / destroy_model_parallel() once per dtype and deallocate option (8 times) in one process. destroy_model_parallel() only sets the group globals to None, so every iteration left its RCCL communicators alive. Each holds about 512 MiB of device memory, and rank 0 failed on its 32nd communicator split. At the end of each iteration, call torch.distributed.destroy_process_group() on every group this rank holds, then destroy_model_parallel() as before. The change is limited to this test; destroy_model_parallel() is unchanged. Co-authored-by: Cursor --- .../test_pipeline_parallel_fwd_bwd.py | 25 ++++++++++++++++++- 1 file changed, 24 insertions(+), 1 deletion(-) diff --git a/tests/L0/run_transformer/test_pipeline_parallel_fwd_bwd.py b/tests/L0/run_transformer/test_pipeline_parallel_fwd_bwd.py index 1b044b845..e0f8aba9f 100644 --- a/tests/L0/run_transformer/test_pipeline_parallel_fwd_bwd.py +++ b/tests/L0/run_transformer/test_pipeline_parallel_fwd_bwd.py @@ -92,6 +92,29 @@ def _get_default_world_sizes_model_parallel_world_size(pipeline_model_parallel_w return tensor_model_parallel_world_size, data_parallel_size, pipeline_model_parallel_world_size +def _destroy_model_parallel() -> None: + # parallel_state.destroy_model_parallel() only drops its references to the groups. + # Each communicator keeps its device buffers until it is destroyed, so running the + # dtype / deallocate loop below in one process runs a 16 GB GPU out of memory. + groups = [] + for name in ( + "_MODEL_PARALLEL_GROUP", + "_TENSOR_MODEL_PARALLEL_GROUP", + "_PIPELINE_MODEL_PARALLEL_GROUP", + "_DATA_PARALLEL_GROUP", + "_EMBEDDING_GROUP", + "_POSITION_EMBEDDING_GROUP", + "_ENCODER_RELATIVE_POSITION_EMBEDDING_GROUP", + "_DECODER_RELATIVE_POSITION_EMBEDDING_GROUP", + ): + group = getattr(parallel_state, name) + if group is not None and all(group is not g for g in groups): + groups.append(group) + for group in groups: + torch.distributed.destroy_process_group(group) + parallel_state.destroy_model_parallel() + + class PipelineParallelForwardBackwardTestBase: GLOBAL_BATCH_SIZE = 16 @@ -228,7 +251,7 @@ def _forward_backward_test_impl( optimizer.step() optimizer.zero_grad(set_to_none=True) - parallel_state.destroy_model_parallel() + _destroy_model_parallel() def test_learning_no_pipelining(self): self._forward_backward_test_impl(False, forward_backward_no_pipelining, 1, None) From 6c85b5a9bb013a64ed750725679dc1afe3188cd0 Mon Sep 17 00:00:00 2001 From: bsarwari Date: Wed, 28 Oct 2026 18:48:24 -0400 Subject: [PATCH 2/2] Destroy the model-parallel process groups in destroy_model_parallel Apex_UT tests that call initialize_model_parallel()/destroy_model_parallel() repeatedly in one process run out of device memory on Radeon hosts on ROCm 10.3.0-20261027 ("HIP failure: 'out of memory'" / "NCCL split failed ... Failed to CUDA calloc 536870912 bytes"): six test_pipeline_parallel_fwd_bwd subtests on 2x gfx1200, and two of those plus test_cross_entropy on 4x gfx1201. destroy_model_parallel() only set the group globals to None, so each iteration's RCCL communicators stayed alive, about 512 MiB each, until the process exited. Call torch.distributed.destroy_process_group() on each group this rank holds before clearing the globals. Groups the rank is not a member of are skipped, and a group object is destroyed only once. This covers every caller, so the pipeline-test-only teardown from the previous commit is removed. Co-authored-by: Cursor --- apex/transformer/parallel_state.py | 25 ++++++++++++++++++- .../test_pipeline_parallel_fwd_bwd.py | 25 +------------------ 2 files changed, 25 insertions(+), 25 deletions(-) diff --git a/apex/transformer/parallel_state.py b/apex/transformer/parallel_state.py index ed834fed6..0214d7bd3 100644 --- a/apex/transformer/parallel_state.py +++ b/apex/transformer/parallel_state.py @@ -640,11 +640,34 @@ def get_data_parallel_rank(): return torch.distributed.get_rank(group=get_data_parallel_group()) +def _destroy_model_parallel_process_groups(): + # Dropping the reference is not enough: each communicator keeps its device + # buffers until it is destroyed, so repeated initialize/destroy cycles run out of memory. + destroyed = [] + for group in ( + _MODEL_PARALLEL_GROUP, + _TENSOR_MODEL_PARALLEL_GROUP, + _PIPELINE_MODEL_PARALLEL_GROUP, + _DATA_PARALLEL_GROUP, + _EMBEDDING_GROUP, + _POSITION_EMBEDDING_GROUP, + _ENCODER_RELATIVE_POSITION_EMBEDDING_GROUP, + _DECODER_RELATIVE_POSITION_EMBEDDING_GROUP, + ): + if group is None or group is torch.distributed.GroupMember.NON_GROUP_MEMBER: + continue + if any(group is g for g in destroyed): + continue + torch.distributed.destroy_process_group(group) + destroyed.append(group) + + # note (mkozuki): `destroy_model_parallel` voids more global variables than Megatron-LM. # Otherwise pipeline parallel forward_backward functions test hangs possibly because # the clean-up of the original is NOT enough. def destroy_model_parallel(): - """Set the groups to none.""" + """Destroy the process groups and set them to None.""" + _destroy_model_parallel_process_groups() global _MODEL_PARALLEL_GROUP _MODEL_PARALLEL_GROUP = None global _TENSOR_MODEL_PARALLEL_GROUP diff --git a/tests/L0/run_transformer/test_pipeline_parallel_fwd_bwd.py b/tests/L0/run_transformer/test_pipeline_parallel_fwd_bwd.py index e0f8aba9f..1b044b845 100644 --- a/tests/L0/run_transformer/test_pipeline_parallel_fwd_bwd.py +++ b/tests/L0/run_transformer/test_pipeline_parallel_fwd_bwd.py @@ -92,29 +92,6 @@ def _get_default_world_sizes_model_parallel_world_size(pipeline_model_parallel_w return tensor_model_parallel_world_size, data_parallel_size, pipeline_model_parallel_world_size -def _destroy_model_parallel() -> None: - # parallel_state.destroy_model_parallel() only drops its references to the groups. - # Each communicator keeps its device buffers until it is destroyed, so running the - # dtype / deallocate loop below in one process runs a 16 GB GPU out of memory. - groups = [] - for name in ( - "_MODEL_PARALLEL_GROUP", - "_TENSOR_MODEL_PARALLEL_GROUP", - "_PIPELINE_MODEL_PARALLEL_GROUP", - "_DATA_PARALLEL_GROUP", - "_EMBEDDING_GROUP", - "_POSITION_EMBEDDING_GROUP", - "_ENCODER_RELATIVE_POSITION_EMBEDDING_GROUP", - "_DECODER_RELATIVE_POSITION_EMBEDDING_GROUP", - ): - group = getattr(parallel_state, name) - if group is not None and all(group is not g for g in groups): - groups.append(group) - for group in groups: - torch.distributed.destroy_process_group(group) - parallel_state.destroy_model_parallel() - - class PipelineParallelForwardBackwardTestBase: GLOBAL_BATCH_SIZE = 16 @@ -251,7 +228,7 @@ def _forward_backward_test_impl( optimizer.step() optimizer.zero_grad(set_to_none=True) - _destroy_model_parallel() + parallel_state.destroy_model_parallel() def test_learning_no_pipelining(self): self._forward_backward_test_impl(False, forward_backward_no_pipelining, 1, None)