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

feat: improve loading strategies with low_cpu_mem_usage and proper DeepSpeed config

Browse Source
This commit is contained in:
Christian Medina
2026-07-01 14:00:41 -04:00
parent 2d13e63c25
commit ce758f7eb5
1 changed files with 46 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

83
training/scripts/train.py
View File

@@ -39,12 +39,19 @@ def train(config_path):
# Try multiple loading strategies in order
print(f"Loading model: {config['base_model']}")
# Strategy 1: Load 4-bit AS-IS
print("\n[1/4] Trying: 4-bit model AS-IS...")
# Strategy 1: Load 4-bit with BitsAndBytes
print("\n[1/4] Trying: 4-bit with BitsAndBytes...")
try:
from transformers import BitsAndBytesConfig
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16,
bnb_4bit_use_double_quant=True,
)
model = AutoModelForCausalLM.from_pretrained(
config["base_model"],
torch_dtype=torch.float16,
quantization_config=bnb_config,
device_map="auto",
trust_remote_code=True,
)
@@ -52,49 +59,67 @@ def train(config_path):
except Exception as e:
print(f"✗ Failed: {e}")
# Strategy 2: Load bf16 to CPU
# Strategy 2: Load bf16 to CPU with low_cpu_mem_usage
print("\n[2/4] 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,
)
print("✓ Success: bf16 model loaded to CPU")
except Exception as e:
print(f"✗ Failed: {e}")
# Strategy 3: Load with accelerate CPU offload
print("\n[3/4] Trying: bf16 with accelerate CPU offload...")
try:
from accelerate import load_checkpoint_and_dispatch
base_model = AutoModelForCausalLM.from_pretrained(
config["base_model"],
torch_dtype=torch.bfloat16,
trust_remote_code=True,
)
model = load_checkpoint_and_dispatch(
base_model,
checkpoint=config["base_model"],
device_map="auto",
dtype=torch.bfloat16,
offload_folder="/tmp/model_offload",
)
print("✓ Success: bf16 with accelerate offload")
except Exception as e:
print(f"✗ Failed: {e}")
# Strategy 4: Use bf16 with DeepSpeed ZeRO-3
print("\n[4/4] Trying: bf16 with DeepSpeed ZeRO-3 CPU offload...")
# 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.bfloat16,
device_map="cpu",
torch_dtype=torch.float16,
device_map="auto",
low_cpu_mem_usage=True,
trust_remote_code=True,
)
print("✓ Success: bf16 model (DeepSpeed will handle offload)")
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:
import deepspeed
model = AutoModelForCausalLM.from_pretrained(
config["base_model"],
torch_dtype=torch.bfloat16,
device_map=None,
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=1e-5)
model, optimizer, _, _ = deepspeed.initialize(
model=model,
model_parameters=model.parameters(),
optimizer=optimizer,
config=ds_config,
)
print("✓ Success: DeepSpeed ZeRO-3 loaded")
except Exception as e:
print(f"✗ Failed: {e}")
raise RuntimeError("All loading strategies failed!")