Strategy 使用指南¶

概述¶

LightRFT 的 strategy 模块是分布式训练能力的核心，为高效的强化学习微调提供了额外优化。Strategy 提供统一的接口来管理：

分布式训练后端：DeepSpeed ZeRO 和 FSDP（完全分片数据并行）
推理引擎集成：vLLM 和 SGLang 用于高吞吐量生成
内存优化：优化器卸载、梯度累积和引擎睡眠模式
序列并行：跨多个 GPU 高效处理长序列

核心 API 扩展¶

LightRFT 为 strategy 接口添加了以下关键方法：

方法	功能
`setup_inference_engine()`	初始化 vLLM 或 SGLang 推理引擎
`update_engine_weights()`	将 actor 模型权重同步到推理引擎
`gather_and_generate()`	分布式生成，自动收集 prompts
`maybe_load_optimizer()`	从 CPU 加载优化器状态（仅 FSDP）
`maybe_offload_optimizer()`	将优化器状态卸载到 CPU（仅 FSDP）
`wakeup_inference_engine()`	从睡眠模式唤醒推理引擎
`maybe_sleep_inference_engine()`	将推理引擎置于睡眠状态以节省内存

创建 Strategy¶

基础设置¶

使用工厂函数 get_strategy() 创建 strategy 实例：

from lightrft.strategy import get_strategy
from lightrft.utils import add_arguments

def train(args):
    # 创建 strategy（根据 args 自动选择 DeepSpeed 或 FSDP）
    strategy = get_strategy(args)

    # 设置用于生成的推理引擎
    strategy.setup_inference_engine(args, engine_type='vllm')

    # 如需要可以访问引擎
    vllm_engine = strategy.inference_engine

    # 创建训练器
    trainer = SPMDPPOTrainer(
        strategy=strategy,
        actor=actor,
        critic=critic,
        reward_model=reward_model,
        initial_model=initial_model,
        ema_model=ema_model,
        actor_optim=actor_optim,
        critic_optim=critic_optim,
        actor_scheduler=actor_scheduler,
        critic_scheduler=critic_scheduler,
        ...
    )

Strategy 选择¶

Strategy 类型由配置参数自动决定：

FSDP：设置 --fsdp 标志
DeepSpeed：不设置 --fsdp 时的默认选项（可通过 --zero_stage 配置）

在 Trainer 中使用 Strategy¶

标准训练操作¶

Strategy 提供标准的分布式训练操作：

# 反向传播
strategy.backward(loss, model, optimizer)

# 优化器步进（带梯度裁剪）
strategy.optimizer_step(optimizer, model, scheduler, name="actor")

# 分布式通信
averaged_value = strategy.all_reduce(local_value, op="mean")
gathered_values = strategy.all_gather(local_value)

内存优化训练¶

对于基于 FSDP 的训练，使用优化器卸载来减少 GPU 内存占用：

def ppo_train(self, global_steps=0):
    torch.cuda.synchronize()
    train_begin = time.time()

    # 从 CPU 加载优化器状态到 GPU（仅 FSDP）
    self.strategy.maybe_load_optimizer(self.actor_optim)

    # 执行训练
    train_ret = super().ppo_train(global_steps)

    # 从 GPU 卸载优化器状态到 CPU（仅 FSDP）
    self.strategy.maybe_offload_optimizer(self.actor_optim)

    torch.cuda.synchronize()
    self.strategy.print(f"PPO Train TIMECOST {time.time() - train_begin}")

    # 将 actor 权重同步到推理引擎
    self.strategy.update_engine_weights(self.actor)

    return train_ret

引擎权重同步¶

训练更新后，将模型权重同步到推理引擎：

# 用最新的 actor 权重更新推理引擎
strategy.update_engine_weights(actor)

这确保推理引擎使用最新的模型参数进行生成。

在 Experience Maker 中使用 Strategy¶

文本生成（LLM）¶

使用 gather_and_generate() 进行分布式文本生成：

