
TRL AsyncDistillationTrainer 异步蒸馏实战指南三终端部署与训练指标分诊【免费下载链接】trlTrain transformer language models with reinforcement learning.项目地址: https://gitcode.com/GitHub_Trending/tr/trlAsyncDistillationTrainer是 TRL 实验模块中的异步 on-policy 蒸馏训练器学生模型在 vLLM 服务器上自己生成样本教师经 vLLM 服务打分生成与梯度更新并发进行。本文覆盖选型、部署、beta/teacher_top_k设置与按问题排查指标。何时选它先看清与同步蒸馏的差异同步版DistillationTrainer把教师加载在本地生成、教师前向、梯度更新在同一进程内顺序执行。异步版把教师完全搬到线上只需一个教师 vLLM 服务器的 URL教师可以跑在独立硬件上规模也可以大到与学生并排装不下。两者对比如下。维度同步DistillationTrainer异步AsyncDistillationTrainer教师加载方式本地加载完整模型永不本地加载仅一个 HTTP URL硬件摆放师生与训练同进程同机教师可独立部署甚至跨机教师权重是否更新不参与训练静态静态trainer 从不向任何教师传权重生成与训练关系交替执行后台 worker 生成打分主进程并行训练on-policy 程度学生自生成学生自生成每个训练样本都由学生产生分布式支持常规后端的通用支持仅 FSDP2不支持 DeepSpeed ZeRO额外依赖无vllm0.22.0 学生侧 NCCL 权重传输选型结论教师与学生必须同机同进程、且追求最简流程时用同步版教师更大、需要独立硬件、或希望训练不等生成时用该 trainer。一个样本的完整旅程先固定五个术语全文沿用rollout一个 prompt 被学生生成一次、被教师打分一次sampleprompt 学生完成结果 教师逐位置的稀疏 top-k 候选是跨进程传输的单元row一个 DP rank 在一次 micro-batch 中前向的内容若干 sample 拼接成单条序列、position_ids在样本边界重置micro-batchworld_size个 row 的集合每个 rank 前向自己的那行optimizer stepgradient_accumulation_steps个 micro-batch 累积而成的一次优化器步。所有指标中的 step 恒指它而非 micro-batch。数据集行 PROMPT消息列表 可选 teacher_id 列 └─ ROLLOUT 1 个 prompt - 1 次学生生成 - 1 次教师打分 ├─ generate 学生 vLLM /v1/completions按 temperature/top_p 采样 └─ score 路由到的教师 /v1/completions带 prompt_logprobsteacher-forced └─ SAMPLE prompt completion 教师稀疏分布 ════════════ 进程边界rollout_buffermp.Queue════════════ └─ SAMPLE 逐个拉取staleness 超过 max_staleness 即丢弃 └─ ROW 规划器分入 dp 个 row 之一按 Σ Lᵢ² 平衡 └─ MICRO-BATCH dp 个 row每 rank 一个 └─ PACKED ROW 拼接成 1 条序列position_ids 逐样本重置 └─ FORWARD compute_lossbs1rank 间填充先被剥掉 └─ OPTIMIZER STEP 由 grad_accum 个 micro-batch 构成蒸馏没有 GRPO 那种需要跨生成求的 advantage 基线所以 prompt 不重复、一次 rollout 恰好产出一个训练样本。三个环节的解耦体现在生成由后台 spawn 子进程负责CUDA_VISIBLE_DEVICES被清空worker 内部跑 asyncio 循环最多max_inflight_tasks个任务在途教师经 HTTP 打分而非本地前向主进程只做拉样本、算损失、更新权重。每weight_sync_steps步更新后的学生权重经 NCCL 推给学生 vLLM 服务器生成始终领先于训练样本可能反映过期策略落后超过max_staleness个权重更新的样本被丢弃并计入sample/dropped_stale_total。三终端如何点亮三个角色必须分卡运行。教师和学生都是普通的vllm serve但启动参数不同trainer 脚本保持最小即可。最小训练脚本train_async_distillation.py# 加载数据集并启动异步蒸馏默认配置即可跑通 from datasets import load_dataset from trl.experimental.async_distillation import AsyncDistillationTrainer dataset load_dataset(trl-lib/DeepMath-103K, splittrain) trainer AsyncDistillationTrainer( modelQwen/Qwen2.5-0.5B-Instruct, train_datasetdataset, ) trainer.train()终端 1教师服务器GPU 0# 教师只被打分、永不更新需要服务端温度与解除 logprob 上限 CUDA_VISIBLE_DEVICES0 vllm serve Qwen/Qwen2.5-1.5B-Instruct \ --port 8001 \ --logprobs-mode processed_logprobs \ --max-logprobs -1终端 2学生 vLLM 服务器GPU 1# 学生只被生成且必须开启 dev 模式与 NCCL 权重传输以接收更新权重 CUDA_VISIBLE_DEVICES1 VLLM_SERVER_DEV_MODE1 vllm serve Qwen/Qwen2.5-0.5B-Instruct \ --port 8000 \ --weight-transfer-config {backend:nccl}终端 3训练GPU 2# trainer 与两个服务器分卡accelerate 单卡启动即可 CUDA_VISIBLE_DEVICES2 accelerate launch train_async_distillation.py两个 flag 各自对应一个机制--logprobs-mode processed_logprobsteacher_temperature由 vLLM 在服务端作用于返回的 logprobs而非客户端重缩放。缺失该 flag 时教师静默返回原始 logprobsteacher_temperature只作用于学生侧。--max-logprobs -1解除 vLLM 默认每 token 最多 20 个 logprob 的上限teacher_top_k超过 20 时必须有它。学生的 dev 模式与 NCCL 配置则解决另一件事trainer 每weight_sync_steps步要把新权重推进生成服务器没有这个通道生成会一直停留在初始策略。约束先装 vLLM 再强制装 transformerspip install vllm0.22.0之后pip install transformers5.2.0 --no-deps。原因当前两包的依赖约束互相冲突直接联装会降级 transformers。损失与两个旋钮beta和teacher_top_k目标是最小化学生与教师在逐位置 token 分布上的广义 Jensen-Shannon 散度。它决定了两个行为。beta决定加权方式也决定在哪些候选上算散度。beta0.0是前向 KLmean-seekingbeta1.0是反向 KLmode-seeking中间值在两者间插值取值必须在[0.0, 1.0]内。支撑集选择随beta变化beta0.0用教师上报的完整teacher_top_k宽支撑外加尾桶。前向 KL 的期望恰好按教师概率加权教师上报的这组候选正好够用。beta ! 0.0支撑收窄到最多两个候选——教师自己的 top-1 与完成结果的实际 token。线上协议在不传输更宽词表的前提下只保证这两个 token 有教师 logprob更宽的支撑只能算概率性覆盖学生可能采样的 token不是保证。特例是beta1.0纯反向 KL 是学生自身加权的期望教师 top-1 对目标无贡献支撑进一步收窄为仅实际 token宽度 1去重后。后果是实际的调teacher_top_k在beta ! 0.0时几乎不改变训练信号它只影响beta0.0的精度与teacher_entropy的下界。teacher_top_k与尾桶决定教师分布的近似质量。线上只传输每个位置teacher_top_k个候选外加 vLLM 总会报告的 realized token即使它在 top-k 之外和add_tail_bucket的尾桶log(1 - Σexp(top_k_logps))候选少时防止散度平凡地接近零。学生侧的散度是精确的完整 logits 本地可得被近似的只有教师侧。默认8是冒烟测试级邻近 RL 框架的 on-policy 蒸馏生产配置在16miles 默认到64EasyOPD之间正式训练上调到16–64合理。内存代价由分块实现控制投影学生隐藏状态过lm_head时按 256 个 token 一块进行(chunk_size, vocab_size)的 logits 是唯一随词表扩展的张量checkpoint 机制在反向时按需重算。峰值 logits 内存是256 × vocab_size而不是全部有效 token 乘词表。参数速查以下每项回答改它改变什么、何时该改。参数默认值改它改变什么 / 何时该改beta0.00前向 KL1反向 KL中间插值同时切换支撑集见上节。MOPD 场景论文用反向 KL需显式设1.0teacher_top_k8beta0.0下决定教师分布近似宽度。冒烟用 8正式训练提至 16–64超过 20 需教师--max-logprobs -1teacher_temperature1.0散度两侧的 softmax 温度服务端计算教师 logprobs 且作用于学生 logits。与采样temperature无关。想平滑教师分布再改add_tail_bucketTrue是否追加尾桶吸收 top-k 之外的概率质量。teacher_top_k很小时关掉会使散度偏小不建议关max_staleness4样本允许落后多少权重更新超过即丢。trainer-bound 时调大换吞吐、代价是 off-policy生成侧慢时它是丢弃率的主要来源max_inflight_tasks-1两个服务器的在途任务上限。自动值 max(max_staleness, 1) × per_device_train_batch_size × gradient_accumulation_steps × num_processes下限 1 防止max_staleness0清零调度导致空队列挂死。想加大并发改它queue_maxsize1024rollout 队列缓冲上限决定样本最多能存多少份。显存/内存紧张或想减小 staleness 上界时调小weight_sync_steps1两次 NCCL 权重同步之间的训练步数。调大省同步开销但生成侧使用的策略更旧token_budgetNone单 row 打包的最大真实 token 数。None时取学生服务器max_model_len启动时查询0切换为固定样本数打包dtypefloat32学生加载精度。默认 float32 因为 training-inference mismatch 度量对 trainer 精度敏感端到端弥合还要求vllm serve --dtype一致。要省显存改bfloat16并同步服务器request_timeout600对任一 vLLM 服务器的单请求超时秒。大 batch 打分慢时调大weight_sync_timeout1800一次权重传输的超时超时 raise 而非挂起。模型大或网络慢时调大heartbeat_stale_after_s300.0worker 心跳超过该秒数视为挂起并中止。长请求阻塞 worker 时调大log_completionsFalse是否周期性记录 (prompt, completion) 对。调试生成质量时开log_completions_steps100按 worker 已打分样本数计数worker 是独立进程看不到global_step不是优化器步num_completions_to_printNone每次打印的完成结果数None为全部另有五项默认值不同于 transformersTrainingArguments容易按旧习惯误判参数该 trainer 默认TrainingArguments默认说明learning_rate1e-65e-5蒸馏更新幅度小学习率相应低两个量级bf16未设fp16时为TrueFalse默认走 bf16gradient_checkpointingTrueFalse默认开启以省显存logging_steps1500每步都记指标密度高ignore_data_skipTrueFalse基础 Trainer 的 skip-and-replay 不适用于实时 rollout 队列见 MOPD 一节的恢复语义指标分诊按四个问题查口径先讲清(numerator, denominator)对按 Σnum/Σden 聚合为比率名字含total的是计数器求和含max/min的取窗口极值其余 gauge 取窗口均值。token 分三种口径generated学生实际生成、forwarded前向处理的全部 tokenprompt生成、trainedcompletion_mask 1覆盖、损失真正计算的子集。注意 trained ≠ generated教师未对某完成位置报告任何候选时该位置在散度中被掩掉但仍参与前向。问题一瓶颈在生成侧还是训练侧rollout 队列相当于生成与训练之间的车间缓冲区四个指标描述它其中两个等待指标互为镜像perf/rollout_wait_s统计队列为空、rollout/backpressure_s统计队列已满同一时刻只可能成立其一因此两者永远不会同时很大。结合队列水位读指标口径含义sample/rollout_queue_sizegauge窗口均值当前排队等待的样本数sample/time_in_queue_sgauge窗口均值单个样本入队到被训练的等待秒数perf/rollout_wait_sgauge窗口均值训练因队列为空阻塞的秒数rollout/backpressure_sgauge窗口均值生成因队列已满被节流的秒数判读队列近空且perf/rollout_wait_s高 → 生成受限再看rollout/generated_tok_s、rollout/inflight、rollout/score_s队列近满且rollout/backpressure_s高 → 训练受限注意sample/staleness_mean随之攀升两者都近零 → 平衡。慢教师直接体现在每次往返上教师调用在每个 rollout 的关键路径里不同于 GRPO 的独立打分循环指标口径含义rollout/duration_sgauge窗口均值一次 rollout 墙钟时间生成加教师调用rollout/score_sgauge窗口均值其中教师/v1/completions耗时teacher_score_s/idgauge仅 MOPD按教师拆分的打分耗时慢专家只拖慢路由给它的 rolloutrollout/generated_tok_s比率窗口生成吞吐停顿会显现rollout/inflightgauge在途 rollout 数正在生成或打分rollout/vllm_retry_totaltotal计数重试过的 vLLM 请求学生或教师服务器退化服务器否则只是莫名的变慢生成侧本身的健康度看完成结果指标口径含义completions/mean_lengthgauge窗口均值每个 rollout 的生成 token 数completions/min_length/completions/max_length极值窗口内最短/最长完成结果completions/clipped_ratio比率未以 EOS 结束、被max_completion_length截断的占比问题二样本在变陈旧吗指标口径含义sample/forwarded_tokens_meangauge单样本 token 数prompt生成sample/forwarded_tokens_max极值单样本最大 token 数sample/trained_tokens_meangauge其中损失覆盖的 token 数sample/staleness_mean/sample/staleness_maxgauge / 极值数据落后当前策略多少个版本jsd展示 off-policy 对损失的效果这里展示原因sample/dropped_stale_totaltotal计数因超过max_staleness被丢弃的样本数问题三学生在收敛还是在坍缩这些是无前缀指标窗口内 trained token 上的均值。指标口径含义jsd窗口均值损失最小化的广义 JSD按配置的beta。下降学生在向教师分布收敛entropy窗口均值学生自身预测熵。jsd下降的同时这里崩塌学生在收窄而非学习teacher_entropy窗口均值教师在所报候选上的熵只有teacher_top_k个候选过线它从下方界定真实值teacher_jsd/id窗口均值仅 MOPD限定在该教师打分 token 上的jsd混合jsd会把不同领域教师的不同速率混为一谈teacher_entropy/id窗口均值仅 MOPDteacher_entropy的同样拆分teacher_token_frac/id比率仅 MOPD该教师打分 token 的占比。路由偏斜否则不可见被饿死的教师仍会报告健康的teacher_jsd/id没有 per-teacher 的entropy学生熵是其自身策略的属性与哪个教师打分无关混合指标已覆盖。问题四算力真正花在训练上的比例吞吐与 MFU 各报两次基于同一次优化器步仅除数不同_fwd_bwd除以perf/fwd_bwd_s纯计算回答有数据时 trainer 跑得多高效低则问题在 trainer_wall_clock除以perf/step_s完整一步含等 rollout回答分配到的算力有多少变成了训练远低于前者则瓶颈大概率在生成。两者之差是perf/rollout_wait_s加优化器与权重同步时间。指标口径含义perf/step_sgauge优化器步之间的墙钟时间含计算、优化器、权重同步、队列等待perf/fwd_bwd_sgauge前向反向对一步的 micro-batch 求和所有_fwd_bwd指标的分母perf/fwd_sgauge其中前向部分fwd_s / fwd_bwd_s接近 1/3 是常见划分更高意味着反向便宜或重计算落在前向perf/optimizer_sgaugeoptimizer.step()耗时perf/rollout_wait_sgauge见问题一perf/weight_sync_sgauge一次完整同步另有_pause_s等 vLLM、_barrier_srank 偏斜、_transfer_s字节传输perf/forwarded_tok_s_fwd_bwd/perf/forwarded_tok_s_wall_clock比率两种口径的每秒前向 token 数perf/trained_tok_s_wall_clock比率同口径只统计损失见过的 tokenperf/mfu_fwd_bwd/perf/mfu_wall_clock比率两种口径的模型 FLOPs 利用率batch 级指标在问题四之外还要会读打包健康度。一步固定含gradient_accumulation_steps × world_size个 row-slot_per_step指标跨全部 rank 求和row_*是行均值两者按samples_per_step ≈ row-slots × samples_per_row校验允许百分之零点几的偏差per-step 是求和、per-row 是均值。指标口径含义batch/samples_per_step求和每优化器步的训练样本数batch/forwarded_tokens_per_step/batch/trained_tokens_per_step求和一步内前向/训练的 token 数batch/microbatches_per_step实测计数实际值非配置读出batch/masked_token_frac比率completion_mask 0的前向 token 占比前向中不产生梯度的部分batch/samples_per_rowgauge规划器每行打包的样本数batch/row_tokens_mean/batch/row_tokens_maxgauge / 极值行内 token 规模batch/row_fill_frac比率行 token 数相对token_budget。低预算没填满。样本很长时这是量化效应而非 bug1 万 token 样本铺 3.2 万预算3 个放得下、4 个永远放不下常只能放 2 个、行约 77% 满1k 样本则几乎完美铺满token_budget是调节杠杆且注意力按序列 O(L²)更满的长行不线性增加内存batch/row_imbalancegauge各行max Σ Lᵢ² / mean Σ Lᵢ²。注意力 O(L²)它预测哪个 rank 拖慢 all-reduce1.0 为完美batch/pad_frac比率rank 间填充只增加广播字节前向之前剥掉batch/dropped_oversize_totaltotal计数超过token_budget被丢弃的样本数两种规划器对应token_budget的两个 regime0走TokenBudgetBatchertoken 预算动态样本数默认0走FixedCountBatcher每 micro-batch 固定per_device_train_batch_size × num_processes个样本。两者都做 Σ Lᵢ² 平衡避免某个 rank 拖尾。多教师 MOPD路由规则与恢复语义MOPD 不是该 trainer 核心目标所在论文2306.13649同步单教师场景的一部分而是独立方法通用 SFT、各域独立 RL 专家训练、MOPD 融合三阶段中的第三阶段。该 trainer 只实现融合阶段——各域专家必须已存在例如分别用GRPOTrainer/RLOOTrainer训练并经 HTTP 提供服务再把teacher_server_urls指向它们。注意论文自己的 Stage 3 用反向 KLbeta1.0而该 trainer 默认beta0.0配置 MOPD 时应显式设置。单个条目所有样本由该教师打分。多个条目每行的teacher_id列选择打分者如数学 prompt 给数学专家、代码 prompt 给代码专家。每个样本只发给匹配的那一个教师绝不跨教师平均或集成。teacher_id缺失或未映射直接 raiseValueError不静默回退。约束每个教师必须与学生共享同一个 tokenizer。原因完成结果以原始 token id 传输教师回报的候选 id 直接索引学生词表词表不同的教师会把学生训练到错误的 token 上除非其词表大于学生否则这种错误是静默的。检查点保存的是已训练位置而非生成器位置首个未被训练的 prompt 索引随每个检查点写入rollout_state.json恢复时 worker 直接快进到该位置。worker 领先训练一个队列深度缓冲中未训练的样本在运行结束时丢失——若从生成器位置恢复会跳过那些已生成但未训练的 prompt。流式数据集IterableDataset无法重新定位其 worker 恢复时从 prompt 0 重启。红线与常见坑依赖装反现象——transformers 被 vLLM 拉到不满足要求的版本。原因——两包依赖约束冲突。对策——先装vllm0.22.0再pip install transformers5.2.0 --no-deps。用 DeepSpeed ZeRO 跑分布式现象——配置不生效或行为异常。原因——该 trainer 分布式只支持 FSDP2。对策——换 FSDP2 配置。teacher_id未映射现象——启动即报ValueError。原因——拒绝静默路由到错误教师。对策——补齐数据集teacher_id列或对应teacher_server_urls条目。teacher_top_k超过 20 教师报错/截断现象——候选数不足。原因——vLLM 默认每 token logprob 上限 20。对策——教师加--max-logprobs -1。超过token_budget的样本消失现象——batch/dropped_oversize_total增长、日志有警告。原因——该样本放不进任何行。对策——调大token_budget或缩短max_completion_length。序列维并行直接抛错现象——cp_size1或sp_size1时__post_init__raise。原因——蒸馏在生成之后才于 trainer 内部构建模型输入context-parallel/Ulysses 的输入分片无法应用于原始生成 batch。对策——cp_sizesp_size1。teacher_temperature不生效现象——只有学生侧变化。原因——教师未带--logprobs-mode processed_logprobs静默返回原始 logprobs。对策——按三终端章节的 flag 重启教师。max_staleness0空队列挂死现象——队列始终为空。原因——max_inflight_tasks自动公式含max_staleness因子下限max(max_staleness, 1)是防呆显式设 0 会清零调度。对策——不要显式压 0或显式给max_inflight_tasks设正值。检查点恢复后数据重复/跳段现象——流式数据集恢复后从 prompt 0 重来。原因——IterableDataset无法重新定位。对策——MOPD/流式场景预期行为可重定位数据集按rollout_state.json快进。边界与延伸阅读该 trainer 刻意保持最小化不打算成长为通用方案需要不支持的功能时官方建议直接克隆仓库改造仓库地址git clone https://gitcode.com/GitHub_Trending/tr/trl。有两个官方注入口测试即通过注入 no-op 实现脱离真实 vLLM 运行RolloutWorkerProtocol可替换的 rollout worker生成打分循环WeightTransferProtocol可替换的权重同步后端。关键源码与示例均为仓库相对路径trl/experimental/async_distillation/async_distillation_config.py参数默认值与启动校验trl/experimental/async_distillation/async_distillation_trainer.py损失、指标规约、打包与检查点trl/experimental/async_distillation/async_rollout_worker.pyspawn 子进程、生成打分循环、RolloutSampletrl/experimental/async_distillation/weight_transfer.pyNCCL 权重传输客户端trl/experimental/async_distillation/vllm_client.pyvLLM HTTP 客户端examples/async_distillation_math/async_distillation_math.py单教师可运行示例GSM8K、max_steps100、learning_rate1e-6examples/async_distillation_math/async_distillation_mopd.py双教师 MOPD 示例数学/代码教师路由显式beta1.0【免费下载链接】trlTrain transformer language models with reinforcement learning.项目地址: https://gitcode.com/GitHub_Trending/tr/trl创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考