chore: import upstream snapshot with attribution

examples/megatron/base_to_chat.sh

@@ -0,0 +1,32 @@
+# 8 * 65GiB
+PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
+NPROC_PER_NODE=8 \
+CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
+megatron sft \
+    --model Qwen/Qwen2.5-14B \
+    --save_safetensors true \
+    --dataset 'liucong/Chinese-DeepSeek-R1-Distill-data-110k-SFT' \
+    --load_from_cache_file true \
+    --split_dataset_ratio 0.01 \
+    --tensor_model_parallel_size 4 \
+    --micro_batch_size 1 \
+    --global_batch_size 16 \
+    --packing true \
+    --recompute_granularity selective \
+    --train_iters 2000 \
+    --eval_iters 50 \
+    --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-14B \
+    --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

examples/megatron/benchmark/deepspeed.sh

@@ -0,0 +1,31 @@
+# 8 * 80GiB
+# Corresponding Megatron-SWIFT script reference:
+# https://github.com/modelscope/ms-swift/tree/main/examples/megatron/base_to_chat.sh
+PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
+NPROC_PER_NODE=8 \
+CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
+swift sft \
+    --model Qwen/Qwen2.5-14B \
+    --tuner_type full \
+    --dataset 'liucong/Chinese-DeepSeek-R1-Distill-data-110k-SFT' \
+    --load_from_cache_file true \
+    --split_dataset_ratio 0.01 \
+    --torch_dtype bfloat16 \
+    --max_steps 2000 \
+    --per_device_train_batch_size 1 \
+    --per_device_eval_batch_size 1 \
+    --learning_rate 1e-5 \
+    --gradient_accumulation_steps 2 \
+    --packing true \
+    --eval_steps 200 \
+    --save_steps 200 \
+    --logging_steps 5 \
+    --max_length 8192 \
+    --warmup_ratio 0.05 \
+    --dataloader_num_workers 8 \
+    --dataset_num_proc 8 \
+    --save_total_limit -1 \
+    --save_only_model true \
+    --output_dir output/Qwen2.5-14B \
+    --deepspeed zero2 \
+    --attn_impl flash_attn

examples/megatron/dense/72b_offload.sh

@@ -0,0 +1,36 @@
+# 8 * 65GiB. 80s/it
+PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
+NPROC_PER_NODE=8 \
+CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
+megatron sft \
+    --model Qwen/Qwen2.5-72B-Instruct \
+    --save_safetensors true \
+    --dataset 'liucong/Chinese-DeepSeek-R1-Distill-data-110k-SFT' \
+    --load_from_cache_file true \
+    --split_dataset_ratio 0.01 \
+    --tensor_model_parallel_size 8 \
+    --micro_batch_size 1 \
+    --global_batch_size 16 \
+    --packing true \
+    --recompute_granularity full \
+    --recompute_method uniform \
+    --recompute_num_layers 1 \
+    --num_train_epochs 5 \
+    --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-72B-Instruct \
+    --eval_steps 500 \
+    --save_steps 500 \
+    --max_length 8192 \
+    --dataloader_num_workers 8 \
+    --dataset_num_proc 8 \
+    --no_save_optim true \
+    --no_save_rng true \
+    --sequence_parallel true \
+    --optimizer_cpu_offload true \
+    --use_precision_aware_optimizer true \
+    --optimizer_offload_fraction 1 \
+    --attention_backend flash

examples/megatron/dense/qwen3_32b.sh

@@ -0,0 +1,33 @@
+# 8 * 80GiB
+PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
+NPROC_PER_NODE=8 \
+CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
+megatron sft \
+    --model Qwen/Qwen3-32B \
+    --save_safetensors true \
+    --dataset 'liucong/Chinese-DeepSeek-R1-Distill-data-110k-SFT' \
+    --load_from_cache_file true \
+    --split_dataset_ratio 0.01 \
+    --tensor_model_parallel_size 8 \
+    --micro_batch_size 1 \
+    --global_batch_size 16 \
+    --packing true \
+    --recompute_granularity full \
+    --recompute_method uniform \
+    --recompute_num_layers 1 \
+    --num_train_epochs 5 \
+    --finetune true \
+    --cross_entropy_loss_fusion true \
+    --lr 1e-5 \
+    --lr_warmup_fraction 0.05 \
+    --min_lr 1e-6 \
+    --output_dir megatron_output/Qwen3-32B \
+    --eval_steps 500 \
+    --save_steps 500 \
+    --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

examples/megatron/embedding/qwen3_emb.sh

@@ -0,0 +1,36 @@
+# 2 * 70GiB
+# For inference code, refer to: examples/infer/demo_embedding.py
+PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
+NPROC_PER_NODE=2 \
+CUDA_VISIBLE_DEVICES=0,1 \
+INFONCE_TEMPERATURE=0.1 \
+megatron sft \
+    --model Qwen/Qwen3-Embedding-8B \
+    --task_type embedding \
+    --save_safetensors true \
+    --tuner_type full \
+    --dataset sentence-transformers/stsb:positive \
+    --load_from_cache_file true \
+    --split_dataset_ratio 0.02 \
+    --tensor_model_parallel_size 2 \
+    --sequence_parallel true \
+    --micro_batch_size 16 \
+    --global_batch_size 16 \
+    --recompute_granularity full \
+    --recompute_method uniform \
+    --recompute_num_layers 1 \
+    --finetune true \
+    --cross_entropy_loss_fusion true \
+    --lr 5e-6 \
+    --lr_warmup_fraction 0.05 \
+    --min_lr 1e-7 \
+    --num_train_epochs 5 \
+    --output_dir megatron_output/Qwen3-Embedding-8B \
+    --save_steps 200 \
+    --eval_steps 100 \
+    --max_length 8192 \
+    --loss_type infonce \
+    --dataloader_num_workers 4 \
+    --no_save_optim true \
+    --no_save_rng true \
+    --dataset_num_proc 4

examples/megatron/embedding/qwen3_vl_emb.sh

@@ -0,0 +1,36 @@
+# 2 * 80GiB
+# For inference code, refer to: examples/infer/demo_embedding.py
+PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
+NPROC_PER_NODE=2 \
+CUDA_VISIBLE_DEVICES=0,1 \
+INFONCE_TEMPERATURE=0.1 \
+megatron sft \
+    --model Qwen/Qwen3-VL-Embedding-8B \
+    --task_type embedding \
+    --save_safetensors true \
+    --tuner_type full \
+    --dataset swift/TextCaps:emb \
+    --load_from_cache_file true \
+    --split_dataset_ratio 0.02 \
+    --tensor_model_parallel_size 2 \
+    --sequence_parallel true \
+    --micro_batch_size 16 \
+    --global_batch_size 16 \
+    --recompute_granularity full \
+    --recompute_method uniform \
+    --recompute_num_layers 1 \
+    --finetune true \
+    --cross_entropy_loss_fusion true \
+    --lr 5e-6 \
+    --lr_warmup_fraction 0.05 \
+    --min_lr 1e-7 \
+    --num_train_epochs 1 \
+    --output_dir megatron_output/Qwen3-VL-Embedding-8B \
+    --save_steps 200 \
+    --eval_steps 100 \
+    --max_length 8192 \
+    --loss_type infonce \
+    --dataloader_num_workers 4 \
+    --no_save_optim true \
+    --no_save_rng true \
+    --dataset_num_proc 4

examples/megatron/export/full.sh

@@ -0,0 +1,23 @@
+# safetensors -> torch_dist
+CUDA_VISIBLE_DEVICES=0,1,2,3 \
+NPROC_PER_NODE=4 \
+megatron export \
+    --model Qwen/Qwen3-30B-A3B-Instruct-2507 \
+    --output_dir Qwen3-30B-A3B-Instruct-2507-mcore \
+    --to_mcore true \
+    --tensor_model_parallel_size 2 \
+    --expert_model_parallel_size 2 \
+    --pipeline_model_parallel_size 2 \
+    --test_convert_precision true
+
+# torch_dist -> safetensors
+CUDA_VISIBLE_DEVICES=0,1,2,3 \
+NPROC_PER_NODE=4 \
+megatron export \
+    --mcore_model Qwen3-30B-A3B-Instruct-2507-mcore \
+    --output_dir Qwen3-30B-A3B-Instruct-2507-hf \
+    --to_hf true \
+    --tensor_model_parallel_size 2 \
+    --expert_model_parallel_size 2 \
+    --pipeline_model_parallel_size 2 \
+    --test_convert_precision true

examples/megatron/export/lora.sh

@@ -0,0 +1,46 @@
+# torch_dist -> safetensors
+# If you need to perform merge-lora and test precision alignment after merge-lora,
+# simply set `--merge_lora true`
+
+# You can also change `--model safetensors-path` to `--mcore_model torch-dist-path`.
+# These two methods are equivalent, and mcore-bridge will handle it automatically.
+CUDA_VISIBLE_DEVICES=0,1,2,3 \
+NPROC_PER_NODE=4 \
+megatron export \
+    --model Qwen/Qwen3-30B-A3B-Instruct-2507 \
+    --mcore_adapter megatron_output/Qwen3-30B-A3B-Instruct-2507/vx-xxx/checkpoint-xxx \
+    --output_dir megatron_output/Qwen3-30B-A3B-Instruct-2507/vx-xxx/checkpoint-xxx-lora \
+    --merge_lora false \
+    --to_hf true \
+    --tensor_model_parallel_size 2 \
+    --expert_model_parallel_size 2 \
+    --pipeline_model_parallel_size 2 \
+    --test_convert_precision true
+
+# safetensors -> torch_dist
+CUDA_VISIBLE_DEVICES=0,1,2,3 \
+NPROC_PER_NODE=4 \
+megatron export \
+    --model Qwen/Qwen3-30B-A3B-Instruct-2507 \
+    --adapters megatron_output/Qwen3-30B-A3B-Instruct-2507/vx-xxx/checkpoint-xxx-lora \
+    --output_dir megatron_output/Qwen3-30B-A3B-Instruct-2507/vx-xxx/checkpoint-xxx-mcore \
+    --merge_lora false \
+    --to_mcore true \
+    --tensor_model_parallel_size 2 \
+    --expert_model_parallel_size 2 \
+    --pipeline_model_parallel_size 2 \
+    --test_convert_precision true
+
+# Merge-LoRA:
+# torch_dist -> torch_dist
+CUDA_VISIBLE_DEVICES=0,1,2,3 \
+NPROC_PER_NODE=4 \
+megatron export \
+    --model Qwen/Qwen3-30B-A3B-Instruct-2507 \
+    --mcore_adapter megatron_output/Qwen3-30B-A3B-Instruct-2507/vx-xxx/checkpoint-xxx \
+    --output_dir megatron_output/Qwen3-30B-A3B-Instruct-2507/vx-xxx/checkpoint-xxx-merged \
+    --merge_lora true \
+    --to_mcore true \
+    --tensor_model_parallel_size 2 \
+    --expert_model_parallel_size 2 \
+    --pipeline_model_parallel_size 2

examples/megatron/export/merge_lora.sh

@@ -0,0 +1,9 @@
+NPROC_PER_NODE=2 \
+CUDA_VISIBLE_DEVICES=0,1 \
+megatron export \
+    --adapters megatron_output/Qwen2.5-7B-Instruct/vx-xxx/checkpoint-xxx \
+    --tensor_model_parallel_size 2 \
+    --to_hf true \
+    --merge_lora true \
+    --torch_dtype bfloat16 \
+    --output_dir megatron_output/Qwen2.5-7B-Instruct/vx-xxx/checkpoint-xxx-merged

examples/megatron/fp8/benchmark.sh

@@ -0,0 +1,39 @@
+# test_env: H20, cuda12.9
+# FP8: 8 * 58GiB 8s/it
+# BF16: 8 * 52GiB 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 sft \
+    --model Qwen/Qwen3-14B-FP8 \
+    --save_safetensors true \
+    --fp8_recipe blockwise \
+    --fp8_format e4m3 \
+    --fp8_param_gather true \
+    --dataset 'swift/Chinese-Qwen3-235B-2507-Distill-data-110k-SFT#20000' \
+    --load_from_cache_file true \
+    --tensor_model_parallel_size 4 \
+    --micro_batch_size 1 \
+    --global_batch_size 16 \
+    --packing true \
+    --recompute_granularity selective \
+    --num_train_epochs 1 \
+    --finetune true \
+    --cross_entropy_loss_fusion true \
+    --cross_entropy_fusion_impl native \
+    --lr 1e-5 \
+    --lr_warmup_fraction 0.05 \
+    --min_lr 1e-6 \
+    --output_dir megatron_output/Qwen3-14B-FP8 \
+    --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 \
+    --use_precision_aware_optimizer true \
+    --exp_avg_dtype bf16 \
+    --exp_avg_sq_dtype bf16 \
+    --attention_backend flash

