Diff Coverage

            raise ValueError(
                f"fsdp_degree must be None or a positive int, got {self.fsdp_degree}"
            )
        if self.pp_degree is not None and self.pp_degree < 1:
            raise ValueError(
                f"pp_degree must be None or a positive int, got {self.pp_degree}"
            )
        if self.pp_microbatch_size < 1:
            raise ValueError(
                f"pp_microbatch_size must be >= 1, got {self.pp_microbatch_size}"
            )

            yield item
        elif isinstance(item, (tuple, list)):
            stack.extend(item)
        elif isinstance(item, dict):
            stack.extend(item.values())


def _tensor_arg_stages(node: fx.Node, node_stage: Dict[fx.Node, int]) -> List[int]:
    """Stages of ``node``'s TENSOR-valued resolved args.

    stages = []
    for arg in _iter_node_args(node):
        stage = node_stage.get(arg)
        if stage is None:
            continue
        val = arg.meta.get("val")
        if val is not None and not isinstance(val, torch.Tensor):
            continue
        stages.append(stage)

        return _even_split(children, pp_degree)

    container_mod = getattr(model, container)
    if isinstance(container_mod, nn.ModuleDict):
        elem_keys = list(container_mod.keys())
    else:
        elem_keys = [str(i) for i in range(len(container_mod))]

    idx = children.index(container)

        model. Copying here would desynchronize the two.
        """
        state_fqns = getattr(graph_module, "state_fqns", None)
        if state_fqns is None:
            raise ValueError(
                "PpPass requires the tracer's state_fqns attribute on the "
                "joint graph (attach it via trace_model_graph)"
            )
        return state_fqns

        """Return the graph's live ``state_is_param`` list, defaulting to
        all-params (matching the tracer's fallback for older traces)."""
        state_is_param = getattr(graph_module, "state_is_param", None)
        if state_is_param is None:
            state_is_param = [True] * len(state_fqns)
            graph_module.state_is_param = state_is_param
        return state_is_param

    @staticmethod
    def _loss_node(graph_module: fx.GraphModule) -> Optional[fx.Node]:

    def _loss_node(graph_module: fx.GraphModule) -> Optional[fx.Node]:
        """The original joint graph's loss node (output index 0)."""
        output = next((n for n in graph_module.graph.nodes if n.op == "output"), None)
        if output is None:
            return None
        returned = output.args[0]
        if isinstance(returned, (list, tuple)) and returned:
            first = returned[0]
            if isinstance(first, fx.Node):

        if isinstance(returned, (list, tuple)) and returned:
            first = returned[0]
            if isinstance(first, fx.Node):
                return first
        return None

    # ------------------------------------------------------------------
    # Group / plan resolution
    # ------------------------------------------------------------------

            seen: Set[str] = set()
            for s, fqns in enumerate(manual):
                for fqn in fqns:
                    if fqn in seen:
                        raise ValueError(
                            f"Module '{fqn}' is assigned to multiple stages "
                            f"(stage {s} repeats it)"
                        )
                    seen.add(fqn)

                        )
                    seen.add(fqn)
            stage_plan = [list(fqns) for fqns in manual]
        else:
            stage_plan = _auto_stage_split(model, pp_degree)

        known = dict(model.named_modules())
        for s, fqns in enumerate(stage_plan):
            for fqn in fqns:

        if not unassigned:
            return
        shown = unassigned[:10]
        if len(unassigned) > len(shown):
            shown.append(f"... {len(unassigned) - len(shown)} more")
        _LOG.warning(
            "No PP stage declares modules %s — their nodes ride with the "
            "dataflow attribution (consumer stage, then last stage). "
            "Declare them via PassPlan.pp_stage() if that is not intended.",

            stage_state_indices,
        )

        if input_ph is None or label_ph is None:
            raise ValueError(
                "Joint graph must expose (state..., input, label) placeholders"
            )
        return (
            sorted(stage_state_indices),

            if idx < num_state_inputs:
                stage = self._state_stage(state_fqns[idx], stage_of_fqn)
                state_phs.append(ph)
                if stage is None:
                    deferred_state.append(idx)
                else:
                    node_stage[ph] = stage
                    if stage == stage_idx:
                        stage_state_indices.append(idx)

            phase = _BWD if node.meta.get("autograd_backward", False) else _FWD
            if phase == _FWD:
                for arg in _iter_node_args(node):
                    if node_phase.get(arg) == _BWD:
                        phase = _BWD
                        break
            node_phase[node] = phase
            if phase == _FWD:
                stage = self._stage_from_stack(node, stage_of_fqn)
                if stage is None:

        Multi-stage consumption is the RoPE-cache-shaped error case and
        fails with an actionable message.
        """
        for idx in deferred_state:
            ph = state_phs[idx]
            consumer_stages = {node_stage[u] for u in ph.users if u.op != "output"}
            if len(consumer_stages) > 1:
                raise ValueError(
                    f"Root-level state '{state_fqns[idx]}' is consumed by "
                    f"stages {sorted(consumer_stages)} — v1 requires state "
                    f"to live under a single stage's module (e.g. move a "
                    f"shared RoPE cache under a per-stage module or split "

                    f"to live under a single stage's module (e.g. move a "
                    f"shared RoPE cache under a per-stage module or split "
                    f"it per stage)"
                )
            stage = next(iter(consumer_stages), 0)
            node_stage[ph] = stage
            if stage == stage_idx:
                stage_state_indices.append(idx)
            _LOG.info(
                "Root-level state '%s' assigned to stage %s by its consumer",
                state_fqns[idx],
                stage,
            )

        for i in range(len(parts) - 1, 0, -1):
            stage = stage_of_fqn.get(".".join(parts[:i]))
            if stage is not None:
                return stage
        return stage_of_fqn.get(fqn)

    def _stage_from_stack(
        self, node: fx.Node, stage_of_fqn: Dict[str, int]
    ) -> Optional[int]:

            if node_phase.get(node) != _BWD:
                continue
            stage = node_stage.get(node)
            if stage is None:
                continue
            for arg in _iter_node_args(node):
                if (
                    isinstance(arg, fx.Node)
                    and node_phase.get(arg) == _BWD

            if node.op == "output":
                continue
            producer = node_stage.get(node)
            if producer is None:
                continue
            for user in node.users:
                if user.op == "output":
                    continue
                consumer = node_stage.get(user)

        activation, replayed from the forward outputs).
        """
        output = next((n for n in graph_module.graph.nodes if n.op == "output"), None)
        if output is None or model is None:
            return
        returned = output.args[0]
        if not isinstance(returned, (list, tuple)):
            return
        state_is_param = getattr(
            graph_module, "state_is_param", [True] * len(state_fqns)
        )
        trainable = self._trainable_state_indices(state_fqns, state_is_param, model)

        )
        trainable = self._trainable_state_indices(state_fqns, state_is_param, model)
        for i, grad_node in enumerate(returned[1:]):
            if i >= len(trainable) or not isinstance(grad_node, fx.Node):
                continue
            pstage = self._state_stage(state_fqns[trainable[i]], stage_of_fqn)
            if pstage is not None:
                node_stage[grad_node] = pstage

                # inputs (e.g. gradient-shape sym_size scalars).
                if node_phase.get(arg) != _FWD and arg not in act_in_set:
                    continue
                if arg.op == "placeholder" and arg not in stage_owned_phs:
                    continue
                seen.add(arg)
                saved.append(arg)
        # Fwd-phase values this stage must ship BACKWARD (consumed by the
        # previous stage's backward) ride the fwd outputs as pass-throughs:

        # computed once in the fwd subgraph, replayed into the bwd
        # subgraph, and appended to its outputs.
        for v in grad_out_list:
            if node_phase[v] == _FWD and v not in seen:
                seen.add(v)
                saved.append(v)
        return act_in_list, act_out_list, grad_in_list, grad_out_list, saved

    # ------------------------------------------------------------------
    # Subgraph construction

        """Return ``name`` (or a suffixed variant) not yet in ``used``."""
        candidate = name
        suffix = 1
        while candidate in used:
            candidate = f"{name}_{suffix}"
            suffix += 1
        used.add(candidate)
        return candidate

    def _copy_node(

        ``env`` maps already-copied same-slice nodes; ``foreign`` maps
        boundary values to their placeholder stand-ins.
        """
        if node.op == "get_attr":
            raise ValueError(
                f"Unexpected get_attr node '{node.name}' — the tracer "
                f"contracts parameters/buffers to be static placeholders"
            )

                if a in env:
                    return env[a]
                if a in foreign:
                    return foreign[a]
                raise ValueError(
                    f"Node '{node.name}' references '{a.name}' which lies "
                    f"outside this stage's slice — boundary detection or "
                    f"the stage plan is inconsistent"
                )
            return a

        new_args = fx.map_arg(node.args, map_fn)
        new_kwargs = fx.map_arg(node.kwargs, map_fn)
        new_node = g.create_node(node.op, node.target, new_args, new_kwargs)

    ) -> List[fx.Node]:
        """Order the fwd subgraph's outputs: boundary values first, then saved."""
        if is_last:
            if loss_node is None:
                raise ValueError("Last PP stage requires a loss node in the output")
            out_orig: List[fx.Node] = [loss_node]
        else:
            if not act_out_list:
                raise ValueError(
                    "Stage ships no boundary values but is not the last "
                    "stage — the stage plan does not cut the graph"
                )
            out_orig = list(act_out_list)

        returned = output.args[0] if output is not None else []
        stage_grad_nodes: List[fx.Node] = []
        for i, grad_node in enumerate(returned[1:]):
            if i >= len(trainable_indices):
                break
            if trainable_indices[i] in stage_state_set:
                stage_grad_nodes.append(grad_node)
        return stage_grad_nodes

            """
            children = child._modules  # pylint: disable=protected-access
            keys = list(children.keys())
            if all(is_kept(f"{cfqn}.{key}") for key in keys):
                return
            for key in keys:
                elem_fqn = f"{cfqn}.{key}"
                if is_kept(elem_fqn):
                    continue
                elem = children[key]
                descendants = [f for f, _ in elem.named_modules() if f]
                if any(is_kept(f"{elem_fqn}.{f}") for f in descendants):
                    prune(elem, elem_fqn + ".")

                # Leaf module (e.g. a bare Linear): keep only if declared.
                if not is_kept(cfqn):
                    setattr(mod, name, None)
                return
            if not any(is_kept(f"{cfqn}.{f}") for f in descendants):
                setattr(mod, name, None)
            else:
                # Mixed subtree (manual plans may cut inside a child).
                prune(child, cfqn + ".")

        def prune(mod: nn.Module, prefix: str) -> None:
            """Recursively prune non-stage children under ``prefix``."""
            for name, child in list(mod.named_children()):

        )
        trainable: List[int] = []
        for idx, fqn in enumerate(state_fqns):
            if idx < len(state_is_param) and not state_is_param[idx]:
                continue
            param = param_lookup.get(fqn)
            if param is not None and param.requires_grad:
                trainable.append(idx)
        return trainable

                # Training uses static shapes per compile, so the hinted
                # concrete values are stable.
                shape = tuple(int(s) for s in val.shape)
                spec.append(("tensor", shape, val.dtype, val.device))
            elif val is None:
                raise ValueError(
                    f"Boundary {kind} value '{v.name}' on stage {stage_idx} "
                    f"has no 'val' meta (traced without FakeTensor "
                    f"metadata) — cannot size the P2P receive buffer"
                )

                    f"has no 'val' meta (traced without FakeTensor "
                    f"metadata) — cannot size the P2P receive buffer"
                )
            else:
                spec.append(("scalar",))
        if nodes and not any(entry[0] == "tensor" for entry in spec):
            raise ValueError(
                f"Boundary {kind} list on stage {stage_idx} carries only "
                f"scalar values — no tensor entry to anchor the P2P "
                f"receive device on"
            )

        label_batch = flat_inputs[-1]

        batch_size = input_batch.shape[0]
        if batch_size % self.microbatch_size != 0:
            raise ValueError(
                f"PP microbatch mismatch: batch size {batch_size} is not "
                f"divisible by pp_microbatch_size {self.microbatch_size}"
            )
        num_microbatches = batch_size // self.microbatch_size

        for value in values:
            if isinstance(value, torch.Tensor):
                tensor = value.contiguous()
            else:
                tensor = torch.tensor(value, dtype=torch.int64, device=anchor_device)
            work = dist.isend(tensor, dst=dst_rank, group=self.pp_group)
            self._pending_sends.append((work, tensor))

    def _recv_values(self, spec: Sequence[Tuple[Any, ...]], src: int) -> List[Any]:

    def _global_rank(self, stage_idx: int) -> int:
        """Map a stage index to its global rank within the PP group."""
        if self.pp_group is None:
            return stage_idx
        return dist.get_global_rank(self.pp_group, stage_idx)


__all__ = ["ScheduleGPipe"]

        # 2. Execution layer: Parallel dimension partitioning
        if getattr(self.config, "fsdp_enabled", False):
            self.passes.append(FSDPPass(pass_plan=self.pass_plan))
        if getattr(self.config, "pp_enabled", False):
            self.passes.append(PpPass(pass_plan=self.pass_plan))

        # 3. Communication-compute overlap optimization
        if getattr(self.config, "enable_overlap", False):
            self.passes.append(AutoOverlapPass())

        merged = PassPlan()
        merged.fsdp_modules = {**self.fsdp_modules, **other.fsdp_modules}
        merged.fsdp_patterns = {**self.fsdp_patterns, **other.fsdp_patterns}
        if other.pp_module_fqns_per_stage is not None:
            merged.pp_module_fqns_per_stage = other.pp_module_fqns_per_stage
        elif self.pp_module_fqns_per_stage is not None:
            merged.pp_module_fqns_per_stage = self.pp_module_fqns_per_stage
        return merged

    def fsdp_wrap(self, module_fqn: str) -> "PassPlan":
        """Mark a specific module for FSDP wrapping (exact match).

    section = config.get(key) or {}
    if not isinstance(section, dict):
        # A present-but-empty section parses to None and is normalized to {}
        # above, so anything landing here is a real scalar/sequence typo.
        raise ValueError(
            f"YAML '{key}' section must be a mapping (e.g. nested keys or an "
            f"empty section); got {type(section).__name__} in {config_path}"
        )
    return section

        if isinstance(stage, dict):
            stage_idx = stage.get("stage", idx)
            module_fqns = list(stage.get("modules", []))
        else:
            stage_idx = idx
            module_fqns = list(stage)
        plan.pp_stage(stage_idx, module_fqns)


def create_simple_sharding_plan() -> PassPlan:

            yield a
        elif isinstance(a, (tuple, list)):
            for x in a:
                if isinstance(x, torch.fx.Node):
                    yield x


def _output_args(fx_g: torch.fx.GraphModule) -> List[torch.fx.Node]:
    """Return the flattened forward arg list of the graph's ``output`` node."""

def _output_args(fx_g: torch.fx.GraphModule) -> List[torch.fx.Node]:
    """Return the flattened forward arg list of the graph's ``output`` node."""
    out_node = next((n for n in fx_g.graph.nodes if n.op == "output"), None)
    if out_node is None:
        return []
    args: List[torch.fx.Node] = []
    for a in out_node.args:
        if isinstance(a, torch.fx.Node):
            args.append(a)
        elif isinstance(a, (tuple, list)):
            for x in a:
                if isinstance(x, torch.fx.Node):
                    args.append(x)

      nodes or the loss source) spans exactly the backward half.
    """
    output = _output_args(fx_g)
    if not output:
        return
    loss_node = output[0]
    grad_seeds = output[1:]

    # Forward loss chain: everything that flows INTO the loss. Backward is

def _copy_fwd_meta(fwd_node: torch.fx.Node, node: torch.fx.Node) -> None:
    """Copy ``custom``/``nn_module_stack``/``stack_trace`` from fwd to bwd."""
    custom = fwd_node.meta.get("custom")
    if custom:
        node.meta.setdefault("custom", {}).update(copy.deepcopy(custom))
    nn_module_stack = fwd_node.meta.get("nn_module_stack")
    if nn_module_stack is not None:
        node.meta["nn_module_stack"] = nn_module_stack.copy()
    stack_trace = fwd_node.meta.get("stack_trace")

    if nn_module_stack is not None:
        node.meta["nn_module_stack"] = nn_module_stack.copy()
    stack_trace = fwd_node.meta.get("stack_trace")
    if stack_trace is not None:
        node.meta["stack_trace"] = stack_trace


def _fakeify_input(fake_mode: FakeTensorMode, x: Any) -> Any:
    """Convert a real tensor input into its fake counterpart."""