| import torch, copy | |
| from ..models.utils import init_weights_on_device | |
| def cast_to(weight, dtype, device): | |
| r = torch.empty_like(weight, dtype=dtype, device=device) | |
| r.copy_(weight) | |
| return r | |
| class AutoWrappedModule(torch.nn.Module): | |
| def __init__(self, module: torch.nn.Module, offload_dtype, offload_device, onload_dtype, onload_device, computation_dtype, computation_device): | |
| super().__init__() | |
| self.module = module.to(dtype=offload_dtype, device=offload_device) | |
| self.offload_dtype = offload_dtype | |
| self.offload_device = offload_device | |
| self.onload_dtype = onload_dtype | |
| self.onload_device = onload_device | |
| self.computation_dtype = computation_dtype | |
| self.computation_device = computation_device | |
| self.state = 0 | |
| def offload(self): | |
| if self.state == 1 and (self.offload_dtype != self.onload_dtype or self.offload_device != self.onload_device): | |
| self.module.to(dtype=self.offload_dtype, device=self.offload_device) | |
| self.state = 0 | |
| def onload(self): | |
| if self.state == 0 and (self.offload_dtype != self.onload_dtype or self.offload_device != self.onload_device): | |
| self.module.to(dtype=self.onload_dtype, device=self.onload_device) | |
| self.state = 1 | |
| def forward(self, *args, **kwargs): | |
| if self.onload_dtype == self.computation_dtype and self.onload_device == self.computation_device: | |
| module = self.module | |
| else: | |
| module = copy.deepcopy(self.module).to(dtype=self.computation_dtype, device=self.computation_device) | |
| return module(*args, **kwargs) | |
| class AutoWrappedLinear(torch.nn.Linear): | |
| def __init__(self, module: torch.nn.Linear, offload_dtype, offload_device, onload_dtype, onload_device, computation_dtype, computation_device): | |
| with init_weights_on_device(device=torch.device("meta")): | |
| super().__init__(in_features=module.in_features, out_features=module.out_features, bias=module.bias is not None, dtype=offload_dtype, device=offload_device) | |
| self.weight = module.weight | |
| self.bias = module.bias | |
| self.offload_dtype = offload_dtype | |
| self.offload_device = offload_device | |
| self.onload_dtype = onload_dtype | |
| self.onload_device = onload_device | |
| self.computation_dtype = computation_dtype | |
| self.computation_device = computation_device | |
| self.state = 0 | |
| def offload(self): | |
| if self.state == 1 and (self.offload_dtype != self.onload_dtype or self.offload_device != self.onload_device): | |
| self.to(dtype=self.offload_dtype, device=self.offload_device) | |
| self.state = 0 | |
| def onload(self): | |
| if self.state == 0 and (self.offload_dtype != self.onload_dtype or self.offload_device != self.onload_device): | |
| self.to(dtype=self.onload_dtype, device=self.onload_device) | |
| self.state = 1 | |
| def forward(self, x, *args, **kwargs): | |
| if self.onload_dtype == self.computation_dtype and self.onload_device == self.computation_device: | |
| weight, bias = self.weight, self.bias | |
| else: | |
| weight = cast_to(self.weight, self.computation_dtype, self.computation_device) | |
| bias = None if self.bias is None else cast_to(self.bias, self.computation_dtype, self.computation_device) | |
| return torch.nn.functional.linear(x, weight, bias) | |
| def enable_vram_management_recursively(model: torch.nn.Module, module_map: dict, module_config: dict, max_num_param=None, overflow_module_config: dict = None, total_num_param=0): | |
| for name, module in model.named_children(): | |
| for source_module, target_module in module_map.items(): | |
| if isinstance(module, source_module): | |
| num_param = sum(p.numel() for p in module.parameters()) | |
| if max_num_param is not None and total_num_param + num_param > max_num_param: | |
| module_config_ = overflow_module_config | |
| else: | |
| module_config_ = module_config | |
| module_ = target_module(module, **module_config_) | |
| setattr(model, name, module_) | |
| total_num_param += num_param | |
| break | |
| else: | |
| total_num_param = enable_vram_management_recursively(module, module_map, module_config, max_num_param, overflow_module_config, total_num_param) | |
| return total_num_param | |
| def enable_vram_management(model: torch.nn.Module, module_map: dict, module_config: dict, max_num_param=None, overflow_module_config: dict = None): | |
| enable_vram_management_recursively(model, module_map, module_config, max_num_param, overflow_module_config, total_num_param=0) | |
| model.vram_management_enabled = True | |