graph_trainer: auto-size GPU_MAX_HW_QUEUES on ROCm

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),
+    )

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()

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)