Quantize model to 4-bit precision using bitsandbytes
(self, model)
| 1005 | self.device_manager = device_manager |
| 1006 | |
| 1007 | def quantize_model(self, model): |
| 1008 | """Quantize model to 4-bit precision using bitsandbytes""" |
| 1009 | def _replace_linear_layers(module): |
| 1010 | for name, child in module.named_children(): |
| 1011 | if isinstance(child, torch.nn.Linear): |
| 1012 | setattr(module, name, bnb.nn.Linear4bit( |
| 1013 | child.in_features, |
| 1014 | child.out_features, |
| 1015 | bias=child.bias is not None, |
| 1016 | compute_dtype=torch.float16 |
| 1017 | )) |
| 1018 | else: |
| 1019 | _replace_linear_layers(child) |
| 1020 | |
| 1021 | _replace_linear_layers(model) |
| 1022 | return model |
| 1023 | |
| 1024 | def load_base_model(self, model_id: str, quantize: bool = True) -> Tuple[AutoModelForCausalLM, AutoTokenizer]: |
| 1025 | def _load_model(): |