cmedina/agenx-lora-training
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: load model to CPU first, let DeepSpeed distribute

Browse Source
This commit is contained in:
Christian Medina
2026-07-01 00:03:24 -04:00
parent 2c1251788d
commit af09d47c6f
1 changed files with 4 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
training/scripts/train.py
View File

@@ -32,13 +32,14 @@ def train(config_path):
print(f"Loading model: {config['base_model']}") print(f"Loading model: {config['base_model']}")
# Load model - preserve NVFP4 quantization # Load model to CPU first, DeepSpeed will distribute
# Use device_map="auto" to distribute across GPUs from the start print("Loading model to CPU...")
model = AutoModelForCausalLM.from_pretrained( model = AutoModelForCausalLM.from_pretrained(
config["base_model"], config["base_model"],
device_map="auto", # Distribute across GPUs device_map="cpu", # Load to CPU first
trust_remote_code=True, trust_remote_code=True,
) )
print(f"Model loaded to CPU. Moving to GPU with DeepSpeed...")
# Add LoRA # Add LoRA
lora_config = LoraConfig( lora_config = LoraConfig(