Repository navigation
Conversation
…st 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 <cursoragent@cursor.com>
Author
|
Validation, commit All 6 |
2 tasks done
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 <cursoragent@cursor.com>
bilalsar
changed the title
Author
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
In short:
destroy_model_parallel()never frees the communication groups it created, so a test that builds them repeatedly in one process runs GPUs with limited VRAM out of memory.test_pipeline_parallel_fwd_bwd.NcclPipelineParallelForwardBackwardTestsubtests fail on ROCm 10.3.0-20261027 (torch2.15.0+rocm10.3.0a20261027) on Radeon hosts withHIP failure: 'out of memory'/NCCL split failed ... Failed to CUDA calloc 536870912 bytes. They passed on 10.0.0-20260729. All runs use apex11867be, therelease/1.15.0pin inrelated_commits. The MI325X and MI350X hosts on the same build do not hit it.ctr-navi4x-aj63-ws13test_inference_no_pipelining,test_inference_pipelining_without_interleaving,test_inference_async_pipelining_without_interleaving,test_learning_no_pipelining,test_learning_pipelining_without_interleaving,test_learning_async_pipelining_without_interleavingctr-navi4x-at48-ws008test_inference_no_pipelining,test_learning_no_pipeliningTechnical Details
In short:
destroy_model_parallel()now destroys the process groups before clearing its references to them.apex/transformer/parallel_state.py:destroy_model_parallel()only set the group globals toNone._forward_backward_test_impl()callsinitialize_model_parallel()/destroy_model_parallel()once per dtype and deallocate option (8 times) in one process, so every iteration left its RCCL communicators alive. Each holds about 512 MiB of device memory; onctr-navi4x-aj63-ws13rank 0 fails on its 32nd communicator split (31 x 512 MiB allocations logged withNCCL_DEBUG=INFO,ALLOC).The new
_destroy_model_parallel_process_groups()callstorch.distributed.destroy_process_group()once on each group this rank holds (the eight_*_GROUPglobals), skippingNoneandGroupMember.NON_GROUP_MEMBERand destroying each group object only once.destroy_model_parallel()calls it first, then clears the globals as before.Effect on other suites: this changes behaviour for every caller of
destroy_model_parallel(): the 13tests/L0/run_transformerfiles that call it,tests/L1/transformer/pipeline_parallel_fwd_bwd_ucc_async.py, and any user code. After the call, process-group handles obtained earlier fromparallel_stateare destroyed and must not be reused. In the full Apex_UT runs below, on 2 and 4 GPUs, every test that passed before still passes and no test hangs.Related: the
test_layersparity errors in the same Apex_UT runs are a separate root cause: #344.Test Result
Full Apex_UT reruns through the frameworks-qa-ci Functional Test Runner with the same image as the nightly (
rocm/frameworks-internal:rocm10.3.0a20261027_ubuntu26.04_py3.14_pytorch_2.15.0_31b8278), apex checked out at this PR's head6c85b5a.11867be(pin)6c85b5actr-navi4x-aj63-ws13jobctr-navi4x-at48-ws008jobtest_inference_no_pipelining,test_learning_no_pipeliningErrors left in the passing runs:
test_layerstest_all_gather_parity(both hosts) andtest_reduce_scatter_parity(at48), fixed by #344, andtest_cross_entropyon at48 (see Not covered). No new errors.apex-pipeline-fwd-bwd-fail-hip-oom.log (run 40688983363, aj63)
apex-destroy-model-parallel-pass-aj63.log (run 40782925437)
apex-destroy-model-parallel-pass-at48.log (run 40782938832)
How the PR commit was run: Apex_UT imports the apex installed in the docker image, not the checkout, so a library change does not reach the tests on its own. The jobs were dispatched from the frameworks-qa-ci branch
test/apex-ut-commit-override, which letsapex_uttake anAPEX_COMMIT_OVERRIDESHA and copies the one file named inAPEX_OVERLAY_FILE(hereapex/transformer/parallel_state.py) from the checkout over the installed package. The matching MD5 lines above show the overlay. Everything else is the same as the nightly job.Not covered:
test_cross_entropy.NcclVocabParallelCrossEntropyTest.test_cross_entropyonctr-navi4x-at48-ws008still fails withFailed to CUDA calloc 536870912 bytesinsideinitialize_model_parallel(), about 4 seconds after the test starts, with this change applied. Its cause is not established and this PR does not address it.related_commitspoints at an apex commit that contains it and the image is rebuilt from that pin (it pinsrelease/1.15.0, so this needs a backport there too).Submission Checklist