| 1000 | self.device_stats[device]['memory_used'] += memory_delta |
| 1001 | |
| 1002 | class ModelManager: |
| 1003 | def __init__(self, cache_manager: CacheManager, device_manager: DeviceManager): |
| 1004 | self.cache_manager = cache_manager |
| 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(): |
| 1026 | logger.info(f"Loading base model: {model_id}") |
| 1027 | |
| 1028 | device = self.device_manager.get_optimal_device() |
| 1029 | logger.info(f"Using device: {device}") |
| 1030 | |
| 1031 | # Load tokenizer |
| 1032 | tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True, token=os.getenv("HF_TOKEN")) |
| 1033 | |
| 1034 | # Base kwargs for model loading |
| 1035 | model_kwargs = { |
| 1036 | "trust_remote_code": True, |
| 1037 | "device_map": "auto" if 'cuda' in device else device |
| 1038 | } |
| 1039 | |
| 1040 | # Configure device-specific optimizations |
| 1041 | if 'cuda' in device: |
| 1042 | compute_capability = torch.cuda.get_device_capability(0) |
| 1043 | if compute_capability[0] >= 8: |
| 1044 | model_kwargs["torch_dtype"] = torch.bfloat16 |
| 1045 | elif compute_capability[0] >= 7: |
| 1046 | model_kwargs["torch_dtype"] = torch.float16 |
| 1047 | |
| 1048 | # Check for flash attention availability |
| 1049 | try: |
| 1050 | import flash_attn |
| 1051 | has_flash_attn = True |
| 1052 | logger.info("Flash Attention 2 is available") |
| 1053 | model_kwargs["attn_implementation"] = "flash_attention_2" |
| 1054 | except ImportError: |
| 1055 | has_flash_attn = False |
| 1056 | logger.info("Flash Attention 2 is not installed - falling back to default attention") |
| 1057 | |
| 1058 | elif 'mps' in device: |
| 1059 | # Special handling for Gemma models which have NaN issues with float16 on MPS |