Diff Coverage

            value = getattr(self, name)
            if not isinstance(value, bool):
                raise ValueError(f"{name} must be a bool")
        if self.fsdp_bucket_modules is not None and not isinstance(self.fsdp_bucket_modules, (list, tuple)):
            raise ValueError("fsdp_bucket_modules must be None or a list/tuple of module groups")
        value = self.max_in_flight_gb
        if (
            isinstance(value, bool)
            or not isinstance(value, (int, float))

        Returns:
            Graph with fused AG/RS, or unchanged graph when disabled/unsupported.
        """
        if not pass_config.fsdp_bucketing or not pass_config.fsdp_enabled or not dist.is_initialized():
            return graph_module
        if dist.get_world_size() == 1:
            return graph_module
        if pass_config.dp_mode != "fsdp" or pass_config.pp_enabled or (pass_config.pp_degree or 1) != 1:
            _LOG.info("Skipped FSDP bucketing: requires dp_mode=fsdp and PP disabled")
            return graph_module
        modules = pass_config.fsdp_bucket_modules
        if modules is not None:
            modules = build_module_bucket_plans(kwargs["model"], module_bucket_plans=modules)
        names = [kwargs.get("fsdp_group_name", "fsdp")]
        if kwargs.get("expert_param_names"):
            names.append(kwargs.get("expert_fsdp_group_name", "fsdp_expert"))
        groups = {name: _resolve_process_group(name) for name in names}
        ranks = {name: tuple(dist.get_process_group_ranks(group)) for name, group in groups.items()}
        for apply in (apply_ag_buckets, apply_rs_buckets):
            plan = plan_fsdp_buckets(graph_module, group_ranks=ranks, module_bucket_plans=modules)
            for name, group in groups.items():
                report = {"group_ranks": {name: ranks[name]}, "group_digests": {name: plan.group_digest(name)}}
                reports = [None] * len(ranks[name])
                dist.all_gather_object(reports, report, group=group)
                validate_group_plans(dict(zip(ranks[name], reports)), name)
            apply(graph_module, plan)
        return graph_module

    Raises:
        ValueError: Types, module paths or group ownership are invalid.
    """
    if not isinstance(model, nn.Module):
        raise ValueError("model must be a torch.nn.Module")
    if (not isinstance(block_module_types, tuple)
            or any(not isinstance(kind, type) or not issubclass(kind, nn.Module) for kind in block_module_types)):
        raise ValueError("block_module_types must be a tuple of torch.nn.Module subclasses")
    explicit = normalize_module_bucket_plans(module_bucket_plans)

    Raises:
        ValueError: Reports are missing ranks or disagree on membership/order/layout.
    """
    if not plans_by_rank:
        raise ValueError("No rank reports supplied")
    reports = {rank: plan.to_dict() if isinstance(plan, BucketPlan) else plan
               for rank, plan in plans_by_rank.items()}
    first = next(iter(reports.values()))
    expected = first["group_ranks"].get(group)

    if not isinstance(source, fx.Node) or not isinstance(group, str) or group not in groups:
        return "unknown_process_group"
    degree = node.args[-2]
    if not _static_int(degree) or degree <= 1 or degree != len(groups[group]):
        return "group_size_mismatch"
    if node.meta.get("comm_group", group) != group:
        return "group_metadata_mismatch"
    name, index = node.meta.get("param_name"), node.meta.get("state_idx")
    if not isinstance(name, str) or not name or not _static_int(index) or index < 0:
        return "missing_parameter_identity"
    return None

    dtype, device = spec.get("dtype"), spec.get("device")
    if not isinstance(dtype, torch.dtype) or dtype not in (torch.float16, torch.bfloat16, torch.float32, torch.float64):
        return "unsupported_dtype"
    if not isinstance(device, str):
        return "unknown_device"
    try:
        torch.device(device)
    except (RuntimeError, ValueError):
        return "unknown_device"

    if reason:
        return None, reason
    waits = list(node.users)
    if len(waits) != 1 or not is_op(waits[0], c10d.wait_tensor):
        return None, "unsupported_wait_users"
    consumers = _consumers(waits[0])
    if not consumers:
        return None, "no_consumer"
    if kind == "fsdp_reduce_scatter":
        if node.args[1] not in ("sum", "avg"):
            return None, "unsupported_reduce_op"
        if any(user.op != "output" for user in consumers):
            return None, "non_output_rs_consumer"
    if module_fqns is None:
        return None, "outside_module_plan"

def _membership_reason(members: list[Candidate], candidate: Candidate, policy: BucketPolicy) -> Optional[str]:
    """Return why compatibility, uniqueness or capacity rejects the candidate."""
    if members[-1].key != candidate.key:
        return "incompatible_key"
    combined = [*members, candidate]
    param_names = {m.param_name for m in combined}
    state_indices = {m.state_idx for m in combined}
    if len(param_names) != len(combined) or len(state_indices) != len(combined):

        if is_collective(node) and node not in owned:
            return "intervening_collective"
        if node.op not in ("placeholder", "get_attr", "output") and (
                bool(node.meta.get("autograd_backward", False)) != candidate.backward):
            return "phase_boundary"
    return None


def _split_reason(members: list[Candidate], candidate: Candidate, nodes: list, order: dict,

    graph.lint()
    if normalized or advanced or sunk:
        graph_module.recompile()
    if _LOG.isEnabledFor(logging.DEBUG):
        after = {node: index for index, node in enumerate(graph.nodes)}
        for match in matches:
            launch, wait = match.launch, match.realization[0]
            _LOG.debug("FSDP gradient overlap completed: bucket=%s launch=%s index_before=%s index_after=%s "
                       "wait=%s wait_before=%s wait_after=%s",
                       launch.meta.get("fsdp_bucket_id"), launch.name, before.get(launch), after[launch],
                       wait.name, before.get(wait), after[wait])
    graph_module.meta["fsdp_grad_overlap"] = {

        launch: Reduce-scatter launch whose input preparation is inspected.
    """
    inputs = launch.all_input_nodes
    if len(inputs) != 1:
        return []
    source = inputs[0]
    if set(source.users) != {launch} or source.meta.get("autograd_backward") is False:
        return []
    owner = source.meta.get("fsdp_reduce_scatter_prepare")

        if node.meta.get("fsdp_bucket_prepare") != bucket_id:
            continue
        schema = getattr(node.target, "_schema", None)
        if schema is None or schema.name not in ("aten::reshape", "aten::cat", "aten::contiguous"):
            return []
        if node not in preparation:
            preparation.add(node)
            stack.extend(node.all_input_nodes)
    if any(set(node.users) - preparation - {launch} for node in preparation):

    node = wait
    while node.op == "call_function" and str(node.target).startswith("_c10d_functional."):
        node.meta.update(param_name=param_name, state_idx=state_idx)
        if not node.args or not isinstance(node.args[0], fx.Node):
            break
        node = node.args[0]


def _annotate_rs_spec(node: fx.Node, degree: int) -> None:

        return
    shape = spec["shape"]
    static = shape and all(isinstance(dim, int) and not isinstance(dim, bool) and dim > 0 for dim in shape)
    if not static or shape[0] % degree:
        return
    node.meta["fsdp_bucket_spec"] = {
        "input_shape": tuple(shape), "output_shape": (shape[0] // degree, *shape[1:]),
        "dtype": spec["dtype"], "device": spec["device"],
    }

def _annotate_ag_spec(node: fx.Node, degree: int, params: dict) -> None:
    """Attach shard/full shapes of one all-gather node from its parameter."""
    param = params.get(node.meta.get("param_name"))
    if param is None:
        return
    shape = tuple(param.shape)
    if not shape or any(not isinstance(dim, int) or isinstance(dim, bool) or dim <= 0 for dim in shape):
        return
    node.meta["fsdp_bucket_spec"] = {
        "input_shape": shape, "output_shape": (shape[0] * degree, *shape[1:]),
        "dtype": param.dtype, "device": str(param.device),
    }

        if kind not in ("fsdp_all_gather", "fsdp_all_gather_backward", "fsdp_reduce_scatter"):
            continue
        degree = node.meta.get("fsdp_degree")
        if (not isinstance(degree, int) or isinstance(degree, bool)) or degree <= 1:
            continue
        if kind == "fsdp_reduce_scatter":
            _annotate_rs_spec(node, degree)
        else:
            _annotate_ag_spec(node, degree, params)

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

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