Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / compile / compiler.py: 91%

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""" 

16Graph Compiler - compile a model into a parallel graph and run fwd+bwd. 

17 

18Compiler surface of the graph-mode stack: 

19 

20- ``compile`` traces the model into a joint forward+backward FX graph and 

21 runs the partitioning passes (FSDP / PP / overlap) on it. 

22- ``forward_backward`` executes the compiled graph against the live model 

23 state and deposits gradients into ``param.grad`` (accumulating, so several 

24 micro-batch calls may run before the caller's optimizer step). 

25 

26Training-policy concerns (optimizer, dataloader loop, logging, batch device 

27placement) are deliberately out of scope: ``GraphTrainer`` composes this 

28class and owns them. 

29""" 

30 

31from __future__ import annotations 

32 

33__all__ = ["GraphCompiler"] 

34 

35import logging 

36import math 

37from typing import Any, Callable, Dict, List, Optional, Tuple 

38 

39import torch 

40import torch.distributed as dist 

41from torch._prims_common import make_contiguous_strides_for 

42from torch.distributed.distributed_c10d import _register_process_group 

43 

44from hyper_parallel.core.dtensor.device_mesh import DeviceMesh, init_device_mesh 

45from hyper_parallel.core.dtensor.layout import Layout 

46from hyper_parallel.core.dtensor.placement_types import Replicate, Shard, StridedShard 

47from hyper_parallel.distributed._builder.fsdp_adapter import FSDP2Manager 

48from hyper_parallel.models.build_options import FSDP2Config 

49from hyper_parallel.models.external_state import get_model_external_state 

50 

51from .graph_parallel_plan import GraphParallelPlan 

52from .pass_config import ( 

53 GradientReduction, 

54 PassConfig, 

55 build_pass_config_from_trainer_config, 

56 normalize_dp_mode, 

57) 

58from .passes.pipeline import PassPipeline 

59from .passes.activation_checkpoint import prepare_activation_checkpoint 

60from .tracer.dynamic_shapes import normalize_dynamic_arg_dims 

61from .tracer.graph_tracer import run_traced_graph, trace_model_graph 

62 

63_LOG = logging.getLogger(__name__) 

64 

65 

66class GraphCompiler: 

67 """ 

68 Graph-mode Compiler 

69 

70 Users provide model code and parallel configuration. 

71 The compiler handles graph capture and all parallel logic 

72 (FSDP / PP / overlap passes), then runs forward+backward per step. 

73 """ 

74 

75 def __init__( 

76 self, 

77 model: torch.nn.Module, 

78 train_fn: Callable, 

79 pass_config: Optional[PassConfig] = None, 

80 parallel_plan: Optional[GraphParallelPlan] = None, 

81 trainer_config: Optional[Any] = None, 

82 device: Optional[torch.device] = None, 

83 mesh_context: Optional[Any] = None, 

84 dynamic: Optional[bool] = None, 

85 dynamic_arg_dims: dict[str, int | list[int]] | None = None, 

86 source_shard_info: Optional[dict] = None, 

87 ) -> None: 

88 """ 

89 Args: 

90 model: Model to compile 

91 train_fn: Training function signature: (model, input, label) -> loss 

92 pass_config: Parallel configuration. When omitted, the compiler 

93 projects ``trainer_config`` onto a graph-mode ``PassConfig``; 

94 if both are omitted, it uses ``PassConfig()``. 

95 parallel_plan: GraphParallelPlan declaring which modules to shard 

96 (optional; enables declarative sharding) 

97 trainer_config: Optional trainer topology used only to infer a 

98 default ``PassConfig`` when ``pass_config`` is omitted. 

99 device: Device to place the model and run the graph on. Defaults 

100 to the NPU device when available, otherwise CPU. 

101 mesh_context: Optional automodel ``MeshContext`` carrying the 

102 pre-built parallel meshes. Model forwards retain their TP/EP 

103 groups; the compiler registers the dense shard sub-mesh under 

104 ``"fsdp"`` and, for basic EP, the expert EDP shard sub-mesh 

105 under ``"fsdp_expert"`` for the graph-mode FSDP pass. 

106 dynamic: Symbolize user input dimensions; inherits trainer_config.compile.dynamic when omitted. 

107 dynamic_arg_dims: Optional dotted input paths to dynamic dimensions, e.g. {"x": [0, 1]}. 

108 Overrides automatic dimension selection and inherits trainer configuration when omitted. 

109 source_shard_info: Source placements and meshes returned by apply_sharding_plan, 

110 before FSDP cuts the localized parameters. 

111 """ 

112 compile_config = getattr(trainer_config, "compile", None) 

113 self.dynamic = ( 

114 getattr(compile_config, "dynamic", False) if dynamic is None else dynamic 

115 ) 

116 if not isinstance(self.dynamic, bool): 

117 raise ValueError("dynamic must be a bool") 

118 if dynamic_arg_dims is None: 

