Fix vp check and add checkpoint configuration support #2020
+74
−3
Conversation
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
- Fix virtual_pipeline_model_parallel_size check from '!= -1' to 'is not None' - Update vp default value from -1 to None for consistency - Add checkpoint save/load configuration handling - Enable save_dir, load_dir, save_interval, most_recent_k options - Wire up checkpoint arguments to recipe configuration Manually applied changes from commits: - 5e78174 (vp check fix) - 4c9881f (vp default value) - 853348f/b25d66d3 (checkpoint functionality)
rsalagame-nvidia
requested review from
a team,
erhoo82 and
malay-nagda
as code owners
added 7 commits
January 21, 2026 17:56
- Add nsys_trace=["cuda"] for CUDA tracing - Add nsys_extra_args with detailed profiling options: - Force overwrite existing profiles - Capture range using CUDA profiler API - Enable CUDA graph tracing at node level - Disable CUDA event trace to reduce overhead - Include NCCL domain in NVTX traces Partially adapted from commit a514e44
- Profile all ranks by default when profiling_ranks is None - Maintains user override capability via --profiling_ranks argument - Matches behavior from original commit a514e44
- Use default directory /nemo_run/code/nemo_experiments/default/checkpoints - Only when save_interval is set but save_dir is not specified - Matches behavior from original checkpoint implementation
- Checkpoint save enabled by default to /nemo_run/code/nemo_experiments/default/checkpoints - Save interval defaults to train_iters (checkpoint at end of training) - Users can override via --save_dir and --save_interval arguments - Removes checkpoint.save = None from common perf overrides
- Remove conditional logic for profiling_ranks - Always profile all GPUs: profile_ranks=list(range(num_gpus)) - Simplifies configuration as all ranks should always be profiled
- When --save_interval is provided, enable checkpointing with default or specified save_dir - Default save_dir to /nemo_run/code/nemo_experiments/default/checkpoints when not specified - Default save_interval to train_iters when only save_dir is provided - Removed manual train_iters=1 setting for checkpoint loading (handled by --max_steps from bash script) - Aligns with ENABLE_CHECKPOINT behavior in bash wrapper script
rsalagame-nvidia
added
performance
performance/release
malay-nagda
reviewed
Jan 22, 2026
| recipe.checkpoint.save = args.save_dir | ||
| logger.info(f"Checkpoint save directory set to: {args.save_dir}") | ||
| else: | ||
| recipe.checkpoint.save = "/nemo_run/code/nemo_experiments/default/checkpoints" |
Contributor
There was a problem hiding this comment.
Suggested change
| recipe.checkpoint.save = "/nemo_run/code/nemo_experiments/default/checkpoints" | |
| recipe.checkpoint.save = "/nemo_run/checkpoints" |
Author
There was a problem hiding this comment.
@malay-nagda We want the default location of the checkpoints to be under the code/nemo_experiments/default folder. This is to keep consistent with other workloads (ex: nemotron4, etc...).
added 3 commits
January 23, 2026 20:27
- Change default from None to -1 in argument_parser.py - Change condition from hasattr() to != -1 in overrides.py - Allows users to explicitly set vp to None when needed - Distinguishes between 'not provided' and 'explicitly None'
- Created _set_checkpoint_overrides() to handle all checkpoint save/load logic - Simplified set_user_overrides() by calling the new method - Improves code organization and reusability - Consistent with other helper methods like _set_megatron_fsdp_overrides()
- Add --nsys_trace CLI arg to customize traced events (defaults to nemo_run defaults) - Add --nsys_extra_args CLI arg to provide additional nsys arguments - Fix profile_ranks to use --profiling_ranks from CLI instead of hardcoding to all ranks - Update NsysPlugin to combine default extra args with user-provided args - Preserve default behavior: profiles rank 0 only when --profiling_ranks not specified This fixes the issue where --profiling_ranks was being ignored and all ranks were always profiled, and allows users to customize nsys trace and extra args.
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.
Manually applied changes from commits: