progressive-llm/config/training_config_8gpu.yaml
2025-07-10 22:25:11 +09:00

110 lines
No EOL
3.2 KiB
YAML

experiment:
name: "progressive_reasoning_8gpu"
base_model: "google/gemma-2-2b-it" # Can scale up to larger models
output_dir: "./outputs"
use_wandb: true
wandb_project: "matsuo-llm-comp-2025"
model:
load_in_4bit: false # Can use FP16/BF16 with multiple GPUs
bnb_4bit_compute_dtype: "bfloat16"
bnb_4bit_use_double_quant: true
device_map: "balanced" # Distribute across all GPUs
gradient_checkpointing: true
use_flash_attention_2: true # Enable if available for speed
use_eager_attention: false
# Multi-GPU specific settings
distributed:
strategy: "ddp" # Distributed Data Parallel
find_unused_parameters: false
gradient_as_bucket_view: true
progressive_stages:
- name: "basic_cot"
description: "Basic Chain-of-Thought reasoning"
dataset_path: "./data/basic_cot/"
adapter_config:
r: 32 # Larger rank since we have more memory
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: 16 # Large batch size per GPU
gradient_accumulation_steps: 1 # No need for accumulation with 8 GPUs
learning_rate: 5e-4
warmup_steps: 100
max_length: 2048
bf16: true
max_grad_norm: 1.0
weight_decay: 0.001
save_steps: 50
logging_steps: 10
dataloader_num_workers: 4 # More workers for data loading
dataloader_pin_memory: true
remove_unused_columns: false
- 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
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: 8 # Reduce for larger model
gradient_accumulation_steps: 2
learning_rate: 3e-4
warmup_steps: 200
max_length: 4096
bf16: true
max_grad_norm: 1.0
weight_decay: 0.001
save_steps: 100
logging_steps: 20
dataloader_num_workers: 4
dataset_config:
streaming: true
max_samples: 100000 # Can process more 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 with multi-GPU
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: 4
gradient_accumulation_steps: 4
learning_rate: 2e-4
warmup_steps: 300
max_length: 8192
bf16: true
max_grad_norm: 1.0
weight_decay: 0.001
save_steps: 200
logging_steps: 50
dataloader_num_workers: 4
dataset_config:
streaming: true
max_samples: 50000
split: "train"
evaluation:
benchmarks:
- "HLE"
- "Do-Not-Answer"
save_results: true
results_dir: "./outputs/evaluation_results"