119 dynamic_arg_dims = getattr(compile_config, "dynamic_arg_dims", None) 

120 self.dynamic_arg_dims = normalize_dynamic_arg_dims(dynamic_arg_dims) 

121 self.model = model 

122 self.train_fn = train_fn 

123 self.parallel_plan = parallel_plan 

124 self.pass_config = self._resolve_pass_config( 

125 pass_config=pass_config, 

126 trainer_config=trainer_config, 

127 ) 

128 self._mesh_context = mesh_context 

129 self.source_shard_info = source_shard_info 

130 self._fsdp_mesh = None 

131 self._expert_fsdp_mesh = None 

132 self._fsdp_config = getattr(trainer_config, 'fsdp_config', None) 

133 self._parameter_meshes = {} 

134 self._parameter_groups = {} 

135 self._parameter_shapes = {} 

136 self._parameter_layouts = None 

137 self._layout_signature = None 

138 self.device = device or ( 

139 torch.device("npu") 

140 if (hasattr(torch, "npu") and torch.npu.is_available()) 

141 else torch.device("cpu") 

142 ) 

143 

144 self.pass_config.validate() 

145 

146 self._joint_graph = None 

147 self._expert_param_names: set[str] = set() 

148 self._pytree_pre_hook: Optional[Callable[[], None]] = None 

149 # Set when the PP hybrid mesh registered a "pp" process group; makes 

150 # ``_build_pass_kwargs`` hand the group name to ``PpPass``. 

151 self._pp_group_name: Optional[str] = None 

152 

153 @staticmethod 

154 def _resolve_pass_config( 

155 *, 

156 pass_config: Optional[PassConfig], 

157 trainer_config: Optional[Any], 

158 ) -> PassConfig: 

159 """Resolve pass configuration from explicit or trainer-level input.""" 

160 if pass_config is not None: 

161 return pass_config 

162 if trainer_config is not None: 

163 return build_pass_config_from_trainer_config(trainer_config) 

164 return PassConfig() 

165 

166 @property 

167 def is_compiled(self) -> bool: 

168 """Whether a joint graph has already been compiled.""" 

169 return self._joint_graph is not None 

170 

171 @property 

172 def parameter_layouts(self) -> Optional[dict]: 

173 """Final parameter-identity layouts, available after lazy compilation. 

174 

175 Cache metadata only: gradients may be replaced on every optimizer step. 

176 Replacing a parameter invalidates the mapping, as does compiling again. 

177 """ 

178 if self._joint_graph is None: 

179 return None 

180 parameters = dict(self.model.named_parameters()) 

181 signature = tuple((id(param), tuple(param.shape), param.dtype, param.device) for param in parameters.values()) 

182 if signature != self._layout_signature: 

183 self._parameter_layouts = self._build_parameter_layouts(parameters) 

184 self._layout_signature = signature 

185 self._joint_graph.graph_module.parameter_layouts = self._parameter_layouts 

186 return self._parameter_layouts 

187 

188 def compile(self, **inputs: Any) -> None: 

189 """ 

190 Compile model into parallel graph 

191 

192 Users can explicitly call this, or it will be automatically compiled 

193 at first forward_backward 

194 

195 Args: 

196 **inputs: Model inputs, forwarded to ``train_fn`` as keyword 

197 arguments and used to trace the joint graph 

198 """ 

199 external = get_model_external_state(self.model) 

200 if external is not None and external.parameters: 

201 raise ValueError("GraphCompiler does not support externally managed parameters yet") 

202 

203 if ( 

204 self.dynamic or self.dynamic_arg_dims is not None 

205 ) and self.pass_config.pp_enabled: 

206 raise ValueError( 

207 "Dynamic shapes with pipeline parallel are not supported yet; disable PP or dynamic shapes" 

208 ) 

209 if self._pytree_pre_hook is not None: 

210 self._pytree_pre_hook() 

211 

212 self._validate_ep_topology() 

213 if ( 

214 isinstance(getattr(self._mesh_context, "tp_size", None), int) 

215 and self._mesh_context.tp_size > 1 

216 and not self.source_shard_info 

217 ): 

218 raise ValueError("Graph TP requires source_shard_info from the sharding plan") 

219 

220 if self.pass_config.fsdp_enabled and dist.is_initialized(): 

221 # Only build the FSDP mesh when distributed is actually up. 

222 # ``FSDPPass`` early-returns when ``world_size == 1``, so a 

223 # single-process run (no dist, or a single rank) compiles and 

224 # runs as plain graph mode without sharding. 

225 self._init_device_mesh(self._mesh_context) 

226 

227 # Activation checkpointing must resolve its scope and run the adapter's 

228 # runtime hook BEFORE tracing: the hook configures state the model's 

229 # forward reads. The pass receives the same scope so the hook runs once. 

230 ac_scope = None 

231 if self.pass_config.ac_mode != "off": 

232 ac_scope = prepare_activation_checkpoint(self.model, self.pass_config) 

