diff --git a/train.py b/train.py index 4f459ce..c50c6d1 100644 --- a/train.py +++ b/train.py @@ -38,24 +38,18 @@ def train(config_path): errors = [] # ------------------------------------------------------------------ - # Strategy 1: QLoRA 4-bit with device_map="auto" (distributed across GPUs, no FSDP) + # Strategy 1: Load already-quantized 4-bit model (distributed across GPUs) # ------------------------------------------------------------------ - print("\n[1/4] Trying: 4-bit QLoRA (distributed across GPUs, no FSDP)...") + print("\n[1/4] Trying: 4-bit model AS-IS (distributed across GPUs)...") try: - bnb_config = BitsAndBytesConfig( - load_in_4bit=True, - bnb_4bit_quant_type="nf4", - bnb_4bit_compute_dtype=torch.bfloat16, - bnb_4bit_use_double_quant=True, - ) model = AutoModelForCausalLM.from_pretrained( config["base_model"], - quantization_config=bnb_config, device_map="auto", # Distribute layers across GPUs automatically + torch_dtype=torch.float16, trust_remote_code=True, low_cpu_mem_usage=True, ) - print("✓ Success: QLoRA 4-bit distributed across GPUs (no FSDP)") + print("✓ Success: 4-bit model distributed across GPUs (no FSDP)") except Exception as e: errors.append(("QLoRA 4-bit", e)) print(f"✗ Failed: {e}")