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: 5e-4
      betas: [0.9, 0.999]
      eps: 1e-8
      weight_decay: 0.001
  
  scheduler:
    type: "WarmupLR"
    params:
      warmup_min_lr: 0
      warmup_max_lr: 5e-4
      warmup_num_steps: 100
  
  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"