Coverage for /home/jenkins/.local/lib/python3.10/site-packages/hyper

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

7# http://www.apache.org/licenses/LICENSE-2.0

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.

18Compiler surface of the graph-mode stack:

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

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

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

28class and owns them.

29"""

31from __future__ import annotations

33__all__ = ["GraphCompiler"]

35import logging

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

38import torch

39import torch.distributed as dist

40from torch.distributed.distributed_c10d import _register_process_group

42from hyper_parallel.core.dtensor.device_mesh import init_device_mesh

44from .pass_config import PassConfig, normalize_dp_mode, build_pass_config_from_trainer_config

45from .graph_parallel_plan import GraphParallelPlan

46from .passes.pipeline import PassPipeline

47from .tracer.dynamic_shapes import normalize_dynamic_arg_dims

48from .tracer.graph_tracer import run_traced_graph, trace_model_graph

50_LOG = logging.getLogger(__name__)

53class GraphCompiler:

54 """

55 Graph-mode Compiler

57 Users provide model code and parallel configuration.

58 The compiler handles graph capture and all parallel logic

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

60 """

62 def __init__(

63 self,

64 model: torch.nn.Module,

65 train_fn: Callable,

66 pass_config: Optional[PassConfig] = None,

67 parallel_plan: Optional[GraphParallelPlan] = None,

68 trainer_config: Optional[Any] = None,

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

70 mesh_context: Optional[Any] = None,

71 dynamic: Optional[bool] = None,

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

73 ) -> None:

74 """

75 Args:

76 model: Model to compile

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

78 pass_config: Parallel configuration. When omitted, the compiler

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

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

81 parallel_plan: GraphParallelPlan declaring which modules to shard

82 (optional; enables declarative sharding)

83 trainer_config: Optional trainer topology used only to infer a

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

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

86 to the NPU device when available, otherwise CPU.

87 mesh_context: Optional automodel ``MeshContext`` carrying a

88 pre-built TP/FSDP mesh. When provided, the TP group is reused

89 as-is (boundary forwards already hold the group object) and

90 only the FSDP shard sub-mesh is registered under ``"fsdp"``.

91 Use this to feed an automodel TP-sharded model into the

92 graph-mode FSDP pass.

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

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

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

96 """

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

98 self.dynamic = getattr(compile_config, "dynamic", False) if dynamic is None else dynamic

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

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

101 if dynamic_arg_dims is None:

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

103 self.dynamic_arg_dims = normalize_dynamic_arg_dims(dynamic_arg_dims)

104 self.model = model

105 self.train_fn = train_fn

106 self.parallel_plan = parallel_plan

107 self.pass_config = self._resolve_pass_config(

108 pass_config=pass_config,

109 trainer_config=trainer_config,

110 )

111 self._mesh_context = mesh_context

112 self.device = device or (

113 torch.device("npu")

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

115 else torch.device("cpu")

116 )

117

118 self.pass_config.validate()

119

120 self._joint_graph = None

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

122

123 @staticmethod

124 def _resolve_pass_config(

125 *,

126 pass_config: Optional[PassConfig],

127 trainer_config: Optional[Any],

128 ) -> PassConfig:

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

130 if pass_config is not None:

131 return pass_config

132 if trainer_config is not None:

133 return build_pass_config_from_trainer_config(trainer_config)

134 return PassConfig()

135

136 @property

137 def is_compiled(self) -> bool:

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

139 return self._joint_graph is not None

140

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

142 """

143 Compile model into parallel graph

144

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

146 at first forward_backward

147

148 Args:

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

150 arguments and used to trace the joint graph

151 """

152 if (self.dynamic or self.dynamic_arg_dims is not None) and self.pass_config.pp_enabled:

153 raise ValueError(

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

155 )

156 if self._pytree_pre_hook is not None:

157 self._pytree_pre_hook()

158

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

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

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

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

163 # runs as plain graph mode without sharding.

164 self._init_device_mesh(self._mesh_context)

165

166 trace_kwargs = {}

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

168 trace_kwargs = {"dynamic": self.dynamic, "dynamic_arg_dims": self.dynamic_arg_dims}

169 joint_graph = trace_model_graph(self.model, self.train_fn, inputs, **trace_kwargs)

170

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

172

173 pass_kwargs = self._build_pass_kwargs()

174

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

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

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

178

179 if joint_graph.input_guards is not None:

180 joint_graph.input_guards.refresh()

181 self._joint_graph = joint_graph

182

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

184 """

185 Execute one compiled forward+backward step.

186

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

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

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

190

191 Args:

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

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

194

195 Returns:

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

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

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