examples/megatron/fp8/llm.sh

@@ -0,0 +1,60 @@
+# 8 * 95GiB
+# "cuda>=12.9"
+# In this example, FP8 training does not provide any speedup.
+PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
+NPROC_PER_NODE=8 \
+CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
+megatron sft \
+    --model Qwen/Qwen3-30B-A3B-Instruct-2507-FP8 \
+    --save_safetensors true \
+    --fp8_recipe blockwise \
+    --fp8_format e4m3 \
+    --fp8_param_gather true \
+    --dataset 'swift/Chinese-Qwen3-235B-2507-Distill-data-110k-SFT#2000' \
+              'swift/self-cognition#1000' \
+    --load_from_cache_file true \
+    --tensor_model_parallel_size 2 \
+    --expert_model_parallel_size 8 \
+    --moe_permute_fusion true \
+    --moe_grouped_gemm true \
+    --moe_shared_expert_overlap true \
+    --moe_aux_loss_coeff 1e-6 \
+    --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-30B-A3B-Instruct-2507-FP8 \
+    --eval_steps 200 \
+    --save_steps 200 \
+    --max_length 2048 \
+    --dataloader_num_workers 8 \
+    --dataset_num_proc 8 \
+    --no_save_optim true \
+    --no_save_rng true \
+    --sequence_parallel true \
+    --moe_expert_capacity_factor 2 \
+    --use_precision_aware_optimizer true \
+    --exp_avg_dtype bf16 \
+    --exp_avg_sq_dtype bf16 \
+    --attention_backend flash \
+    --model_author swift \
+    --model_name swift-robot
+
+# CUDA_VISIBLE_DEVICES=0 \
+# swift infer \
+#     --model megatron_output/Qwen3-30B-A3B-Instruct-2507-FP8/vx-xxx/checkpoint-xxx \
+#     --stream true
+
+# CUDA_VISIBLE_DEVICES=0 \
+# swift infer \
+#     --model megatron_output/Qwen3-30B-A3B-Instruct-2507-FP8/vx-xxx/checkpoint-xxx \
+#     --infer_backend vllm \
+#     --vllm_max_model_len 8192 \
+#     --stream true

examples/megatron/fp8/lora.sh

@@ -0,0 +1,82 @@
+# The generated LoRA delta weights cannot be merged into an FP8 base model via Merge-LoRA.
+# Due to the limited precision of FP8, the LoRA delta will be rounded to 0.
+# However, you can use BF16 weights to perform Merge-LoRA.
+
+# Although the model passed in here is BF16, it will be converted to FP8
+# after being loaded as a Megatron model
+PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
+NPROC_PER_NODE=2 \
+CUDA_VISIBLE_DEVICES=0,1 \
+IMAGE_MAX_TOKEN_NUM=1024 \
+VIDEO_MAX_TOKEN_NUM=128 \
+FPS_MAX_FRAMES=12 \
+megatron sft \
+    --model Qwen/Qwen3.5-4B \
+    --save_safetensors true \
+    --dataset 'AI-ModelScope/alpaca-gpt4-data-zh#500' \
+              'AI-ModelScope/alpaca-gpt4-data-en#500' \
+              'swift/self-cognition#500' \
+              'AI-ModelScope/LaTeX_OCR:human_handwrite#2000' \
+    --model_author swift \
+    --model_name swift-robot \
+    --merge_lora false \
+    --linear_decoupled_in_proj true \
+    --load_from_cache_file true \
+    --add_non_thinking_prefix true \
+    --loss_scale ignore_empty_think \
+    --fp8_recipe blockwise \
+    --fp8_format e4m3 \
+    --fp8_param_gather true \
+    --split_dataset_ratio 0.01 \
+    --tuner_type lora \
+    --lora_rank 16 \
+    --lora_alpha 32 \
+    --tensor_model_parallel_size 2 \
+    --micro_batch_size 1 \
+    --global_batch_size 2 \
+    --recompute_granularity full \
+    --recompute_method uniform \
+    --recompute_num_layers 1 \
+    --num_train_epochs 1 \
+    --packing true \
+    --finetune true \
+    --freeze_llm false \
+    --freeze_vit true \
+    --freeze_aligner true \
+    --cross_entropy_loss_fusion true \
+    --lr 1e-4 \
+    --lr_warmup_fraction 0.05 \
+    --min_lr 1e-5 \
+    --output_dir megatron_output/Qwen3.5-4B \
+    --eval_steps 200 \
+    --save_steps 200 \
+    --max_length 4096 \
+    --dataloader_num_workers 8 \
+    --dataset_num_proc 8 \
+    --no_save_optim true \
+    --no_save_rng true \
+    --sequence_parallel true \
+    --mtp_num_layers 1 \
+    --attention_backend flash
+
+# Merge-LoRA
+# FP8 base model + BF16 LoRA inference requires inference framework support
+# Alternatively, you can use BF16 base model + BF16 LoRA for inference
+CUDA_VISIBLE_DEVICES=0 \
+NPROC_PER_NODE=1 \
+megatron export \
+    --model Qwen/Qwen3.5-4B \
+    --adapters megatron_output/Qwen3.5-4B/vx-xxx/checkpoint-xxx \
+    --output_dir megatron_output/Qwen3.5-4B/vx-xxx/checkpoint-xxx-merged \
+    --to_hf true \
+    --linear_decoupled_in_proj true \
+    --mtp_num_layers 1 \
+    --merge_lora true
+
+
+# Inference with merged weights
+CUDA_VISIBLE_DEVICES=0 \
+swift infer \
+    --model megatron_output/Qwen3.5-4B/vx-xxx/checkpoint-xxx-merged \
+    --stream true \
+    --enable_thinking false

examples/megatron/fp8/quant.sh

@@ -0,0 +1,13 @@
+CUDA_VISIBLE_DEVICES=0,1,2,3 \
+NPROC_PER_NODE=4 \
+megatron export \
+    --model Qwen/Qwen3.5-35B-A3B \
+    --output_dir Qwen3.5-35B-A3B-FP8 \
+    --to_hf true \
+    --fp8_recipe blockwise \
+    --fp8_format e4m3 \
+    --fp8_param_gather true \
+    --mtp_num_layers 1 \
+    --linear_decoupled_in_proj true \
+    --tensor_model_parallel_size 2 \
+    --pipeline_model_parallel_size 2

examples/megatron/fp8/vlm.sh

@@ -0,0 +1,52 @@
+# 8 * 95GiB
+# In this example, FP8 training does not provide any speedup.
+PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
+OMP_NUM_THREADS=14 \
+NPROC_PER_NODE=8 \
+CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
+IMAGE_MAX_TOKEN_NUM=1024 \
+VIDEO_MAX_TOKEN_NUM=128 \
+FPS_MAX_FRAMES=16 \
+megatron sft \
+    --model Qwen/Qwen3-VL-30B-A3B-Instruct-FP8 \
+    --save_safetensors true \
+    --fp8_recipe blockwise \
+    --fp8_format e4m3 \
+    --fp8_param_gather true \
+    --dataset 'AI-ModelScope/alpaca-gpt4-data-zh#10000' \
+              'AI-ModelScope/LaTeX_OCR:human_handwrite#5000' \
+              'swift/VideoChatGPT:Generic#2000' \
+    --load_from_cache_file true \
+    --split_dataset_ratio 0.01 \
+    --moe_permute_fusion true \
+    --tensor_model_parallel_size 2 \
+    --expert_model_parallel_size 8 \
+    --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-5 \
+    --lr_warmup_fraction 0.05 \
+    --min_lr 1e-6 \
+    --output_dir megatron_output/Qwen3-VL-30B-A3B-Instruct \
+    --eval_steps 500 \
+    --save_steps 500 \
+    --max_length 4096 \
+    --packing true \
+    --dataloader_num_workers 8 \
+    --dataset_num_proc 8 \
+    --no_save_optim true \
+    --no_save_rng true \
+    --sequence_parallel true \
+    --moe_expert_capacity_factor 2 \
+    --use_precision_aware_optimizer true \
+    --exp_avg_dtype bf16 \
+    --exp_avg_sq_dtype bf16 \
+    --attention_backend flash

examples/megatron/grpo/dense_colocate.sh

@@ -0,0 +1,64 @@
+# DP size = world_size // (context_parallel_size * tensor_model_parallel_size * pipeline_model_parallel_size)
+#         = 8 // (1 * 1 * 1) = 8
+
+# NOTE: global_batch_size and micro_batch_size are completion-level
+# global_batch_size = micro_batch_size * DP size * gradient_accumulation_steps (128)
+# generation_batch_size = global_batch_size * steps_per_generation (128 * 4 = 512)
+# num_of_prompt_to_rollout = generation_batch_size / num_generations (512 / 8 = 64)
+# num_of_prompt_to_train = generation_batch_size / num_generations (128 / 8 = 16)
+
+CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
+NPROC_PER_NODE=8 \
+MAX_PIXELS=602112 \
+MASTER_PORT=29600 \
+megatron rlhf \
+    --rlhf_type grpo \
+    --model Qwen/Qwen2.5-VL-3B-Instruct \
+    --save_safetensors true \
+    --context_parallel_size 1 \
+    --tensor_model_parallel_size 1 \
+    --pipeline_model_parallel_size 1 \
+    --dataset AI-ModelScope/clevr_cogen_a_train#10000 \
+    --num_train_epochs 1 \
+    --global_batch_size 128 \
+    --micro_batch_size 4 \
+    --steps_per_generation 4 \
+    --num_generations 8 \
+    --external_plugins examples/train/grpo/plugin/plugin.py \
+    --reward_funcs external_r1v_acc format \
+    --use_vllm true \
+    --vllm_mode colocate \
+    --vllm_gpu_memory_utilization 0.7 \
+    --vllm_max_model_len 10240 \
+    --max_length 8192 \
+    --max_completion_length 2048 \
+    --tuner_type full \
+    --lr 1e-6 \
+    --bf16 true \
+    --beta 0.001 \
+    --importance_sampling_level token \
+    --epsilon 0.2 \
+    --epsilon_high 0.2 \
+    --dynamic_sample false \
+    --overlong_filter true \
+    --loss_type grpo \
+    --sleep_level 2 \
+    --offload_model true \
+    --offload_bridge false \
+    --offload_optimizer true \
+    --logging_steps 1 \
+    --recompute_granularity selective \
+    --finetune \
+    --dataloader_num_workers 8 \
+    --dataset_num_proc 8 \
+    --no_save_optim \
+    --no_save_rng \
+    --attention_backend flash \
+    --temperature 1.0 \
+    --system examples/train/grpo/prompt.txt \
+    --padding_free true \
+    --log_completions true \
+    --report_to wandb \
+    --train_iters 100 \
+    --eval_steps 1000 \
+    --save_steps 1000

examples/megatron/grpo/dense_server.sh

@@ -0,0 +1,71 @@
+# MAX_PIXELS=602112 \
+# CUDA_VISIBLE_DEVICES=6,7 \
+# swift rollout \
+#     --model Qwen/Qwen2.5-VL-3B-Instruct \
+#     --vllm_data_parallel_size 2 \
+#     --vllm_max_model_len 10240
+
+# DP size = world_size // (context_parallel_size * tensor_model_parallel_size * pipeline_model_parallel_size)
+#         = 6 // (1 * 1 * 1) = 6
+
+# NOTE: global_batch_size and micro_batch_size are completion-level
+# global_batch_size = micro_batch_size * DP size * gradient_accumulation_steps (96)
+# generation_batch_size = global_batch_size * steps_per_generation (96 * 4 = 384)
+# num_of_prompt_to_rollout = generation_batch_size / num_generations (384 / 8 = 48)
+# num_of_prompt_to_train = generation_batch_size / num_generations (96 / 8 = 12)
+
+CUDA_VISIBLE_DEVICES=0,1,2,3,4,5 \
+NPROC_PER_NODE=6 \
+MAX_PIXELS=602112 \
+MASTER_PORT=29600 \
+megatron rlhf \
+    --rlhf_type grpo \
+    --model Qwen/Qwen2.5-VL-3B-Instruct \
+    --save_safetensors true \
+    --context_parallel_size 1 \
+    --tensor_model_parallel_size 1 \
+    --pipeline_model_parallel_size 1 \
+    --dataset AI-ModelScope/clevr_cogen_a_train#10000 \
+    --num_train_epochs 1 \
+    --global_batch_size 96 \
+    --micro_batch_size 4 \
+    --steps_per_generation 4 \
+    --num_generations 8 \
+    --external_plugins examples/train/grpo/plugin/plugin.py \
+    --reward_funcs external_r1v_acc format \
+    --use_vllm true \
+    --vllm_mode server \
+    --vllm_server_host 127.0.0.1 \
+    --vllm_server_port 8000 \
+    --max_length 8192 \
+    --max_completion_length 2048 \
+    --tuner_type full \
+    --lr 1e-6 \
+    --bf16 true \
+    --beta 0.001 \
+    --importance_sampling_level token \
+    --epsilon 0.2 \
+    --epsilon_high 0.2 \
+    --dynamic_sample false \
+    --overlong_filter true \
+    --loss_type grpo \
+    --sleep_level 2 \
+    --offload_model true \
+    --offload_bridge false \
+    --offload_optimizer true \
+    --logging_steps 1 \
+    --recompute_granularity selective \
+    --finetune \
+    --dataloader_num_workers 8 \
+    --dataset_num_proc 8 \
+    --no_save_optim \
+    --no_save_rng \
+    --attention_backend flash \
+    --temperature 1.0 \
+    --system examples/train/grpo/prompt.txt \
+    --padding_free true \
+    --log_completions true \
+    --report_to wandb \
+    --train_iters 100 \
+    --eval_steps 1000 \
+    --save_steps 1000

