feat: manual 4-bit quantization with PEFT prepare + Linear4bit

train.py

@@ -33,23 +33,48 @@ def train(config_path):
     
     print(f"Loading model: {config['base_model']}")
     
-    # Load model with BnB 4-bit on CPU, then move to GPU
-    print(f"\n[INFO] Loading {config['base_model']} with BnB 4-bit to CPU...")
-    from transformers import BitsAndBytesConfig
-    bnb_config = BitsAndBytesConfig(
-        load_in_4bit=True,
-        bnb_4bit_quant_type="nf4",
-        bnb_4bit_compute_dtype=torch.float16,
-    )
+    # Load bf16 model to CPU
+    print(f"\n[INFO] Loading {config['base_model']} bf16 to CPU...")
     model = AutoModelForCausalLM.from_pretrained(
         config["base_model"],
         device_map="cpu",
-        quantization_config=bnb_config,
-        torch_dtype=torch.float16,
+        torch_dtype=torch.bfloat16,
         trust_remote_code=True,
         low_cpu_mem_usage=True,
     )
-    print(f"   CPU RAM: {torch.cuda.memory_allocated(0) / 1e9:.2f} GB (model on CPU)")
+    print("✓ Model loaded to CPU (~70GB bf16)")
+    
+    # Apply PEFT k-bit training preparation
+    print("   Applying PEFT k-bit preparation...")
+    from peft import prepare_model_for_kbit_training
+    model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=False)
+    print("   ✓ Model prepared for k-bit training")
+    
+    # Manually quantize linear layers
+    print("   Quantizing linear layers to 4-bit...")
+    from bitsandbytes.nn import Linear4bit
+    from torch import nn
+    
+    quantized_count = 0
+    for name, module in model.named_modules():
+        if isinstance(module, nn.Linear) and 'lm_head' not in name:
+            new_module = Linear4bit(
+                module.in_features,
+                module.out_features,
+                bias=module.bias is not None,
+            )
+            new_module.weight = nn.Parameter(module.weight.data.clone())
+            if module.bias is not None:
+                new_module.bias = nn.Parameter(module.bias.data.clone())
+            
+            layers = name.split('.')
+            parent = model
+            for layer in layers[:-1]:
+                parent = getattr(parent, layer)
+            setattr(parent, layers[-1], new_module)
+            quantized_count += 1
+    
+    print(f"   ✓ Quantized {quantized_count} linear layers to 4-bit")
     
     # Move to GPU
     print("   Moving to GPU 0...")