Skip to content

Update shard tensor ring attention example#1665

Merged
coreyjadams merged 3 commits into
2.1.0-rcfrom
fix-shard-tensor-examples
May 22, 2026
Merged

Update shard tensor ring attention example#1665
coreyjadams merged 3 commits into
2.1.0-rcfrom
fix-shard-tensor-examples

Conversation

@coreyjadams
Copy link
Copy Markdown
Collaborator

Now, the generation script will use the expected default results directory.

PhysicsNeMo Pull Request

Description

Checklist

Dependencies

Review Process

All PRs are reviewed by the PhysicsNeMo team before merging.

Depending on which files are changed, GitHub may automatically assign a maintainer for review.

We are also testing AI-based code review tools (e.g., Greptile), which may add automated comments with a confidence score.
This score reflects the AI’s assessment of merge readiness and is not a qualitative judgment of your work, nor is
it an indication that the PR will be accepted / rejected.

AI-generated feedback should be reviewed critically for usefulness.
You are not required to respond to every AI comment, but they are intended to help both authors and reviewers.
Please react to Greptile comments with 👍 or 👎 to provide feedback on their accuracy.

@greptile-apps
Copy link
Copy Markdown
Contributor

greptile-apps Bot commented May 21, 2026

Greptile Summary

This PR updates the ring-attention benchmark script to write JSON results by default (opt-out via --print-only) into an auto-named file under a results/ directory that is resolved relative to the script itself, replacing the previous opt-in --output_file argument. The README is updated to document the new flags.

  • --output_file--results_dir + auto-naming: filenames now encode topology, mode, dtype, and sequence length to match the regex expected by plot_scaling_results.py.
  • Default behaviour flipped: results are saved automatically unless --print-only is passed; all file I/O is correctly guarded inside if dm.rank == 0: to avoid duplicate writes in distributed runs.

Important Files Changed

Filename Overview
examples/minimal/ShardTensorExamples/6_ring_attention/benchmark_sharded_attention.py Replaces --output_file (opt-in, manual path) with --results_dir (opt-out, auto-named output); adds --print-only flag; all file I/O remains correctly guarded under dm.rank == 0.
examples/minimal/ShardTensorExamples/6_ring_attention/README.md Updates argument table and adds description for the new --results_dir and --print-only flags; minor documentation discrepancy about default path relativity.

Reviews (1): Last reviewed commit: "Update shard tensor ring attention to us..." | Re-trigger Greptile

Comment thread examples/minimal/ShardTensorExamples/6_ring_attention/README.md Outdated
@negin513 negin513 self-requested a review May 21, 2026 22:58
),
)
parser.add_argument(
"--print-only",
Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Tiny consistency nitpick: every other flag in this file is underscore-cased (--results_dir, ...) and --print-only is the only one that is - based. Would it make sense to use --print_only instead?

Copy link
Copy Markdown
Member

@negin513 negin513 left a comment

Choose a reason for hiding this comment

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

LGTM! Just a few very minor comments inline.

coreyjadams and others added 2 commits May 21, 2026 19:16
Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
…k_sharded_attention.py

Co-authored-by: Negin Sobhani <negin513@gmail.com>
@coreyjadams coreyjadams merged commit e601904 into 2.1.0-rc May 22, 2026
1 of 2 checks passed
@negin513 negin513 deleted the fix-shard-tensor-examples branch May 22, 2026 00:20
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants