Diff Coverage

        builds the mesh with the repo ``init_device_mesh`` (the same builder the
        eager stack uses) and only supports **pure DP** (``tp_size == 1``, PP
        off); TP/PP topologies must supply a ``MeshContext``.
        """
        dp_mode = normalize_dp_mode(self.pass_config.dp_mode)
        if mesh_context is not None:
            self._init_mesh_from_context(mesh_context, dp_mode)
            return

        self._require_pure_dp_fallback()
        device_type = (
            "npu" if (hasattr(torch, "npu") and torch.npu.is_available()) else "cpu"
        )
        world_size = dist.get_world_size()
        rank_list = tuple(range(world_size))

        if dp_mode == "ddp":
            mesh = init_device_mesh(
                device_type,
                (world_size,),
                mesh_dim_names=("fsdp_replicate",),
                rank_list=rank_list,

                mesh_dim_names=("fsdp_replicate",),
                rank_list=rank_list,
                init_backend=True,
            )
            self._register_mesh_group("dp_replicate", mesh["fsdp_replicate"])
            self.pass_config.dp_replicate_degree = mesh["fsdp_replicate"].size()
            return

        if dp_mode == "hsdp":
            shard, replicate = self._resolve_hsdp_degrees(world_size)
            mesh = init_device_mesh(
                device_type,
                (replicate, shard),
                mesh_dim_names=("fsdp_replicate", "fsdp_shard"),
                rank_list=rank_list,

                mesh_dim_names=("fsdp_replicate", "fsdp_shard"),
                rank_list=rank_list,
                init_backend=True,
            )
            self._register_mesh_group("dp_replicate", mesh["fsdp_replicate"])
            self._register_mesh_group("fsdp", mesh["fsdp_shard"])
            self.pass_config.dp_replicate_degree = mesh["fsdp_replicate"].size()
            self.pass_config.fsdp_degree = mesh["fsdp_shard"].size()
            return

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

            mesh_dim_names=("fsdp_shard",),
            rank_list=rank_list,
            init_backend=True,
        )
        self._register_mesh_group("fsdp", mesh["fsdp_shard"])
        # Back-fill, mirroring the external-mesh branch: FSDPPass resolves the
        # group size from ``fsdp_degree`` (falling back to world_size when
        # ``None``), so setting it here keeps the two paths consistent.
        self.pass_config.fsdp_degree = world_size

        Raises:
            ValueError: When ``tp_size != 1`` or ``pp_enabled``.
        """
        cfg = self.pass_config
        if cfg.tp_size != 1 or cfg.pp_enabled:
            raise ValueError(
                "Graph-mode DP without a MeshContext only supports pure data "
                f"parallelism (tp_size=1, PP disabled); got tp_size={cfg.tp_size}, "
                f"pp_enabled={cfg.pp_enabled}. Pass a MeshContext carrying the DP "
                "sub-mesh for TP/PP hybrids."

        / ``"fsdp"``. The effective mode is first reconciled with the mesh
        topology (see ``_reconcile_mode_with_context``), mirroring the eager
        ``FSDP2Manager`` which picks HSDP from ``dp_replicate_size > 1``.
        """
        fsdp_mesh = (
            getattr(mesh_context, "fsdp_non_moe_mesh", None) or mesh_context.device_mesh
        )
        names = tuple(getattr(fsdp_mesh, "mesh_dim_names", ()) or ())
        dp_mode = self._reconcile_mode_with_context(mesh_context, dp_mode)
        self.pass_config.dp_mode = dp_mode

        if dp_mode == "ddp":
            dim = self._first_dim(names, ("fsdp_replicate", "dp_replicate", "dp"))
            if dim is None:
                raise ValueError(
                    "dp_mode='ddp' needs a replicate axis among "
                    "('fsdp_replicate', 'dp_replicate', 'dp') in the mesh "
                    f"dims {names}"
                )
            sub = fsdp_mesh[dim]
            self._register_mesh_group("dp_replicate", sub)
            self.pass_config.dp_replicate_degree = sub.size()
            return

        if dp_mode == "hsdp":
            rep_dim = self._first_dim(names, ("fsdp_replicate", "dp_replicate"))
            shard_dim = self._first_dim(names, ("fsdp_shard", "fsdp", "dp"))
            if rep_dim is None or shard_dim is None:
                raise ValueError(
                    "dp_mode='hsdp' needs both a replicate axis "
                    "('fsdp_replicate'/'dp_replicate') and a shard axis "
                    f"('fsdp_shard'/'fsdp'/'dp') in the mesh dims {names}"
                )
            rep_sub = fsdp_mesh[rep_dim]
            shard_sub = fsdp_mesh[shard_dim]
            self._register_mesh_group("dp_replicate", rep_sub)
            self._register_mesh_group("fsdp", shard_sub)
            self.pass_config.dp_replicate_degree = rep_sub.size()
            self.pass_config.fsdp_degree = shard_sub.size()
            return

        # dp_mode == "fsdp": automodel's fsdp_non_moe_mesh is
        # ("fsdp_replicate","fsdp_shard","tp"); device_mesh (cp=1) is
        # ("dp","cp","tp") and "dp" is the FSDP axis.
        dim = "fsdp_shard" if "fsdp_shard" in names else "dp"
        sub = fsdp_mesh[dim]
        self._register_mesh_group("fsdp", sub)
        self.pass_config.fsdp_degree = sub.size()

    def _reconcile_mode_with_context(self, mesh_context: Any, dp_mode: str) -> str:
        """Align ``dp_mode`` with the ``MeshContext`` topology (eager parity).

        Raises:
            ValueError: When ``dp_mode='ddp'`` but the mesh shards
                (``dp_shard_size > 1``).
        """
        rep = getattr(mesh_context, "dp_replicate_size", None)
        shard = getattr(mesh_context, "dp_shard_size", None)
        if not isinstance(rep, int):
            # Not a real MeshContext (e.g. a test stub): leave the mode as-is.
            return dp_mode
        if dp_mode == "fsdp" and rep > 1:
            _LOG.warning(
                "dp_mode='fsdp' but MeshContext.dp_replicate_size=%d > 1; "
                "using 'hsdp' to match the mesh",
                rep,
            )
            return "hsdp"
        if dp_mode == "hsdp" and rep <= 1:
            _LOG.warning(
                "dp_mode='hsdp' but MeshContext.dp_replicate_size=%d <= 1; "
                "using 'fsdp'",
                rep,
            )
            return "fsdp"
        if dp_mode == "ddp" and isinstance(shard, int) and shard > 1:
            raise ValueError(
                "dp_mode='ddp' requires a non-sharding mesh "
                f"(dp_shard_size=1), got dp_shard_size={shard}; use 'hsdp'"
            )
        return dp_mode

    @staticmethod
    def _first_dim(names: tuple, candidates: tuple) -> Optional[str]:
        """Return the first candidate axis present in ``names`` (or None)."""
        for cand in candidates:
            if cand in names:
                return cand
        return None

    @staticmethod
    def _register_mesh_group(name: str, sub_mesh: Any) -> None:
        """Register ``sub_mesh``'s process group under ``name``."""
        _register_process_group(name, sub_mesh.get_group())

    def _resolve_hsdp_degrees(self, world_size: int) -> Tuple[int, int]:
        """Resolve ``(shard, replicate)`` for a fallback 2-D HSDP mesh.

        Raises:
            ValueError: When both are ``None``, a degree does not divide
                ``world_size``, or the product does not equal ``world_size``.
        """
        shard = self.pass_config.fsdp_degree
        replicate = self.pass_config.dp_replicate_degree
        if shard is None and replicate is None:
            raise ValueError(
                "dp_mode='hsdp' requires fsdp_degree and/or "
                "dp_replicate_degree to build the 2-D (dp_replicate, fsdp) mesh"
            )
        if shard is None:
            if world_size % replicate != 0:
                raise ValueError(
                    f"world_size ({world_size}) not divisible by "
                    f"dp_replicate_degree ({replicate})"
                )
            shard = world_size // replicate
        if replicate is None:
            if world_size % shard != 0:
                raise ValueError(
                    f"world_size ({world_size}) not divisible by fsdp_degree ({shard})"
                )
            replicate = world_size // shard
        if shard * replicate != world_size:
            raise ValueError(
                f"hsdp degrees fsdp_degree={shard} x "
                f"dp_replicate_degree={replicate} != world_size={world_size}"
            )
        return shard, replicate

    def _build_pass_kwargs(self) -> dict:
        """
        Build kwargs to pass to passes.

        if self.pass_config.fsdp_enabled:
            kwargs["fsdp_group_name"] = "fsdp"
            # DDP / HSDP also reduce gradients on the replicate axis.
            if normalize_dp_mode(self.pass_config.dp_mode) in ("ddp", "hsdp"):
                kwargs["dp_replicate_group_name"] = "dp_replicate"

        return kwargs

    def _run_graph(self, **inputs):

    Raises:
        ValueError: When ``mode`` is not a string or is unknown.
    """
    if not isinstance(mode, str):
        raise ValueError(f"dp_mode must be a string, got {type(mode).__name__}")
    key = mode.strip().lower()
    return _DP_MODE_ALIASES.get(key, key)


@dataclass
class PassConfig:

    pp_microbatch_size: int = 1
    pp_schedule: str = "gpipe"

    def __post_init__(self) -> None:
        self.dp_mode = normalize_dp_mode(self.dp_mode)
        self.validate()

    def validate(self) -> None:
        """Sanity-check invariants; also re-run after manual field mutation.

        if self.fsdp_degree is not None and self.fsdp_degree < 1:
            raise ValueError(
                f"fsdp_degree must be None or a positive int, got {self.fsdp_degree}"
            )
        if self.dp_mode not in DP_MODES:
            raise ValueError(
                f"dp_mode must be one of {DP_MODES} "
                f"(aliases: {tuple(_DP_MODE_ALIASES)}), got {self.dp_mode!r}"
            )
        if self.dp_replicate_degree is not None and self.dp_replicate_degree < 1:
            raise ValueError(
                f"dp_replicate_degree must be None or a positive int, "
                f"got {self.dp_replicate_degree}"
            )
        if self.pp_degree is not None and self.pp_degree < 1:

        # ``run`` fails loudly instead of silently using a wrong default.
        self._fsdp_degree: Optional[int] = None
        # Data-parallel mode ("fsdp" | "ddp" | "hsdp") and the replicate-axis
        # group/degree used by DDP / HSDP gradient reduction.
        self._dp_mode = "fsdp"
        self._dp_replicate_group_name = "dp_replicate"
        self._dp_replicate_degree: Optional[int] = None
        self._processed_params: Set[str] = set()
        self._fsdp_modules: Set[str] = set()
        # Per-param forward all_gather record (the unsharded parameter),
        # populated by ``_insert_all_gather_for_params`` and consumed by the

        if not dist.is_initialized() or dist.get_world_size() == 1:
            _LOG.info("Skipped: distributed not initialized or world_size=1")
            return graph_module

        self._dp_mode = normalize_dp_mode(getattr(pass_config, "dp_mode", "fsdp"))
        # FSDP shard-axis group size: use the explicitly configured degree when
        # present (required for TP+FSDP / HSDP, where the FSDP group is a
        # proper sub-group of the world — using world_size would over-shard
        # along the TP axis).

        # along the TP axis).
        configured = pass_config.fsdp_degree
        self._fsdp_degree = configured if configured else dist.get_world_size()
        self._fsdp_group_name = kwargs.get("fsdp_group_name", self._fsdp_group_name)
        self._dp_replicate_group_name = kwargs.get(
            "dp_replicate_group_name", self._dp_replicate_group_name
        )
        # Replicate-axis degree used by DDP / HSDP gradient reduction.
        replicate_degree = getattr(pass_config, "dp_replicate_degree", None)
        if self._dp_mode == "ddp":
            self._dp_replicate_degree = replicate_degree or dist.get_world_size()
        elif self._dp_mode == "hsdp":
            self._dp_replicate_degree = replicate_degree or (
                dist.get_world_size() // self._fsdp_degree
            )
        else:
            self._dp_replicate_degree = replicate_degree
        self._parallel_plan = kwargs.get("parallel_plan", self._parallel_plan)
        model = kwargs.get("model")
        if model is None:
            raise ValueError(

            return graph_module

        # Per-run state: reset so a reused FSDPPass instance does not carry
        # parameter names from a previous graph (which would skip their gather).
        self._processed_params = set()
        self._fsdp_modules = set()
        self._unsharded_params = {}
        # DDP replicates parameters: no all_gather, no live-model sharding and
        # no reshard. Only the gradient all-reduce path below runs.
        if self._dp_mode != "ddp":
            graph_module = self._insert_all_gather_for_params(
                graph_module, param_nodes, pass_config
            )

        reduce_param_indices = frozenset(node.meta["state_idx"] for node in param_nodes)
        graph_module = self._insert_grad_reduction(
            graph_module,
            reduce_param_indices,
            state_fqns,
            num_state_inputs,

            state_is_param,
            model,
        )

        if self._dp_mode != "ddp":
            # Reshard: free each replicated parameter once forward is done and
            # re-gather + rematerialize it for the backward. Runs after the
            # grad reduction so the grad path is untouched; no-op when disabled.
            graph_module = self._insert_reshard_logic(
                graph_module, param_nodes, pass_config
            )

            # Shard the live model's parameters in place (dim 0, by this rank's

            # Shard the live model's parameters in place (dim 0, by this rank's
            # index in the FSDP group). ``model.parameters()`` then yields the
            # shards, so the trainer's optimizer / grad accumulation stay
            # DP-agnostic; the graph re-gathers each step.
            self._shard_live_model_params(model)

        _LOG.info(
            "Completed dp_mode=%s, processed %d parameters",
            self._dp_mode,
            len(param_nodes),
        )

            # full-size. Inserting an all_gather here would make the graph
            # expect a sharded input while the live model still holds the full
            # tensor (shape mismatch at ``run_traced_graph`` time), so skip them
            # with the same gate the live-model sharder uses.
            param = param_lookup.get(fqn)
            if param is not None and not param.requires_grad:
                _LOG.info("Skip %s: requires_grad=False (frozen)", fqn)
                continue

            # Same gate as ``_shard_live_model_params``: scalar params
            # (empty shape) and non-divisible dim-0 params stay replicated,
            # so the graph and the live model agree on which parameters are

            # (empty shape) and non-divisible dim-0 params stay replicated,
            # so the graph and the live model agree on which parameters are
            # sharded. DDP has no shard axis, so every marked parameter is
            # reduced regardless of divisibility.
            if (
                self._dp_mode != "ddp"
                and param is not None
                and (not param.shape or param.shape[0] % self._fsdp_degree != 0)
            ):

            if not isinstance(grad_node, fx.Node):
                continue

            state_idx = trainable_state_indices[i - 1]
            if state_idx not in reduce_param_indices:
                continue

            with graph.inserting_before(output_node):
                new_returned[i] = self._build_grad_reduction_nodes(graph, grad_node)

        output_node.args = (type(returned)(new_returned),) + tuple(output_node.args[1:])

        return graph_module

        The returned (wait) node replaces ``grad_node`` in the graph output.
        Emits the mode-specific chain (see ``_insert_grad_reduction``).
        """
        reduced = grad_node
        if self._dp_mode in ("fsdp", "hsdp"):
            rs_node = graph.call_function(
                _c10d_functional.reduce_scatter_tensor,
                args=(
                    reduced,
                    _GRAD_REDUCE_OP,

                    self._fsdp_degree,
                    self._fsdp_group_name,
                ),
            )
            rs_node.meta["comm_type"] = "fsdp_reduce_scatter"
            rs_node.meta["comm_group"] = self._fsdp_group_name
            rs_node.meta["fsdp_degree"] = self._fsdp_degree
            reduced = rs_node

        if self._dp_mode in ("ddp", "hsdp"):
            ar_node = graph.call_function(
                _c10d_functional.all_reduce,
                args=(reduced, _GRAD_REDUCE_OP, self._dp_replicate_group_name),
            )
            ar_node.meta["comm_type"] = "dp_all_reduce"
            ar_node.meta["comm_group"] = self._dp_replicate_group_name
            ar_node.meta["dp_replicate_degree"] = self._dp_replicate_degree
            reduced = ar_node

        wait_node = graph.call_function(
            _c10d_functional.wait_tensor,
            args=(reduced,),
        )
        wait_node.meta["wait_for"] = reduced.name
        return wait_node

    def _build_trainable_state_indices(
        self,
        state_fqns: List[str],

            # This gather lives in the backward half; tag it so downstream
            # passes (e.g. PpPass phase classification) do not mistake it and
            # its remat clones for forward nodes — their only arg is a forward
            # state placeholder, which would otherwise imply the forward phase.
            ag_node.meta["autograd_backward"] = True
        with graph.inserting_before(anchor):
            wait_node = graph.call_function(
                _c10d_functional.wait_tensor, args=(ag_node,)
            )

            wait_node = graph.call_function(
                _c10d_functional.wait_tensor, args=(ag_node,)
            )
            wait_node.meta["wait_for"] = ag_node.name
            wait_node.meta["autograd_backward"] = True
        return wait_node

    def _clone_with_remap(
        self,

            clone = graph.call_function(node.target, args=args, kwargs=kwargs)
        self._copy_recreated_meta(node, clone)
        # Rematerialized views serve the backward half, so they carry the
        # backward phase (see ``_insert_backward_gather``).
        clone.meta["autograd_backward"] = True
        recreate[node] = clone
        return clone

    @staticmethod

        with graph.inserting_after(last_reader):
            free_node = graph.call_function(_free_tensor_storage, args=(tensor,))
            free_node.meta["comm_type"] = comm_type
            free_node.meta["param_name"] = record.param_node.meta.get("param_name")
            if backward:
                free_node.meta["autograd_backward"] = True

    def _alias_descendants(self, root: fx.Node) -> Set[fx.Node]:
        """Forward view descendants of ``root`` (nodes aliasing its storage).