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:

    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
        if self._fsdp_mesh is None and context is not None:
            # Plain graph mode still needs metadata, without creating collectives.
            self._init_mesh_from_context(context, self.pass_config.dp_mode, register_groups=False)
        sharded = getattr(self._joint_graph.graph_module, "fsdp_sharded_param_names", frozenset())

        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._expert_fsdp_mesh if name in self._expert_param_names else self._fsdp_mesh
            if mesh is None:
                raise ValueError("Parameter layouts require an initialized FSDP mesh or parameter binding")
            layouts[parameter] = self._build_parameter_layout(
                parameter, mesh, self.source_layout_for_parameter(name), sharded=name in sharded,
                unsharded_shape=self._parameter_shapes.get(name, tuple(parameter.shape)),
            )

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

        if dp_mode == "hsdp":

        )
        if not names:
            return {}
        if self._mesh_context is None:
            raise ValueError("Graph source-layout reduction requires MeshContext")
        context = self._mesh_context
        loss_replicas = 1 if context.loss_parallel else context.tp_size
        reduction = GradientReduction("sum", context.dp_size * context.cp_size * loss_replicas)
        result = {"gradient_reductions": dict.fromkeys(names, reduction)}

        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

            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