Prhub

#7703 [training_utils] fix: honor fused output-head backend

原始 PR 作者 Luosuu 合并时间 2026-09-03 09:59 文件变更 12 提交数 1 评论 1 代码增减 +81 / -18

执行摘要

修复融合线性输出头后端选择逻辑,使配置参数 impl_backend 生效。

PR body 明确指出,此变更修复了 #7461 引入的回归。该变更在 Liger 可导入时便选择 Liger,即使配置后端为 torch,这违反了后端契约。此外,隐式的 Liger 调度绕过了 PyTorch 算子重载,迫使 vexact 项目通过 monkey patch 修改 verl 的私有模块状态以恢复正确性(如关联 Issue#48 所述)。

值得精读,尤其是核心逻辑变更(FusedLinearForPPO)和配置传播设计。此 PR 展示了如何通过显式参数和属性传递来修复一个由隐式依赖引发的架构违规,对处理类似的可选后端集成问题有参考价值。建议关注其配置传播模式(模型实例属性 -> 适配器函数参数 -> 模块初始化参数)以及完整的测试覆盖策略。

讨论亮点

PR 评论区仅有 Codex 自动审查摘要,无实质性人工讨论。核心设计决策和技术权衡已在 PR body 中详尽阐述:

  • 决策结论:采用显式参数控制后端选择(torch 或 liger),替代基于可导入性的隐式检测。默认行为是 torch,需要用户主动配置才能启用 Liger。
  • 关键论据:此修复消除了 impl_backend 配置被忽略的契约违规行为,并移除了 vexact 需要通过 monkey patch 修改模块状态的 workaround,使集成更干净。
  • 未解决疑虑:无。

实现拆解

  1. 核心逻辑修改:在 verl/utils/experimental/torch_functional.py 的 FusedLinearForPPO.__init__ 中,新增 impl_backend 参数(默认为 torch),并在 forward 方法中优先根据 self.impl_backend 选择计算路径,仅当配置为 liger 且 Liger 可用时才调用 Liger 内核,否则回退至原生分块实现。
  2. 配置传播:修改 verl/models/transformers/monkey_patch.py 中的 patch_forward_with_backends 函数,将 fused_kernels_backend(现在支持 liger)存储在模型实例的 _verl_fused_kernels_backend 属性上,并允许 liger 后端使用 torch 路径的前向适配器。
  3. 适配器适配:更新五个模型适配器(dense_common.py, glm4v.py, qwen2_vl.py, qwen3_5.py, qwen3_vl.py)的 forward_with_torch_backend 函数,在实例化 FusedLinearForPPO 时从 self._verl_fused_kernels_backend 读取后端配置并传递给 impl_backend 参数,确保配置穿透。
  4. 测试覆盖:在 tests/utils/test_experimental_torch_functional_on_cpu.py 中新增 test_fused_linear_for_ppo_torch_backend_does_not_dispatch_to_liger 和 test_fused_linear_for_ppo_rejects_unknown_backend 两个测试,验证后端选择逻辑的正确性和参数校验。在 tests/models/test_fused_kernels_ulysses_sp_on_cpu.py 中新增 test_torch_adapter_propagates_fused_linear_backend 测试,验证适配器能正确传播后端配置。
  5. 配置与文档更新:在 verl/trainer/config/model/hf_model.yaml 和 docs/examples/config.rst、docs/perf/perf_tuning.rst 中,将 impl_backend 的可选值从 [triton, torch] 更新为 [triton, torch, liger],并补充说明。
文件 模块 状态 重要度
verl/utils/experimental/torch_functional.py 工具函数 modified 6.49
verl/models/transformers/monkey_patch.py 模型补丁 modified 6.46
tests/utils/test_experimental_torch_functional_on_cpu.py 融合内核测试 modified 6.81

关键符号

FusedLinearForPPO.__init__ FusedLinearForPPO.forward patch_forward_with_backends

关键源码片段

verl/utils/experimental/torch_functional.py core-logic

核心变更文件。修改了 `FusedLinearForPPO` 类,使其接受并遵循 `impl_backend` 参数,是修复回归问题的核心逻辑所在。

class FusedLinearForPPO(torch.nn.Module):
    def __init__(self, chunk_size: int = 512, impl_backend: str = "torch"):
        super().__init__()
        # 显式校验后端参数,确保只有合法的后端被接受
        if impl_backend not in ("torch", "liger"):
            raise ValueError(f"Unsupported FusedLinearForPPO backend: {impl_backend}. Choose 'torch' or 'liger'.")
        self.chunk_size = chunk_size
        # 将后端配置存储为实例属性,供 forward 使用
        self.impl_backend = impl_backend
​
    def forward(
        self,
        hidden_states: torch.FloatTensor,
        vocab_weights: torch.FloatTensor,
        input_ids: torch.LongTensor,
        temperature: float = 1.0,
    ) -> tuple[torch.FloatTensor, torch.FloatTensor]:
        input_ids = input_ids.to(torch.int64)
        # 关键修复:优先检查显式配置的 impl_backend。
        # 仅当配置为 "liger" 且 Liger 可用时,才使用 Liger 路径;
        # 否则(配置为 "torch" 或 Liger 不可用)均使用原生 torch 路径。
        if self.impl_backend == "torch" or _LIGER_FUSED_LINEAR_SCALED_CROSS_ENTROPY is None:
            return FusedLinearForPPOFunction.apply(
                hidden_states,
                vocab_weights,
                input_ids,
                temperature,
                self.chunk_size,
            )
        # ... 后续为 Liger 路径的逻辑 ...
verl/models/transformers/monkey_patch.py data-contract

配置入口点。修改了 `patch_forward_with_backends` 函数,使其接受 `liger` 后端选项,并将后端配置存储在模型实例上,是实现配置传播的关键一步。

