graph_trainer: auto-size GPU_MAX_HW_QUEUES on ROCm by rishisinhanj · Pull Request #4032 · pytorch/torchtitan

rishisinhanj · 2026-07-30T19:56:51Z

Summary

On ROCm, HIP maps logical streams to hardware queues via stream_id % Q. When Q is too small, independent streams alias onto the same queue and lose overlap.

Adds hw_queues.maybe_set_gpu_max_hw_queues(): estimates independent GPU stream count from the graph_trainer config, sets GPU_MAX_HW_QUEUES to the next power of two, and runs in GraphTrainer.__init__ before HIP init.

ROCm-only; no-op on CUDA
Respects an explicit GPU_MAX_HW_QUEUES override

Stream model

lanes = compute + all_reduce
      + (dense AG + RS          if fsdp_ag_rs_overlap else fsdp_comm)
      + expert_comm             if MoE and ep>1
      + tp_all_gather           if tp>1
      + cp_all_gather           if cp>1
      + cudagraph_capture       if cudagraph and not fsdp_ag_rs_overlap
Q = next_pow2(lanes)

EP=2 DeepSeek-V3 with FSDP AG/RS overlap → 5 lanes → Q=8.

Results

EP=2 2-node DeepSeek-V3 (MI355X, gfx950, ROCm 7.2). Throughput (tok/s/GPU):

torchtitan main	q4	q8 (auto)	q16
`fec3e19` (2026-07-14)	9002	9674 (+7.5%)	8357 (-13.6%)
`1ac4653` (2026-07-29, current)	9498	9489 (±0.1%)	9137 (-3.8%)

On current main the win is net-neutral, but the setter lands on the best Q and protects against Q=16 overshoot.

Why this is still worth merging

Q >= #streams is a hard ceiling: if Q is too small, colliding streams share a hardware queue and cannot run concurrently, regardless of scheduler quality. Today upstream scheduling already keeps dense FSDP collectives off the compute queue at Q=4, so the lever is not the binding constraint — but that can change as overlap improves or configs add streams.

This hook is cheap (config-only, before HIP init, no profiling run) and ensures the queue count is sized correctly ahead of time. It also prevents accidental overshoot to Q=16, which regressed throughput even on the current build.

Test plan

PYTHONPATH=. python3 -m torchtitan.experiments.graph_trainer.tests.test_hw_queues
ROCm: unset GPU_MAX_HW_QUEUES, run graph_trainer, confirm auto-set log line
Override: GPU_MAX_HW_QUEUES=4 → no auto-set

pytorch-bot · 2026-07-30T19:56:59Z

Unknown label ciflow/rocm.
Currently recognized labels are

ciflow/8gpu
ciflow/h100.8
ciflow/rl

pytorch-bot · 2026-07-30T19:56:59Z

~~Workflows were awaiting approval.~~ CI has now been triggered for the ciflow labels on this PR.

yuankaichen-amd · 2026-08-03T03:47:19Z

+    lanes = ["compute", "all_reduce"]  # compute + loss/grad all_reduce (distributed)
+    if dp_shard_active:
+        if fsdp_ag_rs_overlap:
+            lanes += ["dense_all_gather", "reduce_scatter"]  # split onto own PGs


why do we need two lanes for "dense_all_gather", "reduce_scatter"?

When enable_fsdp_ag_rs_overlap is on, reassign_collective_pgs_pass puts dense all-gather and reduce-scatter on separate process groups, so each gets its own stream

yuankaichen-amd · 2026-08-03T03:47:48Z

+            lanes.append("fsdp_comm")  # dense AG+RS unsplit -> share one stream
+    if is_moe and ep > 1:
+        lanes.append("expert_all_gather")  # eFSDP mesh AG (its own PG, structural)
+        lanes.append("expert_all_to_all")  # expert-parallel dispatch (distinct EP PG)


"expert_all_gather" and "expert_all_to_all" never overlap. so they can share one lane.

Estimate the compiled step's independent GPU stream count from the training config and set GPU_MAX_HW_QUEUES to the next power of two before HIP init. Ignore if the env var is already set.

Move GPU_MAX_HW_QUEUES setup out of train.py and run it before super().__init__() so it stays before HIP init while keeping the logic scoped to graph_trainer.

sanketpurandare · 2026-08-04T22:26:25Z

@rishisinhanj is this only needed for graph trainer or also needed for the main torchtitan trainer?

rishisinhanj · 2026-08-04T23:21:32Z

@rishisinhanj is this only needed for graph trainer or also needed for the main torchtitan trainer?

@sanketpurandare This PR was intended only for Graphtrainer, as we have only validated the improvement for graph trainer and tested with GT workloads specifically so I've added it as a hook in the graph trainer trainer file. I think it is also valid to add to the regular torchtitan, but that is a bigger change. If you think it is worthwhile, I can change where it is being called as the underlying logic would remain identical.

sanketpurandare · 2026-08-05T23:04:17Z

@rishisinhanj can you check the failing CI tests?

