Coverage for /home/jenkins/.local/lib/python3.10/site-packages/hyper_parallel/compile/passes/inductor

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

16Inductor Compilation Pass - terminal backend compilation of the joint graph.

18The pass lowers the (already partitioned and overlapped) joint FX graph to

19TorchInductor kernels via ``torch.fx.passes.regional_inductor``, which scoops

20nodes annotated with ``meta["custom"]["compile_with_inductor"]`` into regions

21and compiles each region with ``torch._inductor.standalone_compile``. This

22brings kernel fusion, c10d functionalization, PG unboxing and Inductor caching

23to the graph-mode stack without touching eager semantics.

25Modes (``PassConfig.inductor_mode``):

27- ``"off"`` (default): the pass is not added to the pipeline.

28- ``"regional"``: tag functional compute nodes as compile regions while

29 keeping communication (``_c10d_functional`` collectives + ``wait_tensor``)

30 and side-effectful nodes (e.g. the FSDP reshard storage free) eager. This

31 preserves the comm scheduling surface for the overlap passes.

32- ``"full"``: tag every node so the whole graph compiles as one region

33 (c10d functional collectives included), mirroring torchtitan's

34 ``full_inductor_compilation_pass``.

36Pipeline-parallel support: when ``PpPass`` has already split the joint graph,

37the outer graph is a ``call_module`` stub around the pipeline schedule. The

38compute lives in the schedule's per-virtual-stage forward/backward subgraphs,

39so the pass detects the stub and compiles every stage subgraph independently

40(both modes), instead of compiling — or rejecting — the stub. Cross-stage

41fusion is inherently impossible; the schedule's P2P stays outside the

42compiled regions. Unknown ``call_module`` graphs are still rejected. PP

43forward subgraphs export saved values verbatim, which can include multi-result

44ops (``nll_loss_forward``); their nested output is split into ``getitem``

45leaves for compilation and rebuilt by a thin wrapper, because AOTAutograd's

46regional path requires flat outputs.

48Numerics policy: ``inductor.eager_numerics.division_rounding`` is forced on

49(where the running torch provides it, 2.11+) so compiled regions stay

50bitwise-equal to eager execution (same guarantee torchtitan relies on).

52Compatibility: ``torch.fx.passes.regional_inductor`` only exists in newer

53torch builds (it is absent from 2.6 / 2.7 / 2.9). On such builds the pass logs

54a warning, leaves the graph untouched and records ``inductor_compiled =

55"unavailable"`` instead of raising ``ModuleNotFoundError``.

57Contract notes:

59- The pass must run LAST: after compilation the FX graph is a thin wrapper

60 around the compiled callable and is no longer authoritative for

61 graph-level passes. PP stage subgraphs are compiled through the schedule's

62 ``map_stage_graphs`` hook, so the stub remains the outer entry point.

63- ``regional_inductor`` switches a compiled graph module to the boxed calling

64 convention (``forward(self, args_list)``); the pass resets it to the

