Spaces:
Running on Zero
Running on Zero
| # 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) | |
| 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 | |