feat: add fallback to bf16 with CPU offload if 4-bit fails

training/scripts/train.py

@@ -32,30 +32,34 @@ def train(config_path):
     
     print(f"Loading model: {config['base_model']}")
     
-    # Load model (skip broken quantization config)
+    # Load model - try 4-bit first, fall back to bf16 with CPU offload
     print(f"Loading model: {config['base_model']}")
     from transformers import BitsAndBytesConfig, AutoConfig
     
-    # Remove broken quantization config
+    try:
-    model_config = AutoConfig.from_pretrained(config["base_model"])
-    model_config.quantization_config = None
-    
         bnb_config = BitsAndBytesConfig(
             load_in_4bit=True,
             bnb_4bit_quant_type="nf4",
             bnb_4bit_compute_dtype=torch.bfloat16,
             bnb_4bit_use_double_quant=True,
         )
-    
-    # Load with fresh BnB config
         model = AutoModelForCausalLM.from_pretrained(
             config["base_model"],
             quantization_config=bnb_config,
-        torch_dtype=torch.bfloat16,
+            dtype=torch.bfloat16,
             device_map="auto",
             trust_remote_code=True,
         )
         print("Model loaded with QLoRA (4-bit).")
+    except Exception as e:
+        print(f"4-bit failed: {e}, falling back to bf16 with CPU offload")
+        model = AutoModelForCausalLM.from_pretrained(
+            config["base_model"],
+            torch_dtype=torch.bfloat16,
+            device_map="cpu",
+            trust_remote_code=True,
+        )
+        print("Model loaded as bf16 (CPU offload).")
     
     # Prepare model for k-bit training
     from peft import prepare_model_for_kbit_training