def patch_forward_with_backends(
    model: PreTrainedModel,
    use_fused_kernels: bool = False,
    fused_kernels_backend: str = None,
):
    # ...
    # 更新验证条件,允许 "liger" 通过
    if not use_fused_kernels or fused_kernels_backend not in ["triton", "torch", "liger"]:
        print(
            f"Skipping monkey patch for {model.__class__.__name__} as use_fused_kernels is "
            f"{use_fused_kernels} or fused_kernels_backend is {fused_kernels_backend}"
        )
        return
​
    # ... ( 选择模型适配器的逻辑不变 ) ...
​
    # 新增:将后端配置存储到模型实例上,供下游适配器读取
    model._verl_fused_kernels_backend = fused_kernels_backend
    if fused_kernels_backend == "triton":
        model.__class__.forward = forward_with_triton_backend_function
        print(f"Using Triton backend for fused kernels in {model.__class__.__name__}")
    # 更新: "torch" 和 "liger" 后端均使用 torch 路径的适配器
    elif fused_kernels_backend in ("torch", "liger"):
        model.__class__.forward = forward_with_torch_backend_function
        print(f"Using {fused_kernels_backend.capitalize()} backend for fused kernels in {model.__class__.__name__}")
    else:
        raise ValueError(
            f"Unsupported fused_kernels_backend: {fused_kernels_backend}. Choose 'triton', 'torch', or 'liger'."
        )
tests/utils/test_experimental_torch_functional_on_cpu.py test-coverage

关键测试文件。新增了两个直接测试 `FusedLinearForPPO` 后端选择逻辑的测试用例,验证了修复的正确性。

def test_fused_linear_for_ppo_torch_backend_does_not_dispatch_to_liger(monkeypatch):
    # 模拟一个会在被调用时抛出异常的 Liger 后端,以捕获意外调用
    class UnexpectedLigerFusedLinearScaledCrossEntropyFunction:
        @staticmethod
        def apply(*args):
            raise AssertionError("The torch backend must not dispatch to Liger")
​
    monkeypatch.setattr(
        experimental_F,
        "_LIGER_FUSED_LINEAR_SCALED_CROSS_ENTROPY",
        UnexpectedLigerFusedLinearScaledCrossEntropyFunction,
    )
    monkeypatch.setattr(experimental_F, "_FLASH_ATTN_CROSS_ENTROPY_AVAILABLE", False)
​
    # 使用默认的 torch 后端实例化
    # 如果错误地调用了 Liger,上面定义的 apply 方法会抛出异常
    log_probs, entropy = experimental_F.FusedLinearForPPO()(hidden, weight, labels)
​
    # ... ( 后续为验证输出正确性的断言 ) ...
​
​
def test_fused_linear_for_ppo_rejects_unknown_backend():
    # 验证传入非法后端名时,__init__ 会抛出 ValueError
    with pytest.raises(ValueError, match="Unsupported FusedLinearForPPO backend"):
        experimental_F.FusedLinearForPPO(impl_backend="unknown")

评论区精华

显式后端选择 vs 隐式检测 设计

PR body 详细论证了为何需要从基于可导入性的隐式检测(#7461)切换回显式参数控制。核心论点是:隐式检测违反了 `impl_backend` 的配置契约,且绕过了 PyTorch 算子重载,对 vexact 等精确性集成造成了破坏。

结论:采用显式 `impl_backend` 参数进行后端选择,`torch` 为默认值,需要用户主动配置才能使用 Liger。 · 已解决

向后兼容性与配置传播 设计

通过在模型实例上设置 `_verl_fused_kernels_backend` 属性,并在五个模型适配器的 `forward_with_torch_backend` 函数中读取并传递给 `FusedLinearForPPO`,实现了配置从全局补丁点到最终计算单元的穿透,同时保持了默认行为的向后兼容(默认为 `torch`)。

结论:设计实现了清晰的配置传播链,并保持了默认值的安全回退。 · 已解决

风险与影响

  1. 回归风险:默认 impl_backend="torch" 与旧版(引入 #7461 后)行为不一致。在 Liger 已安装的环境中,用户若未显式配置 impl_backend: liger,将回退至 torch 后端,可能带来预期外的性能变化。这是一个有意为之的破坏性行为修正。
  2. 配置传播完整性:需要确认所有可能的模型适配器(dense_common, glm4v, qwen2_vl, qwen3_5, qwen3_vl)均已更新,且 _verl_fused_kernels_backend 属性在所有使用路径上均可被安全访问。当前检查覆盖了五个主要适配器。
  3. 配置验证:monkey_patch.py 中的 fused_kernels_backend 验证逻辑已更新为接受 liger,错误提示也已更新。但 FusedLinearForPPO.__init__ 中的验证与 monkey_patch.py 中的验证是独立的,需保持同步。
  • 用户:用户重新获得对融合内核后端的显式控制权。配置 fused_kernel_options.impl_backend 现在严格生效。需要用户检查并更新现有配置文件,若依赖旧版的隐式 Liger 加速行为,需显式添加 impl_backend: liger。
  • 系统:恢复了清晰的、基于配置的内部调度逻辑,增强了代码的可预测性和可维护性。为 vexact 等外部精确性集成提供了干净的集成点,无需模块状态 hack。
  • 团队:修复了架构约定上的一个技术债务,使得融合内核的后端选择逻辑更易于理解和测试。
显式配置改变默认行为 多文件同步修改需验证完整性

关联 Issue

#7461 [training_utils, env, doc] feat: use Liger fused linear PPO kernel
#48 [ops] fix: keep verl FusedLinearForPPO batch-invariant under Liger

完整报告

参与讨论