chore: import upstream snapshot with attribution

2026-07-13 13:34:58 +08:00
commit a203934033
1368 changed files with 175001 additions and 0 deletions
@@ -0,0 +1,68 @@
+# Ray Megatron GKD multi-turn — colocate mode, colocated teacher.
+# Tests multi-turn GKD on Ray-Megatron backend with MathTipsScheduler.
+rlhf_type: gkd
+
+model: Qwen/Qwen3.5-0.8B
+teacher_model: Qwen/Qwen3.5-2B
+gkd_logits_topk: 64
+enable_thinking: false
+
+dataset:
+  - 'AI-MO/NuminaMath-TIR#2000'
+dataset_num_proc: 4
+split_dataset_ratio: 0
+
+micro_batch_size: 1
+global_batch_size: 16
+num_train_epochs: 1
+logging_steps: 5
+seed: 42
+max_length: 2048
+max_completion_length: 1024
+truncation_strategy: delete
+padding_free: true
+sequence_parallel: true
+attention_backend: flash
+
+lr: 1e-5
+min_lr: 1e-7
+lr_warmup_fraction: 0.05
+
+temperature: 1.0
+lmbda: 1
+beta: 0.5
+sft_alpha: 0.0
+
+# Multi-turn configuration
+multi_turn_scheduler: math_tip_trick
+max_turns: 2
+
+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
+offload_teacher_model: 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_multi_turn_ray
+
+rollout:
+  gpus: 4
+  vllm_tensor_parallel_size: 2
+  vllm_gpu_memory_utilization: 0.4
+  vllm_max_model_len: 4096
@@ -0,0 +1,62 @@
+# Ray Megatron GKD — colocate mode (train + rollout share GPUs), colocated teacher.
+rlhf_type: gkd
+
+model: Qwen/Qwen3.5-2B
+teacher_model: Qwen/Qwen3.5-4B
+gkd_logits_topk: 64  # omit for full-vocab distillation
+
+dataset:
+  - 'AI-ModelScope/alpaca-gpt4-data-zh#2000'
+  - 'AI-ModelScope/alpaca-gpt4-data-en#2000'
+dataset_num_proc: 4
+split_dataset_ratio: 0
+
+micro_batch_size: 1
+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
+offload_teacher_model: true
+sleep_level: 1
+
+train:
+  gpus: 4
+  tuner_type: lora
+  lora_rank: 8
+  lora_alpha: 32
+  tensor_model_parallel_size: 2  # TP2 -> DP2 on 4 GPUs
+  pipeline_model_parallel_size: 1
+  expert_model_parallel_size: 1
+  context_parallel_size: 1
+  output_dir: megatron_output/gkd_rollout_colocate_teacher_colocate
+
+rollout:
+  gpus: 4
+  vllm_tensor_parallel_size: 2
+  vllm_gpu_memory_utilization: 0.4
+  vllm_max_model_len: 16384
@@ -0,0 +1,68 @@
+# 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
@@ -0,0 +1,60 @@
+# Ray Megatron GKD — separate rollout (disjoint GPUs) + colocated teacher.
+# No `colocate_groups` => train and rollout use disjoint GPU sets
+rlhf_type: gkd
+
+model: Qwen/Qwen3.5-2B
+teacher_model: Qwen/Qwen3.5-4B
+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
+
+# No colocate_groups -> train and rollout occupy separate GPU sets
+offload_model: true
+offload_optimizer: true
+offload_teacher_model: true
+
+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_separate_teacher_colocate
+
+rollout:
+  gpus: 4
+  vllm_tensor_parallel_size: 1
+  vllm_gpu_memory_utilization: 0.8
+  vllm_max_model_len: 16384
@@ -0,0 +1,8 @@
+#!/bin/bash
+# Ray Megatron GKD — default example (rollout colocate + colocated teacher).
+# Swap --config for another yaml in this folder for other placements/teacher modes.
+export CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-0,1,2,3}
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+
+megatron rlhf --use_ray true --config "$SCRIPT_DIR/rollout_colocate_teacher_colocate.yaml"