remove auto device map in training strategies

2026-07-01 16:50:48 -04:00
parent cc353b3dea
commit 19f405255c
1 changed files with 2 additions and 2 deletions
--- a/train.py
+++ b/train.py
@@ -41,7 +41,7 @@ def train(config_path):
        model = AutoModelForCausalLM.from_pretrained(
            config["base_model"],
            torch_dtype=torch.float16,
-            device_map="auto",
+            # device_map="auto",
            trust_remote_code=True,
        )
        print("✓ Success: 4-bit AS-IS")
@@ -81,7 +81,7 @@ def train(config_path):
                    model = AutoModelForCausalLM.from_pretrained(
                        config["base_model"],
                        torch_dtype=torch.bfloat16,
-                        device_map="auto",
+                        # device_map="auto",
                        low_cpu_mem_usage=True,
                        trust_remote_code=True,
                    )