diff --git a/train-on-this-server.sh b/train-on-this-server.sh index b744335..34ccce4 100755 --- a/train-on-this-server.sh +++ b/train-on-this-server.sh @@ -18,9 +18,9 @@ pip install --upgrade pip echo "Installing training dependencies..." pip install transformers datasets trl peft accelerate bitsandbytes deepspeed -# Run training with CUDA memory optimization -echo "Starting training with accelerate..." +# Run training with DeepSpeed ZeRO-3 +echo "Starting training with DeepSpeed ZeRO-3..." export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True -accelerate launch --multi_gpu --num_processes 2 train.py --config training/configs/ornith-35b-lora.yaml +torchrun --nproc_per_node=2 train.py --config training/configs/ornith-35b-lora.yaml echo "Training completed!" \ No newline at end of file diff --git a/train.py b/train.py index 46744cb..8f40d20 100644 --- a/train.py +++ b/train.py @@ -37,9 +37,9 @@ def train(config_path): errors = [] # ------------------------------------------------------------------ - # Strategy 1: QLoRA (preferred) + # Strategy 1: QLoRA with DeepSpeed ZeRO-3 (preferred) # ------------------------------------------------------------------ - print("\n[1/4] Trying: 4-bit NF4 on GPU...") + print("\n[1/4] Trying: 4-bit QLoRA (DeepSpeed ZeRO-3)...") try: bnb_config = BitsAndBytesConfig( load_in_4bit=True, @@ -50,11 +50,10 @@ def train(config_path): model = AutoModelForCausalLM.from_pretrained( config["base_model"], quantization_config=bnb_config, - device_map="auto", trust_remote_code=True, low_cpu_mem_usage=True, ) - print("✓ Success: QLoRA 4-bit") + print("✓ Success: QLoRA 4-bit (DeepSpeed will handle placement)") except Exception as e: errors.append(("QLoRA 4-bit", e)) print(f"✗ Failed: {e}") @@ -150,7 +149,7 @@ def train(config_path): }, ) - # Training arguments (max_seq_length removed - passed to SFTTrainer instead) + # Training arguments with DeepSpeed ZeRO-3 training_args = TrainingArguments( output_dir=config["train_params"]["output_dir"], num_train_epochs=config["train_params"]["num_train_epochs"], @@ -165,9 +164,9 @@ def train(config_path): save_total_limit=config["train_params"]["save_total_limit"], eval_strategy=config.get("eval_strategy", "steps"), eval_steps=config.get("eval_steps", 100), - bf16=config.get("mixed_precision", "bf16") == "bf16", - fp16=config.get("mixed_precision", "bf16") == "fp16", + bf16=True, gradient_checkpointing=config.get("gradient_checkpointing", True), + deepspeed=config.get("deepspeed_config_path", "training/configs/ds_zero3.json"), ) # SFT Trainer (DeepSpeed handles distributed training via config) diff --git a/training/configs/ds_zero3.json b/training/configs/ds_zero3.json new file mode 100644 index 0000000..c9c9550 --- /dev/null +++ b/training/configs/ds_zero3.json @@ -0,0 +1,22 @@ +{ + "bf16": { + "enabled": true + }, + + "zero_optimization": { + "stage": 3, + + "overlap_comm": true, + "contiguous_gradients": true, + + "reduce_bucket_size": 50000000, + "stage3_prefetch_bucket_size": 50000000, + "stage3_param_persistence_threshold": 100000 + }, + + "gradient_clipping": 1.0, + + "train_micro_batch_size_per_gpu": 1, + + "steps_per_print": 100 +}