Feat/jax rccl analyser #390

olehtika · 2025-10-16T08:59:45Z

Pull Request Template

Note to AMDers:
This is a public repository. Please do not upload any confidential or customer data. Make sure all such data has been anonymized or removed before making this PR. If you need to attach any private files or links, please insert a Internal OneDrive Link or a Jira Ticket Link instead.

…lyser script.

TraceLens/JAXRcclAnalyser/notes.txt

TraceLens/JAXRcclAnalyser/jax_rccl_analyser.ipynb

jasainio

I'm not sure if JAXRcclAnalyser is the proper place for this code. Why not put this be under NcclAnalyser?
Remove unnecessary files that are not meant to be included in the PR.
rccl_complete_analysis.py should be refactored and only the essential parts should be kept.
Time differences between ranks in different nodes might not be meaningful. These should be investigated further. Probably out of the scope of this work.

TraceLens/JAXRcclAnalyser/bandwidth_help.txt

jasainio · 2025-10-16T09:03:27Z

TraceLens/JAXRcclAnalyser/Jax_trace2tree_example.ipynb

Is this folder the proper place for Jupyter notebooks?

TraceLens/JAXRcclAnalyser/bandwidth_analysis.ipynb

jasainio · 2025-10-16T09:03:48Z

TraceLens/JAXRcclAnalyser/jax_rccl_analyser.ipynb

Is this folder the proper place for Jupyter notebooks?

jasainio · 2025-10-16T09:05:03Z

TraceLens/JAXRcclAnalyser/jax_rccl_analyser.py

Clean this out if not needed.

jasainio · 2025-10-16T09:05:53Z

TraceLens/JAXRcclAnalyser/notes.txt

What is the purpose of this file?

jasainio · 2025-10-16T09:06:18Z

TraceLens/JAXRcclAnalyser/test_1.csv

What is the purpose of this file?

jasainio · 2025-10-16T09:06:26Z

TraceLens/JAXRcclAnalyser/test_2.csv

What is the purpose of this file?

jasainio · 2025-10-16T09:09:52Z

TraceLens/JAXRcclAnalyser/util/rccl_complete_analysis.py

Could you try to assimilate (rewrite/clean/refactor) the essential parts from this file into the current feature (multinode RCCL analysis) and place those in locations that make sense in TraceLense?

jasainio · 2025-10-16T09:16:56Z

TraceLens/JAXRcclAnalyser/jax_rccl_analyser.py

It's probably not correct to compare timestamps between different nodes / processes as they might not be synced. I think this requires further investigations. The safe option would be to remove any calculations / metrics that depend on comparing timestamps on different nodes / processes and add these in a future PR if they make sense.

TraceLens/JAXRcclAnalyser/Jax_trace2tree_example.ipynb

…transpose perm.

TraceLens/JAXRcclAnalyser/jax_rccl_analyser.py

…nization.

examples/jax_nccl_analyser_example.ipynb

anant added 10 commits September 18, 2025 11:53

Initial files for jax-nccl-analyser.

c60c3af

Added replica groups field to hlo-ops dict.

abd5991

Added trace loading function to RCCLAnalyser class.

8bb2429

Added the code needed to extract bytes involved in collective ops.

f58359f

Fixed module not found errors.

69288b1

Corrected typo in JAXRCCLAnalyser folder and added code in JAXRCCLAna…

3a4c3f1

…lyser script.

Updated code related to df analysis.

db77a68

Organizing code dir.

ddf0f4e

Removed a workspace apecific file.

2d0a305

Added code for bandwidth analysis.

afc0768

olehtika marked this pull request as draft October 16, 2025 09:00

olehtika assigned olehtika and amital-amd and unassigned olehtika Oct 16, 2025

olehtika commented Oct 16, 2025

View reviewed changes

TraceLens/JAXRcclAnalyser/notes.txt Outdated Show resolved Hide resolved

olehtika commented Oct 16, 2025

View reviewed changes

TraceLens/JAXRcclAnalyser/jax_rccl_analyser.ipynb Outdated Show resolved Hide resolved

Removed a temporary file.

a290a79

jasainio requested changes Oct 16, 2025

View reviewed changes

olehtika commented Oct 16, 2025

View reviewed changes

TraceLens/JAXRcclAnalyser/Jax_trace2tree_example.ipynb Outdated Show resolved Hide resolved

anant added 3 commits October 16, 2025 11:26

Cleaned the working files from the dir.

de02cd5

Updated logic to extract replica group strings.

1fb13f6

Modified logic to consider replica group string cases which included …

b7c21a6

…transpose perm.

olehtika commented Oct 22, 2025

View reviewed changes

TraceLens/JAXRcclAnalyser/jax_rccl_analyser.py Outdated Show resolved Hide resolved

anant added 2 commits October 22, 2025 11:05

Removed these files from the current location as part of files reorga…

bdf9283

…nization.

Updated code to add bandwidth analysis and added an example notebook.

279b427

amital-amd added the enhancement New feature or request label Oct 22, 2025

olehtika commented Oct 23, 2025

View reviewed changes

examples/jax_nccl_analyser_example.ipynb Outdated Show resolved Hide resolved

anant added 3 commits October 23, 2025 07:12

Added node to protobuf file mapping util file.

d38c892

Updated example usage of util script.

ea4c0c2

Fixed bugs related to replica-groups parsing, missing ops.

df4f292

Updated the import section to fix import errors.

6d2d100

Feat/jax rccl analyser #390

Are you sure you want to change the base?

Feat/jax rccl analyser #390

Uh oh!

Conversation

olehtika commented Oct 16, 2025

Pull Request Template

Uh oh!

Uh oh!

Uh oh!

jasainio left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants