chore: import upstream snapshot with attribution

examples/megatron/rlhf/dpo/dense.sh

@@ -0,0 +1,36 @@
+# 4 * 45GiB
+PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
+NPROC_PER_NODE=4 \
+CUDA_VISIBLE_DEVICES=0,1,2,3 \
+megatron rlhf \
+    --rlhf_type dpo \
+    --model Qwen/Qwen2.5-7B-Instruct \
+    --save_safetensors true \
+    --dataset hjh0119/shareAI-Llama3-DPO-zh-en-emoji \
+    --load_from_cache_file true \
+    --split_dataset_ratio 0.01 \
+    --tensor_model_parallel_size 4 \
+    --micro_batch_size 8 \
+    --global_batch_size 16 \
+    --recompute_granularity full \
+    --recompute_method uniform \
+    --recompute_num_layers 1 \
+    --num_train_epochs 1 \
+    --finetune true \
+    --cross_entropy_loss_fusion true \
+    --lr 1e-5 \
+    --lr_warmup_fraction 0.05 \
+    --min_lr 1e-6 \
+    --output_dir megatron_output/Qwen2.5-7B-Instruct \
+    --eval_steps 100 \
+    --save_steps 100 \
+    --max_length 8192 \
+    --dataloader_num_workers 8 \
+    --dataset_num_proc 8 \
+    --no_save_optim true \
+    --no_save_rng true \
+    --sequence_parallel true \
+    --attention_backend flash \
+    --beta 0.1 \
+    --rpo_alpha 0.1 \
+    --loss_type sigmoid

examples/megatron/rlhf/dpo/group_by_length.sh

@@ -0,0 +1,38 @@
+# 4 * 56GiB
+PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
+NPROC_PER_NODE=4 \
+CUDA_VISIBLE_DEVICES=0,1,2,3 \
+megatron rlhf \
+    --rlhf_type dpo \
+    --model Qwen/Qwen3-4B-Instruct-2507 \
+    --save_safetensors true \
+    --dataset 'AI-ModelScope/orpo-dpo-mix-40k' \
+    --load_from_cache_file true \
+    --split_dataset_ratio 0.01 \
+    --tensor_model_parallel_size 4 \
+    --padding_free false \
+    --group_by_length true \
+    --micro_batch_size 4 \
+    --global_batch_size 16 \
+    --recompute_granularity full \
+    --recompute_method uniform \
+    --recompute_num_layers 1 \
+    --num_train_epochs 1 \
+    --finetune true \
+    --cross_entropy_loss_fusion true \
+    --lr 1e-5 \
+    --lr_warmup_fraction 0.05 \
+    --min_lr 1e-6 \
+    --output_dir megatron_output/Qwen3-4B-Instruct-2507 \
+    --eval_steps 200 \
+    --save_steps 200 \
+    --max_length 8192 \
+    --dataloader_num_workers 8 \
+    --dataset_num_proc 8 \
+    --no_save_optim true \
+    --no_save_rng true \
+    --sequence_parallel true \
+    --attention_backend flash \
+    --beta 0.1 \
+    --rpo_alpha 0.1 \
+    --loss_type sigmoid

examples/megatron/rlhf/dpo/moe.sh

@@ -0,0 +1,46 @@
+# 8 * 46GiB; 13s/it
+PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
+NPROC_PER_NODE=8 \
+CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
+megatron rlhf \
+    --rlhf_type dpo \
+    --model Qwen/Qwen3-30B-A3B-Instruct-2507 \
+    --save_safetensors true \
+    --dataset AI-ModelScope/orpo-dpo-mix-40k \
+    --load_from_cache_file true \
+    --split_dataset_ratio 0.01 \
+    --packing true \
+    --tensor_model_parallel_size 4 \
+    --expert_tensor_parallel_size 1 \
+    --expert_model_parallel_size 8 \
+    --moe_permute_fusion true \
+    --moe_grouped_gemm true \
+    --moe_shared_expert_overlap true \
+    --moe_aux_loss_coeff 1e-3 \
+    --micro_batch_size 1 \
+    --global_batch_size 4 \
+    --recompute_granularity full \
+    --recompute_method uniform \
+    --recompute_num_layers 1 \
+    --num_train_epochs 1 \
+    --finetune true \
+    --cross_entropy_loss_fusion true \
+    --lr 1e-5 \
+    --lr_warmup_fraction 0.05 \
+    --min_lr 1e-6 \
+    --output_dir megatron_output/Qwen3-30B-A3B-Instruct-2507 \
+    --eval_steps 200 \
+    --save_steps 200 \
+    --max_length 8192 \
+    --dataloader_num_workers 8 \
+    --dataset_num_proc 8 \
+    --no_save_optim true \
+    --no_save_rng true \
+    --sequence_parallel true \
+    --attention_backend flash \
+    --optimizer_cpu_offload true \
+    --use_precision_aware_optimizer true \
+    --optimizer_offload_fraction 1 \
+    --beta 0.1 \
+    --rpo_alpha 0.1 \
+    --loss_type sigmoid

