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