执行摘要

• 一句话：新增仅保存 LoRA 适配器权重的 checkpoint 支持
• 推荐动作：此项变更设计清晰，向后兼容，建议 LoRA 训练用户启用该功能。值得关注的设计决策是键模式检测（lora_/.adapter_）和向后兼容策略，类似模式可复用于其他 adapter 场景。

功能与动机

在 LoRA 微调场景下，完整保存模型权重（~54 GiB）造成巨大的存储和传输开销。该 PR 旨在只保存 LoRA 适配器权重（~150 MiB），大幅降低 checkpoint 体积。PR #6409 仅保存元数据 JSON，未实现权重过滤。

实现拆解

1. 新增配置项：在 verl/trainer/config/config.py 的 CheckpointConfig 中添加 save_lora_only: bool = False 字段，并更新文档注释。
2. 基类支持：在 BaseCheckpointManager（verl/utils/checkpoint/checkpoint_manager.py）中添加 should_save_lora_only 属性（从 config 读取）和静态方法 is_lora_only_state_dict（检测 state dict 中所有键是否包含 lora_ 或 .adapter_ 关键字）。
3. FSDP 保存逻辑：在 FSDPCheckpointManager.save_checkpoint 中，当 should_save_lora_only 为 True 且模型包含 peft_config 时，过滤 state dict 仅保留 LoRA 键；计算并日志记录 LoRA 参数大小；若过滤后无键则抛出 ValueError。
4. FSDP 加载逻辑：在 FSDPCheckpointManager.load_checkpoint 中，加载模型 state dict 后调用 is_lora_only_state_dict 判断是否为 LoRA-only checkpoint；若是，使用 strict=False 加载，并检查 unexpected_keys 确保适配器配置匹配，防止静默失败。
5. 测试与文档：新增包含 13 个单元测试的测试文件，覆盖默认值、配置开关、各种 state dict 模式（full/lora-only/mixed）的检测、以及 FSDP 的 _has_lora、加载异常路径等。同时更新了两份文档说明使用方式和效果。

关键文件：

• tests/checkpoint_engine/test_fsdp_lora_only_checkpoint_on_cpu.py（模块 单元测试；类别 test；类型 test-coverage；符号 TestBaseCheckpointManagerLoraOnly, _patch_dist, test_should_save_lora_only_default, test_should_save_lora_only_false）: 新增 13 个单元测试，覆盖 save_lora_only 的配置、state dict 检测、保存过滤和加载异常路径。
• verl/utils/checkpoint/fsdp_checkpoint_manager.py（模块 检查点实现；类别 source；类型 core-logic；符号 _has_lora, load_checkpoint, save_checkpoint）: 核心实现文件：在保存时过滤 LoRA 权重，在加载时自动检测 LoRA-only checkpoint 并合并；新增 _has_lora 方法。
• verl/utils/checkpoint/checkpoint_manager.py（模块 检查点基类；类别 source；类型 core-logic；符号 should_save_lora_only, is_lora_only_state_dict）: 基类新增 should_save_lora_only 属性和 is_lora_only_state_dict 静态方法，为所有 checkpoint 管理器提供统一接口。
• verl/trainer/config/config.py（模块 训练配置；类别 source；类型 configuration；符号 save_lora_only）: 在 CheckpointConfig 中添加 save_lora_only 布尔字段，控制是否仅保存 LoRA 权重。
• docs/advance/checkpoint.rst（模块 开发者文档；类别 docs；类型 documentation）: 新增 LoRA-only Checkpoints 章节，说明使用方式和效果。
• docs/examples/config.rst（模块 示例文档；类别 docs；类型 documentation）: 添加 save_lora_only 配置选项的示例说明。

关键符号：BaseCheckpointManager.should_save_lora_only, BaseCheckpointManager.is_lora_only_state_dict, FSDPCheckpointManager._has_lora, FSDPCheckpointManager.load_checkpoint, FSDPCheckpointManager.save_checkpoint