65 default positional convention so ``graph_module(*flat_inputs)`` (and the

66 schedule's direct ``fwd_gm(*args)`` / ``bwd_gm(*args)`` calls) keep

67 working unchanged.

68- The graph must be traced with ``make_fx`` (placeholders carry fake

69 ``meta["val"]``). Nodes inserted by earlier passes lack fake values, so the

70 pass first re-propagates them with ``FakeTensorProp`` using the graph's own

71 ``FakeTensorMode``.

72"""

74__all__ = ["InductorPass"]

76import importlib

77import logging

78import operator

79from typing import Any, Dict, Optional, Sequence

81import torch

82from torch import fx, nn

83from torch._subclasses import FakeTensor

84from torch.fx.graph import CodeGen

85from torch.fx.node import _side_effectful_functions

86from torch.fx.passes.fake_tensor_prop import FakeTensorProp

88from ..pass_config import PassConfig

89from .base import GraphPass

91_LOG = logging.getLogger(__name__)

93# Communication nodes stay outside compiled regions in "regional" mode. The

94# launch/wait pair is the scheduling surface the overlap passes reorder; a

95# compiled region would fuse a collective into surrounding compute and remove

96# the explicit reorder point. "full" mode does compile them (Inductor's c10d

97# functionalization supports that), matching torchtitan's full pass.

98# ``str(OpOverload)`` has no ``torch.ops.`` prefix, so match the bare namespace.

99_C10D_FUNCTIONAL_NS = "_c10d_functional."

100

101# Meta key recording that (and how) a graph module was compiled.

102_INDUCTOR_COMPILED_META = "inductor_compiled"

103

104# Meta value when the running torch predates regional_inductor (<= 2.9.x).

105_INDUCTOR_UNAVAILABLE = "unavailable"

106

107

108def _regional_inductor_available() -> bool:

109 """Whether this torch build ships ``torch.fx.passes.regional_inductor``.

110

111 The module first appeared after torch 2.9, while this project still

112 supports 2.6 / 2.7 / 2.9. Importing it lazily here lets the pass degrade to

113 a documented no-op instead of raising ``ModuleNotFoundError`` on those

114 builds.

115 """

116 try:

117 importlib.import_module("torch.fx.passes.regional_inductor")

118 except ImportError:

119 return False

120 return True

121

122

123def _get_fake_mode_from_gm(gm: fx.GraphModule) -> Optional[Any]:

124 """Return the ``FakeTensorMode`` shared by the graph's placeholder vals."""

125 for node in gm.graph.nodes:

126 if node.op == "placeholder":

127 val = node.meta.get("val")

128 if isinstance(val, FakeTensor):

129 return val.fake_mode

130 return None

131

132

133def _placeholder_fake_args(gm: fx.GraphModule) -> list:

134 """Return the placeholder values in order, requiring fake ``meta["val"]``."""

135 args = []

136 for node in gm.graph.find_nodes(op="placeholder"):

137 val = node.meta.get("val")

138 if val is None:

139 raise ValueError(

140 f"InductorPass requires a fake meta['val'] on placeholder "

141 f"'{node.name}' to size compiled regions; the graph must be "

142 f"make_fx-traced (PP stage subgraphs copy their placeholders' "

143 f"meta from the joint graph)"

144 )

145 args.append(val)

146 return args

147

148

149def _propagate_fake_vals(gm: fx.GraphModule) -> fx.GraphModule:

150 """Fill ``meta["val"]`` on nodes inserted after tracing (FSDP collectives,

151 pad views, reshard frees) so ``regional_inductor`` can size each region.

152

153 Propagation reuses the graph's own ``FakeTensorMode``: placeholder vals

154 already belong to it and mixing modes is unsupported by ``FakeTensorProp``.

155 Side-effectful free nodes return ``None`` and legitimately carry no val.

156 """

157 fake_mode = _get_fake_mode_from_gm(gm)

158 if fake_mode is None:

159 raise ValueError(

160 "InductorPass requires a make_fx-traced graph whose placeholders "

161 "carry fake meta['val'] tensors"

162 )

163 fake_args = _placeholder_fake_args(gm)

164 FakeTensorProp(gm, fake_mode).propagate_dont_convert_inputs(*fake_args)

165 return gm

166

167

168def _node_is_comm(node: fx.Node) -> bool:

169 """Whether the node is a functional collective or its wait."""

170 return str(node.target).startswith(_C10D_FUNCTIONAL_NS)

171

172

173def _node_is_side_effectful(node: fx.Node) -> bool:

174 """Whether the node targets a function registered as side-effectful.

175

176 ``FSDPPass`` registers the reshard storage-free op there so DCE keeps it;

177 compiling such a node would let Inductor drop or defer the free. The

178 ``ao::`` offload/reload/wait ops are also kept eager: they own a transfer

179 stream and an event registry outside the compiled region, and rewriting

180 them as fused kernels would break the host/device synchronization.

181 """

182 if node.op != "call_function":

183 return False

184 if node.target in _side_effectful_functions:

185 return True

186 return str(node.target).startswith("ao.")

187

188

189def _flatten_nested_outputs(

190 graph_module: fx.GraphModule,

191) -> tuple:

192 """Rewrite tuple/list graph outputs into a flat tuple of leaf values.

193

194 AOTAutograd's regional compilation expects flat tensor outputs. The joint

195 graph is flat, but PP stage forward subgraphs export every value the

196 backward references verbatim — including multi-result ops such as

197 ``nll_loss_forward`` — so their output tuple can contain a nested

198 ``(loss, total_weight)`` segment, which trips AOT's output-spec check.

199 Splitting nested values into ``getitem`` leaves keeps every region output

200 flat; the returned spec lets the caller rebuild the original structure

201 for the schedule / backward subgraph, whose contract is unchanged.

202

203 Returns:

204 ``(spec, graph_module)`` where ``spec`` is ``None`` when the output is

205 already flat, otherwise a nested marker tree consumed by

206 ``_rebuild_nested_output``.

207 """

208 graph = graph_module.graph

209 output = next((n for n in graph.nodes if n.op == "output"), None)

210 returned = output.args[0] if output is not None else None

211 if not isinstance(returned, (tuple, list)):

212 return None, graph_module

213 if not any(

214 isinstance(node.meta.get("val"), (tuple, list))

215 for node in returned

216 if isinstance(node, fx.Node)

217 ):

218 return None, graph_module

219

220 flat_nodes: list = []

221

222 def build(node: fx.Node) -> tuple:

223 """Emit ``getitem`` leaves for a nested value, else record the leaf."""

224 val = node.meta.get("val")

225 if isinstance(val, (tuple, list)) and val:

226 children = []

227 for index, element in enumerate(val):

228 sub = graph.call_function(operator.getitem, (node, index))

229 sub.meta["val"] = element

230 children.append(build(sub))

231 kind = "tuple" if isinstance(val, tuple) else "list"

232 return kind, tuple(children)

233 flat_nodes.append(node)

234 return "leaf", len(flat_nodes) - 1

235

236 with graph.inserting_before(output):

237 spec = tuple(build(node) for node in returned)

238 output.args = (tuple(flat_nodes),)

239 graph.lint()

240 graph_module.recompile()

241 return spec, graph_module

242

243

244def _rebuild_nested_output(spec: tuple, flat: Sequence[Any]) -> list:

245 """Rebuild the original output structure from ``spec`` and flat leaves."""

246 leaves = iter(flat)

247 rebuilt = []

248 for kind, payload in spec:

249 if kind == "leaf":

250 rebuilt.append(next(leaves))

251 else:

252 children = _rebuild_nested_output(payload, leaves)

253 rebuilt.append(tuple(children) if kind == "tuple" else children)

254 return rebuilt

255

256

257class _NestedOutputWrapper(nn.Module):

258 """Thin wrapper restoring a compiled graph's nested output structure."""

259

260 def __init__(

261 self,

262 inner: nn.Module,

263 spec: tuple,

264 meta: Optional[Dict[str, Any]] = None,

265 ) -> None:

266 """Store the compiled inner module and its output marker tree."""

267 super().__init__()

268 self.inner = inner

269 self._output_spec = spec

270 # Mirrors GraphModule.meta so callers can introspect compilation mode.

271 self.meta = dict(meta or {})

272

273 def forward(self, *args: Any, **kwargs: Any) -> tuple:

274 """Run the flat compiled graph and rebuild the original structure."""

275 flat = self.inner(*args, **kwargs)

276 return tuple(_rebuild_nested_output(self._output_spec, flat))

277

278

279def _tag_node(node: fx.Node, inductor_configs: Optional[Dict[str, Any]]) -> None:

280 """Annotate ``node`` for regional_inductor with the user's Inductor config."""

281 node.meta.setdefault("custom", {})["compile_with_inductor"] = {

282 "inductor_configs": dict(inductor_configs or {}),

283 }

284

285

286class InductorPass(GraphPass):

287 """Terminal backend compilation pass (TorchInductor).

288

289 See the module docstring for the mode semantics, the numerics policy and

290 the calling-convention contract.

291 """

292

293 @property

294 def name(self) -> str:

295 """Pass name."""

296 return "inductor_compilation"

297

298 def run(

299 self,

300 graph_module: fx.GraphModule,

301 pass_config: PassConfig,

302 example_inputs: Optional[tuple] = None,

303 **kwargs: Any,

304 ) -> fx.GraphModule:

305 """Compile the joint graph (or its PP stages) with Inductor.

306

307 Args:

308 graph_module: The joint FX graph produced by the tracer and

309 transformed by the earlier partitioning / overlap passes. When

310 ``PpPass`` already ran, this is the schedule stub and its

311 stage subgraphs are compiled instead.

312 pass_config: Graph-mode configuration; ``inductor_mode`` selects

313 the behavior, ``inductor_configs`` overrides Inductor options.

314 example_inputs: Fake inputs matching the joint graph's

315 placeholders in order (the tracer's

316 ``JointGraph.example_inputs``). Ignored for PP stage

317 subgraphs, whose inputs are derived from each subgraph's

318 placeholder metas.

319

320 Returns:

321 The same graph module with compiled regions spliced in and the

322 positional calling convention restored. On torch builds without

323 regional Inductor support it is returned untouched, logging a

324 warning and recording ``"unavailable"`` in its meta.

325

326 Raises:

327 ValueError: When ``example_inputs`` is missing for a non-PP graph,

328 or a ``call_module`` stub is not a pipeline schedule (unknown

329 split).

330 """

331 mode = pass_config.inductor_mode

332 if mode == "off":

333 return graph_module

334

335 compiled = False

336 if graph_module.graph.find_nodes(op="call_module"):

337 compiled = self._compile_pipeline_stages(graph_module, pass_config)

338 else:

339 if example_inputs is None:

340 raise ValueError(

341 "InductorPass requires example_inputs (the tracer's "

342 "JointGraph.example_inputs) to size compiled regions"

343 )

344 if _regional_inductor_available():

345 self._compile_graph(graph_module, pass_config, example_inputs)

346 compiled = True

347 else:

348 self._warn_unavailable(mode)

349

350 graph_module.meta[_INDUCTOR_COMPILED_META] = (

351 mode if compiled else _INDUCTOR_UNAVAILABLE

352 )

353 return graph_module

354

355 @staticmethod

356 def _warn_unavailable(mode: str) -> None:

357 """Warn that the running torch predates regional Inductor compilation."""

358 _LOG.warning(

359 "InductorPass: torch %s does not provide "

360 "torch.fx.passes.regional_inductor; skipping %r compilation "

361 "(requires a newer torch build)",

362 torch.__version__,

363 mode,

364 )

365

366 def _compile_pipeline_stages(

367 self, graph_module: fx.GraphModule, pass_config: PassConfig

368 ) -> bool:

369 """Compile every stage subgraph installed by ``PpPass``.

370

371 ``PpPass`` replaces the joint graph with a ``call_module`` stub around

372 the pipeline schedule; the FX compute lives in the schedule's

373 per-virtual-stage forward/backward subgraphs (see

374 ``PipelineScheduleBase.map_stage_graphs``). Compiling the stub itself

375 would be a no-op and the stub must remain the outer entry point.

376

377 Returns:

378 ``True`` when the stage graphs were compiled, ``False`` when the

379 running torch does not provide regional Inductor compilation (the

380 pass then no-ops after ``_warn_unavailable``).

381 """

382 schedules = []

383 for node in graph_module.graph.find_nodes(op="call_module"):

384 submodule = graph_module.get_submodule(node.target)

385 if not callable(getattr(submodule, "map_stage_graphs", None)):

386 raise ValueError(

387 f"InductorPass cannot compile call_module '{node.target}': "

388 f"only PpPass's pipeline schedule exposes stage subgraphs "

389 f"(map_stage_graphs) for backend compilation"

390 )

391 schedules.append(submodule)

392

393 if not _regional_inductor_available():

394 self._warn_unavailable(pass_config.inductor_mode)

395 return False

396

397 for schedule in schedules:

398 schedule.map_stage_graphs(

399 lambda stage_gm: self._compile_graph(

400 stage_gm, pass_config, example_inputs=None

401 )

402 )

403 _LOG.info(

404 "InductorPass: compiled %d PP schedule(s) in %r mode",

405 len(schedules),

406 pass_config.inductor_mode,

407 )

408 return True

409

410 def _compile_graph(

411 self,

412 graph_module: fx.GraphModule,

413 pass_config: PassConfig,

414 example_inputs: Optional[Sequence[Any]],

415 ) -> fx.GraphModule:

416 """Compile one FX graph (joint graph or one PP stage subgraph).

417

418 ``example_inputs`` is ``None`` for PP stage subgraphs, whose inputs

419 are derived from their own placeholder metas (the joint graph's

420 inputs do not match a stage's placeholder list).

421 """

422 mode = pass_config.inductor_mode

423 output_spec, graph_module = _flatten_nested_outputs(graph_module)

424 _propagate_fake_vals(graph_module)

425 if example_inputs is None:

426 example_inputs = _placeholder_fake_args(graph_module)

427 _LOG.info(

428 "InductorPass: compiling graph in %r mode (%d placeholders, %d nodes)",

429 mode,

430 len(graph_module.graph.find_nodes(op="placeholder")),

431 len(graph_module.graph.nodes),

432 )

433

434 if mode == "full":

435 self._tag_full(graph_module, pass_config.inductor_configs)

436 else:

437 self._tag_regional(graph_module, pass_config.inductor_configs)

438

439 # Imported lazily so this module does not pull in the (heavy)

440 # TorchInductor stack unless compilation actually runs.

441 import torch._inductor.config as inductor_config # pylint: disable=import-outside-toplevel

442 from torch._guards import ( # pylint: disable=import-outside-toplevel

443 TracingContext,

444 tracing,

445 )

446 from torch.fx.passes.regional_inductor import ( # pylint: disable=import-outside-toplevel

447 regional_inductor,

448 )

449

450 # Bitwise-equal numerics vs eager (torchtitan's regional_inductor

451 # policy). emulate_precision_casts / disable_ftz stay off: enabling

452 # them is a numerics-changing opt-in, not this pass's default. Scope

453 # the override so it cannot leak into other Inductor compiles. The

454 # knob only exists in torch >= 2.11; older builds keep their defaults.

455 numerics = {}

456 if hasattr(inductor_config, "eager_numerics"):

457 numerics["eager_numerics.division_rounding"] = True

458

459 fake_mode = _get_fake_mode_from_gm(graph_module)

460 tracing_ctx = TracingContext(fake_mode)

461 with (

462 inductor_config.patch( # pylint: disable=no-member

463 numerics

464 ),

465 tracing(tracing_ctx),

466 ):

467 # ``regional_inductor`` consumes these on torch releases whose

468 # signature expects example inputs; newer releases size regions

469 # from the nodes' fake ``meta["val"]`` instead.

470 regional_inductor(graph_module, example_inputs)

471

472 # regional_inductor compiles regions with the boxed calling

473 # convention (one mutable args list); reset to the default positional

474 # convention so run_traced_graph's graph_module(*flat_inputs) (and the

475 # PP schedule's fwd_gm(*args) / bwd_gm(*args)) keeps working. Setting

476 # the default unconditionally is a no-op when the region compiler

477 # already used positional codegen and avoids peeking at private state.

478 graph_module.graph.set_codegen(CodeGen())

479 graph_module.recompile()

480 graph_module.meta[_INDUCTOR_COMPILED_META] = mode

481 if output_spec is not None:

482 # Rebuild the stage's original (possibly nested) output structure

483 # for the schedule / backward subgraph; only the compiled graph is

484 # flat.

485 return _NestedOutputWrapper(graph_module, output_spec, graph_module.meta)

486 return graph_module

487

488 @staticmethod

489 def _tag_full(

490 graph_module: fx.GraphModule, inductor_configs: Optional[Dict[str, Any]]

491 ) -> None:

492 """Tag every node so the whole graph compiles as one region.

493

494 Mirrors torchtitan's ``full_inductor_compilation_pass``: Inductor's

495 ``reorder_for_peak_memory`` restores the fwd/bwd interleaving that

496 AOTAutograd flattening would otherwise drop. Side-effectful nodes

497 (the FSDP reshard storage free) stay untagged: AOTAutograd's

498 functionalization would re-execute them on functional tensors and

499 lose the storage release.

500 """

501 configs = dict(inductor_configs or {})

502 configs.setdefault("reorder_for_peak_memory", True)

503 # Flat graph (joint graph or one PP stage subgraph): tag its nodes

504 # directly (no nested ``call_module`` submodules to recurse into).

505 for node in graph_module.graph.nodes:

506 if node.op in ("placeholder", "output"):

507 continue

508 if _node_is_side_effectful(node):

509 continue

510 _tag_node(node, configs)

511

512 @staticmethod

513 def _tag_regional(

514 graph_module: fx.GraphModule, inductor_configs: Optional[Dict[str, Any]]

515 ) -> None:

516 """Tag functional compute nodes, keeping comms and side effects eager.

517

518 Untagged nodes become region boundaries: contiguous runs of tagged

519 compute between collectives / frees are each fused and compiled, and

520 the collective launch/wait nodes remain plain FX nodes so the overlap

521 passes' scheduling decisions stay observable at runtime.

522 """

523 for node in graph_module.graph.nodes:

524 if node.op not in ("call_function", "call_method"):

525 continue

526 if _node_is_comm(node) or _node_is_side_effectful(node):

527 continue

528 _tag_node(node, inductor_configs)

Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / compile / passes / inductor_pass.py: 94%

180 statements