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

fix: use DeepSpeed ZeRO-3 directly (skip CPU loading)

Browse Source
This commit is contained in:
Christian Medina
2026-07-01 14:41:48 -04:00
parent 3760aaf316
commit 3fa2cfe4cc
1 changed files with 27 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files

34
training/scripts/train.py
View File

@@ -59,36 +59,8 @@ def train(config_path):
except Exception as e: except Exception as e:
print(f"✗ Failed: {e}") print(f"✗ Failed: {e}")
# Strategy 2: Load bf16 to CPU with low_cpu_mem_usage # Strategy 2: DeepSpeed ZeRO-3 (distribute across GPUs)
print("\n[2/4] Trying: bf16 model to CPU...") print("\n[2/4] Trying: DeepSpeed ZeRO-3...")
try:
model = AutoModelForCausalLM.from_pretrained(
config["base_model"],
torch_dtype=torch.bfloat16,
device_map="cpu",
low_cpu_mem_usage=True,
trust_remote_code=True,
)
print("✓ Success: bf16 model loaded to CPU")
except Exception as e:
print(f"✗ Failed: {e}")
# Strategy 3: Load fp16 with auto placement
print("\n[3/4] Trying: fp16 model with auto placement...")
try:
model = AutoModelForCausalLM.from_pretrained(
config["base_model"],
torch_dtype=torch.float16,
device_map="auto",
low_cpu_mem_usage=True,
trust_remote_code=True,
)
print("✓ Success: fp16 model loaded")
except Exception as e:
print(f"✗ Failed: {e}")
# Strategy 4: DeepSpeed ZeRO-3 with CPU offload
print("\n[4/4] Trying: DeepSpeed ZeRO-3 CPU offload...")
try: try:
import deepspeed import deepspeed
model = AutoModelForCausalLM.from_pretrained( model = AutoModelForCausalLM.from_pretrained(
@@ -112,7 +84,7 @@ def train(config_path):
"bf16": {"enabled": True}, "bf16": {"enabled": True},
} }
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-5) optimizer = torch.optim.AdamW(model.parameters(), lr=float(config["train_params"]["learning_rate"]))
model, optimizer, _, _ = deepspeed.initialize( model, optimizer, _, _ = deepspeed.initialize(
model=model, model=model,
model_parameters=model.parameters(), model_parameters=model.parameters(),