Spaces:
Running on Zero
Running on Zero
File size: 6,289 Bytes
f0395ef | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 | # Copyright (c) 2025 Hansheng Chen
import importlib
import torch
import torch.nn as nn
import torch.distributed as dist
try:
from torch.distributed.fsdp import fully_shard, MixedPrecisionPolicy
except:
pass
from mmcv.parallel.scatter_gather import scatter_kwargs
from mmcv.parallel import MODULE_WRAPPERS
def get_module_object(path):
module_path, attribute = path.rsplit('.', 1)
module = importlib.import_module(module_path)
return getattr(module, attribute)
@MODULE_WRAPPERS.register_module()
class FSDP2Wrapper(nn.Module):
def __init__(
self,
module,
wrap_frozen_modules=False,
ignore_frozen_parameters=False,
param_dtype='bfloat16',
reduce_dtype='float32',
fsdp_modules=None,
exclude_keys=(),
hybrid_sharding=True,
**kwargs):
super().__init__()
self.module = module
fsdp_kwargs = kwargs
self.param_dtype = getattr(torch, param_dtype)
self.reduce_dtype = getattr(torch, reduce_dtype)
if hybrid_sharding:
global_world_size = dist.get_world_size()
num_devices_per_node = torch.cuda.device_count()
mesh = dist.init_device_mesh(
'cuda',
(global_world_size // num_devices_per_node, num_devices_per_node),
mesh_dim_names=('replicate', 'shard'))
fsdp_kwargs.update(mesh=mesh)
if fsdp_modules is not None:
assert isinstance(fsdp_modules, (list, tuple))
fsdp_modules = tuple([get_module_object(m) for m in fsdp_modules])
self.to_fsdp(
wrap_frozen_modules,
ignore_frozen_parameters,
exclude_keys,
fsdp_modules=fsdp_modules,
**fsdp_kwargs)
def to_fsdp(self,
wrap_frozen_modules=False,
ignore_frozen_parameters=False,
exclude_keys=(),
fsdp_modules=(),
**kwargs):
for name, module in self.module._modules.items():
if name in exclude_keys or next(module.parameters(), None) is None:
module = module.cuda()
elif all(not p.requires_grad for p in module.parameters()):
if wrap_frozen_modules:
for submodule in module.modules():
if isinstance(submodule, fsdp_modules):
fsdp_kwargs = kwargs.copy()
fsdp_kwargs.update(
mp_policy=MixedPrecisionPolicy(
param_dtype=self.param_dtype,
reduce_dtype=self.reduce_dtype))
fully_shard(submodule, **fsdp_kwargs)
fsdp_kwargs = kwargs.copy()
fsdp_kwargs.update(
mp_policy=MixedPrecisionPolicy(
param_dtype=self.param_dtype,
reduce_dtype=self.reduce_dtype,
cast_forward_inputs=False))
fully_shard(module, **fsdp_kwargs)
else:
module = module.cuda()
else:
if ignore_frozen_parameters:
ignored_params = []
for p in module.parameters():
if not p.requires_grad:
p.data = p.data.cuda()
ignored_params.append(p)
else:
ignored_params = None
for submodule in module.modules():
if isinstance(submodule, fsdp_modules):
fsdp_kwargs = kwargs.copy()
fsdp_kwargs.update(
mp_policy=MixedPrecisionPolicy(
param_dtype=self.param_dtype,
reduce_dtype=self.reduce_dtype))
if ignored_params is not None: # requires torch >= 2.7
fsdp_kwargs.update(ignored_params=ignored_params)
fully_shard(submodule, **fsdp_kwargs)
fsdp_kwargs = kwargs.copy()
fsdp_kwargs.update(
mp_policy=MixedPrecisionPolicy(
param_dtype=self.param_dtype,
reduce_dtype=self.reduce_dtype,
cast_forward_inputs=False))
if ignored_params is not None: # requires torch >= 2.7
fsdp_kwargs.update(ignored_params=ignored_params)
fully_shard(module, **fsdp_kwargs)
self.module._modules[name] = module
def scatter(self, inputs, kwargs, device_ids):
"""Scatter function.
Args:
inputs (Tensor): Input Tensor.
kwargs (dict): Args for
``mmcv.parallel.scatter_gather.scatter_kwargs``.
device_ids (int): Device id.
"""
return scatter_kwargs(inputs, kwargs, device_ids)
def forward(self, *inputs, **kwargs):
"""Forward function.
Args:
inputs (tuple): Input data.
kwargs (dict): Args for
``mmcv.parallel.scatter_gather.scatter_kwargs``.
"""
inputs, kwargs = self.scatter(inputs, kwargs, [torch.cuda.current_device()])
return self.module(*inputs[0], **kwargs[0])
def train_step(self, *inputs, **kwargs):
"""Train step function.
Args:
inputs (Tensor): Input Tensor.
kwargs (dict): Args for
``mmcv.parallel.scatter_gather.scatter_kwargs``.
"""
inputs, kwargs = self.scatter(inputs, kwargs, [torch.cuda.current_device()])
output = self.module.train_step(*inputs[0], **kwargs[0])
return output
def val_step(self, *inputs, **kwargs):
"""Validation step function.
Args:
inputs (tuple): Input data.
kwargs (dict): Args for ``scatter_kwargs``.
"""
inputs, kwargs = self.scatter(inputs, kwargs, [torch.cuda.current_device()])
output = self.module.val_step(*inputs[0], **kwargs[0])
return output
|