examples/megatron/grpo/moe_colocate_full.sh

@@ -0,0 +1,54 @@
+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 grpo \
+    --model Qwen/Qwen3-30B-A3B-Instruct-2507 \
+    --save_safetensors true \
+    --context_parallel_size 1 \
+    --tensor_model_parallel_size 4 \
+    --expert_model_parallel_size 4 \
+    --pipeline_model_parallel_size 2 \
+    --dataset open-r1/DAPO-Math-17k-Processed \
+    --num_train_epochs 1 \
+    --global_batch_size 8 \
+    --micro_batch_size 1 \
+    --steps_per_generation 1 \
+    --num_generations 8 \
+    --reward_funcs accuracy format \
+    --use_vllm true \
+    --vllm_mode colocate \
+    --vllm_gpu_memory_utilization 0.4 \
+    --vllm_tensor_parallel_size 8 \
+    --vllm_max_model_len 16384 \
+    --max_length 8192 \
+    --max_completion_length 8192 \
+    --tuner_type full \
+    --lr 1e-6 \
+    --bf16 true \
+    --beta 0.00 \
+    --importance_sampling_level sequence \
+    --epsilon 3e-4 \
+    --epsilon_high 4e-4 \
+    --dynamic_sample false \
+    --overlong_filter true \
+    --loss_type grpo \
+    --sleep_level 2 \
+    --offload_model true \
+    --offload_bridge false \
+    --offload_optimizer true \
+    --optimizer_cpu_offload true \
+    --use_precision_aware_optimizer \
+    --logging_steps 1 \
+    --recompute_granularity selective \
+    --finetune \
+    --dataloader_num_workers 8 \
+    --dataset_num_proc 8 \
+    --no_save_optim \
+    --no_save_rng \
+    --attention_backend flash \
+    --temperature 1.0 \
+    --padding_free true \
+    --sequence_parallel true \
+    --log_completions true \
+    --report_to wandb

examples/megatron/grpo/moe_colocate_lora.sh

@@ -0,0 +1,53 @@
+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 grpo \
+    --model Qwen/Qwen3-30B-A3B-Instruct-2507 \
+    --save_safetensors true \
+    --merge_lora false \
+    --context_parallel_size 2 \
+    --tensor_model_parallel_size 2 \
+    --expert_model_parallel_size 4 \
+    --pipeline_model_parallel_size 2 \
+    --dataset open-r1/DAPO-Math-17k-Processed \
+    --num_train_epochs 1 \
+    --global_batch_size 64 \
+    --micro_batch_size 2 \
+    --steps_per_generation 2 \
+    --num_generations 8 \
+    --reward_funcs accuracy format \
+    --use_vllm true \
+    --vllm_mode colocate \
+    --vllm_gpu_memory_utilization 0.3 \
+    --vllm_tensor_parallel_size 4 \
+    --vllm_max_model_len 16384 \
+    --max_length 8192 \
+    --max_completion_length 8192 \
+    --tuner_type lora \
+    --lr 5e-5 \
+    --bf16 true \
+    --beta 0.00 \
+    --importance_sampling_level sequence \
+    --epsilon 3e-4 \
+    --epsilon_high 4e-4 \
+    --dynamic_sample false \
+    --overlong_filter true \
+    --loss_type grpo \
+    --sleep_level 2 \
+    --offload_model true \
+    --offload_bridge false \
+    --offload_optimizer true \
+    --logging_steps 1 \
+    --recompute_granularity selective \
+    --finetune \
+    --dataloader_num_workers 8 \
+    --dataset_num_proc 8 \
+    --no_save_optim \
+    --no_save_rng \
+    --attention_backend flash \
+    --temperature 1.0 \
+    --padding_free true \
+    --sequence_parallel true \
+    --log_completions true \
+    --report_to wandb

examples/megatron/grpo/multi_turn/frozen_lake.jsonl

