Load quantized model
(model, load_path)
| 249 | |
| 250 | |
| 251 | def load_quantized_model(model, load_path): |
| 252 | """Load quantized model""" |
| 253 | |
| 254 | for layer_index in range(len(model.model.layers)): |
| 255 | model.model.layers[layer_index] = torch.load( |
| 256 | os.path.join(load_path, str(layer_index) + ".pth"), |
| 257 | map_location=model.model.layers[layer_index].input_layernorm.weight.device, |
| 258 | ) |
| 259 | for module in model.model.layers[layer_index].modules(): |
| 260 | if isinstance(module, QuantizedWeight): |
| 261 | if not hasattr(module, "codes_storage"): |
| 262 | module.codes_storage = None # backwards compatibility |
| 263 | |
| 264 | model.load_state_dict(torch.load(os.path.join(load_path, "not_quantized_weights.pt")), strict=False) |
| 265 | return model |
| 266 | |
| 267 | |
| 268 | def save_not_quantized_weights(model: nn.Module, save_dir: str): |
no test coverage detected