233 

234 self._bind_parameter_meshes() 

235 self._parameter_shapes = {name: tuple(param.shape) for name, param in self.model.named_parameters()} 

236 self._layout_signature = None 

237 

238 trace_kwargs = {} 

239 if self.dynamic or self.dynamic_arg_dims is not None: 

240 trace_kwargs = { 

241 "dynamic": self.dynamic, 

242 "dynamic_arg_dims": self.dynamic_arg_dims, 

243 } 

244 joint_graph = trace_model_graph( 

245 self.model, self.train_fn, inputs, **trace_kwargs 

246 ) 

247 

248 pipeline = PassPipeline.from_config(self.pass_config, self.parallel_plan) 

249 

250 pass_kwargs = self._build_pass_kwargs() 

251 # Terminal backend compilation (InductorPass) sizes its regions from 

252 # the tracer's fake example inputs; every other pass ignores them. 

253 pass_kwargs["example_inputs"] = joint_graph.example_inputs 

254 if ac_scope is not None: 

255 pass_kwargs["ac_scope"] = ac_scope 

256 

257 # Passes mutate ``graph_module`` in place and return it, so the 

258 # transformed graph lives on ``joint_graph`` for ``forward_backward``. 

259 pipeline.run(joint_graph.graph_module, **pass_kwargs) 

260 

261 if joint_graph.input_guards is not None: 

262 joint_graph.input_guards.refresh() 

263 self._joint_graph = joint_graph 

264 

265 def forward_backward(self, **inputs: Any) -> Any: 

266 """ 

267 Execute one compiled forward+backward step. 

268 

269 Compiles lazily on the first call. Gradients are ACCUMULATED into 

270 ``param.grad`` (not overwritten), so several micro-batch calls may 

271 run before the caller's optimizer step / zero_grad. 

272 

273 Args: 

274 **inputs: Model inputs, forwarded to ``train_fn`` as keyword 

275 arguments (must live on the compiler's device) 

276 

277 Returns: 

278 tuple: ``(loss, loss_dict)`` — the backward loss and the named 

279 loss values emitted by ``train_fn`` as extra graph outputs after 

280 the gradients (empty when ``train_fn`` returned a bare loss 

281 tensor). 

282 """ 

283 if self._joint_graph is None: 

284 self.compile(**inputs) 

285 

286 loss, grads, loss_dict = self._run_graph(**inputs) 

287 

288 self._accumulate_grads(grads) 

289 

290 return loss, loss_dict 

291 

292 def set_pytree_pre_hook(self, hook: Callable[[], None]) -> "GraphCompiler": 

293 """Register a no-arg hook run immediately before first compilation.""" 

294 self._pytree_pre_hook = hook 

295 return self 

296 

297 def to(self, device: torch.device) -> "GraphCompiler": 

298 """Move the model to ``device`` and remember it for graph execution.""" 

299 self.device = torch.device(device) 

300 self.model = self.model.to(self.device) 

301 return self 

302 

303 def source_layout_for_parameter(self, name: str) -> Optional[tuple]: 

304 """Resolve pre-FSDP placement, including tied and replicated parameters. 

305 

306 Args: 

307 name: Parameter FQN in the live model. 

308 

309 Returns: 

310 The source ``(placements, mesh)`` pair, or None for a plain local model. 

311 """ 

312 if self.source_shard_info is not None: 

313 if name in self.source_shard_info: 

314 return self.source_shard_info[name] 

315 parameters = dict(self.model.named_parameters(remove_duplicate=False)) 

316 for alias, layout in self.source_shard_info.items(): 

317 if parameters.get(alias) is parameters.get(name): 

318 return layout 

319 context = self._mesh_context 

320 if context is None: 

321 return None 

322 if name in self._expert_param_names: 

323 return ((Shard(0),), context.fsdp_moe_mesh["ep"]) 

324 mesh = context.fsdp_non_moe_mesh 

325 if mesh is not None and "tp" in mesh.mesh_dim_names: 

326 return ((Replicate(),), mesh["tp"]) 

327 return None 

328 

329 def _build_parameter_layouts(self, parameters: dict) -> Optional[dict]: 

330 """Describe actual local storage using common FSDP layout construction.""" 

331 context = self._mesh_context 

332 if context is None and self._fsdp_mesh is None: 

333 return None 

334 if self._fsdp_mesh is None and context is not None: 

335 # Plain graph mode still needs metadata, without creating collectives. 

336 self._init_mesh_from_context(context, self.pass_config.dp_mode, register_groups=False) 

337 sharded = getattr(self._joint_graph.graph_module, "fsdp_sharded_param_names", frozenset()) 

338 layouts = {} 

339 for name, parameter in parameters.items(): 

340 if not parameter.requires_grad: 

341 continue 

342 mesh = self._parameter_meshes.get(name) 

343 if mesh is None: 

344 mesh = self._expert_fsdp_mesh if name in self._expert_param_names else self._fsdp_mesh 