# 对 prompts 进行分词（为提高效率不使用 padding）
all_prompt_token_ids = self.tokenize_fn(
    all_prompts,
    self.prompt_max_len,
    padding=False
)["input_ids"]

# 自动分布式生成响应
all_outputs = self.strategy.gather_and_generate(
    sampling_params=sampling_params,
    all_prompt_token_ids=all_prompt_token_ids,
    sleep_engine=True  # 生成后自动使引擎进入睡眠状态
)

if dist.get_rank(self.vllm_mp_group) == 0:
    self.strategy.print(f"生成了 {len(all_outputs)} 个输出")

多模态生成（VLM）¶

对于带有图像的视觉-语言模型：

# 使用多模态输入生成
all_outputs = self.strategy.gather_and_generate(
    sampling_params=sampling_params,
    all_prompts=all_prompts,        # 文本 prompts
    all_images=all_images,          # 图像数据
    images_num=images_num,          # 每个 prompt 的图像数量
    sleep_engine=True
)

`gather_and_generate()` 工作原理¶

该方法执行以下操作：

收集（Gather）：将张量并行组内所有 rank 的 prompts 收集到 rank 0
- 示例：当 world_size=8 且 engine_tp_size=4 时，ranks [0,1,2,3] 收集到 rank 0，ranks [4,5,6,7] 收集到 rank 4
生成（Generate）：使用 vLLM/SGLang 引擎对收集的 prompts 执行推理
分发（Distribute）：将生成的输出按相同顺序分散回原始 ranks
睡眠管理（Sleep Management）：根据 sleep_engine 参数自动处理引擎的睡眠/唤醒周期

注意：使用此接口时，用户无需手动管理引擎睡眠状态。

必需参数¶

将 LightRFT 特定的参数添加到你的参数解析器：

from lightrft.utils import add_arguments
import argparse

# 创建解析器
parser = argparse.ArgumentParser()

# 添加 LightRFT 参数
add_arguments(parser)

# 解析参数
args = parser.parse_args()

关键参数¶

推理引擎配置：

--engine_tp_size 4              # 推理引擎的张量并行大小
--engine_mem_util 0.85          # KV 缓存的 GPU 内存利用率（0.0-1.0）
--engine_type vllm              # 引擎类型：'vllm' 或 'sglang'
--enable_engine_sleep           # 启用引擎睡眠模式（默认：True）
--disable_engine_sleep          # 禁用引擎睡眠模式

分布式训练：

--fsdp                          # 使用 FSDP 而非 DeepSpeed
--zero_stage 2                  # DeepSpeed ZeRO 阶段（1、2 或 3）
--fsdp_cpu_offload              # 将 FSDP 优化器状态卸载到 CPU
--adam_offload                  # 卸载 Adam 优化器状态
--sp_size 2                     # 序列并行大小

训练优化：

--packing_samples               # 将多个样本打包到序列中
--use_mp_opt                    # 使用混合精度优化器（FSDP）
--fused_linear_logprob          # 融合线性层和 logprob 计算
--chunk_size 4096               # 融合操作的块大小

监控：

--log_dir ./logs                # 日志和可视化的目录
--plot_every 10                 # 每 N 步绘制生成长度分布

Strategy 实现细节¶

可用的 Strategy¶

LightRFT 提供两种主要的 strategy 实现：

DeepspeedStrategy（默认）
- 使用 DeepSpeed ZeRO 进行内存高效训练
- 可配置 ZeRO 阶段（1、2 或 3）
- 支持梯度累积和混合精度
- 最适合：通用 RLHF 训练、成熟的工作流程
FSDPV2Strategy（设置 --fsdp 时）
- 使用 PyTorch 的完全分片数据并行
- 支持优化器状态的 CPU 卸载
- 原生 PyTorch 实现，集成更好
- 最适合：最大内存效率、PyTorch 原生工作流程

Strategy 选择逻辑¶

# 在 get_strategy() 函数中
if args.fsdp:
    strategy = FSDPV2Strategy(...)
else:
    strategy = DeepspeedStrategy(...)