examples/megatron/rlhf/dpo/packing.sh

@@ -0,0 +1,37 @@
+# 4 * 28GiB; 3.4s/it
+PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
+NPROC_PER_NODE=4 \
+CUDA_VISIBLE_DEVICES=0,1,2,3 \
+megatron rlhf \
+    --rlhf_type dpo \
+    --model Qwen/Qwen3-4B-Instruct-2507 \
+    --save_safetensors true \
+    --dataset 'AI-ModelScope/orpo-dpo-mix-40k' \
+    --load_from_cache_file true \
+    --split_dataset_ratio 0.01 \
+    --tensor_model_parallel_size 4 \
+    --packing true \
+    --micro_batch_size 1 \
+    --global_batch_size 4 \
+    --recompute_granularity full \
+    --recompute_method uniform \
+    --recompute_num_layers 1 \
+    --num_train_epochs 1 \
+    --finetune true \
+    --cross_entropy_loss_fusion true \
+    --lr 1e-5 \
+    --lr_warmup_fraction 0.05 \
+    --min_lr 1e-6 \
+    --output_dir megatron_output/Qwen3-4B-Instruct-2507 \
+    --eval_steps 200 \
+    --save_steps 200 \
+    --max_length 8192 \
+    --dataloader_num_workers 8 \
+    --dataset_num_proc 8 \
+    --no_save_optim true \
+    --no_save_rng true \
+    --sequence_parallel true \
+    --attention_backend flash \
+    --beta 0.1 \
+    --rpo_alpha 0.1 \
+    --loss_type sigmoid

examples/megatron/rlhf/gkd/dense.sh

@@ -0,0 +1,38 @@
+CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
+NPROC_PER_NODE=8 \
+PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
+megatron rlhf \
+    --rlhf_type gkd \
+    --model Qwen/Qwen3-8B-Base \
+    --teacher_model Qwen/Qwen3-32B \
+    --tuner_type lora \
+    --dataset AI-ModelScope/alpaca-gpt4-data-en#2000 AI-ModelScope/alpaca-gpt4-data-zh#2000 \
+    --tensor_model_parallel_size 2 \
+    --expert_model_parallel_size 1 \
+    --pipeline_model_parallel_size 2 \
+    --context_parallel_size 2 \
+    --lmbda 1 \
+    --beta 1 \
+    --torch_dtype bfloat16 \
+    --micro_batch_size 2 \
+    --global_batch_size 16 \
+    --num_train_epochs 1 \
+    --lr 5e-6 \
+    --logging_steps 1 \
+    --max_length 8192 \
+    --max_completion_length 8192 \
+    --attention_backend flash \
+    --use_vllm true \
+    --vllm_mode colocate \
+    --vllm_gpu_memory_utilization 0.5 \
+    --vllm_tensor_parallel_size 1 \
+    --vllm_max_model_len 16384 \
+    --sleep_level 1 \
+    --offload_teacher_model true \
+    --recompute_granularity selective \
+    --finetune \
+    --no_save_optim \
+    --no_save_rng \
+    --temperature 1.0 \
+    --padding_free true \
+    --sequence_parallel true

examples/megatron/rlhf/gkd/multi_turn.sh

