# @package _global_
# Qwen3-30B-A3B multi-teacher on-policy distillation (GKD) on a single 8-GPU H20 node,
# with two data sources (BCP + DAPO math) each routed to its own teacher.
#
# Assumes two teachers are already available, e.g. a math expert and a BCP expert.
# For how to prepare the data and produce them, see Task 1 and Task 2 in docs/en/run-guide.md.
#
# Launch with one teacher per data source. hf_path is always required and only
# supplies the architecture via its config.json; when dist_ckpt_path is set the
# weights come from there instead, and both must describe the same architecture:
#   bash examples/start.sh qwen3_30b_a3b/mopd_h20_1node \
#     hf_model_path=/root/Qwen3-30B-A3B \
#     data_sources.0.dataset.prompt_data_path=/root/browsecomp-plus/raw \
#     data_sources.1.dataset.prompt_data_path=/root/DAPO-Math-17k-unique \
#     opd.teachers.0.hf_path=/root/Qwen3-30B-A3B \
#     opd.teachers.0.dist_ckpt_path=/root/ckpt/bcp_expert/train_step_100/dist_ckpt \
#     opd.teachers.1.hf_path=/root/Qwen3-30B-A3B \
#     opd.teachers.1.dist_ckpt_path=/root/ckpt/math_expert/train_step_100/dist_ckpt
defaults:
  - /default
  - _self_

checkpoint_path: "./checkpoints"
hf_model_path: ???
colocate: true
total_steps: 300

data_sources:
  - dataset:
      prompt_data_path: ???
      data_pre_processor: bcp
      input_key: prompt
      label_key: answer
      shuffle: false
    num_trajectories_per_prompt: 8
    num_prompts_per_step: 32
    teacher_name: "bcp_expert"
    agent:
      name: bcp
      retrieval_service_url: http://127.0.0.1:9000
      search_topk: 3
      max_queries_per_call: 1
      topk_chunks: 3
      chunk_size: 512
      chunk_overlap: 32
      max_turns: 15
      tool_timeout: 30.0
      search_timeout: 300.0
    reward:
      name: bcp
    max_response_len_per_trajectory: 10240
  - dataset:
      prompt_data_path: ???
      input_key: prompt
      label_key: reward_model
    reward:
      name: dapo_math
      overlong_penalty_length: 4096
    teacher_name: "math_expert"
    num_prompts_per_step: 32
    num_trajectories_per_prompt: 8
    max_response_len_per_trajectory: 20480

rollout:
  num_gpus_per_node: 8
  sglang_args:
    mem_fraction_static: 0.6
    disable_cuda_graph: false
  sglang_replicas:
    regular:
      num_nodes: 1
      num_gpus_per_replica: 4
  sampler:
    num_oversample: 0

agentflow:
  dump_trajectory_path: "log/trajectories"
  router:
    proxy_timeout_seconds: 600
  tokenizer:
    generation_prompt_kwargs:
      enable_thinking: true

trainer:
  num_gpus_per_node: 8
  num_nodes: 1
  mini_batch_size: 128
  max_tokens_per_gpu: 16384
  use_dynamic_batch_size: true
  use_rollout_routing_replay: true
  use_fp32_lm_head: false

megatron:
  model:
    tensor_model_parallel_size: 4
    pipeline_model_parallel_size: 2
    virtual_pipeline_model_parallel_size: null
    context_parallel_size: 1
    expert_model_parallel_size: 4
    expert_tensor_parallel_size: 1
    recompute_granularity: full # selective|full
    recompute_method: uniform
    recompute_num_layers: 1
  optimizer:
    lr: 2.0e-6

opd:
  enable: true
  pg_ratio: 0
  gkd_ratio: 1
  pg_kl_method: k1
  gkd_kl_method: full_kl
  topk: 32
  teacher_nodes: 1
  teacher_gpus_per_node: 8
  teachers:
    - name: "bcp_expert"
      hf_path: ???
      # Optional: one concrete dist_ckpt dir, preferred over hf_path's weights.
      dist_ckpt_path: null
    - name: "math_expert"
      hf_path: ???
      dist_ckpt_path: null
  model:
    tensor_model_parallel_size: 4
    pipeline_model_parallel_size: 2
    virtual_pipeline_model_parallel_size: null
    context_parallel_size: 1
    expert_model_parallel_size: 4
    expert_tensor_parallel_size: 1

algorithm:
  advantage_estimator: grpo
  advantage_norm_mode: group_zscore
  policy_loss: grpo
  loss_agg_mode: token-mean
  is_correction:
    enable: false
    action: mask
    level: token
    lower_bound: 0.5
    upper_bound: 2