Skip to content

feat: add NCCL flight recorder configuration support#3806

Open
sbak5 wants to merge 3 commits intoNVIDIA:mainfrom
sbak5:sbak/fr_enable
Open

feat: add NCCL flight recorder configuration support#3806
sbak5 wants to merge 3 commits intoNVIDIA:mainfrom
sbak5:sbak/fr_enable

Conversation

@sbak5
Copy link
Contributor

@sbak5 sbak5 commented Mar 11, 2026

Add flight recorder config fields to DistributedInitConfig and set the corresponding env vars (TORCH_FR_DUMP_TEMP_FILE, TORCH_NCCL_TRACE_BUFFER_SIZE, etc.) before distributed init when flight_recorder_dump_path is specified.

What does this PR do ?

⚠️ For major changes (either in lines of code or in its impact), please make sure to first share a design doc with the team. If you're unsure what's the best way to do so, contact the @mcore-oncall.

Contribution process

Pre-checks

  • I have added relevant unit tests
  • I have added relevant functional tests
  • I have added proper typing to my code Typing guidelines
  • I have added relevant documentation
  • I have run the autoformatter.sh on my PR

Code review

Feel free to message or comment the @mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!

All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.

Step 1: Mark PR as "Ready for Review"

  1. When your PR is ready, click Ready for Review.
  2. An oncall reviewer is auto-assigned and expert reviewers are notified based on your changes.
    • Some PRs may jump straight to step 2. This is determined by .github/CODEOWNERS.

⚠️ Only mark as ready once merge-conflicts are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

Step 2: Final Review

For PRs that change megatron/core, once all expert reviewers have approved, the Final Review label is applied automatically and final reviewers are assigned.

For PRs outside megatron/core, this step is skipped.

Step 3: Approved

Once all required reviewers have approved, the Approved label is applied automatically.

Merge

Any member of mcore-engineers will be able to merge your PR.

For MRs into `dev` branch The proposed review process for `dev` branch is under active discussion.

MRs are mergable after one approval by either eharper@nvidia.com or zijiey@nvidia.com.

Add flight recorder config fields to DistributedInitConfig and set the
corresponding env vars (TORCH_FR_DUMP_TEMP_FILE, TORCH_NCCL_TRACE_BUFFER_SIZE,
etc.) before distributed init when flight_recorder_dump_path is specified.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
@copy-pr-bot
Copy link

copy-pr-bot bot commented Mar 11, 2026

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@svcnvidia-nemo-ci svcnvidia-nemo-ci marked this pull request as draft March 11, 2026 17:26
@github-actions
Copy link
Contributor

This PR has been automatically converted to draft because all PRs must start as drafts.

When you are ready for review, click Ready for Review to begin the review process. This will:

  1. Add the oncall reviewer (optional reviewer)
  2. Add required review teams based on your changes

See the contribution guide for more details.

@sbak5 sbak5 requested a review from maanug-nv March 11, 2026 17:27
@sbak5
Copy link
Contributor Author

sbak5 commented Mar 11, 2026

/claude review

@sbak5 sbak5 requested review from a team March 11, 2026 17:27
…log effective values

Use os.environ.setdefault semantics for optional flight recorder env vars
(TORCH_NCCL_TRACE_BUFFER_SIZE, TORCH_NCCL_DUMP_ON_TIMEOUT, etc.) so that
values already set in the environment are respected. Emit a warning via
warn_rank_0 for each var that is skipped, and print all effective flight
recorder env var values via print_rank_0 once at startup.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
@sbak5 sbak5 requested a review from a team March 11, 2026 17:34
@sbak5
Copy link
Contributor Author

sbak5 commented Mar 11, 2026

/claude review

Copy link
Contributor

@claude claude bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks good overall — clean implementation with nice setdefault semantics for the optional env vars.

One issue: TORCH_FR_DUMP_TEMP_FILE and TORCH_NCCL_DEBUG_INFO_TEMP_FILE are unconditionally overwritten (lines 342-343), unlike the other vars which warn when a pre-existing value is found. Left an inline comment with a suggestion to at least warn before overriding.

… environment

Apply setdefault semantics to TORCH_FR_DUMP_TEMP_FILE and
TORCH_NCCL_DEBUG_INFO_TEMP_FILE consistently with all other flight
recorder vars. The block now also fires when either path var is already
present in the environment (without the MLM arg), so the remaining
defaults (buffer size, dump-on-timeout, etc.) are applied in all cases.

Priority order: pre-existing environment variable > MLM argument.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
@sbak5
Copy link
Contributor Author

sbak5 commented Mar 11, 2026

/claude review

Copy link
Contributor

@claude claude bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@sbak5 sbak5 marked this pull request as ready for review March 11, 2026 18:13
@sbak5 sbak5 enabled auto-merge March 11, 2026 23:25
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants