Add Nemotron Nano pytorch profiler TOML and chakra tracing docs - #983
Add Nemotron Nano pytorch profiler TOML and chakra tracing docs#983saivishal1999 wants to merge 9 commits into
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds a GB200 Megatron-Bridge test configuration with PyTorch Kineto and optional Chakra profiling, plus documentation and workload index entries describing the required settings and generated outputs. ChangesProfiler tracing support
Estimated code review effort: 2 (Simple) | ~10 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 4✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 5
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In
`@conf/experimental/megatron_bridge/test/gb200/megatron_bridge_nemotron_nano_pytorch_profiler.toml`:
- Line 38: Update the test configuration’s hf_token entry to use the supported
user-provided secret mechanism so validation receives a non-empty value without
committing credentials. If this public-model test should not require
authentication, adjust the relevant validation contract instead, while
preserving the profiler test behavior.
In `@doc/workloads/chakra_tracing.rst`:
- Line 66: Update all prose references to pytorch_profiler_collect_chakra in the
background section and Chakra Trace instructions to use the fully qualified
profiling.pytorch_profiler_collect_chakra name, matching the TOML examples and
parameter table; leave unrelated configuration references unchanged.
In `@src/cloudai/workloads/megatron_bridge/megatron_bridge.py`:
- Around line 592-599: Escape experiment_name before interpolating it into the
glob pattern in the surrounding workload method. Apply glob.escape() only to the
experiment-name path component so wildcard characters in wandb_experiment_name
are treated literally while preserving the existing directory structure and
matching behavior.
- Around line 580-582: Update the method containing the
`_copy_profiler_traces(tr)` call to catch `OSError` and `shutil.Error` around
profiler trace collection, then continue returning
`JobStatusResult(is_successful=True)` when copying fails.
- Around line 603-608: Update the torch_profile copy logic around src, dst, and
shutil.copytree so retries refresh the existing destination instead of skipping
it. Remove the prior dst tree or otherwise overwrite it before copying the
newest src selected by os.path.getmtime, while preserving the current logging
behavior.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Enterprise
Run ID: fc8b5e91-d471-47f2-a618-ca0131d3ca7a
📒 Files selected for processing (4)
conf/experimental/megatron_bridge/test/gb200/megatron_bridge_nemotron_nano_pytorch_profiler.tomldoc/workloads/chakra_tracing.rstdoc/workloads/index.rstsrc/cloudai/workloads/megatron_bridge/megatron_bridge.py
| domain = "llm" | ||
| task = "pretrain" | ||
| compute_dtype = "fp8_mx" | ||
| hf_token = "" |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Provide a non-empty hf_token or update the validation contract.
Line 38 leaves the required token empty, so this test can fail validation before the profiler run starts. Supply it through the supported user-provided mechanism without committing a real secret, or explicitly adjust validation for this public-model scenario.
Based on learnings: Megatron-Bridge test TOMLs under conf/experimental/megatron_bridge/test/ require a non-empty hf_token.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In
`@conf/experimental/megatron_bridge/test/gb200/megatron_bridge_nemotron_nano_pytorch_profiler.toml`
at line 38, Update the test configuration’s hf_token entry to use the supported
user-provided secret mechanism so validation receives a non-empty value without
committing credentials. If this public-model test should not require
authentication, adjust the relevant validation contract instead, while
preserving the profiler test behavior.
Source: Learnings
| # Take the most recently modified in case of retries | ||
| src = max(matches, key=os.path.getmtime) | ||
| dst = tr.output_path / "torch_profile" | ||
| if not dst.exists(): | ||
| shutil.copytree(src, dst) | ||
| logging.info(f"Copied pytorch profiler traces from {src} to {dst}") |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Map the target file and inspect the relevant section with line numbers.
sed -n '540,650p' src/cloudai/workloads/megatron_bridge/megatron_bridge.py | cat -n
# Find other references to torch_profile and output_path in the repo.
rg -n --hidden --glob '!**/.git/**' 'torch_profile|output_path' src/cloudai/workloads/megatron_bridge src | sed -n '1,200p'Repository: NVIDIA/cloudai
Length of output: 30385
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Show the surrounding function and any retry-related control flow in the target file.
rg -n -A 8 -B 8 'torch_profile|retry|attempt|output_path|copytree' src/cloudai/workloads/megatron_bridge/megatron_bridge.py | sed -n '1,240p'
# If the file is not too large, inspect nearby definitions for tr/output path creation.
python3 - <<'PY'
from pathlib import Path
p = Path('src/cloudai/workloads/megatron_bridge/megatron_bridge.py')
lines = p.read_text().splitlines()
for i, line in enumerate(lines, 1):
if 'torch_profile' in line or 'output_path' in line or 'copytree' in line:
start = max(1, i-12)
end = min(len(lines), i+12)
print(f"\n--- lines {start}-{end} around {i} ---")
for j in range(start, end+1):
print(f"{j:4}: {lines[j-1]}")
PYRepository: NVIDIA/cloudai
Length of output: 14177
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the path-generation logic for TestRun output directories and any retry handling.
sed -n '180,240p' src/cloudai/_core/base_runner.py | cat -n
sed -n '1,260p' src/cloudai/systems/slurm/single_sbatch_runner.py | cat -n | sed -n '1,260p'
sed -n '150,230p' src/cloudai/configurator/cloudai_gym.py | cat -nRepository: NVIDIA/cloudai
Length of output: 18909
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the TestRun fields and where current_iteration/step are mutated.
sed -n '1,220p' src/cloudai/_core/test_scenario.py | cat -n
sed -n '1,260p' src/cloudai/_core/base_runner.py | cat -n | sed -n '1,180p'
rg -n -A 4 -B 4 'current_iteration|step\s*=|retry|resubmit|iteration' src/cloudai | sed -n '1,240p'Repository: NVIDIA/cloudai
Length of output: 34508
Refresh torch_profile on retries. tr.output_path is reused for the same iteration/step, so the if not dst.exists() guard leaves any prior torch_profile tree in place. A retried run can therefore publish stale or partial traces; remove the destination before copying or overwrite it in place.
🧰 Tools
🪛 Ruff (0.16.0)
[warning] 608-608: Logging statement uses f-string
(G004)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/cloudai/workloads/megatron_bridge/megatron_bridge.py` around lines 603 -
608, Update the torch_profile copy logic around src, dst, and shutil.copytree so
retries refresh the existing destination instead of skipping it. Remove the
prior dst tree or otherwise overwrite it before copying the newest src selected
by os.path.getmtime, while preserving the current logging behavior.
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
src/cloudai/workloads/megatron_bridge/megatron_bridge.py (2)
603-604: 🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick winEscape and validate
wandb_experiment_namebefore globbing.The value is interpolated into the pattern twice, so wildcard characters or path components can match unintended experiment directories and copy the wrong traces. Enforce a path-safe experiment-name contract and apply
glob.escape()to the literal name before constructing the pattern.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/cloudai/workloads/megatron_bridge/megatron_bridge.py` around lines 603 - 604, Update the experiment-directory discovery logic around wandb_experiment_name and base_pattern to validate that the experiment name is path-safe, rejecting wildcard characters and path components. Escape the validated literal with glob.escape() before interpolating it into the pattern, including both occurrences, while preserving the existing directory filtering.
615-617: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winRefresh existing trace destinations on retries.
if not dst.exists()preserves stale or partially copiedtorch_profile/chakratrees, causing a successful retry to publish traces from an earlier attempt. Remove or atomically replace the destination before copying the newest source directory.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/cloudai/workloads/megatron_bridge/megatron_bridge.py` around lines 615 - 617, Update the trace-copy logic around dst in the Megatron bridge to refresh existing destinations on every retry: remove or atomically replace dst before copying src with shutil.copytree. Ensure stale or partially copied torch_profile/chakra trees cannot remain, while preserving the newest source directory as the published result.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@src/cloudai/workloads/megatron_bridge/megatron_bridge.py`:
- Around line 603-604: Update the experiment-directory discovery logic around
wandb_experiment_name and base_pattern to validate that the experiment name is
path-safe, rejecting wildcard characters and path components. Escape the
validated literal with glob.escape() before interpolating it into the pattern,
including both occurrences, while preserving the existing directory filtering.
- Around line 615-617: Update the trace-copy logic around dst in the Megatron
bridge to refresh existing destinations on every retry: remove or atomically
replace dst before copying src with shutil.copytree. Ensure stale or partially
copied torch_profile/chakra trees cannot remain, while preserving the newest
source directory as the published result.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Enterprise
Run ID: d3408d60-1d4b-48d6-97f1-f4c360281374
📒 Files selected for processing (1)
src/cloudai/workloads/megatron_bridge/megatron_bridge.py
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@doc/workloads/chakra_tracing.rst`:
- Around line 34-43: Update the profiler output examples in the Chakra tracing
documentation to reflect that the revised implementation stores profiles under
the nested NemoRun experiment directory rather than directly under
<cloudai_output>. Remove or revise the outdated torch_profile/ and chakra/
direct-path claims, and document the pytorch_profile/ to tensorboard/ rename
consistently with the implementation.
In `@src/cloudai/workloads/megatron_bridge/megatron_bridge.py`:
- Around line 604-608: Update the task directory loop to wrap each
os.rename(src, dst) operation in its own try/except, logging the failure for
that specific task_dir while continuing to process remaining directories.
Preserve the existing directory checks and ensure later valid pytorch_profile
directories are still renamed.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Enterprise
Run ID: 8c15c1dd-9bf2-4476-af38-4817e1bec6df
📒 Files selected for processing (2)
doc/workloads/chakra_tracing.rstsrc/cloudai/workloads/megatron_bridge/megatron_bridge.py
| for task_dir in task_dirs: | ||
| src = os.path.join(task_dir, "pytorch_profile") | ||
| dst = os.path.join(task_dir, "tensorboard") | ||
| if os.path.isdir(src) and not os.path.exists(dst): | ||
| os.rename(src, dst) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Locate the relevant file and inspect the surrounding implementation.
git ls-files src/cloudai/workloads/megatron_bridge/megatron_bridge.py
wc -l src/cloudai/workloads/megatron_bridge/megatron_bridge.py
sed -n '540,680p' src/cloudai/workloads/megatron_bridge/megatron_bridge.py
# Find the success flag / error handling referenced in the comment.
rg -n "was_run_successful|task_dirs|pytorch_profile|tensorboard|os\.rename|rename\(" src/cloudai/workloads/megatron_bridge/megatron_bridge.pyRepository: NVIDIA/cloudai
Length of output: 5035
Handle rename failures per task directory
Wrap each os.rename() in a per-directory try/except; the current outer handler logs the first failure and still returns success, so later valid pytorch_profile directories are left untouched.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/cloudai/workloads/megatron_bridge/megatron_bridge.py` around lines 604 -
608, Update the task directory loop to wrap each os.rename(src, dst) operation
in its own try/except, logging the failure for that specific task_dir while
continuing to process remaining directories. Preserve the existing directory
checks and ensure later valid pytorch_profile directories are still renamed.
| ":doc:`ai_dynamo`", "✅", "✅", "❌", "❌" | ||
| ":doc:`bash_cmd`", "✅", "❌", "❌", "❌" | ||
| ":doc:`chakra_replay`", "✅", "❌", "❌", "❌" | ||
| ":doc:`chakra_tracing`", "✅", "❌", "❌", "❌" |
There was a problem hiding this comment.
this is not a separate workload but rather guidance on MBridge = should not be in this list as a separate item
There was a problem hiding this comment.
this is not a separate workload but rather guidance on MBridge = it should belong to the MBridge docus
There was a problem hiding this comment.
@podkidyshev True that we are adding a capability in MBridge. But would be good to call out how to collect tracing. It might be expanded to others in future. So if Workloads is not the right place, then lets add it in its own pillar.
There was a problem hiding this comment.
@podkidyshev @srivatsankrishnan Moved the docs to doc/workloads/megatron_bridge.rst
| return JobStatusResult(is_successful=False, error_message="\n".join(log_data.splitlines()[-40:])) | ||
|
|
||
| try: | ||
| self._rename_tensorboard_dir(tr) |
There was a problem hiding this comment.
this is excessive I believe as we already parse the tensorboard dir via src/cloudai/report_generator/training/report_generation_strategy.py. It prepares structured training_report.json
if you really want to keep the original folder in some predefined space:
- you don't need a constant wandb_experiment_name. simply iterate recursively
experimentsfolder, there will be always one experiment folder there - don't move, use copy
however I'd be really happy if you could utilize training_report.json
There was a problem hiding this comment.
also the right place for such operation is a custom reporter, not was_run_successful
There was a problem hiding this comment.
Reverted the renaming logic
- Add conf/experimental/megatron_bridge/test/gb200/megatron_bridge_nemotron_nano_pytorch_profiler.toml: Nemotron 3 Nano 3B (2 nodes, 8 GPUs) with PyTorch kineto profiling via [extra_cmd_args] Hydra overrides at steps 45-50. Validated end-to-end on Lyris GB200: profiler fires, rank-0.json.gz and memory_profile.pickle produced. - Add doc/workloads/chakra_tracing.rst: Documents how to enable Kineto and Chakra tracing in MegatronBridge, how to disable nsys (incompatible with Kineto), and full parameter reference. - Modify megatron_bridge.py: Add _copy_profiler_traces() to was_run_successful(). NemoRun places traces at experiments/<name>/<run_id>/<task>/torch_profile/ internally; this method copies them to <cloudai_output>/torch_profile/ after job completion so MegatronBridge and MegatronRun produce traces at a consistent location.
- _copy_profiler_traces now copies both torch_profile/ and chakra/ from NemoRun's nested experiment dir to <cloudai_output>/. chakra/ is only copied if present (produced when pytorch_profiler_collect_chakra=true). - Update chakra_tracing.rst to clarify that kineto trace is always produced and pytorch_profiler_collect_chakra adds a separate Chakra ExecutionTrace.
Copy failures (e.g. disk full, permissions) should not mark a successful training job as failed.
Per reviewer feedback: copying traces to the output root risks overwriting on retries and goes against NemoRun's directory structure. Instead, rename pytorch_profile/ to tensorboard/ within the experiment directory to match MegatronRun's naming convention (pytorch_profile/ contains TensorBoard events, not profiling traces).
…; drop expensive profiling flags from example TOML
f7ec65d to
13ad2c1
Compare
3 participants
Add conf/experimental/megatron_bridge/test/gb200/megatron_bridge_nemotron_nano_pytorch_profiler.toml: Nemotron 3 Nano 3B (2 nodes, 8 GPUs) with PyTorch kineto profiling via [extra_cmd_args] Hydra overrides at steps 45-50. Validated end-to-end on Lyris GB200: profiler fires, rank-0.json.gz and memory_profile.pickle produced.
Add doc/workloads/chakra_tracing.rst: Documents how to enable Kineto and Chakra tracing in MegatronBridge, how to disable nsys (incompatible with Kineto), and full parameter reference.
Modify megatron_bridge.py: Add _copy_profiler_traces() to was_run_successful(). NemoRun places traces at experiments//<run_id>//torch_profile/ internally; this method copies them to <cloudai_output>/torch_profile/ after job completion so MegatronBridge and MegatronRun produce traces at a consistent location.
Test Plan
Workload: MegatronBridge — Nemotron 3 Nano 3B (2 nodes, 8 GPUs, TP=1, EP=8)
Test TOML (added in this PR):
conf/experimental/megatron_bridge/test/gb200/megatron_bridge_nemotron_nano_pytorch_profiler.tomlRun Command