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

fix: use device_map=auto to distribute model across GPUs

Browse Source
This commit is contained in:
Christian Medina
2026-06-30 23:45:57 -04:00
parent 8d4cdca3bb
commit 114e25fcfb
1 changed files with 2 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
training/scripts/train.py
View File

@@ -34,10 +34,11 @@ def train(config_path):
# Load model in bf16 (fp8 not supported for training) # Load model in bf16 (fp8 not supported for training)
# Disable quantization config to train full precision # Disable quantization config to train full precision
# Use device_map="auto" to distribute across GPUs from the start
model = AutoModelForCausalLM.from_pretrained( model = AutoModelForCausalLM.from_pretrained(
config["base_model"], config["base_model"],
torch_dtype=torch.bfloat16, torch_dtype=torch.bfloat16,
device_map="cpu", device_map="auto", # Distribute across GPUs
quantization_config=None, # Override any fp8 quantization quantization_config=None, # Override any fp8 quantization
trust_remote_code=True, trust_remote_code=True,
) )