SWIFT Megatron 版 GKD 指南JSD 蒸馏损失、Teacher 部署与 On-Policy 生成实战【免费下载链接】swiftUse PEFT or Full-parameter to CPT/SFT/DPO/GRPO 600 LLMs (Qwen3.6, DeepSeek-V4, GLM-5.1, InternLM3, Llama4, ...) and 300 MLLMs (Qwen3-VL, Qwen3-Omni, InternVL3.5, Ovis2.5, GLM4.5v, Gemma4, Llava, Phi4, ...) (AAAI 2025).项目地址: https://gitcode.com/GitHub_Trending/swift1/swift本文围绕 SWIFTms-swift仓库中的 Megatron-SWIFT GKDGeneralized Knowledge Distillation广义知识蒸馏文档展开系统讲解如何在 Megatron 后端上用megatron rlhf --rlhf_type gkd完成教师→学生分布对齐训练。读完本文你将掌握 GKD 全部专有参数beta、lmbda、sft_alpha、gkd_logits_topk 等的语义与默认值、本地教师模型 / Teacher API / 多教师路由三种接入方式的差异以及 CP/PP/TP/EP 并行、vLLM on-policy 生成与多轮蒸馏Multi-turn GKD的落地配置并能直接复制仓库中 gkd 示例脚本 跑通完整训练。GKD 是什么用 JSD 损失对齐教师与学生分布GKD 是一种把教师模型teacher的知识迁移给学生模型student的训练方法对同一输入序列分别计算学生与教师的输出分布并以** Jensen-Shannon 散度JSD**作为损失来拉近两者。SWIFT 仓库同时提供 HF 后端与 Megatron 后端的 GKD 实现本文聚焦 Megatron-SWIFT 版本其核心入口是 MegatronGKDTrainer继承MegatronRolloutMixin与MegatronRLHFTrainer。如果你是 GKD / OPD-RL 的初次接触者建议先阅读 蒸馏文档 了解整体概念。Megatron 版 GKD 当前支持的特性见 GKD 官方文档训练模式全参数训练full与 LoRA 微调--tuner_type lora并行策略Context ParallelCP、Pipeline ParallelPP、Tensor ParallelTP、Expert ParallelEP模型支持兼容 Megatron-SWIFT 中的 LLM 与 MLLM教师卸载--offload_teacher_model支持把教师模型卸载到 CPU 以节省 GPU 显存在线生成支持用 vLLM 对学生模型做 on-policy 生成多轮训练通过--multi_turn_scheduler支持多轮 GKD与 GRPO 共享同一套MultiTurnScheduler基础设施核心参数详解GKD 专有参数参数类型默认值说明--teacher_modelstr-教师模型的路径或模型 ID使用teacher_model_server时可省略--teacher_model_serverstrNone教师 API URL支持单 URL 或多教师 JSON 配置--teacher_tag_keystrdataset多教师路由时用于将样本标签匹配到教师tags的列名--gkd_logits_topkintNoneTop-K logit 数量使用外部 API 教师时必填--betafloat0.5JSD 散度插值系数0.0 前向 KL0.5 对称 JSD1.0 反向 KL--lmbdafloat0.5On-Policy 学习概率0.0 纯 Off-Policy1.0 纯 On-Policy--temperaturefloat0.9采样与损失计算使用的温度--sft_alphafloat0混入一定比例的 SFT 损失仅作用于非学生生成的完成--max_completion_lengthint512生成完成的最大 token 数这些参数在 Megatron 参数定义 中声明并在 参数校验逻辑 中强制执行若干约束gkd_logits_topk与teacher_model_server强绑定走 API 路径时教师只返回 top-k logprobs因此一旦设置了teacher_model_server而未设置gkd_logits_topk会直接抛出ValueError(gkd_logits_topk is required when using teacher_model_server)且gkd_logits_topk必须为正整数。教师两种接入方式互斥teacher_model与teacher_model_server不能同时设置解析阶段即拒绝。Ray 模式限制use_ray与teacher_model_server组合不被支持API 教师在 Ray 后端下的请求路由尚未打通从参数校验的注释可以确认该限制。teacher_model_server既接受单个 URL 字符串也接受多教师 JSON 配置解析逻辑复用 parse_teacher_model_server。批量相关参数与 Megatron SFT 一致用以下参数控制批量大小参数说明--micro_batch_size每个 DP 组的训练 batch size--global_batch_size全局 batch sizemicro_batch_size × dp_size × gradient_accumulation_steps其余完整参数请参考 Megatron-SWIFT 命令行参数文档。损失函数原理JSD 插值与 Top-K 模式GKD 损失统一实现在 gkd_loss.py该模块注释明确shared across HF / Megatron / Ray backends——即 HF、Megatron、Ray 三个后端共用同一份数学核心。JSD 插值jsd_lossbeta0退化为前向 KLKL(teacher || student)beta1退化为反向 KLKL(student || teacher)0beta1构造混合分布m beta·teacher (1-beta)·student源码中以logsumexp在对数域计算损失为beta·KL(m||teacher) (1-beta)·KL(m||student)计算按chunk_size512分块进行控制大词表下的显存峰值当 CP 某个 rank 分片内没有有效 tokenN0时直接返回零避免cp_reduce挂死。Top-K 模式TeacherOutput教师输出被统一封装为TeacherOutput数据类支持full_logits全词表与topk_logprobs topk_indicestop-k两种形态。在 gkd_loss 主流程 中top-k 模式下用gather_fn只收集学生在教师 top-k 索引处的 logit 后做 JSD全词表模式下先经 _align_vocab 对师生词表做 padding 对齐再走 TP 感知的log_softmax/kl_div。Megatron 侧注入的正是vocab_parallel_log_softmax与vocab_parallel_kl_div见 gkd_trainer.py 的loss_func使 JSD 计算在词表并行切分下保持正确。SFT 损失混入当sft_alpha 0且当前批次来自数据集off-policy时loss_func 会额外计算标准 causal LM 损失并按loss jsd_loss sft_alpha · sft_loss加权CP1 时 SFT 统计量会先做 context-parallel all-reduce 再平均。学生自生成的批次on-policy不叠加 SFT 损失这与参数表中applied to non-student-generated completions的描述一致。On-Policy 机制lmbda 概率调度与 vLLM 生成--lmbda控制每个训练步的数据来源。从 _determine_data_source 可以看到判定逻辑训练器用隔离的random.Random(seed step)实例生成一个在所有分布式进程上确定一致的随机数保证各 rank 决策相同随机数小于lmbda时进入 On-Policy 分支此时要求--use_vllm true学生模型经 vLLM 在线生成完成若未启用 vLLM会打印一次性告警并回退到数据集响应否则使用数据集自带的响应Off-Policy。生成结果按micro_batch_size切分为若干微批次编码教师 logits 计算本地前向或 API 请求在 _generate_and_score_completions 中统一完成。另外当--truncation_strategy deleteGKD 的默认截断策略见 MegatronRLHFArguments时超长样本会用独立种子seed1的数据迭代器重采样补齐见 _init_resample_data_iterator。自蒸馏OPSD特例当teacher_model与model指向同一模型且学生为 LoRA 时训练器进入自蒸馏模式——教师通过disable_adapter()上下文复用学生权重见 prepare_model避免加载第二份模型此时教师 logits 会在每个训练步重算而非每个生成周期以确保教师始终跟踪最新的 LoRA 权重_on_train_step_batch 中的注释解释了这一设计。训练脚本实战仓库在 examples/megatron/rlhf/gkd/ 下提供四个可直接参考的脚本分别覆盖四种典型场景。场景一本地大教师 vLLM on-policydense.shdense.sh 是标准 8 卡配置用Qwen/Qwen3-8B-Base作学生LoRAQwen/Qwen3-32B作本地教师CUDA_VISIBLE_DEVICES0,1,2,3,4,5,6,7 \ NPROC_PER_NODE8 \ PYTORCH_CUDA_ALLOC_CONFexpandable_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要点解读--lmbda 1表示每个训练步都走 on-policy学生生成--beta 1取反向 KL--offload_teacher_model true把 32B 教师平时挂在 CPU 上仅在计算教师 logits 时加载配合--sleep_level 1vLLM 休眠释放显存实现 32B 教师与 8B 学生共享 8 卡教师前向在 _compute_teacher_logits_local 中以torch.no_grad()执行并按gkd_logits_topk若设置做vocab_parallel_topk压缩--padding_free true与--sequence_parallel true用于变长序列的显存与效率优化。场景二Teacher ServerAPI 教师teacher_server.shteacher_server.sh 演示教师不占训练卡的解耦部署教师单独用swift deploy起 vLLM 服务训练端只通过--teacher_model_server http://localhost:8000拉取 logprobstop_k64 max_prompt_length2048 max_completion_length2048 max_total_length$((max_prompt_length max_completion_length)) # 教师服务需先启动脚本中为注释示例 # CUDA_VISIBLE_DEVICES0 \ # 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_DEVICES1,2 \ NPROC_PER_NODE2 \ PYTORCH_CUDA_ALLOC_CONFexpandable_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 \ --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 \ --max_length $max_prompt_length \ --max_completion_length $max_completion_length \ --use_vllm true \ --vllm_mode colocate \ --sleep_level 1 \ --padding_free true \ --recompute_granularity selective \ ...这里有两个硬性前提值得强调教师服务必须通过swift deploy --max_logprobs $top_k启动使推理端真正返回 top-k logprobs否则gkd_logits_topk无数据可用训练端必须传--gkd_logits_topk否则会触发上文提到的参数校验报错。从源码看API 路径下教师 logits 的获取由 fetch_teacher_parsed_by_routing 完成请求的聚合gather→推理→切回scatter多教师时按--teacher_tag_key默认取样本的dataset列路由到不同教师详见 蒸馏文档的多教师路由章节。场景三多轮 GKDmulti_turn.shmulti_turn.sh 展示与学生模型多轮交互的蒸馏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 \ --lmbda 1 \ --beta 0.5 \ --multi_turn_scheduler math_tip_trick \ --max_turns 2 \ --loss_scale last_round \ --truncation_strategy delete \ --vllm_server_pass_dataset true \ --remove_unused_columns false \ ...--multi_turn_scheduler math_tip_trick启用数学提示型多轮调度器与 GRPO 共用同一套MultiTurnScheduler基础设施调度器注册见 swift/rollout/multi_turn.py--max_turns 2限制交互轮数--loss_scale last_round表示只对最后一轮回复计损失GKD 在 Megatron 侧的默认loss_scale即last_round见 rlhf_args.py--vllm_server_pass_dataset true让数据集字段透传给 vLLM server--remove_unused_columns false保证多轮插件需要的原始列不被裁掉。场景四OPSD 自蒸馏opsd.shopsd.sh 复现 Self-Distilled Reasoner 的固定教师自蒸馏模式--teacher_model与--model同为Qwen/Qwen3-4B教师即关闭 LoRA adapter 后的基础模型脚本注释与 _teacher_use_disable_adapter 的日志均可印证。关键超参lmbda1, beta0.5, temperature1.2, lora_rank64, lora_alpha128, effective batch32并挂 external_plugins 处理样本构建build_opsd_samples。并行策略与分布式细节Megatron 版 GKD 对四种并行的适配体现在损失与数据两条链路TP词表并行loss_func注入tp_gather_topk、vocab_parallel_log_softmax、vocab_parallel_kl_div见 vocab_parallel_utilsJSD 在词表切分下正确归约CP上下文并行JSD 总和经 cp_reduce 在 CP 组内 all-reduce 后按有效 token 数平均最终损失还会除以 CP world sizePP流水线并行非末段 stage 以TeacherOutput()占位符兜底extract_active 中处理自蒸馏模式的教师重算特意放在 pipeline schedule 之外执行因此 PP1 同样受支持_on_train_step_batch 注释EPMoE 模型通过--expert_model_parallel_size配置示例脚本中均已显式声明。延伸阅读蒸馏概念与多教师路由docs/source_en/Instruction/Distillation.mdMegatron-SWIFT 完整参数表docs/source_en/Megatron-SWIFT/Command-line-parameters.md全部 GKD 训练脚本examples/megatron/rlhf/gkd损失数学核心HF/Megatron/Ray 共用swift/rlhf_trainers/gkd_loss.py训练器实现swift/megatron/trainers/gkd_trainer.pyRay 后端的 GKD 实现swift/ray/megatron/gkd_trainer.py 与 swift/ray/megatron/loss/gkd.py【免费下载链接】swiftUse PEFT or Full-parameter to CPT/SFT/DPO/GRPO 600 LLMs (Qwen3.6, DeepSeek-V4, GLM-5.1, InternLM3, Llama4, ...) and 300 MLLMs (Qwen3-VL, Qwen3-Omni, InternVL3.5, Ovis2.5, GLM4.5v, Gemma4, Llava, Phi4, ...) (AAAI 2025).项目地址: https://gitcode.com/GitHub_Trending/swift1/swift创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
SWIFT Megatron 版 GKD 指南JSD 蒸馏损失、Teacher 部署与 On-Policy 生成实战【免费下载链接】swiftUse PEFT or Full-parameter to CPT/SFT/DPO/GRPO 600 LLMs (Qwen3.6, DeepSeek-V4, GLM-5.1, InternLM3, Llama4, ...) and 300 MLLMs (Qwen3-VL, Qwen3-Omni, InternVL3.5, Ovis2.5, GLM4.5v, Gemma4, Llava, Phi4, ...) (AAAI 2025).项目地址: https://gitcode.com/GitHub_Trending/swift1/swift本文围绕 SWIFTms-swift仓库中的 Megatron-SWIFT GKDGeneralized Knowledge Distillation广义知识蒸馏文档展开系统讲解如何在 Megatron 后端上用megatron rlhf --rlhf_type gkd完成教师→学生分布对齐训练。读完本文你将掌握 GKD 全部专有参数beta、lmbda、sft_alpha、gkd_logits_topk 等的语义与默认值、本地教师模型 / Teacher API / 多教师路由三种接入方式的差异以及 CP/PP/TP/EP 并行、vLLM on-policy 生成与多轮蒸馏Multi-turn GKD的落地配置并能直接复制仓库中 gkd 示例脚本 跑通完整训练。GKD 是什么用 JSD 损失对齐教师与学生分布GKD 是一种把教师模型teacher的知识迁移给学生模型student的训练方法对同一输入序列分别计算学生与教师的输出分布并以** Jensen-Shannon 散度JSD**作为损失来拉近两者。SWIFT 仓库同时提供 HF 后端与 Megatron 后端的 GKD 实现本文聚焦 Megatron-SWIFT 版本其核心入口是 MegatronGKDTrainer继承MegatronRolloutMixin与MegatronRLHFTrainer。如果你是 GKD / OPD-RL 的初次接触者建议先阅读 蒸馏文档 了解整体概念。Megatron 版 GKD 当前支持的特性见 GKD 官方文档训练模式全参数训练full与 LoRA 微调--tuner_type lora并行策略Context ParallelCP、Pipeline ParallelPP、Tensor ParallelTP、Expert ParallelEP模型支持兼容 Megatron-SWIFT 中的 LLM 与 MLLM教师卸载--offload_teacher_model支持把教师模型卸载到 CPU 以节省 GPU 显存在线生成支持用 vLLM 对学生模型做 on-policy 生成多轮训练通过--multi_turn_scheduler支持多轮 GKD与 GRPO 共享同一套MultiTurnScheduler基础设施核心参数详解GKD 专有参数参数类型默认值说明--teacher_modelstr-教师模型的路径或模型 ID使用teacher_model_server时可省略--teacher_model_serverstrNone教师 API URL支持单 URL 或多教师 JSON 配置--teacher_tag_keystrdataset多教师路由时用于将样本标签匹配到教师tags的列名--gkd_logits_topkintNoneTop-K logit 数量使用外部 API 教师时必填--betafloat0.5JSD 散度插值系数0.0 前向 KL0.5 对称 JSD1.0 反向 KL--lmbdafloat0.5On-Policy 学习概率0.0 纯 Off-Policy1.0 纯 On-Policy--temperaturefloat0.9采样与损失计算使用的温度--sft_alphafloat0混入一定比例的 SFT 损失仅作用于非学生生成的完成--max_completion_lengthint512生成完成的最大 token 数这些参数在 Megatron 参数定义 中声明并在 参数校验逻辑 中强制执行若干约束gkd_logits_topk与teacher_model_server强绑定走 API 路径时教师只返回 top-k logprobs因此一旦设置了teacher_model_server而未设置gkd_logits_topk会直接抛出ValueError(gkd_logits_topk is required when using teacher_model_server)且gkd_logits_topk必须为正整数。教师两种接入方式互斥teacher_model与teacher_model_server不能同时设置解析阶段即拒绝。Ray 模式限制use_ray与teacher_model_server组合不被支持API 教师在 Ray 后端下的请求路由尚未打通从参数校验的注释可以确认该限制。teacher_model_server既接受单个 URL 字符串也接受多教师 JSON 配置解析逻辑复用 parse_teacher_model_server。批量相关参数与 Megatron SFT 一致用以下参数控制批量大小参数说明--micro_batch_size每个 DP 组的训练 batch size--global_batch_size全局 batch sizemicro_batch_size × dp_size × gradient_accumulation_steps其余完整参数请参考 Megatron-SWIFT 命令行参数文档。损失函数原理JSD 插值与 Top-K 模式GKD 损失统一实现在 gkd_loss.py该模块注释明确shared across HF / Megatron / Ray backends——即 HF、Megatron、Ray 三个后端共用同一份数学核心。JSD 插值jsd_lossbeta0退化为前向 KLKL(teacher || student)beta1退化为反向 KLKL(student || teacher)0beta1构造混合分布m beta·teacher (1-beta)·student源码中以logsumexp在对数域计算损失为beta·KL(m||teacher) (1-beta)·KL(m||student)计算按chunk_size512分块进行控制大词表下的显存峰值当 CP 某个 rank 分片内没有有效 tokenN0时直接返回零避免cp_reduce挂死。Top-K 模式TeacherOutput教师输出被统一封装为TeacherOutput数据类支持full_logits全词表与topk_logprobs topk_indicestop-k两种形态。在 gkd_loss 主流程 中top-k 模式下用gather_fn只收集学生在教师 top-k 索引处的 logit 后做 JSD全词表模式下先经 _align_vocab 对师生词表做 padding 对齐再走 TP 感知的log_softmax/kl_div。Megatron 侧注入的正是vocab_parallel_log_softmax与vocab_parallel_kl_div见 gkd_trainer.py 的loss_func使 JSD 计算在词表并行切分下保持正确。SFT 损失混入当sft_alpha 0且当前批次来自数据集off-policy时loss_func 会额外计算标准 causal LM 损失并按loss jsd_loss sft_alpha · sft_loss加权CP1 时 SFT 统计量会先做 context-parallel all-reduce 再平均。学生自生成的批次on-policy不叠加 SFT 损失这与参数表中applied to non-student-generated completions的描述一致。On-Policy 机制lmbda 概率调度与 vLLM 生成--lmbda控制每个训练步的数据来源。从 _determine_data_source 可以看到判定逻辑训练器用隔离的random.Random(seed step)实例生成一个在所有分布式进程上确定一致的随机数保证各 rank 决策相同随机数小于lmbda时进入 On-Policy 分支此时要求--use_vllm true学生模型经 vLLM 在线生成完成若未启用 vLLM会打印一次性告警并回退到数据集响应否则使用数据集自带的响应Off-Policy。生成结果按micro_batch_size切分为若干微批次编码教师 logits 计算本地前向或 API 请求在 _generate_and_score_completions 中统一完成。另外当--truncation_strategy deleteGKD 的默认截断策略见 MegatronRLHFArguments时超长样本会用独立种子seed1的数据迭代器重采样补齐见 _init_resample_data_iterator。自蒸馏OPSD特例当teacher_model与model指向同一模型且学生为 LoRA 时训练器进入自蒸馏模式——教师通过disable_adapter()上下文复用学生权重见 prepare_model避免加载第二份模型此时教师 logits 会在每个训练步重算而非每个生成周期以确保教师始终跟踪最新的 LoRA 权重_on_train_step_batch 中的注释解释了这一设计。训练脚本实战仓库在 examples/megatron/rlhf/gkd/ 下提供四个可直接参考的脚本分别覆盖四种典型场景。场景一本地大教师 vLLM on-policydense.shdense.sh 是标准 8 卡配置用Qwen/Qwen3-8B-Base作学生LoRAQwen/Qwen3-32B作本地教师CUDA_VISIBLE_DEVICES0,1,2,3,4,5,6,7 \ NPROC_PER_NODE8 \ PYTORCH_CUDA_ALLOC_CONFexpandable_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要点解读--lmbda 1表示每个训练步都走 on-policy学生生成--beta 1取反向 KL--offload_teacher_model true把 32B 教师平时挂在 CPU 上仅在计算教师 logits 时加载配合--sleep_level 1vLLM 休眠释放显存实现 32B 教师与 8B 学生共享 8 卡教师前向在 _compute_teacher_logits_local 中以torch.no_grad()执行并按gkd_logits_topk若设置做vocab_parallel_topk压缩--padding_free true与--sequence_parallel true用于变长序列的显存与效率优化。场景二Teacher ServerAPI 教师teacher_server.shteacher_server.sh 演示教师不占训练卡的解耦部署教师单独用swift deploy起 vLLM 服务训练端只通过--teacher_model_server http://localhost:8000拉取 logprobstop_k64 max_prompt_length2048 max_completion_length2048 max_total_length$((max_prompt_length max_completion_length)) # 教师服务需先启动脚本中为注释示例 # CUDA_VISIBLE_DEVICES0 \ # 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_DEVICES1,2 \ NPROC_PER_NODE2 \ PYTORCH_CUDA_ALLOC_CONFexpandable_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 \ --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 \ --max_length $max_prompt_length \ --max_completion_length $max_completion_length \ --use_vllm true \ --vllm_mode colocate \ --sleep_level 1 \ --padding_free true \ --recompute_granularity selective \ ...这里有两个硬性前提值得强调教师服务必须通过swift deploy --max_logprobs $top_k启动使推理端真正返回 top-k logprobs否则gkd_logits_topk无数据可用训练端必须传--gkd_logits_topk否则会触发上文提到的参数校验报错。从源码看API 路径下教师 logits 的获取由 fetch_teacher_parsed_by_routing 完成请求的聚合gather→推理→切回scatter多教师时按--teacher_tag_key默认取样本的dataset列路由到不同教师详见 蒸馏文档的多教师路由章节。场景三多轮 GKDmulti_turn.shmulti_turn.sh 展示与学生模型多轮交互的蒸馏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 \ --lmbda 1 \ --beta 0.5 \ --multi_turn_scheduler math_tip_trick \ --max_turns 2 \ --loss_scale last_round \ --truncation_strategy delete \ --vllm_server_pass_dataset true \ --remove_unused_columns false \ ...--multi_turn_scheduler math_tip_trick启用数学提示型多轮调度器与 GRPO 共用同一套MultiTurnScheduler基础设施调度器注册见 swift/rollout/multi_turn.py--max_turns 2限制交互轮数--loss_scale last_round表示只对最后一轮回复计损失GKD 在 Megatron 侧的默认loss_scale即last_round见 rlhf_args.py--vllm_server_pass_dataset true让数据集字段透传给 vLLM server--remove_unused_columns false保证多轮插件需要的原始列不被裁掉。场景四OPSD 自蒸馏opsd.shopsd.sh 复现 Self-Distilled Reasoner 的固定教师自蒸馏模式--teacher_model与--model同为Qwen/Qwen3-4B教师即关闭 LoRA adapter 后的基础模型脚本注释与 _teacher_use_disable_adapter 的日志均可印证。关键超参lmbda1, beta0.5, temperature1.2, lora_rank64, lora_alpha128, effective batch32并挂 external_plugins 处理样本构建build_opsd_samples。并行策略与分布式细节Megatron 版 GKD 对四种并行的适配体现在损失与数据两条链路TP词表并行loss_func注入tp_gather_topk、vocab_parallel_log_softmax、vocab_parallel_kl_div见 vocab_parallel_utilsJSD 在词表切分下正确归约CP上下文并行JSD 总和经 cp_reduce 在 CP 组内 all-reduce 后按有效 token 数平均最终损失还会除以 CP world sizePP流水线并行非末段 stage 以TeacherOutput()占位符兜底extract_active 中处理自蒸馏模式的教师重算特意放在 pipeline schedule 之外执行因此 PP1 同样受支持_on_train_step_batch 注释EPMoE 模型通过--expert_model_parallel_size配置示例脚本中均已显式声明。延伸阅读蒸馏概念与多教师路由docs/source_en/Instruction/Distillation.mdMegatron-SWIFT 完整参数表docs/source_en/Megatron-SWIFT/Command-line-parameters.md全部 GKD 训练脚本examples/megatron/rlhf/gkd损失数学核心HF/Megatron/Ray 共用swift/rlhf_trainers/gkd_loss.py训练器实现swift/megatron/trainers/gkd_trainer.pyRay 后端的 GKD 实现swift/ray/megatron/gkd_trainer.py 与 swift/ray/megatron/loss/gkd.py【免费下载链接】swiftUse PEFT or Full-parameter to CPT/SFT/DPO/GRPO 600 LLMs (Qwen3.6, DeepSeek-V4, GLM-5.1, InternLM3, Llama4, ...) and 300 MLLMs (Qwen3-VL, Qwen3-Omni, InternVL3.5, Ovis2.5, GLM4.5v, Gemma4, Llava, Phi4, ...) (AAAI 2025).项目地址: https://gitcode.com/GitHub_Trending/swift1/swift创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
