139 lines
No EOL
3.9 KiB
YAML
139 lines
No EOL
3.9 KiB
YAML
experiment:
|
|
name: "progressive_reasoning_gemma3_1b_8gpu_deepspeed"
|
|
base_model: "google/gemma-3-1b-pt"
|
|
output_dir: "./outputs"
|
|
use_wandb: true
|
|
wandb_project: "matsuo-llm-comp-2025"
|
|
|
|
model:
|
|
load_in_4bit: false # Disable quantization for DeepSpeed
|
|
device_map: null # Let DeepSpeed handle device placement
|
|
gradient_checkpointing: true # Enable for memory efficiency
|
|
use_flash_attention_2: false
|
|
use_eager_attention: true
|
|
|
|
# DeepSpeed Configuration
|
|
deepspeed:
|
|
zero_optimization:
|
|
stage: 2 # ZeRO Stage 2 (partition optimizer states and gradients)
|
|
allgather_partitions: true
|
|
allgather_bucket_size: 500000000 # 500MB buckets
|
|
overlap_comm: true
|
|
reduce_scatter: true
|
|
reduce_bucket_size: 500000000
|
|
contiguous_gradients: true
|
|
cpu_offload: false # Keep on GPU for speed with small model
|
|
|
|
optimizer:
|
|
type: "AdamW"
|
|
params:
|
|
lr: "auto"
|
|
betas: "auto"
|
|
eps: "auto"
|
|
weight_decay: "auto"
|
|
|
|
scheduler:
|
|
type: "WarmupLR"
|
|
params:
|
|
warmup_min_lr: "auto"
|
|
warmup_max_lr: "auto"
|
|
warmup_num_steps: "auto"
|
|
|
|
fp16:
|
|
enabled: false
|
|
|
|
bf16:
|
|
enabled: true
|
|
|
|
gradient_clipping: 1.0
|
|
|
|
train_batch_size: 512 # Total batch size across all GPUs
|
|
train_micro_batch_size_per_gpu: 64 # Per-GPU batch size
|
|
|
|
progressive_stages:
|
|
- name: "basic_cot"
|
|
description: "Basic Chain-of-Thought reasoning"
|
|
dataset_path: "./data/basic_cot/"
|
|
adapter_config:
|
|
r: 32 # Larger rank with 8 GPUs
|
|
lora_alpha: 64
|
|
lora_dropout: 0.1
|
|
target_modules: ["q_proj", "k_proj", "v_proj", "o_proj"]
|
|
init_lora_weights: true
|
|
training:
|
|
num_epochs: 2
|
|
per_device_batch_size: 64 # Large batch with DeepSpeed
|
|
gradient_accumulation_steps: 1 # No accumulation needed
|
|
learning_rate: 5e-4
|
|
warmup_steps: 100
|
|
max_length: 1024
|
|
bf16: true
|
|
max_grad_norm: 1.0
|
|
weight_decay: 0.001
|
|
save_steps: 50
|
|
logging_steps: 10
|
|
dataloader_num_workers: 8
|
|
dataloader_pin_memory: true
|
|
|
|
- name: "math_reasoning"
|
|
description: "Mathematical reasoning with OpenR1-Math-220k dataset"
|
|
dataset_path: "open-r1/OpenR1-Math-220k"
|
|
inherit_from: "basic_cot"
|
|
adapter_config:
|
|
r: 64 # Larger rank for math reasoning
|
|
lora_alpha: 128
|
|
lora_dropout: 0.1
|
|
target_modules: ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]
|
|
init_lora_weights: true
|
|
training:
|
|
num_epochs: 1
|
|
per_device_batch_size: 32 # Reduce for longer sequences
|
|
gradient_accumulation_steps: 1
|
|
learning_rate: 3e-4
|
|
warmup_steps: 200
|
|
max_length: 2048
|
|
bf16: true
|
|
max_grad_norm: 1.0
|
|
weight_decay: 0.001
|
|
save_steps: 100
|
|
logging_steps: 20
|
|
dataloader_num_workers: 8
|
|
dataset_config:
|
|
streaming: true
|
|
max_samples: 500000 # Process more data with 8 GPUs
|
|
split: "train"
|
|
|
|
- name: "complex_reasoning"
|
|
description: "Complex multi-step reasoning with Mixture-of-Thoughts"
|
|
dataset_path: "open-r1/Mixture-of-Thoughts"
|
|
inherit_from: "math_reasoning"
|
|
adapter_config:
|
|
r: 128 # Maximum rank for complex reasoning
|
|
lora_alpha: 256
|
|
lora_dropout: 0.1
|
|
target_modules: ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]
|
|
init_lora_weights: true
|
|
training:
|
|
num_epochs: 1
|
|
per_device_batch_size: 16 # Reduce for very long sequences
|
|
gradient_accumulation_steps: 2
|
|
learning_rate: 2e-4
|
|
warmup_steps: 300
|
|
max_length: 4096
|
|
bf16: true
|
|
max_grad_norm: 1.0
|
|
weight_decay: 0.001
|
|
save_steps: 200
|
|
logging_steps: 50
|
|
dataloader_num_workers: 8
|
|
dataset_config:
|
|
streaming: true
|
|
max_samples: 800000 # Process even more data
|
|
split: "train"
|
|
|
|
evaluation:
|
|
benchmarks:
|
|
- "HLE"
|
|
- "Do-Not-Answer"
|
|
save_results: true
|
|
results_dir: "./outputs/evaluation_results" |