fix: use device_map=auto to distribute model across GPUs

training/scripts/train.py

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