TRL documentation
加速训练
加速训练
本指南介绍了在 TRL 中加速训练的多种方法。每种技术都附有最小示例，并链接到更全面的文档。
用于在线方法中快速生成的 vLLM
诸如 GRPO 或 Online DPO 等在线方法需要模型生成补全内容，这通常是一个缓慢的过程，并可能显著影响训练时间。 为了加速生成，你可以使用 vLLM，这是一个通过 PagedAttention 等技术实现快速生成的库。TRL 的在线训练器支持 vLLM，可大幅提升训练速度。更多详情，请参阅 vLLM 集成。
要使用 vLLM，请先使用以下命令进行安装：
pip install trl[vllm]
首先，运行以下命令启动 vLLM 服务器：
VLLM_SERVER_DEV_MODE=1 vllm serve <model_name> \
--weight-transfer-config '{"backend": "nccl"}' \
--logprobs-mode processed_logprobs \
--max-logprobs -1然后，运行训练脚本，并在训练参数中传入 use_vllm=True。
from trl.experimental.online_dpo import OnlineDPOConfig
training_args = OnlineDPOConfig(..., use_vllm=True, vllm_mode="server")首先，运行以下命令启动 vLLM 服务器：
VLLM_SERVER_DEV_MODE=1 vllm serve <model_name> \
--weight-transfer-config '{"backend": "nccl"}' \
--logprobs-mode processed_logprobs \
--max-logprobs -1然后，运行训练脚本，并在训练参数中传入 use_vllm=True。
from trl import GRPOConfig
training_args = GRPOConfig(..., use_vllm=True, vllm_mode="server")你可以通过传递额外参数来自定义服务器配置。有关更多信息，请参阅 vLLM 集成。
使用 vLLM 时，请确保分配给训练和生成的 GPU 相互独立，以避免资源冲突。例如，如果你计划使用 4 个 GPU 进行训练，另外 4 个用于 vLLM 生成，可以通过
CUDA_VISIBLE_DEVICES来指定 GPU 分配。为 vLLM 生成设置 GPU 0-3：
CUDA_VISIBLE_DEVICES=0,1,2,3 VLLM_SERVER_DEV_MODE=1 vllm serve <model_name> \ --weight-transfer-config '{"backend": "nccl"}' \ --logprobs-mode processed_logprobs \ --max-logprobs -1以及用于训练的 GPU 4-7：
CUDA_VISIBLE_DEVICES=4,5,6,7 torchrun --nproc_per_node 4 train.py
首先，运行以下命令启动 vLLM 服务器：
VLLM_SERVER_DEV_MODE=1 vllm serve <model_name> \
--weight-transfer-config '{"backend": "nccl"}' \
--logprobs-mode processed_logprobs \
--max-logprobs -1然后，运行训练脚本，并在训练参数中传入 use_vllm=True。
from trl import RLOOConfig
training_args = RLOOConfig(..., use_vllm=True, vllm_mode="server")你可以通过传递额外参数来自定义服务器配置。有关更多信息，请参阅 vLLM 集成。
使用 vLLM 时，请确保分配给训练和生成的 GPU 相互独立，以避免资源冲突。例如，如果你计划使用 4 个 GPU 进行训练，另外 4 个用于 vLLM 生成，可以通过
CUDA_VISIBLE_DEVICES来指定 GPU 分配。为 vLLM 生成设置 GPU 0-3：
CUDA_VISIBLE_DEVICES=0,1,2,3 VLLM_SERVER_DEV_MODE=1 vllm serve <model_name> \ --weight-transfer-config '{"backend": "nccl"}' \ --logprobs-mode processed_logprobs \ --max-logprobs -1以及用于训练的 GPU 4-7：
CUDA_VISIBLE_DEVICES=4,5,6,7 torchrun --nproc_per_node 4 train.py
优化的注意力实现
TRL 支持多种优化后的注意力实现，可在降低内存占用的同时显著加快训练速度。您可以直接使用 Kernels Hub 中预优化的内核，也可以使用手动构建的注意力后端。
您可以直接使用 Hub 上预优化的注意力内核，而无需手动编译：
from trl import SFTConfig
training_args = SFTConfig(..., model_init_kwargs={"attn_implementation": "kernels-community/flash-attn2"})其他选项包括 kernels-community/vllm-flash-attn3 和 kernels-community/paged-attention。
优化后的注意力机制适用于所有 TRL 训练器。有关更多详细信息，请参阅 Kernels Hub 集成。
手动构建优化的注意力后端既复杂又耗时。除非绝对必要，否则不建议这样做。请考虑使用上一节中描述的来自 Hub 的 Kernels。
如果您已手动安装了 Flash Attention 2 等优化注意力后端，可以在训练参数中指定它：
from trl import SFTConfig
training_args = SFTConfig(..., model_init_kwargs={"attn_implementation": "flash_attention_2"})混合精度训练
使用 bf16 或 fp16 进行混合精度训练可以加快训练速度并降低内存占用，同时对模型质量的影响微乎其微。
from trl import SFTConfig
training_args = SFTConfig(..., bf16=True) # or fp16=True for older GPUs对于 Ampere 架构（A100、RTX 30xx）或更新的 GPU，请使用 bf16=True；对于较旧的 GPU，请使用 fp16=True。所有 TRL 训练器均支持混合精度训练。