345 if mesh is None: 

346 raise ValueError("Parameter layouts require an initialized FSDP mesh or parameter binding") 

347 layouts[parameter] = self._build_parameter_layout( 

348 parameter, mesh, self.source_layout_for_parameter(name), sharded=name in sharded, 

349 unsharded_shape=self._parameter_shapes.get(name, tuple(parameter.shape)), 

350 ) 

351 return layouts 

352 

353 @staticmethod 

354 def _build_parameter_layout( 

355 parameter: torch.nn.Parameter, fsdp_mesh: DeviceMesh, source: Optional[tuple], *, 

356 sharded: bool, unsharded_shape: Tuple[int, ...], 

357 ) -> Layout: 

358 """Describe actual local storage without changing parameters or source metadata. 

359 

360 Args: 

361 parameter: Live parameter after physical FSDP slicing. 

362 fsdp_mesh: Selected data-parallel child mesh. 

363 source: Source TP/EP placements and mesh before FSDP. 

364 sharded: Whether the pass physically sliced this parameter. 

365 unsharded_shape: Local shape before FSDP padding and slicing. 

366 """ 

367 mesh = fsdp_mesh 

368 placements = [Replicate() for _ in range(mesh.ndim)] 

369 shape = tuple(unsharded_shape) 

370 axis = None 

371 if sharded: 

372 axis = next((index for index, name in enumerate(mesh.mesh_dim_names or ()) 

373 if name in ("fsdp_shard", "edp_shard", "fsdp", "dp", "dp_shard")), None) 

374 if axis is None or parameter.ndim == 0: 

375 raise ValueError("A physically sharded parameter requires a shard mesh axis and a non-scalar shape") 

376 shape = (parameter.shape[0] * mesh.size(axis), *shape[1:]) 

377 if source is not None and source[1] is not None: 

378 source_placements, source_mesh = source 

379 source_layout = Layout.from_device_mesh(source_mesh) 

380 source_layout.set_placements(tuple(source_placements)) 

381 source_layout.placement_to_tensor_map(len(shape)) 

382 shape = source_layout.get_global_shape(shape) 

383 mesh = DeviceMesh.concatenate([fsdp_mesh, source_mesh]) 

384 placements.extend(source_placements) 

385 if axis is not None: 

386 split_factor = math.prod(mesh.mesh_shape[index] for index, value in enumerate(placements) 

387 if index != axis and value.is_shard(0)) 

388 placements[axis] = StridedShard(0, split_factor) if split_factor > 1 else Shard(0) 

389 layout = Layout.from_device_mesh(mesh) 

390 layout.set_placements(tuple(placements)) 

391 layout.placement_to_tensor_map(len(shape)) 

392 layout.set_tensor_meta(tuple(shape), make_contiguous_strides_for(shape), parameter.dtype) 

393 return layout 

394 

395 def _bind_parameter_meshes(self) -> None: 

396 """Reuse eager unit discovery and managed identities without eager hooks.""" 

397 self._parameter_meshes = {} 

398 self._parameter_groups = {} 

399 if not isinstance(self._fsdp_config, FSDP2Config) or self._mesh_context is None: 

400 return 

401 manager = FSDP2Manager(self._fsdp_config, self._mesh_context) 

402 bindings = manager.parameter_mesh_bindings(self.model, self.source_shard_info) 

403 registered = {} 

404 for name, parameter in self.model.named_parameters(remove_duplicate=False): 

405 mesh = bindings.get(parameter) 

406 if mesh is None: 

407 continue 

408 self._parameter_meshes[name] = mesh 

409 shard = mesh[mesh.mesh_dim_names[-1]] 

410 ranks = tuple(shard.rank_list) 

411 if ranks not in registered: 

412 group = f"fsdp_unit_{len(registered)}" 

413 self._register_mesh_group(group, shard) 

414 registered[ranks] = group 

415 self._parameter_groups[name] = (registered[ranks], shard.size()) 

416 

417 def _init_device_mesh(self, mesh_context: Optional[Any] = None): 

418 """Initialize the DP process group(s) for the configured ``dp_mode``. 

419 

420 Four modes (see ``PassConfig.dp_mode``), matching simplefsdp's 

421 ``data_parallel`` and the eager ``fully_shard`` conventions: 

422 

423 * ``"fsdp"`` (default): a 1-D ``("fsdp_shard",)`` mesh; the group is 

424 registered under the graph-mode name ``"fsdp"``. 

425 * ``"ddp"``: a 1-D ``("fsdp_replicate",)`` mesh; registered under 

426 ``"dp_replicate"`` (no parameter sharding, grads all-reduced). 

427 * ``"hsdp"``: a 2-D ``("fsdp_replicate", "fsdp_shard")`` mesh 

428 (replicate, shard); both groups registered. 

429 * PP hybrid (``pp_enabled``): a 2-D ``(pp, DP)`` mesh — see 

430 ``_init_pp_dp_mesh``; the intra-stage DP group and the cross-stage 

431 ``"pp"`` group are both registered. 

432 

433 An external automodel ``MeshContext`` is reused when provided — its 

434 ``fsdp_replicate`` / ``fsdp_shard`` sub-meshes are registered under the 

435 graph-mode names, mirroring the eager ``FSDP2Manager``. The fallback 

436 builds the mesh with the repo ``init_device_mesh`` (the same builder the 

437 eager stack uses) and only supports pure DP (``tp_size == 1``) plus the 

438 PP hybrid; TP topologies must supply a ``MeshContext``. 

439 """ 

