Repository navigation
Conversation
torch.cuda.empty_cache() was added as a workaround for out-of-memory errors on some models. With async-shm checkpointing, some of the memory pressure is reduced. Remaining memory fragmentation and pressure can be reduced by properly configuring the torch allocator. Signed-off-by: Skand Hurkat <shurkat@nvidia.com>
- Update checkpointing to use nvrx with shm. - Provide PYTORCH_ALLOC_CONF environment variable to initiate garbage collection. Use larger blocks to avoid fragmentation due to splitting. Signed-off-by: Skand Hurkat <shurkat@nvidia.com>
Instead of removing the checkpoint-time free_overlap_buffers() and torch.cuda.empty_cache() calls outright, put both behind a new --ckpt-free-memory argument that defaults to off, so runs that rely on the extra GPU headroom for the async checkpoint worker can opt back in. The two Nemotron3 Super GB200 release configs set the flag and therefore no longer need the PYTORCH_ALLOC_CONF override, which is reverted to its previous commented-out state. Signed-off-by: Deepak Narayanan <dnarayanan@nvidia.com>
The previous name parsed as a noun phrase, so it read like it described an amount of free memory rather than an action. Release is unambiguously a verb and gpu says which memory is reclaimed. Signed-off-by: Deepak Narayanan <dnarayanan@nvidia.com>
deepakn94
changed the title
shurkat-nvidia
force-pushed
the
remove-torch-cuda-empty-cache
branch
from
August 10, 2026 21:01
debb98a to
c7830f1
Compare
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
3 participants
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.
Stacked on top of NVIDIA#5678 (targets
shurkat-nvidia:remove-torch-cuda-empty-cache).What does this PR do?
Rather than dropping the checkpoint-time memory reclaim outright, this hides both
free_overlap_buffers()andtorch.cuda.empty_cache()insave_checkpoint_and_time()behind a new--ckpt-release-gpu-memoryargument that defaults to off. Nobody pays the multi-second cost per save by default, but runs that depend on the extra GPU headroom for the async checkpoint worker (the reason NVIDIA#5170 was reverted by NVIDIA#5366) can opt back in without a code change.The two Nemotron3 Super GB200 release configs set
--ckpt-release-gpu-memory: true, so they keep today's behavior and thePYTORCH_ALLOC_CONFoverride is reverted to its previous commented-out state. The--async-strategy: nvrx/--async-ckpt-use-cpu-shm: trueadditions from NVIDIA#5678 are kept.Changes
megatron/training/config/training_config.py: newckpt_release_gpu_memoryfield onCheckpointConfig(defaultFalse), with a pointer to allocator tuning as the cheaper alternative.megatron/training/training.py: gate the two calls onargs.ckpt_release_gpu_memory.tests/functional_tests/test_cases/nemotron/nemotron3_super_release_gb200{,_sm}/model_config.yaml: set--ckpt-release-gpu-memory: true, revertPYTORCH_ALLOC_CONF.Pre-checks