durgappc/infinitetalk2

1# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.2import gc3from functools import partial4 5import torch6from torch.distributed.fsdp import FullyShardedDataParallel as FSDP7from torch.distributed.fsdp import MixedPrecision, ShardingStrategy8from torch.distributed.fsdp.wrap import lambda_auto_wrap_policy9from torch.distributed.utils import _free_storage10 11 12def shard_model(13    model,14    device_id,15    param_dtype=torch.bfloat16,16    reduce_dtype=torch.float32,17    buffer_dtype=torch.float32,18    process_group=None,19    sharding_strategy=ShardingStrategy.FULL_SHARD,20    sync_module_states=True,21):22    model = FSDP(23        module=model,24        process_group=process_group,25        sharding_strategy=sharding_strategy,26        auto_wrap_policy=partial(27            lambda_auto_wrap_policy, lambda_fn=lambda m: m in model.blocks),28        # mixed_precision=MixedPrecision(29        #     param_dtype=param_dtype,30        #     reduce_dtype=reduce_dtype,31        #     buffer_dtype=buffer_dtype),32        device_id=device_id,33        sync_module_states=sync_module_states)34    return model35 36 37def free_model(model):38    for m in model.modules():39        if isinstance(m, FSDP):40            _free_storage(m._handle.flat_param.data)41    del model42    gc.collect()43    torch.cuda.empty_cache()44