440 dp_mode = normalize_dp_mode(self.pass_config.dp_mode) 

441 if self.pass_config.pp_enabled: 

442 self._init_pp_dp_mesh(dp_mode, mesh_context) 

443 return 

444 

445 if mesh_context is not None: 

446 self._init_mesh_from_context(mesh_context, dp_mode) 

447 self._init_expert_dp(mesh_context) 

448 return 

449 

450 device_type = self._mesh_device_type() 

451 world_size = dist.get_world_size() 

452 rank_list = tuple(range(world_size)) 

453 

454 if dp_mode == "ddp": 

455 mesh = init_device_mesh( 

456 device_type, 

457 (world_size,), 

458 mesh_dim_names=("fsdp_replicate",), 

459 rank_list=rank_list, 

460 init_backend=True, 

461 ) 

462 self._fsdp_mesh = mesh 

463 self._register_mesh_group("dp_replicate", mesh["fsdp_replicate"]) 

464 self.pass_config.dp_replicate_degree = mesh["fsdp_replicate"].size() 

465 return 

466 

467 if dp_mode == "hsdp": 

468 shard, replicate = self._resolve_hsdp_degrees(world_size) 

469 mesh = init_device_mesh( 

470 device_type, 

471 (replicate, shard), 

472 mesh_dim_names=("fsdp_replicate", "fsdp_shard"), 

473 rank_list=rank_list, 

474 init_backend=True, 

475 ) 

476 self._fsdp_mesh = mesh 

477 self._register_mesh_group("dp_replicate", mesh["fsdp_replicate"]) 

478 self._register_mesh_group("fsdp", mesh["fsdp_shard"]) 

479 self.pass_config.dp_replicate_degree = mesh["fsdp_replicate"].size() 

480 self.pass_config.fsdp_degree = mesh["fsdp_shard"].size() 

481 return 

482 

483 # dp_mode == "fsdp" (default): 1-D shard mesh over the whole world. 

484 mesh = init_device_mesh( 

485 device_type, 

486 (world_size,), 

487 mesh_dim_names=("fsdp_shard",), 

488 rank_list=rank_list, 

489 init_backend=True, 

490 ) 

491 self._fsdp_mesh = mesh 

492 self._register_mesh_group("fsdp", mesh["fsdp_shard"]) 

493 # Back-fill, mirroring the external-mesh branch: FSDPPass resolves the 

494 # group size from ``fsdp_degree`` (falling back to world_size when 

495 # ``None``), so setting it here keeps the two paths consistent. 

496 self.pass_config.fsdp_degree = world_size 

497 

498 def _validate_ep_topology(self) -> None: 

499 """Limit graph EP to the synchronous, non-pipelined basic topology.""" 

500 context = self._mesh_context 

501 if context is None or getattr(context, "ep_size", 1) <= 1: 

502 return 

503 if any( 

504 getattr(context, axis, 1) != 1 for axis in ("cp_size", "pp_size") 

505 ): 

506 raise ValueError("Graph EP currently requires CP=PP=1") 

507 if self.pass_config.enable_overlap: 

508 raise ValueError("Graph EP currently requires enable_overlap=False") 

509 

510 def _init_expert_dp(self, context: Any) -> None: 

511 """Use EDP for the exact parameters owned by bound local experts.""" 

512 if getattr(context, "ep_size", 1) <= 1: 

513 return 

514 mesh = context.fsdp_moe_mesh 

515 if mesh is None or "edp_shard" not in mesh.mesh_dim_names: 

516 raise ValueError("Graph EP requires the expert data-parallel shard mesh") 

517 if self.pass_config.dp_mode != "fsdp": 

518 raise ValueError( 

519 "Graph EP currently supports sharded data parallelism without replication" 

520 ) 

521 if "edp_replicate" in mesh.mesh_dim_names and mesh["edp_replicate"].size() > 1: 

522 raise ValueError( 

523 "Graph EP currently does not support expert data-parallel replication" 

524 ) 

525 parameters = { 

526 id(parameter) 

527 for module in self.model.modules() 

528 if hasattr(module, "local_expert_count") 

529 for parameter in module.parameters(recurse=False) 

530 } 

531 self._expert_param_names = { 

532 name 

533 for name, parameter in self.model.named_parameters(remove_duplicate=False) 

534 if id(parameter) in parameters 

535 } 

