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

remove auto device map in training strategies

Browse Source
This commit is contained in:
Christian Medina
2026-07-01 16:50:48 -04:00
parent cc353b3dea
commit 19f405255c
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
train.py
View File

@@ -41,7 +41,7 @@ def train(config_path):
model = AutoModelForCausalLM.from_pretrained( model = AutoModelForCausalLM.from_pretrained(
config["base_model"], config["base_model"],
torch_dtype=torch.float16, torch_dtype=torch.float16,
device_map="auto", # device_map="auto",
trust_remote_code=True, trust_remote_code=True,
) )
print("✓ Success: 4-bit AS-IS") print("✓ Success: 4-bit AS-IS")
@@ -81,7 +81,7 @@ def train(config_path):
model = AutoModelForCausalLM.from_pretrained( model = AutoModelForCausalLM.from_pretrained(
config["base_model"], config["base_model"],
torch_dtype=torch.bfloat16, torch_dtype=torch.bfloat16,
device_map="auto", # device_map="auto",
low_cpu_mem_usage=True, low_cpu_mem_usage=True,
trust_remote_code=True, trust_remote_code=True,
) )