diff --git a/training/scripts/train.py b/training/scripts/train.py index 329bdf0..7231ad9 100755 --- a/training/scripts/train.py +++ b/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") - 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).") + print(f"4-bit failed: {e}") + try: + # Try 8-bit with CPU offload + print("Trying 8-bit with CPU offload...") + bnb_config_8bit = BitsAndBytesConfig( + load_in_8bit=True, + llm_int8_enable_fp32_cpu_offload=True, + ) + 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