536 if not self._expert_param_names: 

537 raise ValueError("Graph EP requires bound local expert parameters") 

538 self._expert_fsdp_mesh = mesh["edp_shard"] 

539 self._register_mesh_group("fsdp_expert", self._expert_fsdp_mesh) 

540 

541 def _init_pp_dp_mesh( 

542 self, dp_mode: str, mesh_context: Optional[Any] = None 

543 ) -> None: 

544 """Initialize the 2-D ``(pp, DP)`` mesh for the PP+FSDP/DDP hybrid. 

545 

546 Builds the mesh over the identity rank list (``rank = pp_idx * 

547 dp_size + dp_idx``, PP-outer) and registers the intra-stage DP group 

548 plus the cross-stage ``"pp"`` group that ``PpPass`` pipelines over. 

549 

550 Args: 

551 dp_mode: Normalized DP mode (``"fsdp"`` or ``"ddp"``). 

552 mesh_context: Must be ``None``; an external automodel 

553 ``MeshContext`` composed with PP is not supported yet. 

554 

555 Raises: 

556 ValueError: On a ``mesh_context``, ``"hsdp"``, missing degrees, or 

557 degrees that do not tile the world. 

558 """ 

559 if mesh_context is not None: 

560 raise ValueError( 

561 "PP composed with an external automodel MeshContext is not " 

562 "supported yet; run the PP hybrid with the self-built mesh " 

563 "(omit mesh_context and set pp_degree plus fsdp_degree / " 

564 "dp_replicate_degree)" 

565 ) 

566 if dp_mode == "hsdp": 

567 raise ValueError( 

568 "PP + HSDP requires a 3-D (pp, dp_replicate, fsdp_shard) " 

569 "mesh, which is not supported yet; run PP with dp_mode='fsdp' " 

570 "or 'ddp'" 

571 ) 

572 pp_degree, dp_degree, dp_mesh_name = self._resolve_pp_dp_degrees(dp_mode) 

573 world_size = dist.get_world_size() 

574 mesh = init_device_mesh( 

575 self._mesh_device_type(), 

576 (pp_degree, dp_degree), 

577 mesh_dim_names=("pp", dp_mesh_name), 

578 rank_list=tuple(range(world_size)), 

579 init_backend=True, 

580 ) 

581 self._register_pp_dp_groups(mesh, dp_mode, dp_mesh_name, pp_degree, dp_degree) 

582 _LOG.info( 

583 "PP+DP mesh: pp_degree=%s dp_mode=%s dp_degree=%s world_size=%s", 

584 pp_degree, 

585 dp_mode, 

586 dp_degree, 

587 world_size, 

588 ) 

589 

590 def _resolve_pp_dp_degrees(self, dp_mode: str) -> Tuple[int, int, str]: 

591 """Resolve and validate the ``(pp, DP)`` degrees and DP dim name.""" 

592 world_size = dist.get_world_size() 

593 pp_degree = self.pass_config.pp_degree 

594 if dp_mode == "ddp": 

595 dp_degree = self.pass_config.dp_replicate_degree 

596 dp_mesh_name = "dp_replicate" 

597 else: 

598 dp_degree = self.pass_config.fsdp_degree 

599 dp_mesh_name = "fsdp_shard" 

600 if pp_degree is None and dp_degree is None: 

601 raise ValueError( 

602 "PP + DP hybrid requires pp_degree together with (or instead " 

603 "of) fsdp_degree / dp_replicate_degree so the 2-D mesh can be " 

604 "sized; got both None" 

605 ) 

606 if pp_degree is None: 

607 if world_size % dp_degree != 0: 

608 raise ValueError( 

609 f"world_size ({world_size}) not divisible by the DP degree " 

610 f"({dp_degree}); cannot derive pp_degree" 

611 ) 

612 pp_degree = world_size // dp_degree 

613 if dp_degree is None: 

614 if world_size % pp_degree != 0: 

615 raise ValueError( 

616 f"world_size ({world_size}) not divisible by pp_degree " 

617 f"({pp_degree}); cannot derive the DP degree" 

618 ) 

619 dp_degree = world_size // pp_degree 

620 if pp_degree * dp_degree != world_size: 

621 raise ValueError( 

622 f"PP+DP degrees must tile the world: pp_degree ({pp_degree}) " 

623 f"x dp_degree ({dp_degree}) != world_size ({world_size})" 

624 ) 

625 return pp_degree, dp_degree, dp_mesh_name 

626 

627 def _register_pp_dp_groups( 

628 self, mesh: Any, dp_mode: str, dp_mesh_name: str, pp_degree: int, dp_degree: int 

629 ) -> None: 

630 """Register both hybrid groups and back-fill the pass-config degrees.""" 

631 self._fsdp_mesh = mesh[dp_mesh_name] 

632 self._register_mesh_group("pp", mesh["pp"]) 

633 self._register_mesh_group( 

634 "fsdp" if dp_mesh_name == "fsdp_shard" else "dp_replicate", 

635 mesh[dp_mesh_name], 

636 ) 

