Repository: huggingface/diffusers
File: src/diffusers/training_utils.py (Python)

Code before:
            best_bucket_idx = bucket_idx
    return best_bucket_idx


def get_fsdp_kwargs_from_accelerator(accelerator) -> dict:
    """
    Extract and convert FSDP config from Accelerator into PyTorch FSDP kwargs.
    """

    kwargs = {}
    fsdp_plugin = accelerator.state.fsdp_plugin

    if fsdp_plugin is None:
        # FSDP not enabled in Accelerator
        kwargs["sharding_strategy"] = ShardingStrategy.FULL_SHARD
    else:
        # FSDP is enabled → use plugin's strategy, or default if None
        kwargs["sharding_strategy"] = (
            fsdp_plugin.sharding_strategy or ShardingStrategy.FULL_SHARD
        )

    return kwargs


def wrap_with_fsdp(
        model: torch.nn.Module,
        device: Union[str, torch.device],
        offload: bool = True,
        use_orig_params: bool = True,
        limit_all_gathers: bool = True,
        fsdp_kwargs: Optional[Dict[str, Any]] = None,
        transformer_layer_cls: Optional[Set[Type[torch.nn.Module]]] = None,
) -> FSDP:
    """
    Wrap a model with FSDP using common defaults and optional transformer auto-wrapping.

    Args:
        model: Model to wrap
        device: Target device (e.g., accelerator.device)
        offload: Whether to enable CPU parameter offloading
        use_orig_params: Whether to use original parameters
        limit_all_gathers: Whether to limit all gathers
        fsdp_kwargs: FSDP arguments (sharding_strategy, etc.) — usually from Accelerate config
        transformer_layer_cls: Classes for auto-wrapping (if not using policy from fsdp_kwargs)

    Returns:
        FSDP-wrapped model
    """

    if transformer_layer_cls is None:
        # Set the default layers if transformer_layer_cls is not provided

Review comment (suggestion):
```suggestion
    fsdp_plugin = getattr(accelerator.state, "fsdp_plugin", None)
    if fsdp_plugin is None:
        raise ValueError("Accelerate isn't configured to handle FSDP. Please update your installation.")
```

Code after:
    min_metric = float("inf")
    best_bucket_idx = None
    for bucket_idx, (bucket_h, bucket_w) in enumerate(bucket_options):
        metric = abs(h * bucket_w - w * bucket_h)
        if metric <= min_metric:
            min_metric = metric
            best_bucket_idx = bucket_idx
    return best_bucket_idx


def _to_cpu_contiguous(state_dicts) -> dict:
    return {k: v.detach().cpu().contiguous() if isinstance(v, torch.Tensor) else v for k, v in state_dicts.items()}


def get_fsdp_kwargs_from_accelerator(accelerator) -> dict:
    """
    Extract and convert FSDP config from Accelerator into PyTorch FSDP kwargs.
    """

    kwargs = {}
    fsdp_state = getattr(accelerator.state, "fsdp_plugin", None)

    if fsdp_state is None:
        raise ValueError("Accelerate isn't configured to handle FSDP. Please update your installation.")

    fsdp_plugin = accelerator.state.fsdp_plugin

    if fsdp_plugin is None:
        # FSDP not enabled in Accelerator
        kwargs["sharding_strategy"] = ShardingStrategy.FULL_SHARD
    else:
        # FSDP is enabled → use plugin's strategy, or default if None
        kwargs["sharding_strategy"] = fsdp_plugin.sharding_strategy or ShardingStrategy.FULL_SHARD

    return kwargs


def wrap_with_fsdp(
    model: torch.nn.Module,
    device: Union[str, torch.device],
    offload: bool = True,
    use_orig_params: bool = True,
    limit_all_gathers: bool = True,
    fsdp_kwargs: Optional[Dict[str, Any]] = None,
    transformer_layer_cls: Optional[Set[Type[torch.nn.Module]]] = None,
) -> FSDP:
    """
    Wrap a model with FSDP using common defaults and optional transformer auto-wrapping.

    Args:
        model: Model to wrap