Diff Coverage

Diff: origin/master...HEAD, staged and unstaged changes

Source File Diff Coverage (%) Missing Lines
hyper_parallel/compile/compiler.py 100%  
hyper_parallel/compile/pass_config.py 100%  
hyper_parallel/compile/passes/inductor_pass.py 94.9% 118-119,130,139,159,289,387-388,450
hyper_parallel/compile/passes/parallel/fsdp_pass.py 100%  
hyper_parallel/compile/passes/parallel/pp_schedule.py 71.4% 1131-1132
hyper_parallel/compile/passes/pipeline.py 100%  
hyper_parallel/compile/passes/inductor_pass.py
114
115
116
117
118
119
120
121
122
123
    builds.
    """
    try:
        importlib.import_module("torch.fx.passes.regional_inductor")
    except ImportError:
        return False
    return True


def _get_fake_mode_from_gm(gm: fx.GraphModule) -> Optional[Any]:
126
127
128
129
130
131
132
133
134
        if node.op == "placeholder":
            val = node.meta.get("val")
            if isinstance(val, FakeTensor):
                return val.fake_mode
    return None


def _placeholder_fake_args(gm: fx.GraphModule) -> list:
    """Return the placeholder values in order, requiring fake ``meta["val"]``."""
135
136
137
138
139
140
141
142
143
    args = []
    for node in gm.graph.find_nodes(op="placeholder"):
        val = node.meta.get("val")
        if val is None:
            raise ValueError(
                f"InductorPass requires a fake meta['val'] on placeholder "
                f"'{node.name}' to size compiled regions; the graph must be "
                f"make_fx-traced (PP stage subgraphs copy their placeholders' "
                f"meta from the joint graph)"
155
156
157
158
159
160
161
162
163
    Side-effectful free nodes return ``None`` and legitimately carry no val.
    """
    fake_mode = _get_fake_mode_from_gm(gm)
    if fake_mode is None:
        raise ValueError(
            "InductorPass requires a make_fx-traced graph whose placeholders "
            "carry fake meta['val'] tensors"
        )
    fake_args = _placeholder_fake_args(gm)
285
286
287
288
289
290
291
292
293

    @property
    def name(self) -> str:
        """Pass name."""
        return "inductor_compilation"

    def run(
        self,
        graph_module: fx.GraphModule,
383
384
385
386
387
388
389
390
391
392
                )
            schedules.append(submodule)

        if not _regional_inductor_available():
            self._warn_unavailable(pass_config.inductor_mode)
            return False

        for schedule in schedules:
            schedule.map_stage_graphs(
                lambda stage_gm: self._compile_graph(
446
447
448
449
450
451
452
453
454
        # the override so it cannot leak into other Inductor compiles. The
        # knob only exists in torch >= 2.11; older builds keep their defaults.
        numerics = {}
        if hasattr(inductor_config, "eager_numerics"):
            numerics["eager_numerics.division_rounding"] = True

        fake_mode = _get_fake_mode_from_gm(graph_module)
        tracing_ctx = TracingContext(fake_mode)
        with (
hyper_parallel/compile/passes/parallel/pp_schedule.py
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
        self._pending_sends: Dict[Any, List[Tuple[Any, torch.Tensor]]] = {}

    def map_stage_graphs(self, fn: Callable[[nn.Module], nn.Module]) -> None:
        """Apply ``fn`` to every virtual stage's fwd and bwd subgraph."""
        self.fwd_gms = nn.ModuleList([fn(gm) for gm in self.fwd_gms])
        self.bwd_gms = nn.ModuleList([fn(gm) for gm in self.bwd_gms])

    def _resolve_num_microbatches(self, user_inputs: Sequence[Any]) -> int:
        """Derive the microbatch count from the first tensor model input."""
        batch_size = next(