feat: add 8-bit CPU offload fallback between 4-bit and bf16

training/scripts/train.py

@@ -37,6 +37,7 @@ def train(config_path):
     from transformers import BitsAndBytesConfig, AutoConfig
     
     try:
+        # Try 4-bit QLoRA first
         bnb_config = BitsAndBytesConfig(
             load_in_4bit=True,
             bnb_4bit_quant_type="nf4",
@@ -52,14 +53,30 @@ def train(config_path):
         )
         print("Model loaded with QLoRA (4-bit).")
     except Exception as e:
-        print(f"4-bit failed: {e}, falling back to bf16 with CPU offload")
+        print(f"4-bit failed: {e}")
-        model = AutoModelForCausalLM.from_pretrained(
+        try:
-            config["base_model"],
+            # Try 8-bit with CPU offload
-            torch_dtype=torch.bfloat16,
+            print("Trying 8-bit with CPU offload...")
-            device_map="cpu",
+            bnb_config_8bit = BitsAndBytesConfig(
-            trust_remote_code=True,
+                load_in_8bit=True,
-        )
+                llm_int8_enable_fp32_cpu_offload=True,
-        print("Model loaded as bf16 (CPU offload).")
+            )
+            model = AutoModelForCausalLM.from_pretrained(
+                config["base_model"],
+                quantization_config=bnb_config_8bit,
+                device_map="auto",
+                trust_remote_code=True,
+            )
+            print("Model loaded with 8-bit CPU offload.")
+        except Exception as e2:
+            print(f"8-bit failed: {e2}, falling back to bf16 with DeepSpeed 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 (DeepSpeed CPU offload).")
     
     # Prepare model for k-bit training
     from peft import prepare_model_for_kbit_training