@@ -0,0 +1,45 @@
+CUDA_VISIBLE_DEVICES=0,1,2,3 \
+NPROC_PER_NODE=4 \
+PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
+megatron rlhf \
+    --rlhf_type gkd \
+    --model Qwen/Qwen3.5-2B \
+    --teacher_model Qwen/Qwen3.5-4B \
+    --tuner_type lora \
+    --dataset 'AI-ModelScope/NuminaMath-TIR#2000' \
+    --tensor_model_parallel_size 2 \
+    --lmbda 1 \
+    --beta 0.5 \
+    --temperature 1.0 \
+    --torch_dtype bfloat16 \
+    --micro_batch_size 2 \
+    --global_batch_size 16 \
+    --num_train_epochs 1 \
+    --lr 5e-6 \
+    --logging_steps 5 \
+    --max_length 2048 \
+    --max_completion_length 1024 \
+    --attention_backend flash \
+    --use_vllm true \
+    --vllm_mode colocate \
+    --vllm_gpu_memory_utilization 0.3 \
+    --vllm_tensor_parallel_size 1 \
+    --vllm_max_model_len 4096 \
+    --sleep_level 1 \
+    --offload_teacher_model true \
+    --recompute_granularity selective \
+    --finetune \
+    --no_save_optim \
+    --no_save_rng \
+    --enable_thinking false \
+    --loss_scale last_round \
+    --multi_turn_scheduler math_tip_trick \
+    --max_turns 2 \
+    --truncation_strategy delete \
+    --padding_free true \
+    --sequence_parallel true \
+    --output_dir output \
+    --save_steps 200 \
+    --save_total_limit 2 \
+    --vllm_server_pass_dataset true \
+    --remove_unused_columns false

examples/megatron/rlhf/gkd/opsd.sh

@@ -0,0 +1,47 @@
+# OPSD Fixed Teacher Mode (Self-Distillation) - Megatron
+# Paper: Self-Distilled Reasoner (arXiv:2601.18734)
+# Teacher = base model (disable_adapter), Student = LoRA-adapted model
+# Dataset: open-r1/OpenThoughts-114k-math
+# Model: Qwen3-4B
+#
+# Hyperparameters aligned with paper's run_opsd.sh:
+#   lr=2e-5, lora_r=64, lora_alpha=128, temp=1.2, beta=0.5, lmbda=1, effective batch=32
+
+CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
+NPROC_PER_NODE=8 \
+PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
+megatron rlhf \
+    --rlhf_type gkd \
+    --model Qwen/Qwen3-4B \
+    --teacher_model Qwen/Qwen3-4B \
+    --external_plugins examples/train/rlhf/opsd/opsd_plugin.py \
+    --dataset 'open-r1/OpenThoughts-114k-math' \
+    --use_vllm true \
+    --vllm_mode colocate \
+    --vllm_gpu_memory_utilization 0.6 \
+    --vllm_max_model_len 10240 \
+    --tuner_type lora \
+    --lora_rank 64 \
+    --lora_alpha 128 \
+    --sleep_level 1 \
+    --lmbda 1.0 \
+    --beta 0.5 \
+    --temperature 1.2 \
+    --sft_alpha 0 \
+    --torch_dtype bfloat16 \
+    --micro_batch_size 2 \
+    --global_batch_size 32 \
+    --train_iters 1000 \
+    --lr 2e-5 \
+    --save_steps 100 \
+    --save_total_limit 10 \
+    --logging_steps 1 \
+    --max_length 8192 \
+    --max_completion_length 2048 \
+    --tensor_model_parallel_size 1 \
+    --pipeline_model_parallel_size 1 \
+    --attention_backend flash \
+    --recompute_granularity selective \
+    --finetune \
+    --no_save_optim \
+    --no_save_rng

examples/megatron/rlhf/gkd/teacher_server.sh