关键源码片段

tests/checkpoint_engine/test_fsdp_lora_only_checkpoint_on_cpu.py

新增 13 个单元测试，覆盖 save_lora_only 的配置、state dict 检测、保存过滤和加载异常路径。

import pytest
from verl.utils.checkpoint.checkpoint_manager import BaseCheckpointManager
​
​
class TestBaseCheckpointManagerLoraOnly:
    """Tests for LoRA-only checkpoint properties on BaseCheckpointManager."""
​
    @pytest.fixture(autouse=True)
    def _patch_dist(self, monkeypatch):
        # 模拟分布式环境，避免需要真实初始化
        import torch.distributed
        monkeypatch.setattr(torch.distributed, "get_rank", lambda: 0)
        monkeypatch.setattr(torch.distributed, "get_world_size", lambda: 1)
        monkeypatch.setattr(torch.distributed, "barrier", lambda: None)
​
    def test_should_save_lora_only_true(self):
        # 当 checkpoint_config 中的 save_lora_only 为 True 时，属性应返回 True
        mgr = self._make_manager(checkpoint_config={"save_lora_only": True})
        assert mgr.should_save_lora_only is True
​
    def test_is_lora_only_state_dict_lora(self):
        # 当所有键包含 "lora_" 时，应判定为 LoRA-only 状态
        state_dict = {
            "base_model.model.model.layers.0.self_attn.q_proj.lora_A.weight": None,
            "base_model.model.model.layers.0.self_attn.q_proj.lora_B.weight": None,
        }
        assert BaseCheckpointManager.is_lora_only_state_dict(state_dict) is True
​
    @staticmethod
    def _make_manager(checkpoint_config):
        # 辅助方法，创建带有 mock 的 BaseCheckpointManager 实例
        class MockModel:
            pass
        class MockOptimizer:
            pass
        return BaseCheckpointManager(
            model=MockModel(), optimizer=MockOptimizer(),
            lr_scheduler=None, processing_class=None,
            checkpoint_config=checkpoint_config,
        )

verl/utils/checkpoint/fsdp_checkpoint_manager.py

核心实现文件：在保存时过滤 LoRA 权重，在加载时自动检测 LoRA-only checkpoint 并合并；新增 _has_lora 方法。

def save_checkpoint(self, local_path: str, hdfs_path: str = None, global_step: int = -1):
    # ... 省略其他部分
    if self.should_save_model:
        model_state_dict = self.model.state_dict()
        # 当开启 save_lora_only 且模型包含 peft_config 时，过滤 state dict
        if self.should_save_lora_only and self._has_lora():
            n_total = len(model_state_dict)
            # 只保留键中包含 "lora_" 或 ".adapter_" 的参数（LoRA 权重）
            model_state_dict = {
                k: v for k, v in model_state_dict.items()
                if "lora_" in k or ".adapter_" in k
            }
            # 若过滤后 state dict 为空，则抛出异常，避免保存空文件
            if not model_state_dict:
                raise ValueError(
                    f"save_lora_only is True and the model has a peft_config, "
                    f"but no LoRA/adapter parameters were found in the state dict. "
                    f"Total params checked: {n_total}."
                )
            # 计算参数总字节数（兼容 ShardedTensor）
            lora_bytes = 0
            for v in model_state_dict.values():
                if hasattr(v, "numel") and hasattr(v, "element_size"):
                    lora_bytes += v.numel() * v.element_size()
                elif hasattr(v, "local_shards"):
                    # 兼容 PyTorch < 2.1 中的 ShardedTensor
                    for shard in v.local_shards():
                        lora_bytes += shard.tensor.numel() * shard.tensor.element_size()
            lora_mib = lora_bytes / 1024**2
            log_with_rank(
                f"LoRA-only save: {len(model_state_dict)}/{n_total} params ({lora_mib:.1f} MiB)",
                rank=self.rank, logger=logger, log_only_rank_0=True,
            )
        torch.save(model_state_dict, model_path)
        # ... 后续日志