引擎睡眠/唤醒机制¶

Strategy 通过引擎睡眠模式提供自动内存管理：

# 引擎生命周期管理
strategy.setup_inference_engine(args, engine_type='vllm')  # 创建并唤醒引擎
strategy.maybe_sleep_inference_engine()                     # 睡眠以节省内存
strategy.wakeup_inference_engine()                          # 唤醒以进行生成

自动管理：当使用 gather_and_generate() 并设置 sleep_engine=True 时，睡眠/唤醒周期会自动处理。

配置示例¶

高吞吐量设置（8 GPU，DeepSpeed）¶

# 使用 DeepSpeed ZeRO-2 和大张量并行
python train.py \
    --zero_stage 2 \
    --engine_tp_size 4 \
    --engine_mem_util 0.9 \
    --enable_engine_sleep \
    --micro_train_batch_size 1 \
    --train_batch_size 128

内存高效设置（8 GPU，FSDP + CPU 卸载）¶

# 使用 FSDP + CPU 卸载以实现最大内存效率
python train.py \
    --fsdp \
    --fsdp_cpu_offload \
    --use_mp_opt \
    --engine_tp_size 2 \
    --engine_mem_util 0.85 \
    --enable_engine_sleep \
    --micro_train_batch_size 1 \
    --train_batch_size 64

视觉-语言模型设置¶

# 使用多模态数据训练 VLM
python train_vl.py \
    --fsdp \
    --engine_tp_size 4 \
    --mixed_mm_data \
    --packing_samples \
    --enable_engine_sleep \
    --plot_every 20

最佳实践¶

1. 张量并行配置¶

根据模型大小和 GPU 数量设置 engine_tp_size
7B 模型：engine_tp_size=1 或 2
13B-70B 模型：engine_tp_size=4 或 8
确保 world_size % engine_tp_size == 0

2. 内存管理¶

对于内存受限的设置启用引擎睡眠模式：--enable_engine_sleep
根据可用内存调整 engine_mem_util（0.5-0.9）
使用 FSDP + CPU 卸载以最大化节省内存：--fsdp --fsdp_cpu_offload

3. 性能优化¶

对于不同长度的序列使用 --packing_samples
对于大词汇表模型启用 --fused_linear_logprob
设置适当的 micro_train_batch_size 以充分利用 GPU

4. 调试和监控¶

使用 --plot_every 和 --log_dir 跟踪生成长度分布
使用 strategy.report_memory(prefix="checkpoint_name") 监控内存
使用 strategy.inference_engine_status 检查引擎状态

高级功能¶

序列并行¶

为超长序列启用序列并行：

# 在参数中
--sp_size 4  # 将序列分割到 4 个 GPU

Strategy 会自动创建序列并行组并处理通信。

自定义奖励模型¶

支持多个奖励模型或远程奖励 API：

# 多个奖励模型
reward_models = [reward_model_1, reward_model_2, reward_model_3]
strategy = get_strategy(args)

# 模型自动分片到各个 GPU
prepared_rms = [strategy.prepare_model(rm, shard_size=8) for rm in reward_models]

混合精度训练¶

控制混合精度行为：

# 启用 BF16 训练
--bf16

# 使用混合精度优化器（FSDP）
--use_mp_opt

故障排除¶

常见问题¶

问题：生成时内存不足

解决方案：降低 engine_mem_util 或增加 engine_tp_size

问题：引擎未使用新权重更新

解决方案：确保训练后调用 update_engine_weights()

问题：生成速度慢

解决方案：增加 micro_rollout_batch_size 或减少 engine_tp_size

问题：FSDP 优化器卸载错误

解决方案：验证是否使用 FSDP strategy（--fsdp）并成对调用 offload/load

API 参考¶

详细的 API 文档请参阅：

lightrft.strategy.strategy_base.StrategyBase - 基础 strategy 类
lightrft.strategy.get_strategy() - Strategy 工厂函数
lightrft.utils.add_arguments() - 参数配置