@@ -0,0 +1,57 @@
+top_k=64
+max_prompt_length=2048
+max_completion_length=2048
+max_total_length=$((max_prompt_length + max_completion_length))
+
+export IMAGE_MAX_TOKEN_NUM=1024
+
+# Teacher server must be running first:
+
+# CUDA_VISIBLE_DEVICES=0 \
+# swift deploy \
+#     --model Qwen/Qwen3.5-4B \
+#     --infer_backend vllm \
+#     --port 8000 \
+#     --max_logprobs $top_k \
+#     --max_length $max_total_length \
+#     --vllm_max_model_len $max_total_length
+
+CUDA_VISIBLE_DEVICES=1,2 \
+NPROC_PER_NODE=2 \
+PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
+megatron rlhf \
+    --rlhf_type gkd \
+    --model Qwen/Qwen3.5-4B \
+    --teacher_model_server http://localhost:8000 \
+    --gkd_logits_topk $top_k \
+    --dataset 'AI-ModelScope/clevr_cogen_a_train' \
+    --tensor_model_parallel_size 1 \
+    --pipeline_model_parallel_size 1 \
+    --context_parallel_size 1 \
+    --expert_model_parallel_size 1 \
+    --lmbda 1 \
+    --beta 0.5 \
+    --torch_dtype bfloat16 \
+    --micro_batch_size 2 \
+    --global_batch_size 32 \
+    --train_iters 500 \
+    --lr 5e-5 \
+    --lr_warmup_fraction 0.1 \
+    --logging_steps 1 \
+    --save_steps 100 \
+    --save_total_limit 10 \
+    --max_length $max_prompt_length \
+    --max_completion_length $max_completion_length \
+    --attention_backend flash \
+    --use_vllm true \
+    --vllm_mode colocate \
+    --vllm_gpu_memory_utilization 0.5 \
+    --vllm_tensor_parallel_size 1 \
+    --vllm_max_model_len $max_total_length \
+    --sleep_level 1 \
+    --finetune \
+    --no_save_optim \
+    --no_save_rng \
+    --temperature 1.0 \
+    --padding_free true \
+    --recompute_granularity selective

examples/megatron/rlhf/kto/dense.sh

@@ -0,0 +1,37 @@
+# 4 * 43GiB
+PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
+NPROC_PER_NODE=4 \
+CUDA_VISIBLE_DEVICES=0,1,2,3 \
+megatron rlhf \
+    --rlhf_type kto \
+    --model Qwen/Qwen2.5-7B-Instruct \
+    --save_safetensors true \
+    --dataset 'AI-ModelScope/ultrafeedback-binarized-preferences-cleaned-kto#20000' \
+    --load_from_cache_file true \
+    --split_dataset_ratio 0.01 \
+    --tensor_model_parallel_size 4 \
+    --packing true \
+    --micro_batch_size 1 \
+    --global_batch_size 4 \
+    --recompute_granularity full \
+    --recompute_method uniform \
+    --recompute_num_layers 1 \
+    --num_train_epochs 1 \
+    --finetune true \
+    --cross_entropy_loss_fusion true \
+    --lr 1e-5 \
+    --lr_warmup_fraction 0.05 \
+    --min_lr 1e-6 \
+    --output_dir megatron_output/Qwen2.5-7B-Instruct \
+    --eval_steps 200 \
+    --save_steps 200 \
+    --max_length 8192 \
+    --dataloader_num_workers 8 \
+    --dataset_num_proc 8 \
+    --no_save_optim true \
+    --no_save_rng true \
+    --sequence_parallel true \
+    --attention_backend flash \
+    --beta 0.1 \
+    --desirable_weight 1 \
+    --undesirable_weight 1

examples/megatron/rlhf/kto/moe.sh

@@ -0,0 +1,46 @@
+# 2 * 48GiB
+PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
+NPROC_PER_NODE=2 \
+CUDA_VISIBLE_DEVICES=0,1 \
+megatron rlhf \
+    --rlhf_type kto \
+    --model Qwen/Qwen3-30B-A3B-Instruct-2507 \
+    --save_safetensors true \
+    --merge_lora false \
+    --dataset 'AI-ModelScope/ultrafeedback-binarized-preferences-cleaned-kto#20000' \
+    --load_from_cache_file true \
+    --packing true \
+    --tuner_type lora \
+    --lora_rank 8 \
+    --lora_alpha 32 \
+    --target_modules all-linear \
+    --split_dataset_ratio 0.01 \
+    --expert_model_parallel_size 2 \
+    --moe_permute_fusion true \
+    --moe_grouped_gemm true \
+    --moe_shared_expert_overlap true \
+    --moe_aux_loss_coeff 1e-3 \
+    --micro_batch_size 1 \
+    --global_batch_size 4 \
+    --recompute_granularity full \
+    --recompute_method uniform \
+    --recompute_num_layers 1 \
+    --num_train_epochs 1 \
+    --finetune true \
+    --cross_entropy_loss_fusion true \
+    --lr 1e-4 \
+    --lr_warmup_fraction 0.05 \
+    --min_lr 1e-5 \
+    --output_dir megatron_output/Qwen3-30B-A3B-Instruct-2507 \
+    --eval_steps 100 \
+    --save_steps 100 \
+    --max_length 8192 \
+    --dataloader_num_workers 8 \
+    --dataset_num_proc 8 \
+    --no_save_optim true \
+    --no_save_rng true \
+    --sequence_parallel true \
+    --attention_backend flash \
+    --beta 0.1 \
+    --desirable_weight 1 \
+    --undesirable_weight 1