199 tensor).

200 """

201 if self._joint_graph is None:

202 self.compile(**inputs)

203

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

205

206 self._accumulate_grads(grads)

207

208 return loss, loss_dict

209

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

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

212 self._pytree_pre_hook = hook

213 return self

214

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

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

217 self.device = torch.device(device)

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

219 return self

220

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

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

223

224 Three modes (see ``PassConfig.dp_mode``), matching simplefsdp's

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

226

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

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

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

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

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

232 (replicate, shard); both groups registered.

233

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

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

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

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

238 eager stack uses) and only supports **pure DP** (``tp_size == 1``, PP

239 off); TP/PP topologies must supply a ``MeshContext``.

240 """

241 dp_mode = normalize_dp_mode(self.pass_config.dp_mode)

242 if mesh_context is not None:

243 self._init_mesh_from_context(mesh_context, dp_mode)

244 return

245

246 self._require_pure_dp_fallback()

247 device_type = (

248 "npu" if (hasattr(torch, "npu") and torch.npu.is_available()) else "cpu"

249 )

250 world_size = dist.get_world_size()

251 rank_list = tuple(range(world_size))

252

253 if dp_mode == "ddp":

254 mesh = init_device_mesh(

255 device_type,

256 (world_size,),

257 mesh_dim_names=("fsdp_replicate",),

258 rank_list=rank_list,

259 init_backend=True,

260 )

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

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

263 return

264

265 if dp_mode == "hsdp":

266 shard, replicate = self._resolve_hsdp_degrees(world_size)

267 mesh = init_device_mesh(

268 device_type,

269 (replicate, shard),

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

271 rank_list=rank_list,

272 init_backend=True,

273 )

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

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

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

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

278 return

279

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

281 mesh = init_device_mesh(

282 device_type,

283 (world_size,),

284 mesh_dim_names=("fsdp_shard",),

285 rank_list=rank_list,

286 init_backend=True,

287 )

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

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

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

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

292 self.pass_config.fsdp_degree = world_size

293

294 def _require_pure_dp_fallback(self) -> None:

295 """Reject TP/PP topologies for the self-built DP mesh.

296

297 The fallback mesh spans the whole world; with a TP or PP axis it would

298 shard/reduce across the wrong ranks. Those topologies must supply an

299 automodel ``MeshContext`` (which carries the proper DP sub-mesh).

300

301 Raises:

302 ValueError: When ``tp_size != 1`` or ``pp_enabled``.

303 """

304 cfg = self.pass_config

305 if cfg.tp_size != 1 or cfg.pp_enabled:

306 raise ValueError(

307 "Graph-mode DP without a MeshContext only supports pure data "

308 f"parallelism (tp_size=1, PP disabled); got tp_size={cfg.tp_size}, "

309 f"pp_enabled={cfg.pp_enabled}. Pass a MeshContext carrying the DP "

310 "sub-mesh for TP/PP hybrids."

311 )

312

313 def _init_mesh_from_context(self, mesh_context: Any, dp_mode: str) -> None:

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

315

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

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

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

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

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

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

322 """

323 fsdp_mesh = (

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

325 )

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

327 dp_mode = self._reconcile_mode_with_context(mesh_context, dp_mode)

328 self.pass_config.dp_mode = dp_mode

329

330 if dp_mode == "ddp":

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

332 if dim is None:

333 raise ValueError(

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

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

336 f"dims {names}"

337 )

338 sub = fsdp_mesh[dim]

339 self._register_mesh_group("dp_replicate", sub)

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

341 return

342

343 if dp_mode == "hsdp":

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

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

346 if rep_dim is None or shard_dim is None:

347 raise ValueError(

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

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

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

351 )

352 rep_sub = fsdp_mesh[rep_dim]

353 shard_sub = fsdp_mesh[shard_dim]

354 self._register_mesh_group("dp_replicate", rep_sub)

355 self._register_mesh_group("fsdp", shard_sub)

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

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

358 return

359

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

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

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

363 dim = "fsdp_shard" if "fsdp_shard" in names else "dp"

364 sub = fsdp_mesh[dim]

365 self._register_mesh_group("fsdp", sub)

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

367

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

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

370

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

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

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

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

375

376 Raises:

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

378 (``dp_shard_size > 1``).

379 """

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

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

382 if not isinstance(rep, int):

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

384 return dp_mode

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

386 _LOG.warning(

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

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

389 rep,

390 )

391 return "hsdp"

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

393 _LOG.warning(

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

395 "using 'fsdp'",

396 rep,

397 )

398 return "fsdp"

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

400 raise ValueError(

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

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

403 )

