Zengtudor/Wan2.1
1
0
Fork 0
mirror of https://github.com/Wan-Video/Wan2.1.git synced 2025-06-05 14:54:54 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

add free stroage api for FSDP (#277)

Browse Source 
Signed-off-by: wenting.swt <wenting.swt@alibaba-inc.com>
This commit is contained in:
Wenting Shen 2025-03-26 15:49:37 +08:00 committed by GitHub
parent 11a6a00822
commit bc3249d61c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 10 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
wan/distributed/fsdp.py
View File

@ -1,11 +1,12 @@
# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. # Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
import gc
from functools import partial from functools import partial
import torch import torch
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp import MixedPrecision, ShardingStrategy from torch.distributed.fsdp import MixedPrecision, ShardingStrategy
from torch.distributed.fsdp.wrap import lambda_auto_wrap_policy from torch.distributed.fsdp.wrap import lambda_auto_wrap_policy
from torch.distributed.utils import _free_storage
def shard_model( def shard_model(
model, model,
@ -30,3 +31,11 @@ def shard_model(
device_id=device_id, device_id=device_id,
sync_module_states=sync_module_states) sync_module_states=sync_module_states)
return model return model
def free_model(model):
for m in model.modules():
if isinstance(m, FSDP):
_free_storage(m._handle.flat_param.data)
del model
gc.collect()
torch.cuda.empty_cache()