fix: load QLoRA model to CPU first, FSDP shards later
This commit is contained in:
4
train.py
4
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}")
|
||||
|
||||
Reference in New Issue
Block a user