# LoRA fine-tuning config for DeepSeek-V3 (transformers v5).
#
# `DeepseekV3NaiveMoe` (v5 patchgen `experts` module) stores fused
# `gate_up_proj` (3D, [num_experts, 2*moe_intermediate_size, hidden]) and
# `down_proj` (3D, [num_experts, hidden, moe_intermediate_size]). Its
# `CheckpointTensorConverter` stacks the per-expert HF checkpoint keys
# (`mlp.experts.{j}.gate_proj.weight` / `up_proj.weight`) into the fused
# `gate_up_proj` at load time so VeOmni MoE-LoRA wraps the same fused
# layout it uses for Qwen3-MoE / Qwen3.5-MoE / Qwen3-VL-MoE / Qwen3-Omni-MoE.
#
# `lora_config` schema (`veomni/trainer/base.py::_setup_lora`):
#   - `lora_modules`         — nn.Linear targets handled by PEFT (q/k/v/o etc.).
#   - `target_parameters`    — 3D nn.Parameter targets on the experts module.
#   - `share_expert_lora`    — when True, `target_parameters` are wrapped by VeOmni's
#                              `LoraSharedExperts` (single 2D LoRA pair shared across
#                              all experts in each layer); when False, `LoraIndependentExperts`
#                              (3D LoRA per-expert tensors, Mode 1, the default).
#
# Note: DeepSeek-V3 also has *shared* experts (`mlp.shared_experts.{gate,up,down}_proj`)
# implemented as plain `nn.Linear`. Add them to `lora_modules` if you want PEFT to
# wrap them too.

model:
  model_path: deepseek-ai/DeepSeek-V3
  ops_implementation:
    attn_implementation: flash_attention_2
    moe_implementation: eager
    cross_entropy_loss_implementation: eager
    rms_norm_implementation: eager
    swiglu_mlp_implementation: eager
    rotary_pos_emb_implementation: eager
  lora_config:
    rank: 16
    alpha: 32
    use_rslora: false
    lora_modules:
      - q_a_proj
      - q_b_proj
      - kv_a_proj_with_mqa
      - kv_b_proj
      - o_proj
    target_parameters:
      - model.layers.*.mlp.experts.gate_up_proj
      - model.layers.*.mlp.experts.down_proj
    share_expert_lora: true

data:
  train_path: tulu-3-sft-mixture
  data_type: conversation
  chat_template: chatml
  max_seq_len: 2048
  train_size: 750000
  text_keys: messages
  dyn_bsz_buffer_size: 200

train:
  accelerator:
    ulysses_size: 1
    ep_size: 8
    fsdp_config:
      fsdp_mode: fsdp2
  gradient_checkpointing:
    enable: true
  global_batch_size: 8
  micro_batch_size: 1
  bsz_warmup_ratio: 0.007
  optimizer:
    type: adamw
    lr: 1.0e-4
    lr_decay_style: cosine
    max_grad_norm: 1.0
  num_train_epochs: 1
  init_device: meta
  enable_full_determinism: false
  empty_cache_steps: 500
  checkpoint:
    output_dir: deepseek_v3-lora-tulu
    manager: dcp
    save_steps: 0
    save_epochs: 1
    save_hf_weights: true
  wandb:
    project: VeOmni
    name: deepseek_v3-lora-tulu