637 self._pp_group_name = "pp" 

638 self.pass_config.pp_degree = pp_degree 

639 if dp_mode == "ddp": 

640 self.pass_config.dp_replicate_degree = dp_degree 

641 else: 

642 self.pass_config.fsdp_degree = dp_degree 

643 

644 @staticmethod 

645 def _mesh_device_type() -> str: 

646 """Return the backend device type for the self-built mesh.""" 

647 if hasattr(torch, "npu") and torch.npu.is_available(): 

648 return "npu" 

649 return "cpu" 

650 

651 def _init_mesh_from_context(self, mesh_context: Any, dp_mode: str, *, register_groups: bool = True) -> None: 

652 """Resolve and register the DP group(s) from an automodel ``MeshContext``. 

653 

654 Reuses automodel's already-built mesh (boundary forwards hold the TP 

655 group object directly) and registers the ``fsdp_replicate`` / 

656 ``fsdp_shard`` sub-meshes under the graph-mode names ``"dp_replicate"`` 

657 / ``"fsdp"``. The effective mode is first reconciled with the mesh 

658 topology (see ``_reconcile_mode_with_context``), mirroring the eager 

659 ``FSDP2Manager`` which picks HSDP from ``dp_replicate_size > 1``. 

660 

661 Args: 

662 mesh_context: Existing data/model parallel topology. 

663 dp_mode: Requested data-parallel mode. 

664 register_groups: False for metadata-only access in plain graph mode. 

665 """ 

666 fsdp_mesh = ( 

667 getattr(mesh_context, "fsdp_non_moe_mesh", None) or mesh_context.device_mesh 

668 ) 

669 names = tuple(getattr(fsdp_mesh, "mesh_dim_names", ()) or ()) 

670 dp_mode = self._reconcile_mode_with_context(mesh_context, dp_mode) 

671 self.pass_config.dp_mode = dp_mode 

672 

673 if dp_mode == "ddp": 

674 dim = self._first_dim(names, ("fsdp_replicate", "dp_replicate", "dp")) 

675 if dim is None: 

676 raise ValueError( 

677 "dp_mode='ddp' needs a replicate axis among " 

678 "('fsdp_replicate', 'dp_replicate', 'dp') in the mesh " 

679 f"dims {names}" 

680 ) 

681 sub = fsdp_mesh[dim] 

682 self._fsdp_mesh = sub 

683 if register_groups: 

684 self._register_mesh_group("dp_replicate", sub) 

685 self.pass_config.dp_replicate_degree = sub.size() 

686 return 

687 

688 if dp_mode == "hsdp": 

689 rep_dim = self._first_dim(names, ("fsdp_replicate", "dp_replicate")) 

690 shard_dim = self._first_dim(names, ("fsdp_shard", "fsdp", "dp_shard", "dp")) 

691 if rep_dim is None or shard_dim is None: 

692 raise ValueError( 

693 "dp_mode='hsdp' needs both a replicate axis " 

694 "('fsdp_replicate'/'dp_replicate') and a shard axis " 

695 f"('fsdp_shard'/'fsdp'/'dp') in the mesh dims {names}" 

696 ) 

697 rep_sub = fsdp_mesh[rep_dim] 

698 shard_sub = fsdp_mesh[shard_dim] 

699 if shard_dim == "dp_shard" and "cp" in names: 

700 shard_sub = fsdp_mesh[(shard_dim, "cp")].flatten("fsdp_shard") 

701 self._fsdp_mesh = DeviceMesh.concatenate([rep_sub, shard_sub]) 

702 else: 

703 self._fsdp_mesh = fsdp_mesh[(rep_dim, shard_dim)] 

704 if register_groups: 

705 self._register_mesh_group("dp_replicate", rep_sub) 

706 self._register_mesh_group("fsdp", shard_sub) 

707 self.pass_config.dp_replicate_degree = rep_sub.size() 

708 self.pass_config.fsdp_degree = shard_sub.size() 

709 return 

710 

711 # dp_mode == "fsdp": automodel's fsdp_non_moe_mesh is 

712 # ("fsdp_replicate","fsdp_shard","tp"); device_mesh (cp=1) is 

713 # ("dp","cp","tp") and "dp" is the FSDP axis. 

714 dim = self._first_dim(names, ("fsdp_shard", "dp_shard", "dp")) 

715 sub = fsdp_mesh[dim] 

716 if dim == "dp_shard" and "cp" in names: 

717 sub = fsdp_mesh[(dim, "cp")].flatten("fsdp_shard") 

718 self._fsdp_mesh = sub 

719 if register_groups: 

720 self._register_mesh_group("fsdp", sub) 

721 self.pass_config.fsdp_degree = sub.size() 

722 

723 def _reconcile_mode_with_context(self, mesh_context: Any, dp_mode: str) -> str: 

