durgappc/infinitetalk

1import copy2 3import torch4 5from src.utils import init_weights_on_device6import optimum.quanto.nn.qlinear as qlinear7 8def cast_to(weight, dtype, device):9    r = torch.empty_like(weight, dtype=dtype, device=device)10    r.copy_(weight)11    return r12 13def cast_to_device(weight, device):14    if hasattr(weight, '__class__') and 'optimum.quanto' in str(weight.__class__):  15        return weight.to(device)16    else:17        r = torch.empty_like(weight, device=device)18        r.copy_(weight)19        return r20 21class AutoWrappedModule(torch.nn.Module):22    def __init__(23        self,24        module: torch.nn.Module,25        offload_dtype,26        offload_device,27        onload_dtype,28        onload_device,29        computation_dtype,30        computation_device,31    ):32        super().__init__()33        self.module = module.to(dtype=offload_dtype, device=offload_device)34        self.offload_dtype = offload_dtype35        self.offload_device = offload_device36        self.onload_dtype = onload_dtype37        self.onload_device = onload_device38        self.computation_dtype = computation_dtype39        self.computation_device = computation_device40        self.state = 041 42    def offload(self):43        if self.state == 1 and (44            self.offload_dtype != self.onload_dtype45            or self.offload_device != self.onload_device46        ):47            self.module.to(dtype=self.offload_dtype, device=self.offload_device)48            self.state = 049 50    def onload(self):51        if self.state == 0 and (52            self.offload_dtype != self.onload_dtype53            or self.offload_device != self.onload_device54        ):55            self.module.to(dtype=self.onload_dtype, device=self.onload_device)56            self.state = 157 58    def forward(self, *args, **kwargs):59        if (60            self.onload_dtype == self.computation_dtype61            and self.onload_device == self.computation_device62        ):63            module = self.module64        else:65            module = copy.deepcopy(self.module).to(66                dtype=self.computation_dtype, device=self.computation_device67            )68        return module(*args, **kwargs)69 70 71 72class AutoWrappedQLinear(qlinear.QLinear):73    def __init__(74        self,75        module: qlinear.QLinear,76        offload_dtype,77        offload_device,78        onload_dtype,79        onload_device,80        computation_dtype,81        computation_device,82    ):83        with init_weights_on_device(device=torch.device("meta")):84            super().__init__(85                in_features=module.in_features,86                out_features=module.out_features,87                bias=module.bias is not None,88                device=offload_device,89            )90        self.weight = module.weight91        self.bias = module.bias92        self.offload_device = offload_device93 94        self.onload_device = onload_device95        self.computation_device = computation_device96        self.state = 097 98    def offload(self):99        if self.state == 1 and (100             self.offload_device != self.onload_device101        ):102            self.to(device=self.offload_device)103            self.state = 0104 105    def onload(self):106        if self.state == 0 and (107            self.offload_device != self.onload_device108        ):109            self.to(device=self.onload_device)110            self.state = 1111 112    def forward(self, x, *args, **kwargs):113        if (114            self.onload_device == self.computation_device115        ):116            117            return torch.nn.functional.linear(x, self.weight, bias=self.bias)118        else:119            120            qweight = cast_to_device(self.weight, self.computation_device)121            bias = (122                None123                if self.bias is None124                else cast_to_device(self.bias, self.computation_device)125            )126            return torch.nn.functional.linear(x, qweight, bias)127 128class AutoWrappedLinear(torch.nn.Linear):129    def __init__(130        self,131        module: torch.nn.Linear,132        offload_dtype,133        offload_device,134        onload_dtype,135        onload_device,136        computation_dtype,137        computation_device,138    ):139        with init_weights_on_device(device=torch.device("meta")):140            super().__init__(141                in_features=module.in_features,142                out_features=module.out_features,143                bias=module.bias is not None,144                dtype=offload_dtype,145                device=offload_device,146            )147        self.weight = module.weight148        self.bias = module.bias149        self.offload_dtype = offload_dtype150        self.offload_device = offload_device151        self.onload_dtype = onload_dtype152        self.onload_device = onload_device153        self.computation_dtype = computation_dtype154        self.computation_device = computation_device155        self.state = 0156 157    def offload(self):158        if self.state == 1 and (159            self.offload_dtype != self.onload_dtype160            or self.offload_device != self.onload_device161        ):162            self.to(dtype=self.offload_dtype, device=self.offload_device)163            self.state = 0164 165    def onload(self):166        if self.state == 0 and (167            self.offload_dtype != self.onload_dtype168            or self.offload_device != self.onload_device169        ):170            self.to(dtype=self.onload_dtype, device=self.onload_device)171            self.state = 1172 173    def forward(self, x, *args, **kwargs):174        if (175            self.onload_dtype == self.computation_dtype176            and self.onload_device == self.computation_device177        ):178            weight, bias = self.weight, self.bias179        else:180            weight = cast_to(181                self.weight, self.computation_dtype, self.computation_device182            )183            bias = (184                None185                if self.bias is None186                else cast_to(self.bias, self.computation_dtype, self.computation_device)187            )188        return torch.nn.functional.linear(x, weight, bias)189 190 191def enable_vram_management_recursively(192    model: torch.nn.Module,193    module_map: dict,194    module_config: dict,195    max_num_param=None,196    overflow_module_config: dict = None,197    total_num_param=0,198):199    for name, module in model.named_children():200        for source_module, target_module in module_map.items():201            if isinstance(module, source_module):202                num_param = sum(p.numel() for p in module.parameters())203                # print(str(module) + ':' + str(num_param))204                if (205                    max_num_param is not None206                    and total_num_param + num_param > max_num_param207                ):208                    # print(str(module) + '-->\t\t num:' + str(num_param) + "\t total:" + str(total_num_param))209                    module_config_ = overflow_module_config210                else:211                    module_config_ = module_config212                module_ = target_module(module, **module_config_)213                setattr(model, name, module_)214                total_num_param += num_param215                break216        else:217            total_num_param = enable_vram_management_recursively(218                module,219                module_map,220                module_config,221                max_num_param,222                overflow_module_config,223                total_num_param,224            )225    return total_num_param226 227 228def enable_vram_management(229    model: torch.nn.Module,230    module_map: dict,231    module_config: dict,232    max_num_param=None,233    overflow_module_config: dict = None,234):235    enable_vram_management_recursively(236        model,237        module_map,238        module_config,239        max_num_param,240        overflow_module_config,241        total_num_param=0,242    )243    model.vram_management_enabled = True244