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"])
+        bnb_config = BitsAndBytesConfig(
-    model_config.quantization_config = None
+            load_in_4bit=True,
-    
+            bnb_4bit_quant_type="nf4",
-    bnb_config = BitsAndBytesConfig(
+            bnb_4bit_compute_dtype=torch.bfloat16,
-        load_in_4bit=True,
+            bnb_4bit_use_double_quant=True,
-        bnb_4bit_quant_type="nf4",
+        )
-        bnb_4bit_compute_dtype=torch.bfloat16,
+        model = AutoModelForCausalLM.from_pretrained(
-        bnb_4bit_use_double_quant=True,
+            config["base_model"],
-    )
+            quantization_config=bnb_config,
-    
+            dtype=torch.bfloat16,
-    # Load with fresh BnB config
+            device_map="auto",
-    model = AutoModelForCausalLM.from_pretrained(
+            trust_remote_code=True,
-        config["base_model"],
+        )
-        quantization_config=bnb_config,
+        print("Model loaded with QLoRA (4-bit).")
-        torch_dtype=torch.bfloat16,
+    except Exception as e:
-        device_map="auto",
+        print(f"4-bit failed: {e}, falling back to bf16 with CPU offload")
-        trust_remote_code=True,
+        model = AutoModelForCausalLM.from_pretrained(
-    )
+            config["base_model"],
-    print("Model loaded with QLoRA (4-bit).")
+            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