diff --git a/train.py b/train.py index c7d3175..a6ef3ce 100644 --- a/train.py +++ b/train.py @@ -50,11 +50,11 @@ def train(config_path): model = AutoModelForCausalLM.from_pretrained( config["base_model"], quantization_config=bnb_config, - device_map="auto", + device_map="cpu", # Load to CPU, FSDP shards later trust_remote_code=True, low_cpu_mem_usage=True, ) - print("✓ Success: QLoRA 4-bit") + print("✓ Success: QLoRA 4-bit loaded to CPU") except Exception as e: errors.append(("QLoRA 4-bit", e)) print(f"✗ Failed: {e}")