aigc-apps-VideoX-Fun/videox_fun/dist/fsdp.py

# Copied from https://github.com/Wan-Video/Wan2.1/blob/main/wan/distributed/fsdp.py
# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
import gc
from functools import partial

import torch
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp import MixedPrecision, ShardingStrategy
from torch.distributed.fsdp.wrap import (lambda_auto_wrap_policy,
                                         transformer_auto_wrap_policy)
from torch.distributed.utils import _free_storage


def find_classes_in_model(model, class_names):
    """
    Recursively find unique module classes in the model that match the given class names.
    
    Args:
        model: The PyTorch model to traverse.
        class_names: A list of class name strings to look for.
        
    Returns:
        A set of matched class types.
    """
    found_classes = set()
    class_names_set = set(class_names)
    
    def traverse(module):
        if module.__class__.__name__ in class_names_set:
            found_classes.add(module.__class__)
        for child in module.children():
            traverse(child)
    
    traverse(model)
    print(f"Found transformer classes: {found_classes}")
    return found_classes


def create_transformer_auto_wrap_policy(
    model,
    transformer_layer_cls_to_wrap,
):
    """
    Creates an auto wrap policy that only wraps modules belonging to the specified 
    transformer layer classes.
    
    Args:
        model: The PyTorch model to analyze for class types.
        transformer_layer_cls_to_wrap: A list of class name strings to wrap.
        
    Returns:
        A callable auto wrap policy function.
    """
    # Dynamically find the actual class types corresponding to the provided names
    transformer_classes = find_classes_in_model(model, transformer_layer_cls_to_wrap)
    
    if not transformer_classes:
        raise ValueError(
            f"No modules found with class names {transformer_layer_cls_to_wrap}. "
            "Please check the class names or the model structure."
        )
    
    def transformer_policy(module, recurse, nonwrapped_numel=None, **kwargs):
        # Use the standard transformer auto wrap policy with the discovered classes
        policy_kwargs = dict(module=module, recurse=recurse, transformer_layer_cls=transformer_classes)
        if nonwrapped_numel is not None:
            policy_kwargs["nonwrapped_numel"] = nonwrapped_numel
        return transformer_auto_wrap_policy(**policy_kwargs)
    
    return transformer_policy


def shard_model(
    model,
    device_id,
    param_dtype=torch.bfloat16,
    reduce_dtype=torch.float32,
    buffer_dtype=torch.float32,
    process_group=None,
    sharding_strategy=ShardingStrategy.FULL_SHARD,
    sync_module_states=True,
    module_to_wrapper=None,
    transformer_layer_cls_to_wrap=None,
    cast_dtype=True,
    ignored_modules=None,
):  
    """
    Wraps the model with FSDP using the specified configuration.
    
    Args:
        model: The PyTorch model to shard.
        device_id: The CUDA device ID.
        param_dtype: Data type for parameters.
        reduce_dtype: Data type for gradient reduction.
        buffer_dtype: Data type for buffers.
        process_group: The process group for distributed training.
        sharding_strategy: The FSDP sharding strategy.
        sync_module_states: Whether to sync module states across ranks.
        module_to_wrapper: Specific modules to wrap if using lambda policy.
        transformer_layer_cls_to_wrap: List of class names to wrap using transformer policy.
        cast_dtype: Whether to cast the managed parameters to `param_dtype` before wrapping. Set False to
            keep a pre-quantized storage dtype (e.g. float8) — MixedPrecision still computes in `param_dtype`.
        ignored_modules: Modules excluded from FSDP (kept replicated in their own dtype via ignored_states).
            Use for modules a mixed-precision checkpoint pins to float32: casting them into the shard dtype
            makes AdaLN/output-head rounding accumulate coherently over the denoising trajectory (flicker).
        
    Returns:
        The FSDP-wrapped model.
    """
    if transformer_layer_cls_to_wrap is not None:
        # Create policy based strictly on transformer layer classes
        auto_wrap_policy = create_transformer_auto_wrap_policy(
            model=model,
            transformer_layer_cls_to_wrap=transformer_layer_cls_to_wrap,
        )
    else:
        # Fallback to lambda policy if no transformer classes are specified
        auto_wrap_policy = partial(
            lambda_auto_wrap_policy, 
            lambda_fn=lambda m: m in (model.blocks if module_to_wrapper is None else module_to_wrapper)
        )

    # FSDP flattens each wrap unit's parameters into one flat buffer and requires a uniform dtype inside it.
    # Models that pin a few modules to fp32 for numerical precision (e.g. MiniMax-H3's embedders/output heads)
    # would otherwise fail with "Must flatten tensors with uniform dtype"; MixedPrecision computes in
    # `param_dtype` anyway, so cast the managed params up front. With `cast_dtype=False` the caller
    # intentionally keeps a different storage dtype (e.g. a pre-applied float8 quantization); FSDP keeps one
    # flat buffer per dtype and MixedPrecision casts to `param_dtype` for the compute. The ignored modules
    # keep their own dtype (typically float32) and stay replicated — a blanket model.to() on them would round
    # the AdaLN modulation and accumulate coherently over the sampling trajectory into temporal flicker.
    ignored_modules = list(ignored_modules) if ignored_modules else []
    ignored_param_ids = {id(p) for m in ignored_modules for p in m.parameters()}
    if cast_dtype and param_dtype is not None:
        for p in model.parameters():
            # A pre-applied float8 quantization stays as the storage dtype (MixedPrecision casts it to
            # `param_dtype` for the compute, matching the non-FSDP qfloat8 dequant wrapper numerics); only
            # the remaining dtypes (e.g. fp32 heads) are homogenized into `param_dtype`.
            if (p.dtype != param_dtype and id(p) not in ignored_param_ids
                    and p.dtype not in (torch.float8_e4m3fn, torch.float8_e5m2)):
                p.data = p.data.to(param_dtype)

    model = FSDP(
        module=model,
        process_group=process_group,
        sharding_strategy=sharding_strategy,
        auto_wrap_policy=auto_wrap_policy,
        mixed_precision=MixedPrecision(
            param_dtype=param_dtype,
            reduce_dtype=reduce_dtype,
            buffer_dtype=buffer_dtype),
        device_id=device_id,
        sync_module_states=sync_module_states,
        ignored_states=ignored_modules if ignored_modules else None)

    # `device_id`/`sync_module_states` only manage the sharded params; the ignored modules must be placed on
    # the device explicitly.
    for m in ignored_modules:
        m.to(device_id)

    return model


def free_model(model):
    """
    Frees memory associated with the FSDP model.
    
    Args:
        model: The FSDP-wrapped model to free.
    """
    for m in model.modules():
        if isinstance(m, FSDP):
            _free_storage(m._handle.flat_param.data)
    del model
    gc.collect()
    torch.cuda.empty_cache()