rishisinhanj · 2026-08-05T23:11:25Z

@sanketpurandare We have observed that failure in our PR to reenable the ROCm CI in the past. I have linked an issue here with our findings: #4064 -> it enters a deadlock due to an SPMD race

It's unrelated to this PR and any changes it makes so shouldn't affect merge status.

tianyu-l · 2026-08-05T23:39:28Z

+    return 1 << (n - 1).bit_length()
+
+
+def maybe_set_gpu_max_hw_queues(config: JobConfig) -> None:


if it's for rocm only, at least the function name and file name should reflect that? o/w it could be confusing to cuda users what this is doing.

sanketpurandare

We had some internal discussion about this, we understand that the abstraction that rocm has requires that we need to know the number of unique streams so that the hw queues can be set accordingly. The part that we are unhappy is that we are introducing hw dependent if-else style code in torchtitan. I would recommend instead adding a section in README on how to set the env variable GPU_MAX_HW_QUEUES based on parallelism configs. Does that sound okay?

SherlockNoMad

we don't have Rocm CI for graph trainer, and ROCM CI is failing for main for a long time...

yuankaichen-amd · 2026-08-06T02:46:14Z

We had some internal discussion about this, we understand that the abstraction that rocm has requires that we need to know the number of unique streams so that the hw queues can be set accordingly. The part that we are unhappy is that we are introducing hw dependent if-else style code in torchtitan. I would recommend instead adding a section in README on how to set the env variable GPU_MAX_HW_QUEUES based on parallelism configs. Does that sound okay?

Thanks for the review! Moving the whole logic to README is messy, as it is not as concise and direct as it is run by the code.

In many places of the pytorch ecosystem, we have hip/rocm specific code. I think it is fine as long as the function names clearly indicate it is for AMD platforms (as Tianyu commented in this PR).

pytorch-bot Bot added ciflow/8gpu ciflow/rocm module: rocm labels Jul 30, 2026

meta-cla Bot added the CLA Signed This label is managed by the Meta Open Source bot. label Jul 30, 2026

yuankaichen-amd reviewed Aug 3, 2026

View reviewed changes

graph_trainer: auto-size GPU_MAX_HW_QUEUES gated only on ROCm

a2fd5c3

Estimate the compiled step's independent GPU stream count from the training config and set GPU_MAX_HW_QUEUES to the next power of two before HIP init. Ignore if the env var is already set.

rishisinhanj force-pushed the graphtrainer/rocm-hw-queues-autoset branch from bb4406b to a2fd5c3 Compare August 3, 2026 22:12

rishisinhanj added 2 commits August 3, 2026 22:20

graph_trainer: remove verify_hw_queues dev script

934a19a

graph_trainer: call hw_queues hook from GraphTrainer.__init__

5f9f622

Move GPU_MAX_HW_QUEUES setup out of train.py and run it before super().__init__() so it stays before HIP init while keeping the logic scoped to graph_trainer.

rishisinhanj marked this pull request as ready for review August 3, 2026 23:10

rishisinhanj requested review from IvanKobzarev, SherlockNoMad, aditvenk, sanketpurandare, tianyu-l and xmfan as code owners August 3, 2026 23:10

sanketpurandare approved these changes Aug 5, 2026

View reviewed changes

tianyu-l reviewed Aug 5, 2026

View reviewed changes

sanketpurandare requested changes Aug 6, 2026

View reviewed changes

SherlockNoMad requested changes Aug 6, 2026

View reviewed changes

		return 1 << (n - 1).bit_length()


		def maybe_set_gpu_max_hw_queues(config: JobConfig) -> None:

Uh oh!

Conversation

rishisinhanj commented Jul 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Stream model

Results

Why this is still worth merging

Test plan

Uh oh!

pytorch-bot Bot commented Jul 30, 2026

Uh oh!

pytorch-bot Bot commented Jul 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

yuankaichen-amd Aug 3, 2026

Choose a reason for hiding this comment

Uh oh!

rishisinhanj Aug 3, 2026

Choose a reason for hiding this comment

Uh oh!

yuankaichen-amd Aug 3, 2026

Choose a reason for hiding this comment

Uh oh!

rishisinhanj Aug 3, 2026

Choose a reason for hiding this comment

Uh oh!

sanketpurandare commented Aug 4, 2026

Uh oh!

rishisinhanj commented Aug 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

sanketpurandare commented Aug 5, 2026

Uh oh!

rishisinhanj commented Aug 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

tianyu-l Aug 5, 2026

Choose a reason for hiding this comment

Uh oh!

sanketpurandare left a comment

Choose a reason for hiding this comment

Uh oh!

SherlockNoMad left a comment

Choose a reason for hiding this comment

Uh oh!

yuankaichen-amd commented Aug 6, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

rishisinhanj commented Jul 30, 2026 •

edited

Loading

pytorch-bot Bot commented Jul 30, 2026 •

edited

Loading

rishisinhanj commented Aug 4, 2026 •

edited

Loading

rishisinhanj commented Aug 5, 2026 •

edited

Loading