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

refactor: switch from DeepSpeed ZeRO-3 to FSDP for QLoRA compatibility

Browse Source
This commit is contained in:
Christian Medina
2026-07-01 22:33:06 -04:00
parent a655040db2
commit e64515e3ca
2 changed files with 14 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
train-on-this-server.sh
View File

@@ -18,9 +18,9 @@ pip install --upgrade pip
echo "Installing training dependencies..." echo "Installing training dependencies..."
pip install transformers datasets trl peft accelerate bitsandbytes deepspeed pip install transformers datasets trl peft accelerate bitsandbytes deepspeed
# Run training with DeepSpeed ZeRO-3 # Run training with FSDP
echo "Starting training with DeepSpeed ZeRO-3..." echo "Starting training with FSDP..."
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
torchrun --nproc_per_node=2 train.py --config training/configs/ornith-35b-lora.yaml accelerate launch --multi_gpu --num_processes 2 train.py --config training/configs/ornith-35b-lora.yaml
echo "Training completed!" echo "Training completed!"

16
train.py
View File

@@ -37,9 +37,9 @@ def train(config_path):
errors = [] errors = []
# ------------------------------------------------------------------ # ------------------------------------------------------------------
# Strategy 1: QLoRA with DeepSpeed ZeRO-3 (preferred) # Strategy 1: QLoRA with FSDP (preferred)
# ------------------------------------------------------------------ # ------------------------------------------------------------------
print("\n[1/4] Trying: 4-bit QLoRA (DeepSpeed ZeRO-3)...") print("\n[1/4] Trying: 4-bit QLoRA (FSDP)...")
try: try:
bnb_config = BitsAndBytesConfig( bnb_config = BitsAndBytesConfig(
load_in_4bit=True, load_in_4bit=True,
@@ -50,10 +50,11 @@ def train(config_path):
model = AutoModelForCausalLM.from_pretrained( model = AutoModelForCausalLM.from_pretrained(
config["base_model"], config["base_model"],
quantization_config=bnb_config, quantization_config=bnb_config,
device_map="auto",
trust_remote_code=True, trust_remote_code=True,
low_cpu_mem_usage=True, low_cpu_mem_usage=True,
) )
print("✓ Success: QLoRA 4-bit (DeepSpeed will handle placement)") print("✓ Success: QLoRA 4-bit")
except Exception as e: except Exception as e:
errors.append(("QLoRA 4-bit", e)) errors.append(("QLoRA 4-bit", e))
print(f"✗ Failed: {e}") print(f"✗ Failed: {e}")
@@ -149,7 +150,7 @@ def train(config_path):
}, },
) )
# Training arguments with DeepSpeed ZeRO-3 # Training arguments with FSDP
training_args = TrainingArguments( training_args = TrainingArguments(
output_dir=config["train_params"]["output_dir"], output_dir=config["train_params"]["output_dir"],
num_train_epochs=config["train_params"]["num_train_epochs"], num_train_epochs=config["train_params"]["num_train_epochs"],
@@ -166,7 +167,12 @@ def train(config_path):
eval_steps=config.get("eval_steps", 100), eval_steps=config.get("eval_steps", 100),
bf16=True, bf16=True,
gradient_checkpointing=config.get("gradient_checkpointing", True), gradient_checkpointing=config.get("gradient_checkpointing", True),
deepspeed=config.get("deepspeed_config_path", "training/configs/ds_zero3.json"), fsdp="full_shard auto_wrap",
fsdp_config={
"backward_prefetch": "backward_pre",
"forward_prefetch": "true",
"cpu_offload": "false",
},
) )
# SFT Trainer (DeepSpeed handles distributed training via config) # SFT Trainer (DeepSpeed handles distributed training via config)