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

feat: add 6 loading variants with different offload strategies

Browse Source
This commit is contained in:
Christian Medina
2026-07-01 14:45:29 -04:00
parent 22d1bfd573
commit 7de54746a2
1 changed files with 53 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files

108
training/scripts/train.py
View File

@@ -59,75 +59,71 @@ def train(config_path):
except Exception as e: except Exception as e:
print(f"✗ Failed: {e}") print(f"✗ Failed: {e}")
# Strategy 1: Load 4-bit model AS-IS with DeepSpeed ZeRO-3 # Strategy 1: 4-bit model variants
print("\n[1/3] Trying: 4-bit model with DeepSpeed ZeRO-3...") print("\n[1/6] Trying: 4-bit model AS-IS with DeepSpeed...")
try: try:
import deepspeed import deepspeed
model = AutoModelForCausalLM.from_pretrained( model = AutoModelForCausalLM.from_pretrained(
config["base_model"], config["base_model"],
torch_dtype=torch.float16, # Load as fp16 (already 4-bit) torch_dtype=torch.float16,
device_map="auto", device_map="auto",
trust_remote_code=True, trust_remote_code=True,
) )
ds_config = {"zero_optimization": {"stage": 3, "offload_optimizer": {"device": "cpu"}, "offload_param": {"device": "cpu"}}, "fp16": {"enabled": True}}
ds_config = {
"train_micro_batch_size_per_gpu": 1,
"gradient_accumulation_steps": 1,
"zero_optimization": {
"stage": 3,
"contiguous_gradients": True,
"overlap_comm": True,
"offload_optimizer": {"device": "cpu", "pin_memory": True},
"offload_param": {"device": "cpu", "pin_memory": True},
},
"fp16": {"enabled": True},
}
optimizer = torch.optim.AdamW(model.parameters(), lr=float(config["train_params"]["learning_rate"])) optimizer = torch.optim.AdamW(model.parameters(), lr=float(config["train_params"]["learning_rate"]))
model, optimizer, _, _ = deepspeed.initialize( model, optimizer, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), optimizer=optimizer, config=ds_config)
model=model, print("✓ Success: 4-bit model AS-IS")
model_parameters=model.parameters(),
optimizer=optimizer,
config=ds_config,
)
print("✓ Success: 4-bit model with DeepSpeed ZeRO-3")
except Exception as e: except Exception as e:
print(f"✗ Failed: {e}") print(f"✗ Failed: {e}")
# Strategy 2: bf16 with DeepSpeed ZeRO-3 + CPU offload print("\n[2/6] Trying: 4-bit model with bf16 compute...")
print("\n[2/3] Trying: bf16 with DeepSpeed ZeRO-3 CPU offload...")
try: try:
model = AutoModelForCausalLM.from_pretrained( from transformers import BitsAndBytesConfig
config["base_model"], bnb_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16, bnb_4bit_use_double_quant=True)
torch_dtype=torch.bfloat16, model = AutoModelForCausalLM.from_pretrained(config["base_model"], quantization_config=bnb_config, device_map="auto", trust_remote_code=True)
device_map="cpu", ds_config = {"zero_optimization": {"stage": 3, "offload_optimizer": {"device": "cpu"}, "offload_param": {"device": "cpu"}}, "bf16": {"enabled": True}}
low_cpu_mem_usage=True,
trust_remote_code=True,
)
ds_config = {
"train_micro_batch_size_per_gpu": 1,
"gradient_accumulation_steps": 1,
"zero_optimization": {
"stage": 3,
"contiguous_gradients": True,
"overlap_comm": True,
"offload_optimizer": {"device": "cpu", "pin_memory": True},
"offload_param": {"device": "cpu", "pin_memory": True},
},
"bf16": {"enabled": True},
}
optimizer = torch.optim.AdamW(model.parameters(), lr=float(config["train_params"]["learning_rate"])) optimizer = torch.optim.AdamW(model.parameters(), lr=float(config["train_params"]["learning_rate"]))
model, optimizer, _, _ = deepspeed.initialize( model, optimizer, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), optimizer=optimizer, config=ds_config)
model=model, print("✓ Success: 4-bit with bf16 compute")
model_parameters=model.parameters(), except Exception as e:
optimizer=optimizer, print(f"✗ Failed: {e}")
config=ds_config,
) print("\n[3/6] Trying: 4-bit model to CPU then DeepSpeed...")
print("✓ Success: bf16 with DeepSpeed ZeRO-3 CPU offload") try:
except Exception as e2: model = AutoModelForCausalLM.from_pretrained(config["base_model"], torch_dtype=torch.float16, device_map="cpu", trust_remote_code=True)
print(f"✗ Failed: {e2}") ds_config = {"zero_optimization": {"stage": 3, "offload_optimizer": {"device": "cpu"}, "offload_param": {"device": "cpu"}}, "fp16": {"enabled": True}}
optimizer = torch.optim.AdamW(model.parameters(), lr=float(config["train_params"]["learning_rate"]))
model, optimizer, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), optimizer=optimizer, config=ds_config)
print("✓ Success: 4-bit to CPU")
except Exception as e:
print(f"✗ Failed: {e}")
print("\n[4/6] Trying: 4-bit model fp32...")
try:
model = AutoModelForCausalLM.from_pretrained(config["base_model"], torch_dtype=torch.float32, device_map="auto", trust_remote_code=True)
ds_config = {"zero_optimization": {"stage": 3, "offload_optimizer": {"device": "cpu"}, "offload_param": {"device": "cpu"}}, "fp32": {"enabled": True}}
optimizer = torch.optim.AdamW(model.parameters(), lr=float(config["train_params"]["learning_rate"]))
model, optimizer, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), optimizer=optimizer, config=ds_config)
print("✓ Success: 4-bit fp32")
except Exception as e:
print(f"✗ Failed: {e}")
print("\n[5/6] Trying: bf16 model to CPU...")
try:
model = AutoModelForCausalLM.from_pretrained(config["base_model"], torch_dtype=torch.bfloat16, device_map="cpu", low_cpu_mem_usage=True, trust_remote_code=True)
ds_config = {"zero_optimization": {"stage": 3, "offload_optimizer": {"device": "cpu"}, "offload_param": {"device": "cpu"}}, "bf16": {"enabled": True}}
optimizer = torch.optim.AdamW(model.parameters(), lr=float(config["train_params"]["learning_rate"]))
model, optimizer, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), optimizer=optimizer, config=ds_config)
print("✓ Success: bf16 to CPU")
except Exception as e:
print(f"✗ Failed: {e}")
print("\n[6/6] Trying: bf16 model auto placement...")
try:
model = AutoModelForCausalLM.from_pretrained(config["base_model"], torch_dtype=torch.bfloat16, device_map="auto", low_cpu_mem_usage=True, trust_remote_code=True)
print("✓ Success: bf16 auto placement")
except Exception as e:
print(f"✗ Failed: {e}")
raise RuntimeError("All loading strategies failed!") raise RuntimeError("All loading strategies failed!")
# Prepare model for k-bit training # Prepare model for k-bit training