Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / compile / pass_config.py: 94%
51 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""
16Parallel Configuration - Graph-mode parallel configuration.
18Pure configuration dataclass: no environment probes. Whether FSDP actually
19runs is decided by ``fsdp_enabled`` here (the user's intent) AND the runtime
20distributed guard inside ``FSDPPass`` (which still early-returns on
21``world_size == 1``). The previous ``fsdp_enabled`` property returned
22``dist.is_initialized() and world_size > 1`` — that turned on FSDP whenever
23distributed was initialized, leaving no way to run pure-TP / pure-PP graph
24mode. The explicit field fixes that. Pipeline-parallel (``pp_enabled``)
25follows the same contract: intent here, runtime guard in ``PpPass``.
26"""
28__all__ = ["PassConfig", "PP_SCHEDULES", "DP_MODES", "normalize_dp_mode"]
30from dataclasses import dataclass
31from typing import Any, Optional
34# Pipeline schedule names accepted by ``pp_schedule``. Kept torch-free here so
35# the config dataclass stays importable anywhere; ``pp_schedule``'s registry
36# (which does import torch) must expose exactly these keys.
37PP_SCHEDULES = ("gpipe", "1f1b")
39# Data-parallel modes, mirroring simplefsdp's ``data_parallel`` modes:
40# "fsdp" -> fully_shard (Shard(0) on the fsdp axis)
41# "ddp" -> replicate (Replicate; grads all-reduced)
42# "hsdp" -> hybrid_shard (Replicate x Shard(0); rs + all-reduce)
43DP_MODES = ("fsdp", "ddp", "hsdp")
45# simplefsdp / fully_shard spellings accepted as aliases.
46_DP_MODE_ALIASES = {
47 "fully_shard": "fsdp",
48 "replicate": "ddp",
49 "hybrid_shard": "hsdp",
50}
53def normalize_dp_mode(mode: str) -> str:
54 """Return the canonical DP mode for ``mode`` (accepts simplefsdp aliases).
56 Raises:
57 ValueError: When ``mode`` is not a string or is unknown.
58 """
59 if not isinstance(mode, str):
60 raise ValueError(f"dp_mode must be a string, got {type(mode).__name__}")
61 key = mode.strip().lower()
62 return _DP_MODE_ALIASES.get(key, key)
65@dataclass
66class PassConfig:
67 """Parallel configuration for graph-mode FSDP (+ optional TP / SP / PP) training.
69 Attributes:
70 enable_overlap: Drive ``AutoOverlapPass`` to move ``wait_tensor`` for
71 communication/compute overlap.
72 fsdp_enabled: Drive ``FSDPPass`` (data-parallel partitioning: parameter
73 all_gather + gradient reduction + live-model sharding, selected by
74 ``dp_mode``). ``False`` skips data parallelism entirely — set this
75 for pure-TP / pure-PP graph-mode runs. ``FSDPPass`` itself still
76 early-returns when distributed is not initialized or
77 ``world_size == 1``, so single-card runs are a no-op regardless.
78 fsdp_reshard_after_forward: Drive the FSDP reshard optimization.
79 When ``True`` (default, matching FSDP's ``reshard_after_forward``)
80 ``FSDPPass`` sinks each param's all_gather to its first forward
81 use, frees the replicated full parameter once its last forward
82 reader is done, and — when the backward still needs it — inserts a
83 fresh all_gather plus a rematerialization of the saved forward
84 views just before the first backward consumer. Peak memory then
85 tracks the forward working set instead of the sum of all
86 replicated parameters. Set ``False`` to keep every gathered
87 parameter resident for the whole joint graph (previous behavior).
88 fsdp_degree: Size of the FSDP group. ``None`` (default) means
89 "resolve at runtime": the trainer back-fills it from the
90 automodel ``MeshContext`` (TP+FSDP hybrid, where the FSDP group
91 is a proper sub-group of the world), and ``FSDPPass`` falls back
92 to ``world_size`` for the FSDP-only path. Mutating this after
93 construction is supported but discouraged — prefer passing the
94 resolved degree at construction time (see ``GraphTrainer``).
95 For ``dp_mode="hsdp"`` this is the *shard* axis degree.
96 dp_mode: Data-parallel mode, one of ``DP_MODES``: ``"fsdp"``
97 (fully_shard, the default), ``"ddp"`` (replicate; no parameter
98 sharding, gradients all-reduced), or ``"hsdp"`` (hybrid_shard;
99 parameters replicated on the ``dp_replicate`` axis and sharded on
100 the ``fsdp`` axis, gradients reduce-scattered then all-reduced).
101 simplefsdp spellings ``"fully_shard"`` / ``"replicate"`` /
102 ``"hybrid_shard"`` are accepted as aliases and normalized.
103 dp_replicate_degree: Size of the replicate axis for ``"ddp"`` /
104 ``"hsdp"``. ``None`` means "resolve at runtime" (from the mesh,
105 or ``world_size // fsdp_degree`` for HSDP / ``world_size`` for
106 DDP). Kept separate from ``fsdp_degree`` (the shard axis) so an
107 HSDP mesh is unambiguous.
108 tp_size: Tensor-parallel degree. Informational today (TP collectives
109 live inside boundary forwards baked by automodel, not in the
110 graph-mode passes); kept so a future TP-aware pass can read it
111 without API churn.
112 sequence_parallel: Enable sequence parallel (SP) on the TP axis.
113 loss_parallel: Enable loss parallel (LP) on the TP axis.
114 pp_enabled: Drive ``PpPass`` (pipeline-parallel graph split). When
115 ``True`` the joint fwd+bwd graph is sliced to this rank's stage
116 along module-FQN boundaries declared in ``GraphParallelPlan``, cross-stage
117 activations are exchanged via P2P ``isend``/``irecv``, and a
118 self-contained GPipe schedule is installed as a ``call_module``
119 stub inside the rewritten graph, so the trainer needs no PP
120 wiring. ``PpPass`` early-returns when distributed is not
121 initialized or ``world_size == 1``.
122 pp_degree: Number of pipeline-parallel ranks (== number of stages
123 for the single-virtual-stage schedule shipped here). ``None``
124 means "resolve at runtime" from ``world_size`` (the pure-PP
125 path); for a PP+FSDP hybrid the FSDP group is a proper
126 sub-group of the world and the trainer back-fills this from
127 the mesh, exactly like ``fsdp_degree``.
128 pp_microbatch_size: Number of samples per microbatch for the
129 pipeline schedule. The per-step batch size (the leading dim of
130 the trainer's input tensor) must be divisible by this. ``1`` means
131 one sample per microbatch (maximum pipeline concurrency, most
132 P2P traffic); larger values trade bubble size for fewer P2P
133 round-trips.
134 pp_schedule: Pipeline schedule name, one of ``PP_SCHEDULES``.
135 ``"gpipe"`` (default) runs all forwards before all backwards
136 (simplest ordering, largest bubble). ``"1f1b"`` interleaves one
137 backward per forward in steady state, shrinking the pipeline
138 bubble and freeing each microbatch's activations as soon as its
139 backward runs (GPipe retains all of them until the sweep ends);
140 it requires ``num_microbatches >= pp_degree``.
141 ``PpPass`` resolves the name to a schedule class via
142 ``pp_schedule.get_schedule_class``.
144 Note:
145 ``fsdp_enabled`` / ``pp_enabled`` no longer probe
146 ``torch.distributed``. The distributed-initialized check moved into
147 the respective passes' ``run`` (their original location) so this
148 dataclass stays torch-free and importable anywhere.
149 """
151 enable_overlap: bool = True
152 fsdp_enabled: bool = True
153 fsdp_reshard_after_forward: bool = True
154 fsdp_degree: Optional[int] = None
155 dp_mode: str = "fsdp"
156 dp_replicate_degree: Optional[int] = None
157 tp_size: int = 1
158 sequence_parallel: bool = False
159 loss_parallel: bool = False
160 pp_enabled: bool = False
161 pp_degree: Optional[int] = None
162 pp_microbatch_size: int = 1
163 pp_schedule: str = "gpipe"
165 def __post_init__(self) -> None:
166 self.dp_mode = normalize_dp_mode(self.dp_mode)
167 self.validate()
169 def validate(self) -> None:
170 """Sanity-check invariants; also re-run after manual field mutation.
172 Raises:
173 ValueError: On a negative ``tp_size``, a non-positive explicit
174 ``fsdp_degree`` / ``pp_degree``, a non-positive
175 ``pp_microbatch_size``, an unknown ``pp_schedule`` /
176 ``dp_mode``, or a non-positive ``dp_replicate_degree``.
177 """
178 if self.tp_size < 1:
179 raise ValueError(f"tp_size must be >= 1, got {self.tp_size}")
180 if self.fsdp_degree is not None and self.fsdp_degree < 1:
181 raise ValueError(
182 f"fsdp_degree must be None or a positive int, got {self.fsdp_degree}"
183 )
184 if self.dp_mode not in DP_MODES:
185 raise ValueError(
186 f"dp_mode must be one of {DP_MODES} "
187 f"(aliases: {tuple(_DP_MODE_ALIASES)}), got {self.dp_mode!r}"
188 )
189 if self.dp_replicate_degree is not None and self.dp_replicate_degree < 1:
190 raise ValueError(
191 f"dp_replicate_degree must be None or a positive int, "
192 f"got {self.dp_replicate_degree}"
193 )
194 if self.pp_degree is not None and self.pp_degree < 1:
195 raise ValueError(
196 f"pp_degree must be None or a positive int, got {self.pp_degree}"
197 )
198 if self.pp_microbatch_size < 1:
199 raise ValueError(
200 f"pp_microbatch_size must be >= 1, got {self.pp_microbatch_size}"
201 )
202 if not isinstance(self.pp_schedule, str) or (
203 self.pp_schedule.lower() not in PP_SCHEDULES
204 ):
205 raise ValueError(
206 f"pp_schedule must be one of {PP_SCHEDULES}, got {self.pp_schedule!r}"
207 )
210def build_pass_config_from_trainer_config(
211 config: Any,
212 *,
213 fsdp_enabled: Optional[bool] = None,
214 enable_overlap: bool = False,
215) -> PassConfig:
216 """Project trainer topology intent onto graph-mode pass config.
218 This helper stays in the pass-config module because it computes the
219 graph-pass surface from the trainer topology without altering runtime
220 objects. Accepting ``Any`` keeps this module importable without the
221 trainer package.
222 """
223 accelerator = config.accelerator
224 if fsdp_enabled is None:
225 fsdp_enabled = (
226 config.fsdp_config.dp_shard_size > 1
227 or config.fsdp_config.edp_shard_size > 1
228 )
229 fsdp_degree = config.fsdp_config.dp_shard_size if fsdp_enabled else None
230 return PassConfig(
231 enable_overlap=enable_overlap,
232 fsdp_enabled=fsdp_enabled,
233 fsdp_degree=fsdp_degree,
234 tp_size=accelerator.tp_size,
235 sequence_parallel=accelerator.sequence_parallel,
236 loss_parallel=accelerator.loss_parallel,
237 )
240__all__ = ["PassConfig", "build_pass_config_from_trainer_config"]