From a2fd5c38f5e8027fd35b68729b3e4b6ec1b6df0f Mon Sep 17 00:00:00 2001 From: Rishi Sinha <60558850+rishisinhanj@users.noreply.github.com> Date: Mon, 3 Aug 2026 22:11:30 +0000 Subject: [PATCH 1/3] graph_trainer: auto-size GPU_MAX_HW_QUEUES gated only on ROCm 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. --- tools/verify_hw_queues.py | 71 ++++++++++++ .../experiments/graph_trainer/hw_queues.py | 105 ++++++++++++++++++ .../graph_trainer/tests/test_hw_queues.py | 59 ++++++++++ torchtitan/train.py | 5 + 4 files changed, 240 insertions(+) create mode 100644 tools/verify_hw_queues.py create mode 100644 torchtitan/experiments/graph_trainer/hw_queues.py create mode 100644 torchtitan/experiments/graph_trainer/tests/test_hw_queues.py diff --git a/tools/verify_hw_queues.py b/tools/verify_hw_queues.py new file mode 100644 index 0000000000..118aafd4aa --- /dev/null +++ b/tools/verify_hw_queues.py @@ -0,0 +1,71 @@ +#!/usr/bin/env python3 +"""Offline check: the GPU_MAX_HW_QUEUES formula vs. real profiler traces. + +Counts distinct GPU stream lanes in each trace and asserts the config-computed +Q removes collisions for them (Q >= next_pow2(observed_lanes)). Trace lanes can +be fewer than the logical estimate (HIP-graph capture folds lanes), so the bound +is >=, not ==. Loads hw_queues.py by path so it runs on bare python3 (no torch). + +Usage: python tools/verify_hw_queues.py ... +""" +import gzip +import importlib.util +import json +import sys +from pathlib import Path + +_path = Path(__file__).resolve().parents[1] / "torchtitan/experiments/graph_trainer/hw_queues.py" +_spec = importlib.util.spec_from_file_location("hw_queues", _path) +hwq = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(hwq) # module top is stdlib-only; torch is imported lazily + +# Parallelism config behind the calibration traces: EP=2 2-node DeepSeek-V3 with +# dense FSDP AG/RS overlap on (dedicated PGs -> their own streams). +EP2_DSV3 = dict( + dp_shard_active=True, is_moe=True, ep=2, tp=1, cp=1, + fsdp_ag_rs_overlap=True, cudagraph=True, +) + + +def observed_lanes(path): + opener = None + if path.suffix == ".gz": + opener = gzip.open + else: + opener = open + with opener(path) as f: + d = json.load(f) + events = None + if isinstance(d, dict): + events = d["traceEvents"] + else: + events = d + return len( + { + e["args"]["name"] + for e in events + if e.get("ph") == "M" + and e.get("name") == "thread_name" + and str(e.get("args", {}).get("name", "")).startswith("stream") + } + ) + + +def main(traces): + if not traces: + print("usage: verify_hw_queues.py ...") + return 1 + q = hwq._next_pow2(len(hwq._stream_lanes(**EP2_DSV3))) + print(f"formula (EP=2 DSv3): Q={q}") + rc = 0 + for t in map(Path, traces): + obs = observed_lanes(t) + need = hwq._next_pow2(obs) + ok = q >= need + rc |= not ok + print(f" [{'OK' if ok else 'FAIL'}] {t.name}: observed={obs}, need Q>={need}, have {q}") + return rc + + +if __name__ == "__main__": + sys.exit(main(sys.argv[1:])) diff --git a/torchtitan/experiments/graph_trainer/hw_queues.py b/torchtitan/experiments/graph_trainer/hw_queues.py new file mode 100644 index 0000000000..ad3fcd5f8f --- /dev/null +++ b/torchtitan/experiments/graph_trainer/hw_queues.py @@ -0,0 +1,105 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""ROCm GPU_MAX_HW_QUEUES sizing for graph_trainer. + +HIP maps logical streams to hardware queues via ``stream_id % Q``. When +``Q`` is smaller than the number of concurrently active streams, unrelated +streams alias onto the same queue and lose overlap. This module estimates +the stream count from the training config and sets ``GPU_MAX_HW_QUEUES`` to +the next power of two before HIP init. +""" + +from __future__ import annotations + +import os +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from torchtitan.config import JobConfig + + +def _stream_lanes( + *, + dp_shard_active: bool, + is_moe: bool, + ep: int, + tp: int, + cp: int, + fsdp_ag_rs_overlap: bool, + cudagraph: bool, +) -> list[str]: + """Return labels for each independent GPU stream the compiled step creates.""" + lanes = ["compute", "all_reduce"] + if dp_shard_active: + if fsdp_ag_rs_overlap: + # reassign_collective_pgs_pass gives dense AG and RS dedicated PGs. + lanes += ["dense_all_gather", "reduce_scatter"] + else: + lanes.append("fsdp_comm") + if is_moe and ep > 1: + # Expert AG and a2a never overlap, so they share one stream. + lanes.append("expert_comm") + if tp > 1: + lanes.append("tp_all_gather") + if cp > 1: + lanes.append("cp_all_gather") + # cudagraph_pass is skipped when fsdp_ag_rs_overlap rewrites the graph. + if cudagraph and not fsdp_ag_rs_overlap: + lanes.append("cudagraph_capture") + return lanes + + +def _next_pow2(n: int) -> int: + return 1 << (n - 1).bit_length() + + +def maybe_set_gpu_max_hw_queues(config: JobConfig) -> None: + """Set GPU_MAX_HW_QUEUES on ROCm before HIP init. + + No-op on CUDA and when the env var is already set. + """ + import torch + + from torchtitan.tools.logging import logger + + if torch.version.hip is None: + return + if "GPU_MAX_HW_QUEUES" in os.environ: + return + + from torchtitan.models.common.moe import MoE + + p = config.parallelism + dp_shard = p.data_parallel_shard_degree + dp_shard_active = dp_shard == -1 or dp_shard > 1 + is_moe = ( + config.model_spec is not None + and next(config.model_spec.model.traverse(MoE.Config), None) is not None + ) + compile_cfg = config.compile + cudagraph = getattr(compile_cfg, "enable_passes", False) and ( + "cudagraph_pass" not in getattr(compile_cfg, "disable_passes", []) + ) + fsdp_ag_rs_overlap = getattr(compile_cfg, "enable_fsdp_ag_rs_overlap", False) + + lanes = _stream_lanes( + dp_shard_active=dp_shard_active, + is_moe=is_moe, + ep=p.expert_parallel_degree, + tp=p.tensor_parallel_degree, + cp=p.context_parallel_degree, + fsdp_ag_rs_overlap=fsdp_ag_rs_overlap, + cudagraph=cudagraph, + ) + q = _next_pow2(len(lanes)) + os.environ["GPU_MAX_HW_QUEUES"] = str(q) + logger.info( + "GPU_MAX_HW_QUEUES auto-set to %d (estimated %d GPU streams: %s)", + q, + len(lanes), + ", ".join(lanes), + ) diff --git a/torchtitan/experiments/graph_trainer/tests/test_hw_queues.py b/torchtitan/experiments/graph_trainer/tests/test_hw_queues.py new file mode 100644 index 0000000000..7efa9c5716 --- /dev/null +++ b/torchtitan/experiments/graph_trainer/tests/test_hw_queues.py @@ -0,0 +1,59 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import unittest + +from torchtitan.experiments.graph_trainer.hw_queues import _next_pow2, _stream_lanes + + +class TestHwQueues(unittest.TestCase): + def test_ep2_dsv3_with_fsdp_overlap(self): + lanes = _stream_lanes( + dp_shard_active=True, + is_moe=True, + ep=2, + tp=1, + cp=1, + fsdp_ag_rs_overlap=True, + cudagraph=True, + ) + self.assertNotIn("cudagraph_capture", lanes) + self.assertEqual(len(lanes), 5) + self.assertEqual(_next_pow2(len(lanes)), 8) + + def test_ep2_dsv3_with_cudagraph(self): + lanes = _stream_lanes( + dp_shard_active=True, + is_moe=True, + ep=2, + tp=1, + cp=1, + fsdp_ag_rs_overlap=False, + cudagraph=True, + ) + self.assertIn("cudagraph_capture", lanes) + self.assertEqual(len(lanes), 5) + self.assertEqual(_next_pow2(len(lanes)), 8) + + def test_plain_fsdp(self): + lanes = _stream_lanes( + dp_shard_active=True, + is_moe=False, + ep=1, + tp=1, + cp=1, + fsdp_ag_rs_overlap=False, + cudagraph=False, + ) + self.assertEqual(lanes, ["compute", "all_reduce", "fsdp_comm"]) + self.assertEqual(_next_pow2(len(lanes)), 4) + + def test_next_pow2(self): + self.assertEqual((_next_pow2(7), _next_pow2(8), _next_pow2(9)), (8, 8, 16)) + + +if __name__ == "__main__": + unittest.main() diff --git a/torchtitan/train.py b/torchtitan/train.py index 405f5e2718..11dfbfba48 100644 --- a/torchtitan/train.py +++ b/torchtitan/train.py @@ -12,6 +12,9 @@ from torchtitan.observability import structured_logger as sl from torchtitan.tools.logging import init_logger, logger from torchtitan.trainer import Trainer +from torchtitan.experiments.graph_trainer.hw_queues import ( + maybe_set_gpu_max_hw_queues, +) def main() -> None: @@ -38,6 +41,8 @@ def main() -> None: ) sl.log_trace_instant("structured_logger_started") + maybe_set_gpu_max_hw_queues(config) + trainer: Trainer | None = None try: From 934a19a625a26c958d871f89aa30ce8b8510d54a Mon Sep 17 00:00:00 2001 From: Rishi Sinha <60558850+rishisinhanj@users.noreply.github.com> Date: Mon, 3 Aug 2026 22:20:12 +0000 Subject: [PATCH 2/3] graph_trainer: remove verify_hw_queues dev script --- tools/verify_hw_queues.py | 71 --------------------------------------- 1 file changed, 71 deletions(-) delete mode 100644 tools/verify_hw_queues.py diff --git a/tools/verify_hw_queues.py b/tools/verify_hw_queues.py deleted file mode 100644 index 118aafd4aa..0000000000 --- a/tools/verify_hw_queues.py +++ /dev/null @@ -1,71 +0,0 @@ -#!/usr/bin/env python3 -"""Offline check: the GPU_MAX_HW_QUEUES formula vs. real profiler traces. - -Counts distinct GPU stream lanes in each trace and asserts the config-computed -Q removes collisions for them (Q >= next_pow2(observed_lanes)). Trace lanes can -be fewer than the logical estimate (HIP-graph capture folds lanes), so the bound -is >=, not ==. Loads hw_queues.py by path so it runs on bare python3 (no torch). - -Usage: python tools/verify_hw_queues.py ... -""" -import gzip -import importlib.util -import json -import sys -from pathlib import Path - -_path = Path(__file__).resolve().parents[1] / "torchtitan/experiments/graph_trainer/hw_queues.py" -_spec = importlib.util.spec_from_file_location("hw_queues", _path) -hwq = importlib.util.module_from_spec(_spec) -_spec.loader.exec_module(hwq) # module top is stdlib-only; torch is imported lazily - -# Parallelism config behind the calibration traces: EP=2 2-node DeepSeek-V3 with -# dense FSDP AG/RS overlap on (dedicated PGs -> their own streams). -EP2_DSV3 = dict( - dp_shard_active=True, is_moe=True, ep=2, tp=1, cp=1, - fsdp_ag_rs_overlap=True, cudagraph=True, -) - - -def observed_lanes(path): - opener = None - if path.suffix == ".gz": - opener = gzip.open - else: - opener = open - with opener(path) as f: - d = json.load(f) - events = None - if isinstance(d, dict): - events = d["traceEvents"] - else: - events = d - return len( - { - e["args"]["name"] - for e in events - if e.get("ph") == "M" - and e.get("name") == "thread_name" - and str(e.get("args", {}).get("name", "")).startswith("stream") - } - ) - - -def main(traces): - if not traces: - print("usage: verify_hw_queues.py ...") - return 1 - q = hwq._next_pow2(len(hwq._stream_lanes(**EP2_DSV3))) - print(f"formula (EP=2 DSv3): Q={q}") - rc = 0 - for t in map(Path, traces): - obs = observed_lanes(t) - need = hwq._next_pow2(obs) - ok = q >= need - rc |= not ok - print(f" [{'OK' if ok else 'FAIL'}] {t.name}: observed={obs}, need Q>={need}, have {q}") - return rc - - -if __name__ == "__main__": - sys.exit(main(sys.argv[1:])) From 5f9f622934a3bda52529cd4f5c38ab570c1d41a8 Mon Sep 17 00:00:00 2001 From: Rishi Sinha <60558850+rishisinhanj@users.noreply.github.com> Date: Mon, 3 Aug 2026 22:23:25 +0000 Subject: [PATCH 3/3] graph_trainer: call hw_queues hook from GraphTrainer.__init__ 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. --- torchtitan/experiments/graph_trainer/trainer.py | 2 ++ torchtitan/train.py | 5 ----- 2 files changed, 2 insertions(+), 5 deletions(-) diff --git a/torchtitan/experiments/graph_trainer/trainer.py b/torchtitan/experiments/graph_trainer/trainer.py index 52588b25f0..f13cc4474c 100644 --- a/torchtitan/experiments/graph_trainer/trainer.py +++ b/torchtitan/experiments/graph_trainer/trainer.py @@ -22,6 +22,7 @@ trace_input_preparer_keys, ) from torchtitan.experiments.graph_trainer.cudagraph import cudagraph_teardown +from torchtitan.experiments.graph_trainer.hw_queues import maybe_set_gpu_max_hw_queues from torchtitan.experiments.graph_trainer.make_fx_tracer import ( minimal_fx_tracer, run_traced, @@ -106,6 +107,7 @@ class Config(Trainer.Config): ) def __init__(self, config): + maybe_set_gpu_max_hw_queues(config) super().__init__(config) _maybe_apply_numa_binding(self.device.index, self.device.type) diff --git a/torchtitan/train.py b/torchtitan/train.py index 11dfbfba48..405f5e2718 100644 --- a/torchtitan/train.py +++ b/torchtitan/train.py @@ -12,9 +12,6 @@ from torchtitan.observability import structured_logger as sl from torchtitan.tools.logging import init_logger, logger from torchtitan.trainer import Trainer -from torchtitan.experiments.graph_trainer.hw_queues import ( - maybe_set_gpu_max_hw_queues, -) def main() -> None: @@ -41,8 +38,6 @@ def main() -> None: ) sl.log_trace_instant("structured_logger_started") - maybe_set_gpu_max_hw_queues(config) - trainer: Trainer | None = None try: