# Ray Megatron GKD — standalone vLLM teacher replicas (top-k only).
# gkd_logits_topk is REQUIRED; the teacher group below serves the teacher model and its max_logprobs must be >= gkd_logits_topk.
# Do NOT set a top-level `teacher_model` here.
rlhf_type: gkd

model: Qwen/Qwen3.5-2B
gkd_logits_topk: 64

dataset: AI-ModelScope/alpaca-gpt4-data-en#2000
dataset_num_proc: 4
split_dataset_ratio: 0

micro_batch_size: 2
global_batch_size: 16
num_train_epochs: 1
logging_steps: 1
seed: 42
max_length: 8192
max_completion_length: 8192
padding_free: true
sequence_parallel: true
attention_backend: flash

lr: 1e-6
min_lr: 1e-6
lr_warmup_fraction: 0.0

temperature: 1.0
lmbda: 1
beta: 0.5
sft_alpha: 0.0

finetune: true
no_save_optim: true
no_save_rng: true
recompute_granularity: selective

use_vllm: true

colocate_groups: [[train, rollout]]
offload_model: true
offload_optimizer: true
sleep_level: 1

train:
  gpus: 4
  tuner_type: lora
  lora_rank: 8
  lora_alpha: 32
  tensor_model_parallel_size: 2
  pipeline_model_parallel_size: 1
  expert_model_parallel_size: 1
  context_parallel_size: 1
  output_dir: megatron_output/gkd_rollout_colocate_teacher_standalone

rollout:
  gpus: 4
  vllm_tensor_parallel_size: 1
  vllm_gpu_memory_utilization: 0.4
  vllm_max_model_len: 16384

teacher:
  gpus: 4
  model: Qwen/Qwen3.5-4B          # teacher served by these replicas
  vllm_tensor_parallel_size: 1
  vllm_gpu_memory_utilization: 0.9
  vllm_max_model_len: 16384
  vllm_engine_kwargs: {"max_logprobs": 64}   # must be >= gkd_logits_topk