examples/megatron/rlhf/rm/dense.sh

@@ -0,0 +1,36 @@
+# 2 * 70GiB
+PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
+NPROC_PER_NODE=2 \
+MAX_PIXELS=1003520 \
+CUDA_VISIBLE_DEVICES=0,1 \
+megatron rlhf \
+    --rlhf_type rm \
+    --model Qwen/Qwen2.5-VL-7B-Instruct \
+    --save_safetensors true \
+    --dataset 'swift/RLAIF-V-Dataset#20000' \
+    --load_from_cache_file true \
+    --split_dataset_ratio 0.01 \
+    --tensor_model_parallel_size 2 \
+    --packing true \
+    --micro_batch_size 1 \
+    --global_batch_size 4 \
+    --recompute_granularity full \
+    --recompute_method uniform \
+    --recompute_num_layers 1 \
+    --num_train_epochs 1 \
+    --finetune true \
+    --cross_entropy_loss_fusion true \
+    --lr 1e-5 \
+    --lr_warmup_fraction 0.05 \
+    --min_lr 1e-6 \
+    --output_dir megatron_output/Qwen2.5-VL-7B-Instruct \
+    --eval_steps 200 \
+    --save_steps 200 \
+    --max_length 8192 \
+    --dataloader_num_workers 8 \
+    --dataset_num_proc 8 \
+    --no_save_optim true \
+    --no_save_rng true \
+    --sequence_parallel true \
+    --attention_backend flash \
+    --center_rewards_coefficient 0.01

examples/megatron/rlhf/rm/moe.sh

@@ -0,0 +1,44 @@
+# 2 * 45GiB
+PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
+NPROC_PER_NODE=2 \
+CUDA_VISIBLE_DEVICES=0,1 \
+megatron rlhf \
+    --rlhf_type rm \
+    --model Qwen/Qwen3-30B-A3B-Instruct-2507 \
+    --save_safetensors true \
+    --merge_lora false \
+    --dataset 'AI-ModelScope/orpo-dpo-mix-40k' \
+    --load_from_cache_file true \
+    --packing true \
+    --tuner_type lora \
+    --lora_rank 8 \
+    --lora_alpha 32 \
+    --target_modules all-linear \
+    --split_dataset_ratio 0.01 \
+    --expert_model_parallel_size 2 \
+    --moe_permute_fusion true \
+    --moe_grouped_gemm true \
+    --moe_shared_expert_overlap true \
+    --moe_aux_loss_coeff 1e-6 \
+    --micro_batch_size 1 \
+    --global_batch_size 4 \
+    --recompute_granularity full \
+    --recompute_method uniform \
+    --recompute_num_layers 1 \
+    --num_train_epochs 1 \
+    --finetune true \
+    --cross_entropy_loss_fusion true \
+    --lr 1e-4 \
+    --lr_warmup_fraction 0.05 \
+    --min_lr 1e-5 \
+    --output_dir megatron_output/Qwen3-30B-A3B-Instruct-2507 \
+    --eval_steps 100 \
+    --save_steps 100 \
+    --max_length 8192 \
+    --dataloader_num_workers 8 \
+    --dataset_num_proc 8 \
+    --no_save_optim true \
+    --no_save_rng true \
+    --sequence_parallel true \
+    --attention_backend flash \
+    --center_rewards_coefficient 0.01