# Phase 4: Qwen3-8B LoRA fine-tune config (used by tanglish.llm.qwen_train).
#
# Strategy: full LoRA on bf16 base (NOT QLoRA/4-bit — L40S 48 GB has enough
# VRAM for cleaner training with higher final quality).
#
# Never-lose-model design (per user 2026-07-15):
#   1. All output goes to /workspace/tanglish/checkpoints/qwen_lora_v1/
#      which lives on the NETWORK VOLUME (survives pod eviction).
#   2. A TrainerCallback S3-syncs every checkpoint AS IT'S SAVED
#      (not just at end).
#   3. Post-training merge → S3 sync of merged model.
#   4. Auto-terminate has 6 gates (train exited cleanly, no failed marker,
#      TRAINING COMPLETE marker, best-val LoRA on disk, merged model on disk,
#      S3 has both).
#   5. Even mid-training pod eviction loses at most ~1 val-check-interval
#      of progress.

# -----------------------------------------------------------------------------
# Model
# -----------------------------------------------------------------------------
model:
  base: "Qwen/Qwen3-8B"                       # Qwen3 series has no "-Instruct" suffix; already instruction-tuned by default
  torch_dtype: bfloat16
  attn_implementation: flash_attention_2      # 2-3x faster than sdpa on Qwen3
  gradient_checkpointing: true                # ~30 % VRAM save, ~15 % slower
  # Do NOT quantize — cleaner LoRA training on bf16 base; VRAM headroom exists.

# -----------------------------------------------------------------------------
# LoRA adapter
# -----------------------------------------------------------------------------
lora:
  r: 16
  alpha: 32              # 2x r is the standard heuristic
  dropout: 0.05
  bias: none
  target_modules:
    - q_proj
    - k_proj
    - v_proj
    - o_proj
    - gate_proj
    - up_proj
    - down_proj

# -----------------------------------------------------------------------------
# Data
# -----------------------------------------------------------------------------
data:
  # Downloaded from S3 in setup_qwen_train.sh, decompressed to plain .jsonl
  train_file: /workspace/tanglish/data/pairs/tanglish_v1.train.jsonl
  val_file:   /workspace/tanglish/data/pairs/tanglish_v1.val.jsonl
  max_seq_length: 2048              # dialogues are short; leaves plenty of room
  packing: false                    # sample-level SFT (safer for chat format)

# -----------------------------------------------------------------------------
# Trainer — TRL SFTTrainer args (Hugging Face TrainingArguments compatible)
# -----------------------------------------------------------------------------
training:
  output_dir: /workspace/tanglish/checkpoints/qwen_lora_v1/lora

  num_train_epochs: 3
  per_device_train_batch_size: 4
  per_device_eval_batch_size: 8
  gradient_accumulation_steps: 4          # effective batch = 16

  learning_rate: 2.0e-4                   # LoRA convention
  lr_scheduler_type: cosine
  warmup_ratio: 0.05                      # 5 % warmup
  weight_decay: 0.01
  max_grad_norm: 1.0

  bf16: true
  optim: adamw_torch
  dataloader_num_workers: 4
  seed: 42

  logging_steps: 20                       # wandb + local every 20 steps
  eval_strategy: steps
  eval_steps: 500                         # ~10-15 min on L40S
  save_strategy: steps
  save_steps: 500                         # save with eval so metric aligns
  save_total_limit: 3                     # top-3 by eval_loss + last kept
  metric_for_best_model: eval_loss
  greater_is_better: false
  load_best_model_at_end: true            # so merge picks the actual best

  report_to: wandb                        # step-level curves like Phase 2

# -----------------------------------------------------------------------------
# Early stopping — kills training when val loss plateaus
# -----------------------------------------------------------------------------
early_stopping:
  patience: 4                             # 4 eval checks with no improvement
  threshold: 0.005                        # must improve by ≥ this each time

# -----------------------------------------------------------------------------
# S3 sync callback — hits S3 after EVERY save event (paranoid but cheap)
# -----------------------------------------------------------------------------
s3_sync:
  enabled: true
  s3_prefix: "checkpoints/qwen3-8b-tanglish-v1/lora-ckpts/"
  # Also sync the merged model after training:
  merged_s3_prefix: "checkpoints/qwen3-8b-tanglish-v1/merged/"

# -----------------------------------------------------------------------------
# Wandb
# -----------------------------------------------------------------------------
wandb:
  project: qwen3-8b-tanglish
  run_name: v1                            # override via CLI flag if you want