verl/utils/checkpoint/checkpoint_manager.py

基类新增 should_save_lora_only 属性和 is_lora_only_state_dict 静态方法，为所有 checkpoint 管理器提供统一接口。

@property
def should_save_lora_only(self) -> bool:
    # 如果 checkpoint_config 为 None，则返回 False（安全默认值）
    if not self.checkpoint_config:
        return False
    # 从 config 对象中获取 save_lora_only 字段，默认 False
    return getattr(self.checkpoint_config, "save_lora_only", False)
​
@staticmethod
def is_lora_only_state_dict(state_dict: dict) -> bool:
    # 空 dict 或 None 视为普通 checkpoint
    if not state_dict:
        return False
    # 当且仅当所有键都包含 "lora_" 或 ".adapter_" 时才认为是 LoRA-only
    return all("lora_" in k or ".adapter_" in k for k in state_dict)

评论区精华

• 加载时缺少 unexpected keys 验证：gemini-code-assist[bot] 指出使用 strict=False 加载 LoRA-only checkpoint 时，若目标模型没有 LoRA 适配器，load_state_dict 会静默忽略所有键，导致加载失败无提示。作者随后在 load_checkpoint 中增加了对 result.unexpected_keys 的检查，若有意外键则抛出 ValueError。

• ShardedTensor 属性缺失及空 state dict 风险：gemini-code-assist[bot] 指出在 SHARDED_STATE_DICT 模式下，state dict 值可能是 ShardedTensor，没有 numel/element_size 属性，计算大小时会 AttributeError；同时过滤后 state dict 可能为空，保存空文件造成问题。作者修复了大小计算逻辑（兼容 ShardedTensor 的 local_shards），并增加了空 state dict 时抛出 ValueError 的保护。

  • 加载 LoRA-only checkpoint 时缺少对 unexpected keys 的验证 (correctness): 作者在后续提交中向 load_checkpoint 增加了对 unexpected_keys 的检查，若存在则抛出 ValueError。
  • 计算 ShardedTensor 大小时可能 AttributeError 以及空 state dict 保护 (correctness): 作者修复了大小计算逻辑，使用 hasattr 检查并处理 ShardedTensor 的 local_shards；同时增加了空 state dict 时抛出 ValueError 的保护。

风险与影响

• 风险：
  • 键命名依赖：LoRA 键检测依赖 lora_ 和 .adapter_ 关键字，若用户使用自定义前缀可能导致检测失败，回退到完整保存（安全）。
  • 后端局限：目前仅 FSDPCheckpointManager 实现了 _has_lora 和过滤逻辑，其他后端（如 Megatron）不生效。若用户在其他后端开启 save_lora_only，should_save_lora_only 返回 True 但保存逻辑无改动，可能仍保存全量权重，造成误导。
  • 空 state dict 异常：当 LoRA 键全部被过滤掉时会抛出 ValueError，防止保存无效 checkpoint，但可能中断训练，需用户确认 LoRA 配置。
• 影响：
  • 用户：LoRA 训练用户可显著减少 checkpoint 存储和网络带宽占用，加载时自动合并权重，对训练流程透明。
  • 系统：新增配置项不影响已有配置，旧 checkpoint 不受影响，完全向后兼容。
  • 团队：需维护此功能并扩展至其他 checkpoint 后端；新增 13 个单元测试覆盖主要路径，降低回归风险。
  • 风险标记：LoRA 键命名依赖, 仅 FSDP 后端生效, 空 state dict 异常打断训练

关联脉络

• PR #6409 [ckpt] feat: save lora train meta: 前序工作，仅保存 LoRA 训练的元数据（rank, alpha 等）到 JSON，未实现权重过滤。本 PR 是其实质性扩展，实现了 LoRA 权重的实际保存与加载。