404 return dp_mode

405

406 @staticmethod

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

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

409 for cand in candidates:

410 if cand in names:

411 return cand

412 return None

413

414 @staticmethod

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

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

417 _register_process_group(name, sub_mesh.get_group())

418

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

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

421

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

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

424

425 Raises:

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

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

428 """

429 shard = self.pass_config.fsdp_degree

430 replicate = self.pass_config.dp_replicate_degree

431 if shard is None and replicate is None:

432 raise ValueError(

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

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

435 )

436 if shard is None:

437 if world_size % replicate != 0:

438 raise ValueError(

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

440 f"dp_replicate_degree ({replicate})"

441 )

442 shard = world_size // replicate

443 if replicate is None:

444 if world_size % shard != 0:

445 raise ValueError(

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

447 )

448 replicate = world_size // shard

449 if shard * replicate != world_size:

450 raise ValueError(

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

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

453 )

454 return shard, replicate

455

456 def _build_pass_kwargs(self) -> dict:

457 """

458 Build kwargs to pass to passes.

459 """

460 kwargs = {}

461

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

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

464 kwargs["model"] = self.model

465

466 if self.pass_config.fsdp_enabled:

467 kwargs["fsdp_group_name"] = "fsdp"

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

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

470 kwargs["dp_replicate_group_name"] = "dp_replicate"

471

472 return kwargs

473

474 def _run_graph(self, **inputs):

475 """Execute compiled graph"""

476 if self._joint_graph is None:

477 raise RuntimeError(

478 "Graph not compiled. Call compiler.compile() or "

479 "compiler.forward_backward() first."

480 )

481

482 # The joint graph's parameters/buffers are static inputs: feed the

483 # live (FSDP-sharded) model state in FQN order each step.

484 return run_traced_graph(

485 self._joint_graph,

486 self.model,

487 inputs,

488 )

489

490 def _accumulate_grads(self, grads: List[torch.Tensor]) -> None:

491 """Accumulate graph-computed gradients into the live model's parameters.

492

493 The graph emits gradients in ``state_fqns`` order (trainable

494 parameters only, shared parameters included once per FQN), which

495 diverges from ``model.parameters()`` (deduplicated) when the model

496 ties weights. Mapping by FQN keeps every gradient on the right

497 parameter; the count check refuses to assign on mismatch instead of

498 letting ``zip`` silently truncate.

499

500 Accumulation (not overwrite) keeps ``forward_backward`` composable:

501 several micro-batch steps may run before the caller's optimizer step

502 (which ends with ``zero_grad``), so per-step grads must sum into

503 ``param.grad``.

504 """

505 # ``state_is_param`` is attached to the traced GraphModule by

506 # ``trace_model_graph``, not to the JointGraph dataclass itself.

507 state_is_param = getattr(self._joint_graph.graph_module, "state_is_param", None)

508 fqn_to_param = dict(self.model.named_parameters(remove_duplicate=False))

509 trainable = self._trainable_params_in_state_order(

510 self._joint_graph.state_fqns, state_is_param, fqn_to_param

511 )

512 if len(trainable) != len(grads):

513 raise ValueError(

514 f"Gradient count ({len(grads)}) does not match trainable "

515 f"parameter count ({len(trainable)}). The traced graph and "

516 f"the live model disagree on which parameters are trainable; "

517 f"refusing to assign gradients to avoid silent misalignment."

518 )

519

520 for param, grad in zip(trainable, grads):

521 if param.grad is None:

522 param.grad = grad

523 else:

524 param.grad += grad

525

526 @staticmethod

527 def _trainable_params_in_state_order(

528 state_fqns: List[str],

529 state_is_param: Optional[List[bool]],

530 fqn_to_param: Dict[str, torch.nn.Parameter],

531 ) -> List[torch.nn.Parameter]:

532 """Return the live trainable parameters in the graph's state order.

533

534 Mirrors the tracer's gradient emission order (``state_fqns`` order,

535 parameters only, trainable only, shared parameters kept per FQN), so

536 gradient ``i`` belongs to the returned parameter ``i``. Buffers share

537 ``state_fqns`` but are absent from the parameter lookup; they are

538 skipped explicitly so a missing ``state_is_param`` flag (old traces)

539 degrades to parameter-only instead of raising KeyError.

540 """

541 trainable: List[torch.nn.Parameter] = []

542 for idx, fqn in enumerate(state_fqns):

543 if state_is_param is not None and not state_is_param[idx]:

544 continue

545 param = fqn_to_param.get(fqn)

546 if param is not None and param.requires_grad:

547 trainable.append(param)

548 return trainable

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

201 statements