This (should) be responsible for de-quantizing on the fly
| 91 | return self.tensor_shape |
| 92 | |
| 93 | class GGMLLayer(torch.nn.Module): |
| 94 | """ |
| 95 | This (should) be responsible for de-quantizing on the fly |
| 96 | """ |
| 97 | comfy_cast_weights = True |
| 98 | dequant_dtype = None |
| 99 | patch_dtype = None |
| 100 | largest_layer = False |
| 101 | torch_compatible_tensor_types = {None, gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16} |
| 102 | |
| 103 | def is_ggml_quantized(self, *, weight=None, bias=None): |
| 104 | if weight is None: |
| 105 | weight = self.weight |
| 106 | if bias is None: |
| 107 | bias = self.bias |
| 108 | return is_quantized(weight) or is_quantized(bias) |
| 109 | |
| 110 | def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs): |
| 111 | weight, bias = state_dict.get(f"{prefix}weight"), state_dict.get(f"{prefix}bias") |
| 112 | # NOTE: using modified load for linear due to not initializing on creation, see GGMLOps todo |
| 113 | if self.is_ggml_quantized(weight=weight, bias=bias) or isinstance(self, torch.nn.Linear): |
| 114 | return self.ggml_load_from_state_dict(state_dict, prefix, *args, **kwargs) |
| 115 | # Not strictly required, but fixes embedding shape mismatch. Threshold set in loader.py |
| 116 | if isinstance(self, torch.nn.Embedding) and self.weight.shape[0] >= (64 * 1024): |
| 117 | return self.ggml_load_from_state_dict(state_dict, prefix, *args, **kwargs) |
| 118 | return super()._load_from_state_dict(state_dict, prefix, *args, **kwargs) |
| 119 | |
| 120 | def ggml_load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs): |
| 121 | prefix_len = len(prefix) |
| 122 | for k,v in state_dict.items(): |
| 123 | if k[prefix_len:] == "weight": |
| 124 | self.weight = torch.nn.Parameter(v, requires_grad=False) |
| 125 | elif k[prefix_len:] == "bias" and v is not None: |
| 126 | self.bias = torch.nn.Parameter(v, requires_grad=False) |
| 127 | else: |
| 128 | unexpected_keys.append(k) |
| 129 | |
| 130 | # For Linear layer with missing weight |
| 131 | if self.weight is None and isinstance(self, torch.nn.Linear): |
| 132 | v = torch.zeros(self.in_features, self.out_features) |
| 133 | self.weight = torch.nn.Parameter(v, requires_grad=False) |
| 134 | missing_keys.append(prefix+"weight") |
| 135 | |
| 136 | # for vram estimation (TODO: less fragile logic?) |
| 137 | if getattr(self.weight, "is_largest_weight", False): |
| 138 | self.largest_layer = True |
| 139 | |
| 140 | def _save_to_state_dict(self, *args, **kwargs): |
| 141 | if self.is_ggml_quantized(): |
| 142 | return self.ggml_save_to_state_dict(*args, **kwargs) |
| 143 | return super()._save_to_state_dict(*args, **kwargs) |
| 144 | |
| 145 | def ggml_save_to_state_dict(self, destination, prefix, keep_vars): |
| 146 | # This is a fake state dict for vram estimation |
| 147 | weight = torch.zeros_like(self.weight, device=torch.device("meta")) |
| 148 | destination[prefix + "weight"] = weight |
| 149 | if self.bias is not None: |
| 150 | bias = torch.zeros_like(self.bias, device=torch.device("meta")) |
nothing calls this directly
no outgoing calls
no test coverage detected