@@ -0,0 +1,128 @@
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":0}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":1}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":2}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":3}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":4}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":5}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":6}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":7}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":8}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":9}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":10}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":11}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":12}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":13}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":14}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":15}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":16}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":17}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":18}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":19}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":20}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":21}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":22}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":23}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":24}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":25}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":26}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":27}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":28}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":29}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":30}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":31}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":32}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":33}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":34}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":35}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":36}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":37}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":38}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":39}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":40}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":41}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":42}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":43}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":44}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":45}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":46}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":47}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":48}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":49}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":50}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":51}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":52}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":53}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":54}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":55}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":56}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":57}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":58}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":59}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":60}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":61}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":62}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":63}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":64}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":65}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":66}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":67}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":68}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":69}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":70}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":71}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":72}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":73}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":74}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":75}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":76}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":77}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":78}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":79}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":80}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":81}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":82}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":83}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":84}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":85}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":86}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":87}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":88}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":89}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":90}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":91}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":92}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":93}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":94}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":95}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":96}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":97}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":98}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":99}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":100}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":101}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":102}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":103}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":104}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":105}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":106}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":107}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":108}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":109}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":110}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":111}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":112}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":113}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":114}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":115}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":116}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":117}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":118}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":119}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":120}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":121}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":122}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":123}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":124}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":125}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":126}}
+{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":127}}

examples/megatron/grpo/multi_turn/frozen_lake.sh

@@ -0,0 +1,66 @@
+# Multi-turn GRPO with the FrozenLake env.
+# Env lives in frozen_lake_plugin.py (loaded via --external_plugins);
+# with --use_gym_env true, the env's total_reward is consumed directly — no reward_funcs needed.
+# To prevent excessively long generations, max_completion_length is capped at 512 (per turn);
+# since prompts are short, max_length (first 9 turns + prompt) is capped at 6120.
+# vllm_max_model_len = max_length + last-turn length = 6632
+# reward improves from 0.2 → 0.6 within 120 steps: https://github.com/modelscope/ms-swift/pull/9405
+
+CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
+NPROC_PER_NODE=8 \
+megatron rlhf \
+    --rlhf_type grpo \
+    --model Qwen/Qwen3.5-2B \
+    --enable_thinking false \
+    --save_safetensors true \
+    --context_parallel_size 1 \
+    --tensor_model_parallel_size 1 \
+    --pipeline_model_parallel_size 1 \
+    --dataset 'examples/megatron/grpo/multi_turn/frozen_lake.jsonl#1024' \
+    --load_from_cache_file false \
+    --train_iters 120 \
+    --global_batch_size 64 \
+    --micro_batch_size 1 \
+    --steps_per_generation 4 \
+    --num_generations 8 \
+    --external_plugins examples/megatron/grpo/multi_turn/frozen_lake_plugin.py \
+    --use_vllm true \
+    --vllm_mode colocate \
+    --vllm_gpu_memory_utilization 0.5 \
+    --vllm_max_model_len 6632 \
+    --max_length 6120 \
+    --max_completion_length 512 \
+    --multi_turn_scheduler gym_scheduler \
+    --gym_env frozen_lake \
+    --use_gym_env true \
+    --max_turns 10 \
+    --tuner_type lora \
+    --lr 5e-5 \
+    --bf16 true \
+    --beta 0.001 \
+    --importance_sampling_level token \
+    --epsilon 0.2 \
+    --epsilon_high 0.2 \
+    --dynamic_sample false \
+    --overlong_filter true \
+    --loss_type grpo \
+    --sleep_level 2 \
+    --offload_model true \
+    --offload_bridge false \
+    --offload_optimizer true \
+    --logging_steps 1 \
+    --recompute_granularity selective \
+    --finetune \
+    --dataloader_num_workers 4 \
+    --dataset_num_proc 4 \
+    --no_save_optim \
+    --no_save_rng \
+    --attention_backend flash \
+    --temperature 1.0 \
+    --top_p 1.0 \
+    --top_k 80 \
+    --padding_free true \
+    --log_completions true \
+    --report_to tensorboard swanlab \
+    --eval_steps 1000 \
+    --save_steps 1000

examples/megatron/grpo/multi_turn/frozen_lake_plugin.py

@@ -0,0 +1,194 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+"""Text-based FrozenLake env for multi-turn GRPO training.
+
+Each turn the LLM sees an ASCII grid (S/G/H/F/P) and must reply with a
+single move inside ``<action>...</action>``. Reaching G yields reward 1.0;
+stepping into H ends the episode with 0. The outer step budget comes from
+swift's ``--max_turns`` flag. Per-row ``env_config.seed`` controls map
+generation so all ``num_generations`` rollouts of a row share the same map.
+
+Register via ``--external_plugins`` and select with ``--gym_env frozen_lake``.
+"""
+# code borrowed from ROLL/roll/pipeline/agentic/env/frozen_lake
+
+from __future__ import annotations
+
+import random
+import re
+from typing import Any, Dict, List, Optional, Tuple
+
+from swift.infer_engine.protocol import RolloutInferRequest
+from swift.rollout.gym_env import Env, envs
+from swift.template import Messages
+from swift.utils import get_logger
+
+logger = get_logger()
+
+
+def _is_valid(board: List[List[str]], size: int) -> bool:
+    """BFS from S; return True iff G is reachable through non-hole cells."""
+    start = None
+    for r in range(size):
+        for c in range(size):
+            if board[r][c] == 'S':
+                start = (r, c)
+                break
+        if start is not None:
+            break
+    if start is None:
+        return False
+
+    frontier = [start]
+    discovered = set()
+    while frontier:
+        r, c = frontier.pop()
+        if (r, c) in discovered:
+            continue
+        discovered.add((r, c))
+        for dr, dc in ((1, 0), (0, 1), (-1, 0), (0, -1)):
+            nr, nc = r + dr, c + dc
+            if 0 <= nr < size and 0 <= nc < size:
+                cell = board[nr][nc]
+                if cell == 'G':
+                    return True
+                if cell != 'H':
+                    frontier.append((nr, nc))
+    return False
+
+
+def generate_random_map(size: int = 4, p: float = 0.8, seed: Optional[int] = None) -> List[str]:
+    """Generate a random solvable FrozenLake map.
+
+    Args:
+        size: side length of the square grid.
+        p: probability of a non-hole (F) cell — 1 - p of any cell being H.
+        seed: RNG seed; same seed always yields the same map.
+
+    Returns:
+        List of `size` strings, each of length `size`, using chars S/G/F/H.
+        Both S and G positions are randomised (unlike gym's ``FrozenLake-v1``
+        which pins S to top-left and G to bottom-right); BFS-validated.
+    """
+    rng = random.Random(seed)
+    while True:
+        board = [['F' if rng.random() < p else 'H' for _ in range(size)] for _ in range(size)]
+        start_r, start_c = rng.randrange(size), rng.randrange(size)
+        goal_r, goal_c = rng.randrange(size), rng.randrange(size)
+        if (start_r, start_c) == (goal_r, goal_c):
+            continue
+        board[start_r][start_c] = 'S'
+        board[goal_r][goal_c] = 'G'
+        if _is_valid(board, size):
+            return [''.join(row) for row in board]
+
+
+# (row_delta, col_delta) for each canonical action token.
+ACTIONS: Dict[str, Tuple[int, int]] = {
+    'up': (-1, 0),
+    'down': (1, 0),
+    'left': (0, -1),
+    'right': (0, 1),
+}
+
+SYSTEM_PROMPT = ('You are playing FrozenLake. You see a grid where:\n'
+                 '  P = your current position, S = start, G = goal, H = hole, F = safe ice.\n'
+                 'Move one cell per turn. Reach G to win (+1 reward). Stepping into H ends '
+                 'the episode with 0 reward. Moves that would go off the grid leave you in '
+                 'place.\n\n'
+                 'On every turn, output your move inside <action>...</action>. '
+                 'The action must be exactly one of: up, down, left, right.\n\n'
+                 'Example: <action>down</action>')
+
+_ACTION_TAG_RE = re.compile(r'<action>\s*(up|down|left|right)\s*</action>', re.IGNORECASE)
+_BARE_ACTION_RE = re.compile(r'\b(up|down|left|right)\b', re.IGNORECASE)
+
+
+def _render(grid: List[str], row: int, col: int) -> str:
+    """Render the grid with the player position marked as 'P'."""
+    rendered = []
+    for r, line in enumerate(grid):
+        if r == row:
+            chars = list(line)
+            chars[col] = 'P' if chars[col] != 'G' else '*'  # '*' = player on goal
+            rendered.append(' '.join(chars))
+        else:
+            rendered.append(' '.join(line))
+    return '\n'.join(rendered)
+
+
+def _parse_action(completion: str) -> Optional[str]:
+    """Extract the action from the assistant message. Returns None if missing."""
+    m = _ACTION_TAG_RE.search(completion)
+    if m:
+        return m.group(1).lower()
+    matches = _BARE_ACTION_RE.findall(completion)
+    if matches:
+        return matches[-1].lower()
+    return None
+
+
+class FrozenLakeEnv(Env):
+
+    def __init__(self, env_config: Dict[str, Any]):
+        super().__init__(env_config)
+        self.size: int = int(env_config.get('size', 4))
+        self.p: float = float(env_config.get('p', 0.8))
+        seed = env_config.get('seed')
+        self.seed: Optional[int] = int(seed) if seed is not None else None
+        self.grid: List[str] = []
+        self.row: int = 0
+        self.col: int = 0
+        self.steps: int = 0
+
+    async def reset(self, config: RolloutInferRequest) -> Tuple[str, Dict[str, Any], str]:
+        self.grid = generate_random_map(size=self.size, p=self.p, seed=self.seed)
+        for r, line in enumerate(self.grid):
+            if 'S' in line:
+                self.row, self.col = r, line.index('S')
+                break
+        self.steps = 0
+
+        observation = (f'FrozenLake {self.size}x{self.size}:\n'
+                       f'{_render(self.grid, self.row, self.col)}\n\n'
+                       f'You are at row {self.row}, col {self.col}. Output your first move.')
+        info = {'seed': self.seed, 'size': self.size}
+        return observation, info, SYSTEM_PROMPT
+
+    async def step(self, action: Messages) -> Tuple[str, float, bool, Dict[str, Any]]:
+        completion = action[-1].get('content', '') if action else ''
+        move = _parse_action(completion)
+        self.steps += 1
+
+        info: Dict[str, Any] = {'seed': self.seed, 'step': self.steps, 'parsed_action': move}
+
+        if move is None:
+            obs = (f'Invalid response: could not find a move. Reply with '
+                   f'<action>up|down|left|right</action>.\n\n'
+                   f'{_render(self.grid, self.row, self.col)}')
+            return obs, 0.0, False, {**info, 'status': 'invalid_action'}
+
+        dr, dc = ACTIONS[move]
+        new_row = max(0, min(len(self.grid) - 1, self.row + dr))
+        new_col = max(0, min(len(self.grid[0]) - 1, self.col + dc))
+        cell = self.grid[new_row][new_col]
+        self.row, self.col = new_row, new_col
+
+        if cell == 'G':
+            obs = (f'You moved {move} and reached the goal!\n'
+                   f'{_render(self.grid, self.row, self.col)}')
+            return obs, 1.0, True, {**info, 'status': 'goal'}
+        if cell == 'H':
+            obs = (f'You moved {move} and fell into a hole. Episode over.\n'
+                   f'{_render(self.grid, self.row, self.col)}')
+            return obs, 0.0, True, {**info, 'status': 'hole'}
+
+        obs = (f'You moved {move}. Now at row {self.row}, col {self.col} (step {self.steps}).\n'
+               f'{_render(self.grid, self.row, self.col)}\n'
+               f'Output your next move.')
+        return obs, 0.0, False, {**info, 'status': 'ok'}
+
+    async def close(self):
+        pass
+
+
+envs['frozen_lake'] = FrozenLakeEnv

examples/megatron/grpo/multi_turn/frozen_lake_server.sh

@@ -0,0 +1,71 @@
+# server mode version of frozen_lake.sh
+CUDA_VISIBLE_DEVICES=0,1,2,3 \
+swift rollout \
+    --model Qwen/Qwen3.5-2B \
+    --enable_thinking false \
+    --vllm_data_parallel_size 4 \
+    --multi_turn_scheduler gym_scheduler \
+    --gym_env frozen_lake \
+    --use_gym_env true \
+    --max_turns 10 \
+    --max_length 6120 \
+    --vllm_max_model_len 6632 \
+    --external_plugins examples/megatron/grpo/multi_turn/frozen_lake_plugin.py > frozen_lake_rollout.log 2>&1 &
+
+# Wait for rollout server to be ready
+echo "Waiting for rollout server to start..."
+until curl -s http://127.0.0.1:8000/health/ > /dev/null 2>&1; do
+    sleep 10
+done
+echo "Rollout server is ready!"
+
+CUDA_VISIBLE_DEVICES=4,5,6,7 \
+NPROC_PER_NODE=4 \
+megatron rlhf \
+    --rlhf_type grpo \
+    --model Qwen/Qwen3.5-2B \
+    --save_safetensors true \
+    --context_parallel_size 1 \
+    --tensor_model_parallel_size 1 \
+    --pipeline_model_parallel_size 1 \
+    --train_iters 120 \
+    --dataset 'examples/megatron/grpo/multi_turn/frozen_lake.jsonl#1024' \
+    --load_from_cache_file false \
+    --global_batch_size 32 \
+    --micro_batch_size 1 \
+    --steps_per_generation 4 \
+    --num_generations 8 \
+    --external_plugins examples/megatron/grpo/multi_turn/frozen_lake_plugin.py \
+    --use_vllm true \
+    --vllm_mode server \
+    --vllm_server_host 127.0.0.1 \
+    --vllm_server_port 8000 \
+    --max_length 6120 \
+    --max_completion_length 512 \
+    --enable_thinking false \
+    --tuner_type lora \
+    --lr 5e-5 \
+    --bf16 true \
+    --beta 0.001 \
+    --importance_sampling_level token \
+    --epsilon 0.2 \
+    --epsilon_high 0.2 \
+    --dynamic_sample false \
+    --overlong_filter true \
+    --loss_type grpo \
+    --logging_steps 1 \
+    --recompute_granularity selective \
+    --finetune \
+    --dataloader_num_workers 4 \
+    --dataset_num_proc 4 \
+    --no_save_optim \
+    --no_save_rng \
+    --attention_backend flash \
+    --temperature 1.0 \
+    --top_p 1.0 \
+    --top_k 80 \
+    --padding_free true \
+    --log_completions true \
+    --report_to tensorboard swanlab \
+    --eval_steps 1000 \
+    --save_steps 1000 \

examples/megatron/grpo/opd_rl.sh

@@ -0,0 +1,61 @@
+SYSTEM_PROMPT="""You are a helpful math assistant. Solve the problem step by step and put your final answer within \\boxed{}."""
+
+CUDA_VISIBLE_DEVICES=0,1 \
+NPROC_PER_NODE=2 \
+megatron rlhf \
+    --rlhf_type grpo \
+    --model Qwen/Qwen3.5-2B \
+    --teacher_model Qwen/Qwen3.5-9B \
+    --enable_thinking false \
+    --save_safetensors true \
+    --expert_model_parallel_size 1 \
+    --tensor_model_parallel_size 2 \
+    --pipeline_model_parallel_size 1 \
+    --sequence_parallel true \
+    --system "$SYSTEM_PROMPT" \
+    --dataset modelscope/gsm8k \
+    --num_train_epochs 1 \
+    --global_batch_size 16 \
+    --micro_batch_size 2 \
+    --steps_per_generation 2 \
+    --num_generations 1 \
+    --use_vllm true \
+    --vllm_mode colocate \
+    --vllm_gpu_memory_utilization 0.6 \
+    --vllm_tensor_parallel_size 1 \
+    --vllm_max_model_len 12288 \
+    --sleep_level 1 \
+    --freeze_vit false \
+    --max_length 4096 \
+    --max_completion_length 2048 \
+    --tuner_type lora \
+    --lora_rank 8 \
+    --lora_alpha 32 \
+    --target_modules all-linear \
+    --lr 3e-5 \
+    --bf16 true \
+    --beta 0 \
+    --epsilon 0.2 \
+    --epsilon_high 0.2 \
+    --loss_type grpo \
+    --offload_model true \
+    --offload_bridge false \
+    --offload_optimizer true \
+    --log_rollout_offpolicy_metrics true \
+    --logging_steps 1 \
+    --recompute_granularity selective \
+    --finetune \
+    --dataloader_num_workers 8 \
+    --dataset_num_proc 8 \
+    --no_save_optim \
+    --no_save_rng \
+    --attention_backend flash \
+    --temperature 1 \
+    --padding_free false \
+    --log_completions true \
+    --report_to swanlab \
+    --train_iters 1000 \
+    --eval_steps 1000 \
+    --save_steps 100 \
+    --top_k 0 \
+    --top_p 1

examples/megatron/grpo/real.sh

@@ -0,0 +1,48 @@
+# REAL, https://arxiv.org/abs/2602.05630
+
+CUDA_VISIBLE_DEVICES=2 \
+swift rollout \
+    --model Qwen/Qwen2.5-0.5B-Instruct
+
+CUDA_VISIBLE_DEVICES=0,1 \
+NPROC_PER_NODE=2 \
+MASTER_PORT=29600 \
+megatron rlhf \
+    --rlhf_type grpo \
+    --model Qwen/Qwen2.5-0.5B-Instruct \
+    --save_safetensors true \
+    --context_parallel_size 1 \
+    --tensor_model_parallel_size 1 \
+    --pipeline_model_parallel_size 1 \
+    --dataset 'AI-MO/NuminaMath-TIR#4000' \
+    --num_train_epochs 1 \
+    --micro_batch_size 8 \
+    --global_batch_size 128 \
+    --num_generations 8 \
+    --reward_funcs accuracy \
+    --use_vllm true \
+    --vllm_mode server \
+    --vllm_server_host 127.0.0.1 \
+    --vllm_server_port 8000 \
+    --max_completion_length 2048 \
+    --tuner_type full \
+    --gradient_accumulation_fusion false \
+    --lr 2e-6 \
+    --bf16 true \
+    --beta 0.001 \
+    --dynamic_sample false \
+    --loss_type real \
+    --logging_steps 1 \
+    --recompute_granularity selective \
+    --finetune \
+    --dataloader_num_workers 8 \
+    --dataset_num_proc 8 \
+    --padding_free true \
+    --attention_backend flash \
+    --no_save_optim \
+    --no_save_rng \
+    --temperature 0.6 \
+    --system """You are a helpful math assistant. Solve the problem step by step and put your final answer within \\boxed{}.""" \
+    --log_completions true \
+    --eval_steps 100 \
+    --save_steps 100

examples/megatron/grpo/sapo.sh

@@ -0,0 +1,54 @@
+# SAPO https://arxiv.org/abs/2511.20347
+CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
+NPROC_PER_NODE=8 \
+MAX_PIXELS=602112 \
+megatron rlhf \
+  --rlhf_type grpo \
+  --loss_type sapo \
+  --tau_pos 1 \
+  --tau_neg 1.05 \
+  --model Qwen/Qwen2.5-VL-3B-Instruct \
+  --context_parallel_size 1 \
+  --tensor_model_parallel_size 1 \
+  --pipeline_model_parallel_size 1 \
+  --dataset AI-ModelScope/clevr_cogen_a_train \
+  --save_safetensors true \
+  --external_plugins examples/train/grpo/plugin/plugin.py \
+  --reward_funcs external_r1v_acc format \
+  --dynamic_sample false \
+  --steps_per_generation 4 \
+  --micro_batch_size 2 \
+  --global_batch_size 128 \
+  --num_generations 8 \
+  --use_vllm true \
+  --vllm_mode colocate \
+  --vllm_gpu_memory_utilization 0.7 \
+  --vllm_max_model_len 8192 \
+  --max_length 4096 \
+  --max_completion_length 4096 \
+  --tuner_type full \
+  --bf16 true \
+  --importance_sampling_level token \
+  --epsilon 0.2 \
+  --epsilon_high 0.2 \
+  --overlong_filter true \
+  --num_train_epochs 1 \
+  --eval_steps 1000 \
+  --save_steps 1000 \
+  --sleep_level 2 \
+  --offload_model true \
+  --offload_optimizer true \
+  --logging_steps 1 \
+  --recompute_granularity selective \
+  --finetune \
+  --lr 1e-6 \
+  --dataloader_num_workers 8 \
+  --dataset_num_proc 8 \
+  --no_save_optim \
+  --no_save_rng \
+  --attention_backend flash \
+  --temperature 1.0 \
+  --system examples/train/grpo/prompt.txt \
+  --beta 0.001 \
+  --padding_free true \
+  --report_to wandb

examples/megatron/long_text.sh

@@ -0,0 +1,36 @@
+# Env: 4 * A100
+# Max Length: 32K
+# GPU Memory: 4 * 50GB, Training Speed 23s/it
+PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
+NPROC_PER_NODE=4 \
+CUDA_VISIBLE_DEVICES=0,1,2,3 \
+megatron sft \
+    --model Qwen/Qwen2.5-7B \
+    --save_safetensors true \
+    --dataset 'ZhipuAI/LongWriter-6k' \
+    --load_from_cache_file true \
+    --split_dataset_ratio 0.01 \
+    --tensor_model_parallel_size 4 \
+    --micro_batch_size 1 \
+    --global_batch_size 8 \
+    --packing true \
+    --recompute_granularity full \
+    --recompute_method uniform \
+    --recompute_num_layers 1 \
+    --train_iters 1000 \
+    --eval_iters 50 \
+    --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 \
+    --eval_steps 200 \
+    --save_steps 200 \
+    --max_length 32768 \
+    --dataloader_num_workers 8 \
+    --dataset_num_proc 8 \
+    --no_save_optim true \
+    --no_save_rng true \
+    --sequence_parallel true \
+    --attention_backend flash

examples/megatron/lora/dense.sh

@@ -0,0 +1,39 @@
+# full: 2 * 70GiB 0.61s/it
+# lora: 2 * 14GiB 0.45s/it
+PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
+NPROC_PER_NODE=2 \
+CUDA_VISIBLE_DEVICES=0,1 \
+megatron sft \
+    --model Qwen/Qwen2.5-7B-Instruct \
+    --save_safetensors true \
+    --merge_lora false \
+    --dataset 'AI-ModelScope/alpaca-gpt4-data-zh#500' \
+              'AI-ModelScope/alpaca-gpt4-data-en#500' \
+              'swift/self-cognition#500' \
+    --tuner_type lora \
+    --lora_rank 8 \
+    --lora_alpha 32 \
+    --target_modules all-linear \
+    --tensor_model_parallel_size 2 \
+    --sequence_parallel true \
+    --micro_batch_size 16 \
+    --global_batch_size 16 \
+    --recompute_granularity full \
+    --recompute_method uniform \
+    --recompute_num_layers 1 \
+    --finetune true \
+    --cross_entropy_loss_fusion true \
+    --lr 1e-4 \
+    --lr_warmup_fraction 0.05 \
+    --min_lr 1e-5 \
+    --num_train_epochs 1 \
+    --output_dir megatron_output/Qwen2.5-7B-Instruct \
+    --save_steps 100 \
+    --max_length 2048 \
+    --system 'You are a helpful assistant.' \
+    --dataloader_num_workers 4 \
+    --no_save_optim true \
+    --no_save_rng true \
+    --dataset_num_proc 4 \
+    --model_author swift \
+    --model_name swift-robot

examples/megatron/lora/dpo.sh

@@ -0,0 +1,45 @@
+# 2 * 65GiB; 4.50s/it
+PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
+NPROC_PER_NODE=2 \
+CUDA_VISIBLE_DEVICES=0,1 \
+megatron rlhf \
+    --rlhf_type dpo \
+    --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 \
+    --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 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-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 2048 \
+    --dataloader_num_workers 8 \
+    --dataset_num_proc 8 \
+    --no_save_optim true \
+    --no_save_rng true \
+    --sequence_parallel true \
+    --attention_backend flash \
+    --rpo_alpha 0.1 \
+    --beta 0.1 \
+    --loss_type sigmoid

examples/megatron/lora/loss_scale.sh

@@ -0,0 +1,44 @@
+# 2 * 60GiB, 3.4s/it
+PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
+NPROC_PER_NODE=2 \
+CUDA_VISIBLE_DEVICES=0,1 \
+megatron sft \
+    --model Qwen/Qwen3-30B-A3B-Base \
+    --save_safetensors true \
+    --merge_lora false \
+    --tuner_type lora \
+    --dataset AI-ModelScope/function-calling-chatml#10000 \
+    --load_from_cache_file true \
+    --loss_scale hermes \
+    --agent_template hermes \
+    --lora_rank 8 \
+    --lora_alpha 32 \
+    --target_modules all-linear \
+    --modules_to_save word_embeddings output_layer \
+    --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 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-4 \
+    --lr_warmup_fraction 0.05 \
+    --min_lr 1e-5 \
+    --output_dir megatron_output/Qwen3-30B-A3B-Base \
+    --eval_steps 200 \
+    --save_steps 200 \
+    --max_length 2048 \
+    --dataloader_num_workers 8 \
+    --dataset_num_proc 8 \
+    --no_save_optim true \
+    --no_save_rng true \
+    --sequence_parallel true \
+    --attention_backend flash

examples/megatron/lora/moe.sh

@@ -0,0 +1,44 @@
+# 2 * 62GiB, 5.10s/it
+PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
+NPROC_PER_NODE=2 \
+CUDA_VISIBLE_DEVICES=0,1 \
+megatron sft \
+    --model Qwen/Qwen3-30B-A3B \
+    --save_safetensors true \
+    --merge_lora false \
+    --dataset 'swift/Qwen3-SFT-Mixin#2000' \
+              'swift/self-cognition:empty_think#600' \
+    --loss_scale ignore_empty_think \
+    --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 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-4 \
+    --lr_warmup_fraction 0.05 \
+    --min_lr 1e-5 \
+    --output_dir megatron_output/Qwen3-30B-A3B \
+    --eval_steps 200 \
+    --save_steps 200 \
+    --max_length 2048 \
+    --dataloader_num_workers 8 \
+    --dataset_num_proc 8 \
+    --no_save_optim true \
+    --no_save_rng true \
+    --sequence_parallel true \
+    --attention_backend flash \
+    --model_author swift \
+    --model_name swift-robot

examples/megatron/lora/mtp.sh

@@ -0,0 +1,61 @@
+# demo: thinking -> non-thinking
+# 4 * 70GiB; 40s/it
+PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
+NPROC_PER_NODE=4 \
+CUDA_VISIBLE_DEVICES=0,1,2,3 \
+megatron sft \
+    --model ZhipuAI/GLM-4.5-Air \
+    --save_safetensors true \
+    --merge_lora true \
+    --mtp_num_layers 1 \
+    --dataset 'swift/Chinese-Qwen3-235B-2507-Distill-data-110k-SFT' \
+    --load_from_cache_file true \
+    --tuner_type lora \
+    --lora_rank 32 \
+    --lora_alpha 64 \
+    --target_modules linear_qkv linear_proj \
+    --split_dataset_ratio 0.01 \
+    --moe_permute_fusion true \
+    --tensor_model_parallel_size 4 \
+    --expert_tensor_parallel_size 1 \
+    --expert_model_parallel_size 4 \
+    --moe_grouped_gemm true \
+    --moe_shared_expert_overlap true \
+    --moe_aux_loss_coeff 1e-3 \
+    --micro_batch_size 1 \
+    --global_batch_size 16 \
+    --recompute_granularity full \
+    --recompute_method uniform \
+    --recompute_num_layers 1 \
+    --num_train_epochs 2 \
+    --finetune true \
+    --cross_entropy_loss_fusion true \
+    --lr 1e-4 \
+    --lr_warmup_fraction 0.05 \
+    --min_lr 1e-5 \
+    --output_dir megatron_output/GLM-4.5-Air \
+    --eval_steps 200 \
+    --save_steps 200 \
+    --packing true \
+    --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
+
+# If not using the MTP module, please remove the speculative-related parameters.
+# CUDA_VISIBLE_DEVICES=0,1,2,3 \
+# swift infer \
+#     --model megatron_output/GLM-4.5-Air/vx-xxx/checkpoint-xxx-merged \
+#     --sglang_tp_size 4 \
+#     --infer_backend sglang \
+#     --load_data_args true \
+#     --sglang_context_length 8192 \
+#     --max_new_tokens 2048 \
+#     --sglang_mem_fraction_static 0.7 \
+#     --sglang_speculative_algorithm EAGLE \
+#     --sglang_speculative_eagle_topk 1 \
+#     --sglang_speculative_num_steps 3 \
+#     --sglang_speculative_num_draft_tokens 4

examples/megatron/lora/new_special_tokens.sh

@@ -0,0 +1,45 @@
+# 2 * 60GiB, 2.7s/it
+# Note: The conversion script has no differences.
+# It will read the new_special_tokens parameter from args.json.
+
+PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
+NPROC_PER_NODE=2 \
+CUDA_VISIBLE_DEVICES=0,1 \
+megatron sft \
+    --model Qwen/Qwen3-30B-A3B \
+    --save_safetensors true \
+    --merge_lora false \
+    --dataset 'swift/new_special_tokens' \
+    --new_special_tokens 'examples/train/new_special_tokens/tokens.txt' \
+    --tuner_type lora \
+    --lora_rank 8 \
+    --lora_alpha 32 \
+    --target_modules all-linear \
+    --modules_to_save word_embeddings output_layer \
+    --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 32 \
+    --global_batch_size 64 \
+    --recompute_granularity full \
+    --recompute_method uniform \
+    --recompute_num_layers 1 \
+    --num_train_epochs 5 \
+    --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 \
+    --eval_steps 500 \
+    --save_steps 500 \
+    --max_length 2048 \
+    --dataloader_num_workers 8 \
+    --dataset_num_proc 8 \
+    --no_save_optim true \
+    --no_save_rng true \
+    --sequence_parallel true \
+    --attention_backend flash

examples/megatron/lora/qwen3_235b.sh

@@ -0,0 +1,48 @@
+# 8 * 80GiB, 3.2s/it
+# If you're doing full-parameter training, you'll need 64 × 80 GiB of GPU memory
+
+PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
+NPROC_PER_NODE=8 \
+CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
+megatron sft \
+    --model Qwen/Qwen3-235B-A22B-Instruct-2507 \
+    --dataset 'swift/Chinese-Qwen3-235B-2507-Distill-data-110k-SFT#2000' \
+              'swift/self-cognition#1000' \
+    --save_safetensors true \
+    --merge_lora false \
+    --load_from_cache_file true \
+    --tuner_type lora \
+    --lora_rank 8 \
+    --lora_alpha 32 \
+    --target_modules all-linear \
+    --split_dataset_ratio 0.01 \
+    --moe_permute_fusion true \
+    --tensor_model_parallel_size 4 \
+    --expert_tensor_parallel_size 1 \
+    --expert_model_parallel_size 8 \
+    --moe_grouped_gemm true \
+    --moe_shared_expert_overlap true \
+    --moe_aux_loss_coeff 1e-3 \
+    --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-4 \
+    --lr_warmup_fraction 0.05 \
+    --min_lr 1e-5 \
+    --output_dir megatron_output/Qwen3-235B-A22B-Instruct-2507 \
+    --eval_steps 200 \
+    --save_steps 200 \
+    --max_length 2048 \
+    --dataloader_num_workers 8 \
+    --dataset_num_proc 8 \
+    --no_save_optim true \
+    --no_save_rng true \
+    --sequence_parallel true \
+    --attention_backend flash \
+    --model_author swift \
+    --model_name swift-robot

examples/megatron/mcore_bridge/full/dense.sh

@@ -0,0 +1,47 @@
+# 2 * 76GiB
+PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
+NPROC_PER_NODE=2 \
+IMAGE_MAX_TOKEN_NUM=1024 \
+VIDEO_MAX_TOKEN_NUM=128 \
+FPS_MAX_FRAMES=16 \
+CUDA_VISIBLE_DEVICES=0,1 \
+megatron sft \
+    --model Qwen/Qwen3-VL-8B-Instruct \
+    --save_safetensors true \
+    --dataset 'AI-ModelScope/LaTeX_OCR:human_handwrite#5000' \
+    --load_from_cache_file true \
+    --tensor_model_parallel_size 2 \
+    --sequence_parallel true \
+    --packing true \
+    --freeze_llm false \
+    --freeze_vit true \
+    --freeze_aligner true \
+    --split_dataset_ratio 0.01 \
+    --micro_batch_size 1 \
+    --global_batch_size 4 \
+    --recompute_granularity full \
+    --recompute_method uniform \
+    --recompute_num_layers 1 \
+    --finetune true \
+    --cross_entropy_loss_fusion true \
+    --lr 1e-5 \
+    --lr_warmup_fraction 0.05 \
+    --min_lr 1e-6 \
+    --num_train_epochs 1 \
+    --output_dir megatron_output/Qwen3-VL-8B-Instruct \
+    --save_steps 200 \
+    --max_length 2048 \
+    --dataloader_num_workers 4 \
+    --no_save_optim true \
+    --no_save_rng true \
+    --dataset_num_proc 8
+
+# PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
+# IMAGE_MAX_TOKEN_NUM=1024 \
+# VIDEO_MAX_TOKEN_NUM=128 \
+# FPS_MAX_FRAMES=16 \
+# CUDA_VISIBLE_DEVICES=0 \
+# swift infer \
+#     --model megatron_output/Qwen3-VL-8B-Instruct/vx-xxx/checkpoint-xxx \
+#     --load_data_args true \
+#     --stream true

examples/megatron/mcore_bridge/full/moe.sh

@@ -0,0 +1,41 @@
+# 8 * 76GiB, 3s/it
+PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
+CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
+NPROC_PER_NODE=8 \
+megatron sft \
+    --model Qwen/Qwen3-30B-A3B-Instruct-2507 \
+    --save_safetensors true \
+    --dataset 'swift/Chinese-Qwen3-235B-Thinking-2507-Distill-data-110k-SFT#20000' \
+    --load_from_cache_file true \
+    --split_dataset_ratio 0.01 \
+    --moe_permute_fusion true \
+    --pipeline_model_parallel_size 2 \
+    --decoder_first_pipeline_num_layers 25 \
+    --tensor_model_parallel_size 4 \
+    --expert_model_parallel_size 4 \
+    --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-5 \
+    --lr_warmup_fraction 0.05 \
+    --min_lr 1e-6 \
+    --output_dir megatron_output/Qwen3-30B-A3B-Instruct-2507 \
+    --eval_steps 500 \
+    --save_steps 500 \
+    --max_length 8192 \
+    --packing true \
+    --dataloader_num_workers 8 \
+    --dataset_num_proc 8 \
+    --no_save_optim true \
+    --no_save_rng true \
+    --sequence_parallel true \
+    --moe_expert_capacity_factor 2 \
+    --attention_backend flash

examples/megatron/mcore_bridge/lora/moe.sh

@@ -0,0 +1,52 @@
+# 50GiB
+PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
+NPROC_PER_NODE=2 \
+CUDA_VISIBLE_DEVICES=0,1 \
+megatron sft \
+    --model Qwen/Qwen3-30B-A3B-Instruct-2507 \
+    --save_safetensors true \
+    --merge_lora false \
+    --dataset 'swift/Chinese-Qwen3-235B-2507-Distill-data-110k-SFT#2000' \
+              'swift/self-cognition#1000' \
+    --load_from_cache_file true \
+    --tuner_type lora \
+    --lora_rank 8 \
+    --lora_alpha 32 \
+    --target_modules all-linear \
+    --split_dataset_ratio 0.01 \
+    --moe_permute_fusion true \
+    --expert_model_parallel_size 2 \
+    --moe_grouped_gemm true \
+    --moe_shared_expert_overlap true \
+    --moe_aux_loss_coeff 1e-3 \
+    --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-4 \
+    --lr_warmup_fraction 0.05 \
+    --min_lr 1e-5 \
+    --output_dir megatron_output/Qwen3-30B-A3B-Instruct-2507 \
+    --eval_steps 200 \
+    --save_steps 200 \
+    --max_length 2048 \
+    --dataloader_num_workers 8 \
+    --dataset_num_proc 8 \
+    --no_save_optim true \
+    --no_save_rng true \
+    --sequence_parallel true \
+    --moe_expert_capacity_factor 2 \
+    --attention_backend flash \
+    --model_author swift \
+    --model_name swift-robot
+
+
+# CUDA_VISIBLE_DEVICES=0 \
+# swift infer \
+#     --model Qwen/Qwen3-30B-A3B-Instruct-2507 \
+#     --adapters megatron_output/Qwen3-30B-A3B-Instruct-2507/vx-xxx/checkpoint-xxx \
+#     --stream true

examples/megatron/mcore_bridge/lora/new_special_tokens.sh

@@ -0,0 +1,50 @@
+# 2 * 50GiB
+PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
+NPROC_PER_NODE=2 \
+CUDA_VISIBLE_DEVICES=0,1 \
+megatron sft \
+    --model Qwen/Qwen3-30B-A3B-Instruct-2507 \
+    --save_safetensors true \
+    --merge_lora false \
+    --dataset 'swift/new_special_tokens' \
+    --load_from_cache_file true \
+    --new_special_tokens 'examples/train/new_special_tokens/tokens.txt' \
+    --tuner_type lora \
+    --lora_rank 8 \
+    --lora_alpha 32 \
+    --target_modules all-linear \
+    --modules_to_save word_embeddings output_layer \
+    --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 32 \
+    --global_batch_size 64 \
+    --recompute_granularity full \
+    --recompute_method uniform \
+    --recompute_num_layers 1 \
+    --num_train_epochs 5 \
+    --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 500 \
+    --save_steps 500 \
+    --max_length 2048 \
+    --dataloader_num_workers 8 \
+    --dataset_num_proc 8 \
+    --no_save_optim true \
+    --no_save_rng true \
+    --moe_expert_capacity_factor 2 \
+    --sequence_parallel true \
+    --attention_backend flash
+
+# CUDA_VISIBLE_DEVICES=0,1 \
+# swift infer \
+#     --adapters megatron_output/Qwen3-30B-A3B-Instruct-2507/vx-xxx/checkpoint-xxx \
+#     --load_data_args true \
+#     --stream true

examples/megatron/mcore_bridge/lora/seq_cls.sh

@@ -0,0 +1,54 @@
+# 2 * 15GiB
+PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
+NPROC_PER_NODE=2 \
+IMAGE_MAX_TOKEN_NUM=1024 \
+VIDEO_MAX_TOKEN_NUM=128 \
+FPS_MAX_FRAMES=16 \
+CUDA_VISIBLE_DEVICES=0,1 \
+megatron sft \
+    --model Qwen/Qwen3-VL-8B-Instruct \
+    --save_safetensors true \
+    --merge_lora false \
+    --dataset 'tany0699/garbage265#20000' \
+    --load_from_cache_file true \
+    --tuner_type lora \
+    --lora_rank 8 \
+    --lora_alpha 32 \
+    --target_modules all-linear \
+    --tensor_model_parallel_size 2 \
+    --sequence_parallel true \
+    --packing true \
+    --freeze_llm false \
+    --freeze_vit true \
+    --freeze_aligner true \
+    --split_dataset_ratio 0.01 \
+    --micro_batch_size 1 \
+    --global_batch_size 4 \
+    --recompute_granularity full \
+    --recompute_method uniform \
+    --recompute_num_layers 1 \
+    --finetune true \
+    --cross_entropy_loss_fusion true \
+    --lr 1e-4 \
+    --lr_warmup_fraction 0.05 \
+    --min_lr 1e-5 \
+    --num_train_epochs 1 \
+    --output_dir megatron_output/Qwen3-VL-8B-Instruct \
+    --save_steps 200 \
+    --max_length 2048 \
+    --dataloader_num_workers 4 \
+    --no_save_optim true \
+    --no_save_rng true \
+    --num_labels 265 \
+    --task_type seq_cls \
+    --dataset_num_proc 8
+
+# PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
+# IMAGE_MAX_TOKEN_NUM=1024 \
+# VIDEO_MAX_TOKEN_NUM=128 \
+# FPS_MAX_FRAMES=16 \
+# CUDA_VISIBLE_DEVICES=0 \
+# swift infer \
+#     --adapters megatron_output/Qwen3-VL-8B-Instruct/vx-xxx/checkpoint-xxx \
+#     --load_data_args true \
+#     --stream true

examples/megatron/moe/deepseek_v3.sh

@@ -0,0 +1,41 @@
+# 8 * 60GiB; 9s/it
+# For ease of use, we use moonshotai/Moonlight-16B-A3B-Instruct, which is also based on the DeepseekV3ForCausalLM architecture.
+# https://modelscope.cn/models/moonshotai/Moonlight-16B-A3B-Instruct/file/view/master/config.json?status=1
+PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
+NPROC_PER_NODE=8 \
+CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
+megatron sft \
+    --model moonshotai/Moonlight-16B-A3B-Instruct \
+    --save_safetensors true \
+    --dataset 'liucong/Chinese-DeepSeek-R1-Distill-data-110k-SFT' \
+    --load_from_cache_file true \
+    --split_dataset_ratio 0.01 \
+    --pipeline_model_parallel_size 2 \
+    --decoder_last_pipeline_num_layers 13 \
+    --moe_permute_fusion true \
+    --expert_model_parallel_size 4 \
+    --moe_grouped_gemm true \
+    --moe_shared_expert_overlap true \
+    --moe_aux_loss_coeff 1e-3 \
+    --micro_batch_size 1 \
+    --global_batch_size 16 \
+    --packing true \
+    --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/Moonlight-16B-A3B-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

examples/megatron/moe/moe.sh

@@ -0,0 +1,42 @@
+# 8 * 57GiB, 2.95s/it
+PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
+NPROC_PER_NODE=8 \
+CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
+megatron sft \
+    --model Qwen/Qwen1.5-MoE-A2.7B \
+    --save_safetensors true \
+    --dataset 'liucong/Chinese-DeepSeek-R1-Distill-data-110k-SFT' \
+    --load_from_cache_file true \
+    --split_dataset_ratio 0.01 \
+    --pipeline_model_parallel_size 2 \
+    --decoder_last_pipeline_num_layers 11 \
+    --expert_model_parallel_size 4 \
+    --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 16 \
+    --packing true \
+    --moe_router_dtype fp32 \
+    --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/Qwen1.5-MoE-A2.7B \
+    --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 \
+    --overlap_param_gather true \
+    --overlap_grad_reduce true

examples/megatron/moe/qwen3_moe.sh

@@ -0,0 +1,42 @@
+# ZeRO3: 91.2s/it; 16 * 80GiB
+# Megatron-LM: 9.6s/it; 16 * 60GiB
+# Launch using Alibaba Cloud DLC
+# https://help.aliyun.com/zh/pai/user-guide/general-environment-variables
+# ref: https://github.com/modelscope/ms-swift/blob/main/examples/train/multi-node/dlc/train.sh
+PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
+NNODES=$WORLD_SIZE \
+NODE_RANK=$RANK \
+megatron sft \
+    --model Qwen/Qwen3-30B-A3B-Base \
+    --save_safetensors true \
+    --dataset 'liucong/Chinese-DeepSeek-R1-Distill-data-110k-SFT' \
+    --load_from_cache_file true \
+    --split_dataset_ratio 0.01 \
+    --pipeline_model_parallel_size 2 \
+    --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 16 \
+    --packing true \
+    --recompute_granularity full \
+    --recompute_method uniform \
+    --recompute_num_layers 1 \
+    --num_train_epochs 3 \
+    --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-Base \
+    --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

examples/megatron/moe/qwen3_moe_offload.sh

@@ -0,0 +1,40 @@
+# 28s/it; 4 * 75GiB
+PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
+NPROC_PER_NODE=4 \
+CUDA_VISIBLE_DEVICES=0,1,2,3 \
+megatron sft \
+    --model Qwen/Qwen3-30B-A3B-Base \
+    --save_safetensors true \
+    --dataset 'liucong/Chinese-DeepSeek-R1-Distill-data-110k-SFT' \
+    --load_from_cache_file true \
+    --split_dataset_ratio 0.01 \
+    --expert_model_parallel_size 4 \
+    --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 16 \
+    --packing true \
+    --recompute_granularity full \
+    --recompute_method uniform \
+    --recompute_num_layers 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-Base \
+    --eval_steps 200 \
+    --save_steps 200 \
+    --max_length 8192 \
+    --num_train_epochs 3 \
+    --dataloader_num_workers 8 \
+    --dataset_num_proc 8 \
+    --no_save_optim true \
+    --no_save_rng true \
+    --sequence_parallel true \
+    --optimizer_cpu_offload true \
+    --use_precision_aware_optimizer true \
+    --optimizer_offload_fraction 1 \
+    --attention_backend flash

examples/megatron/multi-node/node1.sh

@@ -0,0 +1,37 @@
+# For more information on multi-node training launch methods, refer to:
+# https://github.com/modelscope/ms-swift/tree/main/examples/train/multi-node
+
+PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
+CUDA_VISIBLE_DEVICES=0,1,2,3 \
+NNODES=2 \
+NODE_RANK=0 \
+MASTER_ADDR=127.0.0.1 \
+MASTER_PORT=29500 \
+NPROC_PER_NODE=4 \
+megatron sft \
+    --model Qwen/Qwen2.5-14B \
+    --save_safetensors true \
+    --dataset 'liucong/Chinese-DeepSeek-R1-Distill-data-110k-SFT' \
+    --load_from_cache_file true \
+    --split_dataset_ratio 0.01 \
+    --tensor_model_parallel_size 4 \
+    --micro_batch_size 1 \
+    --global_batch_size 16 \
+    --packing true \
+    --recompute_granularity selective \
+    --num_train_epochs 3 \
+    --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-14B \
+    --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

examples/megatron/multi-node/node2.sh

@@ -0,0 +1,34 @@
+PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
+CUDA_VISIBLE_DEVICES=0,1,2,3 \
+NNODES=2 \
+NODE_RANK=1 \
+MASTER_ADDR=xxx.xxx.xxx.xxx \
+MASTER_PORT=29500 \
+NPROC_PER_NODE=4 \
+megatron sft \
+    --model Qwen/Qwen2.5-14B \
+    --save_safetensors true \
+    --dataset 'liucong/Chinese-DeepSeek-R1-Distill-data-110k-SFT' \
+    --load_from_cache_file true \
+    --split_dataset_ratio 0.01 \
+    --tensor_model_parallel_size 4 \
+    --micro_batch_size 1 \
+    --global_batch_size 16 \
+    --packing true \
+    --recompute_granularity selective \
+    --num_train_epochs 3 \
+    --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-14B \
+    --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

examples/megatron/multimodal/dense/dpo.sh

@@ -0,0 +1,41 @@
+# 4 * 50GiB 14s/it
+PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
+NPROC_PER_NODE=4 \
+MAX_PIXELS=1003520 \
+CUDA_VISIBLE_DEVICES=0,1,2,3 \
+megatron rlhf \
+    --rlhf_type dpo \
+    --model Qwen/Qwen2.5-VL-7B-Instruct \
+    --save_safetensors true \
+    --dataset 'swift/RLAIF-V-Dataset#20000' \
+    --load_from_cache_file true \
+    --tuner_type full \
+    --tensor_model_parallel_size 4 \
+    --sequence_parallel true \
+    --freeze_llm false \
+    --freeze_vit true \
+    --freeze_aligner true \
+    --packing true \
+    --split_dataset_ratio 0.01 \
+    --micro_batch_size 1 \
+    --global_batch_size 4 \
+    --recompute_granularity full \
+    --recompute_method uniform \
+    --recompute_num_layers 1 \
+    --finetune true \
+    --cross_entropy_loss_fusion true \
+    --lr 1e-5 \
+    --lr_warmup_fraction 0.05 \
+    --min_lr 1e-6 \
+    --num_train_epochs 1 \
+    --output_dir megatron_output/Qwen2.5-VL-7B-Instruct \
+    --save_steps 200 \
+    --max_length 8192 \
+    --dataloader_num_workers 4 \
+    --no_save_optim true \
+    --no_save_rng true \
+    --dataset_num_proc 8 \
+    --attention_backend flash \
+    --rpo_alpha 0.1 \
+    --beta 0.1 \
+    --loss_type sigmoid

examples/megatron/multimodal/dense/full.sh

@@ -0,0 +1,35 @@
+# 2 * 72GiB; 4.1s/it
+PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
+NPROC_PER_NODE=2 \
+MAX_PIXELS=1003520 \
+CUDA_VISIBLE_DEVICES=0,1 \
+megatron sft \
+    --model Qwen/Qwen2.5-VL-7B-Instruct \
+    --save_safetensors true \
+    --dataset 'AI-ModelScope/LaTeX_OCR:human_handwrite#5000' \
+    --load_from_cache_file true \
+    --tensor_model_parallel_size 2 \
+    --sequence_parallel true \
+    --packing true \
+    --freeze_llm false \
+    --freeze_vit true \
+    --freeze_aligner true \
+    --split_dataset_ratio 0.01 \
+    --micro_batch_size 1 \
+    --global_batch_size 4 \
+    --recompute_granularity full \
+    --recompute_method uniform \
+    --recompute_num_layers 1 \
+    --finetune true \
+    --cross_entropy_loss_fusion true \
+    --lr 1e-5 \
+    --lr_warmup_fraction 0.05 \
+    --min_lr 1e-6 \
+    --num_train_epochs 1 \
+    --output_dir megatron_output/Qwen2.5-VL-7B-Instruct \
+    --save_steps 200 \
+    --max_length 2048 \
+    --dataloader_num_workers 4 \
+    --no_save_optim true \
+    --no_save_rng true \
+    --dataset_num_proc 8

examples/megatron/multimodal/dense/lora.sh

@@ -0,0 +1,40 @@
+# 2 * 23GiB; 2.3s/it
+PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
+NPROC_PER_NODE=2 \
+MAX_PIXELS=1003520 \
+CUDA_VISIBLE_DEVICES=0,1 \
+megatron sft \
+    --model Qwen/Qwen2.5-VL-7B-Instruct \
+    --save_safetensors true \
+    --merge_lora false \
+    --dataset 'AI-ModelScope/LaTeX_OCR:human_handwrite#5000' \
+    --load_from_cache_file true \
+    --tuner_type lora \
+    --lora_rank 8 \
+    --lora_alpha 32 \
+    --target_modules all-linear \
+    --tensor_model_parallel_size 1 \
+    --sequence_parallel true \
+    --freeze_llm false \
+    --freeze_vit true \
+    --freeze_aligner true \
+    --packing true \
+    --split_dataset_ratio 0.01 \
+    --micro_batch_size 1 \
+    --global_batch_size 4 \
+    --recompute_granularity full \
+    --recompute_method uniform \
+    --recompute_num_layers 1 \
+    --finetune true \
+    --cross_entropy_loss_fusion true \
+    --lr 1e-4 \
+    --lr_warmup_fraction 0.05 \
+    --min_lr 1e-5 \
+    --num_train_epochs 1 \
+    --output_dir megatron_output/Qwen2.5-VL-7B-Instruct \
+    --save_steps 200 \
+    --max_length 2048 \
+    --dataloader_num_workers 4 \
+    --no_save_optim true \
+    --no_save_rng true \
+    --dataset_num_proc 8

examples/megatron/multimodal/lora_llm_vit_full/sft.sh

@@ -0,0 +1,69 @@
+# Currently, 'lora_llm' must set `--merge_lora true` and `--no_save_optim true`
+
+PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
+NPROC_PER_NODE=4 \
+CUDA_VISIBLE_DEVICES=0,1,2,3 \
+IMAGE_MAX_TOKEN_NUM=1024 \
+VIDEO_MAX_TOKEN_NUM=128 \
+FPS_MAX_FRAMES=12 \
+megatron sft \
+    --model Qwen/Qwen3.5-35B-A3B \
+    --save_safetensors true \
+    --merge_lora true \
+    --dataset 'AI-ModelScope/alpaca-gpt4-data-zh#500' \
+              'AI-ModelScope/alpaca-gpt4-data-en#500' \
+              'swift/self-cognition#500' \
+              'AI-ModelScope/LaTeX_OCR:human_handwrite#2000' \
+    --load_from_cache_file true \
+    --add_non_thinking_prefix true \
+    --loss_scale ignore_empty_think \
+    --split_dataset_ratio 0.01 \
+    --tuner_type lora_llm \
+    --lora_rank 8 \
+    --lora_alpha 32 \
+    --target_modules all-linear \
+    --expert_model_parallel_size 4 \
+    --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 2 \
+    --finetune true \
+    --cross_entropy_loss_fusion true \
+    --lr 1e-4 \
+    --vit_lr 1e-5 \
+    --aligner_lr 2e-5 \
+    --lr_warmup_fraction 0.05 \
+    --min_lr 1e-5 \
+    --output_dir megatron_output/Qwen3.5-35B-A3B \
+    --eval_steps 200 \
+    --save_steps 200 \
+    --max_length 2048 \
+    --dataloader_num_workers 8 \
+    --dataset_num_proc 8 \
+    --no_save_optim true \
+    --no_save_rng true \
+    --sequence_parallel true \
+    --attention_backend flash \
+    --padding_free true \
+    --packing true \
+    --model_author swift \
+    --model_name swift-robot
+
+
+# PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
+# CUDA_VISIBLE_DEVICES=0,1,2,3 \
+# IMAGE_MAX_TOKEN_NUM=1024 \
+# VIDEO_MAX_TOKEN_NUM=128 \
+# FPS_MAX_FRAMES=12 \
+# swift infer \
+#     --model megatron_output/Qwen3.5-35B-A3B/vx-xxx/checkpoint-xxx-merged \
+#     --stream true \
+#     --experts_impl grouped_mm \
+#     --enable_thinking false \
+#     --load_data_args true

examples/megatron/multimodal/moe/full_dpo_offload.sh

@@ -0,0 +1,52 @@
+# 16s/it; 8 * 65GiB
+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 OpenGVLab/InternVL3_5-30B-A3B \
+    --save_safetensors true \
+    --dataset 'swift/RLAIF-V-Dataset#20000' \
+    --load_from_cache_file true \
+    --tuner_type full \
+    --tensor_model_parallel_size 4 \
+    --expert_tensor_parallel_size 1 \
+    --pipeline_model_parallel_size 2 \
+    --decoder_first_pipeline_num_layers 23 \
+    --split_dataset_ratio 0.01 \
+    --expert_model_parallel_size 4 \
+    --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 \
+    --packing true \
+    --recompute_granularity full \
+    --recompute_method uniform \
+    --recompute_num_layers 1 \
+    --finetune true \
+    --cross_entropy_loss_fusion true \
+    --lr 1e-5 \
+    --lr_warmup_fraction 0.05 \
+    --min_lr 1e-6 \
+    --output_dir megatron_output/InternVL3_5-30B-A3B \
+    --eval_steps 500 \
+    --save_steps 500 \
+    --max_length 16384 \
+    --num_train_epochs 1 \
+    --dataloader_num_workers 8 \
+    --dataset_num_proc 8 \
+    --no_save_optim true \
+    --no_save_rng true \
+    --sequence_parallel true \
+    --freeze_llm false \
+    --freeze_vit true \
+    --freeze_aligner true \
+    --optimizer_cpu_offload true \
+    --use_precision_aware_optimizer true \
+    --optimizer_offload_fraction 0.65 \
+    --attention_backend flash \
+    --rpo_alpha 0.1 \
+    --beta 0.1 \
+    --loss_type sigmoid

examples/megatron/multimodal/moe/lora.sh

@@ -0,0 +1,45 @@
+# 2 * 43GiB, 8s/it
+PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
+NPROC_PER_NODE=2 \
+CUDA_VISIBLE_DEVICES=0,1 \
+megatron sft \
+    --model OpenGVLab/InternVL3_5-30B-A3B \
+    --save_safetensors true \
+    --merge_lora false \
+    --dataset 'AI-ModelScope/LaTeX_OCR:human_handwrite#5000' \
+    --load_from_cache_file true \
+    --tuner_type lora \
+    --lora_rank 8 \
+    --lora_alpha 32 \
+    --target_modules all-linear \
+    --sequence_parallel true \
+    --freeze_llm false \
+    --freeze_vit true \
+    --freeze_aligner true \
+    --packing true \
+    --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 \
+    --finetune true \
+    --cross_entropy_loss_fusion true \
+    --lr 1e-4 \
+    --lr_warmup_fraction 0.05 \
+    --min_lr 1e-5 \
+    --num_train_epochs 1 \
+    --output_dir megatron_output/InternVL3_5-30B-A3B \
+    --eval_steps 200 \
+    --save_steps 200 \
+    --max_length 2048 \
+    --dataloader_num_workers 8 \
+    --dataset_num_proc 8 \
+    --no_save_optim true \
+    --no_save_rng true \
+    --attention_backend flash

examples/megatron/multimodal/omni/dense.sh

@@ -0,0 +1,41 @@
+# 2 * 75GiB; 7s/it
+# Supports mixed modalities
+PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
+ENABLE_AUDIO_OUTPUT=0 \
+NPROC_PER_NODE=2 \
+MAX_PIXELS=1003520 \
+VIDEO_MAX_PIXELS=50176 \
+FPS_MAX_FRAMES=12 \
+CUDA_VISIBLE_DEVICES=0,1 \
+megatron sft \
+    --model Qwen/Qwen2.5-Omni-7B \
+    --save_safetensors true \
+    --dataset 'AI-ModelScope/alpaca-gpt4-data-zh#10000' \
+              'AI-ModelScope/LaTeX_OCR:human_handwrite#5000' \
+              'speech_asr/speech_asr_aishell1_trainsets:validation#5000' \
+    --load_from_cache_file true \
+    --tensor_model_parallel_size 2 \
+    --sequence_parallel true \
+    --packing true \
+    --freeze_llm false \
+    --freeze_vit true \
+    --freeze_aligner true \
+    --split_dataset_ratio 0.01 \
+    --micro_batch_size 1 \
+    --global_batch_size 4 \
+    --recompute_granularity full \
+    --recompute_method uniform \
+    --recompute_num_layers 1 \
+    --finetune true \
+    --cross_entropy_loss_fusion true \
+    --lr 1e-5 \
+    --lr_warmup_fraction 0.05 \
+    --min_lr 1e-6 \
+    --num_train_epochs 1 \
+    --output_dir megatron_output/Qwen2.5-Omni-7B \
+    --save_steps 100 \
+    --max_length 4096 \
+    --dataloader_num_workers 4 \
+    --no_save_optim true \
+    --no_save_rng true \
+    --dataset_num_proc 8

examples/megatron/multimodal/omni/moe.sh

@@ -0,0 +1,51 @@
+# 2 * 50GiB
+PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
+ENABLE_AUDIO_OUTPUT=0 \
+NPROC_PER_NODE=2 \
+IMAGE_MAX_TOKEN_NUM=1024 \
+VIDEO_MAX_TOKEN_NUM=128 \
+FPS_MAX_FRAMES=12 \
+CUDA_VISIBLE_DEVICES=0,1 \
+megatron sft \
+    --model Qwen/Qwen3-Omni-30B-A3B-Instruct \
+    --save_safetensors true \
+    --merge_lora false \
+    --dataset 'AI-ModelScope/alpaca-gpt4-data-zh#10000' \
+              'AI-ModelScope/LaTeX_OCR:human_handwrite#5000' \
+              'speech_asr/speech_asr_aishell1_trainsets:validation#5000' \
+    --load_from_cache_file true \
+    --tuner_type lora \
+    --lora_rank 8 \
+    --lora_alpha 32 \
+    --target_modules all-linear \
+    --sequence_parallel true \
+    --packing true \
+    --freeze_llm false \
+    --freeze_vit true \
+    --freeze_aligner true \
+    --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 \
+    --finetune true \
+    --cross_entropy_loss_fusion true \
+    --lr 1e-4 \
+    --lr_warmup_fraction 0.05 \
+    --min_lr 1e-5 \
+    --num_train_epochs 1 \
+    --output_dir megatron_output/Qwen3-Omni-30B-A3B-Instruct \
+    --eval_steps 100 \
+    --save_steps 100 \
+    --max_length 4096 \
+    --dataloader_num_workers 8 \
+    --dataset_num_proc 8 \
+    --no_save_optim true \
+    --no_save_rng true \
+    --attention_backend flash

examples/megatron/muon.sh

@@ -0,0 +1,36 @@
+# muon: 2 * 65GiB, 4m 14s
+# adam(w): 2 * 78GiB, 1m 19s
+# mcore>=0.16
+PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
+NPROC_PER_NODE=2 \
+CUDA_VISIBLE_DEVICES=0,1 \
+megatron sft \
+    --model Qwen/Qwen2.5-7B-Instruct \
+    --save_safetensors true \
+    --dataset 'AI-ModelScope/alpaca-gpt4-data-zh#500' \
+              'AI-ModelScope/alpaca-gpt4-data-en#500' \
+              'swift/self-cognition#500' \
+    --optimizer dist_muon \
+    --tensor_model_parallel_size 2 \
+    --sequence_parallel true \
+    --micro_batch_size 16 \
+    --global_batch_size 16 \
+    --recompute_granularity full \
+    --recompute_method uniform \
+    --recompute_num_layers 1 \
+    --finetune true \
+    --cross_entropy_loss_fusion true \
+    --lr 1e-5 \
+    --lr_warmup_fraction 0.05 \
+    --min_lr 1e-6 \
+    --num_train_epochs 1 \
+    --output_dir megatron_output/Qwen2.5-7B-Instruct \
+    --save_steps 100 \
+    --max_length 2048 \
+    --system 'You are a helpful assistant.' \
+    --dataloader_num_workers 4 \
+    --no_save_optim true \
+    --no_save_rng true \
+    --dataset_num_proc 4 \
+    --model_author swift \
+    --model_name swift-robot

examples/megatron/pretrain.sh

@@ -0,0 +1,30 @@
+# 4 * 80GiB
+PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
+NPROC_PER_NODE=4 \
+CUDA_VISIBLE_DEVICES=0,1,2,3 \
+megatron pt \
+    --model Qwen/Qwen2.5-7B \
+    --save_safetensors true \
+    --dataset swift/chinese-c4 \
+    --streaming true \
+    --packing true \
+    --tensor_model_parallel_size 4 \
+    --micro_batch_size 1 \
+    --global_batch_size 16 \
+    --recompute_granularity selective \
+    --train_iters 10000 \
+    --finetune true \
+    --cross_entropy_loss_fusion true \
+    --lr 1e-5 \
+    --lr_warmup_iters 300 \
+    --min_lr 1e-6 \
+    --output_dir megatron_output/Qwen2.5-7B \
+    --eval_steps 500 \
+    --save_steps 500 \
+    --max_length 8192 \
+    --dataloader_num_workers 4 \
+    --dataset_num_proc 8 \
+    --no_save_optim true \
+    --no_save_rng true \
+    --sequence_parallel true \
+    --attention_backend flash

examples/megatron/reranker/qwen3_reranker.sh

@@ -0,0 +1,35 @@
+# 2 * 80GiB
+# For inference code, refer to: examples/infer/demo_reranker.py
+PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
+NPROC_PER_NODE=2 \
+CUDA_VISIBLE_DEVICES=0,1 \
+megatron sft \
+    --model Qwen/Qwen3-Reranker-8B \
+    --task_type generative_reranker \
+    --save_safetensors true \
+    --tuner_type full \
+    --dataset MTEB/scidocs-reranking \
+    --load_from_cache_file true \
+    --split_dataset_ratio 0.02 \
+    --tensor_model_parallel_size 2 \
+    --sequence_parallel true \
+    --micro_batch_size 2 \
+    --global_batch_size 16 \
+    --recompute_granularity full \
+    --recompute_method uniform \
+    --recompute_num_layers 1 \
+    --finetune true \
+    --cross_entropy_loss_fusion true \
+    --lr 5e-6 \
+    --lr_warmup_fraction 0.05 \
+    --min_lr 1e-7 \
+    --num_train_epochs 1 \
+    --output_dir megatron_output/Qwen3-Reranker-8B \
+    --save_steps 200 \
+    --eval_steps 50 \
+    --max_length 4096 \
+    --loss_type pointwise_reranker \
+    --dataloader_num_workers 4 \
+    --no_save_optim true \
+    --no_save_rng true \
+    --dataset_num_proc 4

examples/megatron/reranker/qwen3_vl_reranker.sh

@@ -0,0 +1,35 @@
+# 4 * 50GiB
+# For inference code, refer to: examples/infer/demo_reranker.py
+PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
+NPROC_PER_NODE=4 \
+CUDA_VISIBLE_DEVICES=0,1,2,3 \
+megatron sft \
+    --model Qwen/Qwen3-VL-Reranker-8B \
+    --task_type generative_reranker \
+    --save_safetensors true \
+    --tuner_type full \
+    --dataset swift/TextCaps:rerank \
+    --load_from_cache_file true \
+    --split_dataset_ratio 0.02 \
+    --tensor_model_parallel_size 4 \
+    --sequence_parallel true \
+    --micro_batch_size 1 \
+    --global_batch_size 16 \
+    --recompute_granularity full \
+    --recompute_method uniform \
+    --recompute_num_layers 1 \
+    --finetune true \
+    --cross_entropy_loss_fusion true \
+    --lr 5e-6 \
+    --lr_warmup_fraction 0.05 \
+    --min_lr 1e-7 \
+    --num_train_epochs 1 \
+    --output_dir megatron_output/Qwen3-VL-Reranker-8B \
+    --save_steps 200 \
+    --eval_steps 50 \
+    --max_length 4096 \
+    --loss_type pointwise_reranker \
+    --dataloader_num_workers 4 \
+    --no_save_optim true \
+    --no_save_rng true \
+    --dataset_num_proc 4

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

examples/megatron/seq_cls/full.sh

@@ -0,0 +1,39 @@
+# 4 * 60GiB; 7.5s/it
+PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
+NPROC_PER_NODE=4 \
+CUDA_VISIBLE_DEVICES=0,1,2,3 \
+megatron sft \
+    --model Qwen/Qwen2.5-VL-7B-Instruct \
+    --save_safetensors true \
+    --dataset 'tany0699/garbage265#20000' \
+    --load_from_cache_file true \
+    --tensor_model_parallel_size 2 \
+    --sequence_parallel true \
+    --packing true \
+    --freeze_llm false \
+    --freeze_vit true \
+    --freeze_aligner true \
+    --split_dataset_ratio 0.01 \
+    --micro_batch_size 2 \
+    --global_batch_size 4 \
+    --model_kwargs '{"max_pixels": 1003520}' \
+    --recompute_granularity full \
+    --recompute_method uniform \
+    --recompute_num_layers 1 \
+    --finetune true \
+    --cross_entropy_loss_fusion true \
+    --lr 1e-5 \
+    --lr_warmup_fraction 0.05 \
+    --min_lr 1e-6 \
+    --num_train_epochs 1 \
+    --output_dir megatron_output/Qwen2.5-VL-7B-Instruct \
+    --save_steps 200 \
+    --eval_steps 200 \
+    --vit_gradient_checkpointing true \
+    --max_length 8192 \
+    --dataloader_num_workers 4 \
+    --no_save_optim true \
+    --no_save_rng true \
+    --dataset_num_proc 8 \
+    --num_labels 265 \
+    --task_type seq_cls \

examples/megatron/seq_cls/lora/infer.sh

@@ -0,0 +1,9 @@
+# acc: 0.9298597194388778
+# 60GiB
+CUDA_VISIBLE_DEVICES=0 \
+swift infer \
+  --model megatron_output/Qwen3-30B-A3B-Instruct-2507/vx-xxx/checkpoint-xxx \
+  --load_data_args true \
+  --max_batch_size 16 \
+  --attn_impl flash_attn \
+  --metric acc

examples/megatron/seq_cls/lora/train.sh

@@ -0,0 +1,44 @@
+# 2 * 40GiB; 5s/it
+# eval_acc: 0.924
+PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
+NPROC_PER_NODE=2 \
+CUDA_VISIBLE_DEVICES=0,1 \
+megatron sft \
+    --model Qwen/Qwen3-30B-A3B-Instruct-2507 \
+    --save_safetensors true \
+    --merge_lora true \
+    --dataset 'DAMO_NLP/jd:cls' \
+    --load_from_cache_file true \
+    --split_dataset_ratio 0.01 \
+    --tuner_type lora \
+    --lora_rank 8 \
+    --lora_alpha 32 \
+    --target_modules all-linear all-router \
+    --packing true \
+    --expert_model_parallel_size 2 \
+    --moe_permute_fusion true \
+    --moe_grouped_gemm true \
+    --moe_shared_expert_overlap true \
+    --moe_aux_loss_coeff 1e-6 \
+    --sequence_parallel true \
+    --micro_batch_size 1 \
+    --global_batch_size 2 \
+    --recompute_granularity full \
+    --recompute_method uniform \
+    --recompute_num_layers 1 \
+    --finetune true \
+    --cross_entropy_loss_fusion true \
+    --lr 1e-4 \
+    --lr_warmup_fraction 0.05 \
+    --min_lr 1e-5 \
+    --num_train_epochs 1 \
+    --output_dir megatron_output/Qwen3-30B-A3B-Instruct-2507 \
+    --eval_steps 200 \
+    --save_steps 200 \
+    --max_length 2048 \
+    --task_type seq_cls \
+    --num_labels 2 \
+    --dataloader_num_workers 4 \
+    --no_save_optim true \
+    --no_save_rng true \
+    --dataset_num_proc 4

examples/megatron/sft.sh

@@ -0,0 +1,33 @@
+# 2 * 80GiB
+PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
+NPROC_PER_NODE=2 \
+CUDA_VISIBLE_DEVICES=0,1 \
+megatron sft \
+    --model Qwen/Qwen2.5-7B-Instruct \
+    --save_safetensors true \
+    --dataset 'AI-ModelScope/alpaca-gpt4-data-zh#500' \
+              'AI-ModelScope/alpaca-gpt4-data-en#500' \
+              'swift/self-cognition#500' \
+    --tensor_model_parallel_size 2 \
+    --sequence_parallel true \
+    --micro_batch_size 16 \
+    --global_batch_size 16 \
+    --recompute_granularity full \
+    --recompute_method uniform \
+    --recompute_num_layers 1 \
+    --finetune true \
+    --cross_entropy_loss_fusion true \
+    --lr 1e-5 \
+    --lr_warmup_fraction 0.05 \
+    --min_lr 1e-6 \
+    --num_train_epochs 1 \
+    --output_dir megatron_output/Qwen2.5-7B-Instruct \
+    --save_steps 100 \
+    --max_length 2048 \
+    --system 'You are a helpful assistant.' \
+    --dataloader_num_workers 4 \
+    --no_save_optim true \
+    --no_save_rng true \
+    --dataset_num_proc 4 \
+    --model_author swift \
+    --model_name swift-robot