724 """Align ``dp_mode`` with the ``MeshContext`` topology (eager parity). 

725 

726 The eager ``fully_shard`` path chooses FSDP vs HSDP from the mesh 

727 (``dp_replicate_size > 1`` means HSDP), not from a config flag. Mirror 

728 that so a mesh carrying a replicate axis is never silently treated as 

729 plain FSDP, and a 1-wide replicate axis is treated as FSDP. 

730 

731 Raises: 

732 ValueError: When ``dp_mode='ddp'`` but the mesh shards 

733 (``dp_shard_size > 1``). 

734 """ 

735 rep = getattr(mesh_context, "dp_replicate_size", None) 

736 shard = getattr(mesh_context, "dp_shard_size", None) 

737 if not isinstance(rep, int): 

738 # Not a real MeshContext (e.g. a test stub): leave the mode as-is. 

739 return dp_mode 

740 if dp_mode == "fsdp" and rep > 1: 

741 _LOG.warning( 

742 "dp_mode='fsdp' but MeshContext.dp_replicate_size=%d > 1; " 

743 "using 'hsdp' to match the mesh", 

744 rep, 

745 ) 

746 return "hsdp" 

747 if dp_mode == "hsdp" and rep <= 1: 

748 _LOG.warning( 

749 "dp_mode='hsdp' but MeshContext.dp_replicate_size=%d <= 1; " 

750 "using 'fsdp'", 

751 rep, 

752 ) 

753 return "fsdp" 

754 if dp_mode == "ddp" and isinstance(shard, int) and shard > 1: 

755 raise ValueError( 

756 "dp_mode='ddp' requires a non-sharding mesh " 

757 f"(dp_shard_size=1), got dp_shard_size={shard}; use 'hsdp'" 

758 ) 

759 return dp_mode 

760 

761 @staticmethod 

762 def _first_dim(names: tuple, candidates: tuple) -> Optional[str]: 

763 """Return the first candidate axis present in ``names`` (or None).""" 

764 for cand in candidates: 

765 if cand in names: 

766 return cand 

767 return None 

768 

769 @staticmethod 

770 def _register_mesh_group(name: str, sub_mesh: Any) -> None: 

771 """Register ``sub_mesh``'s process group under ``name``.""" 

772 _register_process_group(name, sub_mesh.get_group()) 

773 

774 def _resolve_hsdp_degrees(self, world_size: int) -> Tuple[int, int]: 

775 """Resolve ``(shard, replicate)`` for a fallback 2-D HSDP mesh. 

776 

777 Either degree may be left ``None`` and is derived from ``world_size`` 

778 and the other; both must be consistent with ``world_size``. 

779 

780 Raises: 

781 ValueError: When both are ``None``, a degree does not divide 

782 ``world_size``, or the product does not equal ``world_size``. 

783 """ 

784 shard = self.pass_config.fsdp_degree 

785 replicate = self.pass_config.dp_replicate_degree 

786 if shard is None and replicate is None: 

787 raise ValueError( 

788 "dp_mode='hsdp' requires fsdp_degree and/or " 

789 "dp_replicate_degree to build the 2-D (dp_replicate, fsdp) mesh" 

790 ) 

791 if shard is None: 

792 if world_size % replicate != 0: 

793 raise ValueError( 

794 f"world_size ({world_size}) not divisible by " 

795 f"dp_replicate_degree ({replicate})" 

796 ) 

797 shard = world_size // replicate 

798 if replicate is None: 

799 if world_size % shard != 0: 

800 raise ValueError( 

801 f"world_size ({world_size}) not divisible by fsdp_degree ({shard})" 

802 ) 

803 replicate = world_size // shard 

804 if shard * replicate != world_size: 

805 raise ValueError( 

806 f"hsdp degrees fsdp_degree={shard} x " 

807 f"dp_replicate_degree={replicate} != world_size={world_size}" 

808 ) 

809 return shard, replicate 

810 

811 def _build_pass_kwargs(self) -> dict: 

812 """ 

813 Build kwargs to pass to passes. 

814 """ 

815 kwargs = {} 

816 

817 # Live model: partitioning passes (FSDPPass) physically shard 

818 # parameters in place, keeping the compiler FSDP-agnostic. 

819 kwargs["model"] = self.model 

820 

821 if self.pass_config.fsdp_enabled: 

822 kwargs["fsdp_group_name"] = "fsdp" 

823 if self._parameter_groups: 

824 kwargs["parameter_groups"] = self._parameter_groups 

825 if self._expert_param_names: 

826 kwargs["expert_param_names"] = self._expert_param_names 

827 kwargs["expert_fsdp_group_name"] = "fsdp_expert" 

828 kwargs.update(self._source_reduction_kwargs()) 

829 # DDP / HSDP also reduce gradients on the replicate axis. 

830 if normalize_dp_mode(self.pass_config.dp_mode) in ("ddp", "hsdp"): 

831 kwargs["dp_replicate_group_name"] = "dp_replicate"