Diff Coverage

                if parameters.get(alias) is parameters.get(name):
                    return layout
        context = self._mesh_context
        if context is None:
            return None
        if name in self._expert_param_names:
            return ((Shard(0),), context.fsdp_moe_mesh["ep"])
        mesh = context.fsdp_non_moe_mesh
        if mesh is not None and "tp" in mesh.mesh_dim_names:
            return ((Replicate(),), mesh["tp"])
        return None

    def _build_parameter_layouts(self, parameters: dict) -> Optional[dict]:
        """Describe actual local storage using common FSDP layout construction."""
        context = self._mesh_context
        if context is None and self._fsdp_mesh is None:
            return None
        sharded = getattr(self._joint_graph.graph_module, "fsdp_sharded_param_names", frozenset())
        layouts = {}
        for name, parameter in parameters.items():
            if not parameter.requires_grad:
                continue
            mesh = self._parameter_meshes.get(name)
            if mesh is None:
                mesh = (self._fsdp_mesh if self.pass_config.dp_mode == "ddp" or context is None
                        else select_fsdp_mesh(context, name in self._expert_param_names))

                    "('fsdp_replicate', 'dp_replicate', 'dp') in the mesh "
                    f"dims {names}"
                )
            sub = fsdp_mesh[dim]
            self._fsdp_mesh = sub
            self._register_mesh_group("dp_replicate", sub)
            self.pass_config.dp_replicate_degree = sub.size()
            return

        )
        if not names:
            return {}
        if self._mesh_context is None:
            raise ValueError("Graph source-layout reduction requires MeshContext")
        reduction = resolve_gradient_reduction(self._mesh_context, has_source_layout=True)
        result = {"gradient_reductions": dict.fromkeys(names, reduction)}
        if self.source_shard_info is None:
            return result

        groups, registered = {}, {}
        for name in names:
            layout = self.source_layout_for_parameter(name)
            if layout is None:
                continue
            placements, mesh = layout
            if mesh is None or not placements:
                continue
            axes = tuple(axis for index, axis in enumerate(mesh.mesh_dim_names)
                         if placements[index].is_replicate() and mesh.size(index) > 1)
            if not axes:
                continue

    # Functional all-gather only reliably supports gather_dim=0; for other dims we
    # gather on dim 0 then reshape to avoid _maybe_view_chunk_cat view() bugs
    # in Dynamo fake tensor tracing.
    if not -data.ndim <= concat_dim < data.ndim:
        raise IndexError("All-gather concatenation dimension is out of range")
    concat_dim %= data.ndim
    group_size = dist.get_world_size(group)
    all_gather = getattr(_functional_collectives, "all_gather_single_autograd", None)
    if all_gather is None:

    if world_mesh is None:
        raise ValueError("FSDP2 requires MeshContext.fsdp_non_moe_mesh")
    mesh_dim_names = world_mesh.mesh_dim_names
    if mesh_dim_names is None:
        raise ValueError("FSDP2 requires named DeviceMesh dimensions")
    shard_dim_names = tuple(
        dim_name for dim_name in ("dp_shard", "cp") if dim_name in mesh_dim_names
    )
    if not shard_dim_names:
        raise ValueError("FSDP2 requires a dp_shard or cp dimension in the world mesh")
    shard_mesh = (
        world_mesh[shard_dim_names[0]]
        if len(shard_dim_names) == 1
        else world_mesh[shard_dim_names].flatten("fsdp_shard")

        if len(shard_dim_names) == 1
        else world_mesh[shard_dim_names].flatten("fsdp_shard")
    )
    if "dp_replicate" not in mesh_dim_names:
        return shard_mesh
    return DeviceMesh.concatenate([world_mesh["dp_replicate"], shard_mesh])


def select_fsdp_mesh(context: MeshContext, expert: bool = False) -> DeviceMesh:

            parameters = get_managed_modules_parameters((owner,), claimed)
            source = _build_managed_source_shard_info(self, owner, metadata, claimed)
            expert = self._uses_expert_mesh(source)
            if owner is model and expert:
                raise ValueError("Routed expert parameters must belong to a nested experts FSDP unit")
            mesh = self._build_fsdp_actual_mesh(expert)
            bindings.update(dict.fromkeys(parameters, mesh))
            claimed.update(parameters)
        return bindings

    def _log_gradient_scaling(self, gradient_scaled_units: int) -> None:
        """Log the source-layout gradient scaling applied to FSDP units."""
        if not gradient_scaled_units:
            return
        reduction = resolve_gradient_reduction(self.mesh_context, has_source_layout=True)
        logger.info(
            "Applied source-layout SUM gradient scaling: dp_size=%d, cp_size=%d, tp_size=%d, "
            "source_gradient_domain_size=%d, factor=%s, units=%d",
            self.mesh_context.dp_size,