Repository navigation
Conversation
fa02200 to
432c06d
Compare
432c06d to
836b32e
Compare
当前 critic head 加载与初始化逻辑1. Critic 模型支持的加载来源Critic 模型可以从两类 actor 权重初始化：
2. Value head 的随机初始化方式只有当 _load_params() 返回的 loaded_keys 中没有：lm_head.weight，也就是 runtime 上的 critic head 确实没有加载成功时，才会随机初始化。 初始化分布为： 如果 value head 已经通过默认 key 或自定义 key 成功加载，则不会重新初始化，也不会覆盖已加载权重。 3. 自定义 critic head 权重名如果 critic head 默认保存名是： 那么应配置为： 完整转换链路是： 如果 hf_key_mapping 中没有配置 critic head 的映射，系统会使用模型默认的 HF key，例如： value_head.weight。如果 checkpoint 中使用了自定义名字但没有配置 mapping，系统不会猜测该权重是不是 critic head，最终会视为缺失并随机初始化。 4. Runtime 名称与保存名称Critic head 在模型运行时的模块名和参数名仍然是： 但在保存成 HF 格式时，会把它重命名为： 同时保存的 index metadata 中会写入： |
74491ec to
77e4640
Compare
PengchengShi00
changed the title
c259804 to
c07829a
Compare
|
@claude review |
|
Claude: Summary本 PR 新增 critic 另外提示：PR 描述称 Main Flowchart after this PRflowchart TD
A["TransformerConfig.head_type='value_head'"] --> B["config.build_head()"]
B --> C["ValueHead = nn.Linear(hidden, 1)<br/>runtime 模块名仍为 lm_head"]
C -.->|"module 层 import rl 层"| P1["value_head.py:9<br/>import xtuner.v1.rl.loss.critic_loss<br/>⚠ 分层倒置 (W1)"]
A --> D["from_hf()"]
D --> E["_validate_head_checkpoint()"]
E --> F["_load_params()"]
F --> G["_maybe_init_unloaded_critic_head()"]
G --> P2["from_hf 内联 loaded/unloaded/missing_keys 修补<br/>⚠ 抽象层级混杂 (W5)"]
P2 --> H["返回 keys 三元组"]
C --> I["ValueHead.forward(hidden, loss_ctx)"]
I --> J["CriticLossContext.forward"]
J --> P3["仅 eager 路径；mode='chunk' 被静默忽略<br/>⚠ (W2)"]
J --> K["loss_fn: MSE / clipped value loss"]
K --> P4["CI 中无真实执行路径（测试被 env + 8 卡 gate）<br/>⚠ (W3)"]
L["dense.py / moe.py"] --> P5["tie 抑制 + fp32 head 选择重复 3+2 处<br/>⚠ head 策略分散 (W4)"]
style P1 fill:#ffd6d6
style P2 fill:#ffd6d6
style P3 fill:#ffd6d6
style P4 fill:#ffd6d6
style P5 fill:#ffd6d6
style C fill:#d6e8ff
style G fill:#d6e8ff
style J fill:#d6e8ff
核心原理实现与单测
抽象与信息隐藏评估
公开 Interface 的线性业务流程评估
单测建议
其他 Issues本次未发现 Warning 及以上的正确性、安全或性能问题；RL worker 侧对 critic 返回的普通 VerdictCOMMENT |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| from torch.distributed.tensor import DTensor | ||
| from typing_extensions import overload | ||
|
|
||
| from xtuner.v1.rl.loss.critic_loss import CriticLossContext |
There was a problem hiding this comment.
Claude: [设计] 这个顶层 import 让核心 module 层反向依赖 RL 层。链路为：xtuner/v1/module/__init__.py → .head → value_head → 执行 xtuner/v1/rl/loss/__init__.py → grpo_loss → ..utils → ray_accelerator_worker → import ray；而 xtuner/v1/model/__init__.py 又 import 了 xtuner.v1.module.router.greedy，于是任何模型 import 都会连带拉起整个 RL loss 包与 ray。
目前 xtuner/v1/loss/__init__.py → rl_loss → xtuner.v1.rl.utils.misc 已经使 ray 成为事实上的硬依赖，所以这不是新的安装期 break；真实摩擦在分层与环依赖脆弱性：xtuner/v1/rl/loss/* 一旦出现任何顶层 xtuner.v1.model|module import（现在没有），立刻变成硬 import cycle。对比 LMHead 只依赖同层的 xtuner.v1.loss.ce_loss。
critic_loss.py 实际只依赖 xtuner.v1.loss.* 与 xtuner.v1.utils.*，没有任何 RL 专属依赖。建议二选一：
- 把
critic_loss.py移到xtuner/v1/loss/critic_loss.py（与LMHeadLossContext同层），在xtuner.v1.rl.loss继续 re-export 保持兼容； - 至少改为
TYPE_CHECKINGimport（该类型仅用于注解与 overload）：
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from xtuner.v1.rl.loss.critic_loss import CriticLossContext方案 1 更能提升 Locality：head 与其 loss ctx 留在同一层，调用者无需跨层学习接口。
There was a problem hiding this comment.
这个要修改，不应该让sft的运行环境依赖ray
| self._validate_head_checkpoint(Path(hf_path)) | ||
| hf_loader = HFCheckpointLoader(hf_path) | ||
| loaded_keys, unloaded_keys, missing_keys = self._load_params(hf_loader, strict=strict) | ||
| # XTuner runtime critic 模型内部仍然叫：lm_head.weight | ||
| if self._maybe_init_unloaded_critic_head(hf_loader): | ||
| # 如果随机初始化 ValueHead，对 loaded_keys、unloaded_keys、missing_keys 进行更新 | ||
| # missing_keys 是 HF 名，loaded_keys / unloaded_keys 是运行时名 | ||
| loaded_keys.add("lm_head.weight") | ||
| unloaded_keys.discard("lm_head.weight") | ||
| head_load_spec = self.load_spec_mapping.get("lm_head.weight") | ||
| if head_load_spec is not None: | ||
| missing_keys -= set(head_load_spec.global_hf_keys) |
There was a problem hiding this comment.
Claude: [设计] 公开接口 from_hf 的主流程里混进了低层 key 集合修补：loaded_keys.add("lm_head.weight")、unloaded_keys.discard(...)、查 load_spec_mapping、missing_keys -= global_hf_keys，还带着魔法运行时名 "lm_head.weight"。按线性业务流程，高层只应表达“校验 head → 加载参数 → 补齐缺失 critic head”，三类 key 的一致性规则属于实现细节，且目前“决定是否初始化”和“修正 key 集合”这两半规则被拆在两处，新增 head 类型时要同时改调用点和私有方法。
建议收敛为单步，让该 Module 同时拥有决策、初始化与 key 对账：
self._validate_head_checkpoint(Path(hf_path))
hf_loader = HFCheckpointLoader(hf_path)
loaded_keys, unloaded_keys, missing_keys = self._load_params(hf_loader, strict=strict)
return self._init_missing_critic_head(hf_loader, loaded_keys, unloaded_keys, missing_keys)收益：from_hf 保持同一抽象层级；critic head 的命名与 key 对账规则集中在一处，便于单测直接针对该方法验证。
| if self.config.head_type == "value_head": | ||
| head_fp32 = self.fsdp_config.fp32_value_head | ||
| else: | ||
| lm_head_mp_policy = mp_policy | ||
| head_fp32 = self.fsdp_config.fp32_lm_head | ||
| if head_fp32: | ||
| head_mp_policy = MixedPrecisionPolicy(param_dtype=torch.float32, reduce_dtype=torch.float32) | ||
| else: | ||
| head_mp_policy = mp_policy |
There was a problem hiding this comment.
Claude: [设计] [复杂] head 的策略判断散落在多个调用点，缺少 Locality：
dense.py:68、dense.py:272、dense.py:308三处重复config.tie_word_embeddings and config.head_type != "value_head"；- fp32 选择逻辑在
dense.py本段与moe.py:1407-1414逐字重复（head_type == "value_head"?fp32_value_head:fp32_lm_head）； - 落盘换名在
base.py:_to_hf_key_list中靠临时改写共享状态self.config.tie_word_embeddings = False加正则_remap_critic_hf_keys实现。
即“value head 不 tie / 用独立 fp32 开关 / 落盘叫 value_head”这三条本属 head 自身的规则，被要求由每个模型族的 __init__ 与 fully_shard 分别记住；再加一个 head 类型或新模型族时需要同步修改 5+ 处。
建议把策略收敛到 build_head() 所在的 config/head 接缝，例如在 TransformerConfig 上提供：
@property
def tie_head(self) -> bool:
return self.tie_word_embeddings and self.head_type != "value_head"
def head_fp32(self, fsdp_cfg: FSDPConfig) -> bool:
return fsdp_cfg.fp32_value_head if self.head_type == "value_head" else fsdp_cfg.fp32_lm_headdense/moe 只消费 config.tie_head 与 config.head_fp32(fsdp_cfg)；HF 名则由 head 侧声明（而非 BaseModel 临时改 config 再正则替换），避免对共享 config 的瞬时写入。收益：规则与缺陷集中一处，Interface 更小，单测可直接针对 config 行为验证。
|
除了已提出的claude的意见外，我没有其他意见了。另外，base.py 的改动看起来是为支持 critic model 的补丁，但我暂时也没想到更好的实现，base.py 也请 @jayhenry review下 |
c5e0504 to
1a06068
Compare
jayhenry
left a comment
jayhenry
left a comment
There was a problem hiding this comment.
基于 head 1a06068 提交 10 条 inline review：F1–F6、D1、D2、A1、A3，均附建议伪代码。F1/F2 与 build(device="cpu") 参数错误已在同一 head 的 CPU 最小复现中确认；设计候选与待补测试尚未实施。F4/P0 是指定的 SP>1 验证门槛，不表示已证实 SP 数学错误。A3 将最终 load-spec key 集合、head 例外及复制后的严格检查统一留在 _load_params，使 from_hf 保持线性流程。
ProduceBatchResult impact: not affected
RoutedExperts impact: not affected
Ray concurrency impact: not affected
| super().__init__(loss_cfg, loss_kwargs) | ||
|
|
||
| @staticmethod | ||
| def build_batches(loss_ctx_list: list["CriticLossContext"]) -> list["CriticLossContext"]: # type: ignore[override] |
There was a problem hiding this comment.
[F1 / P1] 批量构建 Interface 与现有调用方不兼容。
BaseModel.build_loss_ctx_batch（model/base.py:1348）固定调用 build_batches(contexts, cu_seq_lens_list=..., sp_mesh=...)，这里却只接受列表。配置 lm_loss_cfg=CriticLossConfig() 后调用模型公开批量入口，会抛 TypeError: unexpected keyword argument 'cu_seq_lens_list'。已在当前 head 的真实模块上用 CPU 最小复现确认；现有两个 loss 测试手工调用 build_batches([ctx])，绕过了该入口。
建议新增 context 对齐既有批量协议，保留 token 归一化实现，让模型继续通过多态调用：
@staticmethod
def build_batches(
loss_ctx_list,
cu_seq_lens_list=None,
sp_mesh=None,
):
# critic 按有效 token 全局校准，不使用样本边界；接受统一调用协议。
... # 保留当前 mask、WORLD 分母与归一化。
return loss_ctx_list请补一个从 model.build_loss_ctx_batch([data]) 进入、再 forward/backward 的公开路径回归测试。影响限定为使用 critic 配置的批量入口，普通 CE 训练不受此参数错误影响。
There was a problem hiding this comment.
根因
build_batches 按调用方分三类协议:
| 类别 | 代表 | 调用方 | 签名 |
|---|---|---|---|
| 主 loss | CELossContext(ce_loss.py:125)、CriticLossContext |
BaseModel.build_loss_ctx_batch(base.py:1348),关键字传全三参 |
(loss_ctx_list, cu_seq_lens_list, sp_mesh) |
| 辅助 loss | BalancingLossContext(moe_loss.py:89)、ZLossContext |
_add_auxiliary_loss(base.py:1394)裸调 |
(loss_ctx_list) |
| RL loss | GRPOLossContext(grpo_loss.py:99)、OrealLossContext |
RL worker 裸调(worker.py:1240) |
(loss_ctx_list),继承 CE 却收窄签名(# type: ignore[override]) |
主 loss 需要 cu_seq_lens_list(sample/square reduction 按样本边界算有效 token)和 sp_mesh(SP>1 时 sp_gather/sp_split);辅助/RL loss 只需 _batch_size 或全局 token 分母,不需要序列结构信息。入口对所有挂 lm_loss_cfg 的 loss 无差别传满协议。
CriticLossContext 照搬了 RL loss 的单参数写法,但 RL context 收窄签名安全的前提是"永远走裸调协议";critic 挂在 lm_loss_cfg 上、会被 SFT 公开入口以完整协议调用,于是第一步组 batch 就抛 TypeError: unexpected keyword argument 'cu_seq_lens_list'。critic 同时走 SFT 与 RL 两条链,必须保持与主 loss 一致的完整签名。
Fix 之后
CriticLossContext.build_batches 对齐模型批量协议,新增 cu_seq_lens_list / sp_mesh 两个默认 None 的参数并显式丢弃——critic 按 token 全局校准,不使用样本边界;逐 token 的 ignore 掩码已覆盖 SP padding(sp_split 时 padding 填 -100),因此这两个参数对 critic 无语义,仅为兼容统一调用协议:
@staticmethod
def build_batches( # type: ignore[override]
loss_ctx_list: list["CriticLossContext"],
cu_seq_lens_list: Sequence[torch.IntTensor] | None = None,
sp_mesh: DeviceMesh | None = None,
) -> list["CriticLossContext"]:
# Accepted for the `BaseModel.build_loss_ctx_batch` call protocol; unused because
# critic loss calibrates per token via the ignore mask below.
del cu_seq_lens_list, sp_mesh
...RL worker 侧的裸调 build_batches(step_loss_ctxs) 不受影响(取默认 None),两条调用链共用同一套校准逻辑(逐 token 掩码权重 + dist.all_reduce 全局分母)。
| loss_weight = loss_kwargs.loss_weight | ||
| assert loss_weight is not None, "loss_weight can not be None" | ||
|
|
||
| returns = returns.squeeze(-1) |
There was a problem hiding this comment.
[F2 / P2] B>1,S=1 时这里会把不同样本的 target 交叉广播。
values.squeeze(-1) 是 [B,S]，但二维 returns/weight 为 [B,1] 时再次 squeeze 会变为 [B]；差值随后扩成 [B,B]。MSE/clipped 两分支都会使用这些形状；SP 后本地长度为 1 时也需覆盖。
当前 head 的 CPU 公开路径复现：ValueHead(1) 权重为 1，hidden 为 [[[1.0]], [[3.0]]]，returns 为 [[0.0], [2.0]]，两个 token 权重均为 1/2。实际 loss=3.0、head weight gradient=6.0；正确值分别为 0.5 和 2.0。
建议只去掉确定存在的 head 输出维，监督字段与逐 token 权重统一到预测形状：
values = F.linear(hidden_states, head_weight, head_bias).float()
v = values.squeeze(-1) # [B,S]
returns = loss_kwargs.returns.reshape_as(v)
weight = loss_kwargs.loss_weight.reshape_as(v)
old = loss_kwargs.old_values
old = None if old is None else old.reshape_as(v)
# 以下保留现有 MSE/clipped 公式，用 v、returns、weight、old 计算。请参数化 (B,S)=(1,32),(2,1),(2,3)，比较 MSE/clipped 的数值与梯度。测试 oracle 也有同样的按末维 squeeze 模式，单 token 用例应使用逐 token 手算期望。
| model_cfg: TransformerConfig | BaseComposeConfig | ||
| optim_cfg: OptimConfig | ||
| loss_cfg: BaseRLLossConfig | ||
| loss_cfg: BaseRLLossConfig | CriticLossConfig |
There was a problem hiding this comment.
[F3 / P2] 这里接受 critic config，但当前 worker 的 fit Implementation 仍只支持 actor 数据流。
_pack_one_batch:988 调用 loss_cfg.build(data=loss_inputs, device="cpu")，CriticLossConfig.build 不接受 device（已 CPU 复现 TypeError）。即便补上参数，当前 pack 只准备 advantages/logprobs 等 actor 字段，没有必需的 returns，critic build 仍返回 None。ref/metrics 的两个 isinstance 不能让整个训练 workflow 可替换。
既然本层明确将 critic worker 接入交给 #2186，建议本层先收回 union 与两个守卫，让配置 Interface 与已交付能力一致：
class WorkerConfig(BaseModel):
loss_cfg: BaseRLLossConfig
# TrainingWorker 恢复既有 actor 契约。
if worker_cfg.loss_cfg.use_kl_loss:
... # 构建 ref model。
extra_info_dict = loss_cfg.finalize_metrics(extra_info_dict, DEVICE)
# #2186 在接通 returns/old_values、CPU 物化、SP 和 fit 后扩展 critic workflow。这是提前放宽配置的问题，不代表当前 PR 已接通的 PPO 流程发生回归。若保留 union，应通过完整 worker 数据准备入口证明 critic 可训练；仅让 build(device="cpu") 不报错不够。
There was a problem hiding this comment.
critic 完整 RL 流程在下一个 PR(#2186)接入。device 参数的原因:xtuner/v1/rl/loss 下的 LossConfig.build 均接收 device(base_loss.py:105,worker 在 _pack_one_batch 固定传 device="cpu",即 worker.py:988),而 xtuner/v1/loss 下的均不接收。critic 的 SFT 与 RL 共用同一个 CriticLossConfig,因此对齐 RL 侧协议、在 build 中增加 device 参数(默认 DEVICE)是合理的,已在 xtuner/v1/loss/critic_loss.py 中添加。
| engine.model.train() | ||
| loss_ctx = loss_cfg.build( | ||
| data={"shifted_labels": shifted_labels, "returns": returns}, | ||
| sp_mesh=None, |
There was a problem hiding this comment.
[F4 / P0，验证门槛] 请补齐 SP>1 的真实分布式回归测试。
现有 critic 数值用例均使用 sp_mesh=None，且标准模型批量入口还被 F1 阻断。静态看 critic 的本地 mask→WORLD 分母→autograd WORLD SUM 与 CE token 路径一致；此处 P0 表示要求补齐验证，不表示已证实 SP 数学错误。
至少覆盖 MSE/clipped、SP=2、非整除长度 padding、分片跨 packed sample、某 rank 全 ignore、多 microbatch 共用分母，并与 SP=1 比较 loss、head 与 backbone 的全局梯度：
@pytest.mark.parametrize("loss_kind", ["mse", "clipped"])
@pytest.mark.parametrize("case", [
"odd_length", "cross_sample", "empty_rank", "microbatches",
])
def test_critic_sp_matches_sp1(loss_kind, case):
weights, batch = make_tiny_training_case(loss_kind, case)
ref = run_public_training_case(weights, batch, sp_size=1)
actual = run_public_training_case(weights, batch, sp_size=2)
assert_close(actual.loss, ref.loss)
assert_close(actual.head_grad, ref.head_grad)
assert_close(actual.backbone_grads, ref.backbone_grads)run_public_training_case 应创建真实进程组/mesh/tiny model，并走 build_loss_ctx_batch → train_step → clip_grad_norm(do_clip=False)，在既有梯度归约完成后重建全局梯度，不 mock loss、attention 或 SP。两次运行使用相同全局数据及初始权重。
还需固定字段契约：build(..., sp_mesh=...) 输入完整 returns/old_values，只切一次；若冻结前向已产出本地 old_values，不能再作为完整序列重复切分。
| """ | ||
| cleaned = self._clean_param_name(name) | ||
| if getattr(self.config, "head_type", "lm_head") == "value_head" and cleaned.startswith("lm_head"): | ||
| return self._remap_critic_hf_keys(self.to_hf_key_list(name)) |
There was a problem hiding this comment.
[D1，Interface] 公开 to_hf_key_list 与新增私有 _to_hf_key_list 对同一参数给出不同语义。
Qwen3.5 critic 的公开调用 to_hf_key_list("lm_head.weight") 仍得到 lm_head.weight，私有 wrapper 才得到实际 I/O 使用的 value_head.weight（均在应用 hf_key_mapping 前）。调用者必须知道应绕过旧公开方法。
运行时 _init_load_spec:1250、严格预检 :1970 已切换，但 xtuner/_testing/testcase.py 的 load_params_from_hf:143、xtuner_ckpt_key:159、materialize_submodule:227 仍调用旧公开入口，critic head 会绕过重命名。
建议保留单一公开完整映射 Interface，将模型族规则下沉为私有 hook：
def to_hf_key_list(self, name):
keys = self._model_hf_keys(name) # 模型族 override 迁到这个 hook。
cleaned = self._clean_param_name(name)
if getattr(self.config, "head_type", "lm_head") == "value_head" and cleaned.startswith("lm_head."):
keys = self._remap_critic_hf_keys(keys)
return keys
# I/O、测试工具和外部消费者继续调用公开 to_hf_key_list。若暂时保留 wrapper，则需要完整映射的消费者应切换过去；override、wrapper 内部委派与 InternVL 包装不能机械替换，否则可能递归或改变扩展语义。请增加公开方法与真实保存 key 一致、测试工具能加载 critic head 的回归用例。
| else: | ||
| return self.to_hf_key_list(name) | ||
|
|
||
| def _head_checkpoint_meta(self) -> dict[str, Any]: |
There was a problem hiding this comment.
[D2，设计候选] 在保持当前 runtime lm_head 的前提下，尝试根据实际 checkpoint keys 判断角色，移除 head metadata 协议。
当前 HF 产物已经包含 value_head，所以这项精简不要求运行时改名。可用模型族 head 命名和显式来源别名解析角色，替代保存、compose 转发、读取 index metadata.head_type 的一套逻辑。
present = set(loader.weight_map) # 全局 keys，各 rank 一致。
actor_keys, critic_keys = self._checkpoint_head_keys()
has_actor = bool(present & actor_keys)
has_critic = bool(present & critic_keys)
source_role = (
"mixed" if has_actor and has_critic else
"critic" if has_critic else "actor" if has_actor else "unknown"
)
if target_is_actor and source_role in ("critic", "mixed"):
raise RuntimeError("Cannot load a checkpoint containing a value head into an actor")
... # 继续现有参数校验/加载与缺失 critic head 初始化。
# 同步/异步/compose 停止写 head_type，其他 metadata 与 weight_map 保留。
# 删除 _head_checkpoint_meta 及对应 override/metadata 方向校验。_checkpoint_head_keys 要拥有模型族前缀与明确别名的解析规则；不能简单搜索子串。自定义 critic.weight 必须有来源映射，未知别名不能靠 shape 猜角色；未找到独立 head 时保留 unknown，因为 tied actor 可能没有独立 lm_head 权重。
这是在明确可识别 key 契约后尝试移除 metadata 的候选，不能承诺未知来源仍能被无条件识别。请通过真实无 head metadata 的 checkpoint 覆盖 critic→actor 拒绝（含 strict=False）、actor→critic 初始化、critic round-trip、自定义别名与 tied actor。
| returns = data["returns"] | ||
| old_values = data.get("old_values", None) | ||
|
|
||
| loss_kwargs = CriticLossKwargs( |
There was a problem hiding this comment.
[A1，架构] 建议深化现有 loss 准备 Module，集中形状、设备和 SP 规则，并明确 batch 校准阶段。
目前 BaseModel 在 build 时迁移/切分，worker 在 CPU packing、设备迁移、SP 和校准的不同位置组织 context，测试又手工拼装。Interface 的隐含顺序分散在多个调用方；F1 的参数不匹配与 F2 的形状处理是具体摩擦。
建议现有 Config/Context 统一准备协议，由 critic 构建时将完整监督字段规范为 [B,S]，设备迁移与 SP 切分在 loss Module 内完成一次：
class CriticLossConfig(BaseLossConfig):
def build(self, data, sp_mesh=None, device=None):
labels = data["shifted_labels"]
old = data.get("old_values")
kwargs = CriticLossKwargs(
shifted_labels=labels,
returns=data["returns"].reshape_as(labels),
old_values=None if old is None else old.reshape_as(labels),
).to(DEVICE if device is None else device)
if sp_mesh is not None and sp_mesh.size() > 1:
kwargs.sp_split(sp_mesh)
return CriticLossContext(self, kwargs)
# CPU packing 保留原始监督字段；执行阶段使用相同 build 协议。
contexts = [loss_cfg.build(data, sp_mesh=sp_mesh, device=device) for data in batch]
... # actor 需要时，由 worker 采集并填入已是本地分片的冻结预测。
contexts = loss_cfg.loss_ctx_cls.build_batches(
contexts, cu_seq_lens_list=cu_seq_lens_list, sp_mesh=sp_mesh
)伪代码省略了当前缺字段返回 None 的分支，应保留。critic loss_fn 之后直接消费 [B,S] 的监督字段，不再 squeeze；actor 冻结预测调度仍由上层拥有，不能重复切分已是本地的 old_logprobs/old_values。
这会把 Locality 收回 loss Module，准备路径同时服务模型、worker 和测试；复用现有类即可。回归测试应从模型公开批量入口验证最终 loss/gradient，避免再增加只逐项转发步骤的浅 wrapper。
| def _load_params(self, checkpoint_loader: HFCheckpointLoader, strict=True) -> tuple: | ||
| matched_hf_keys: set[str] = set(checkpoint_loader.weight_map) | ||
| expected_hf_keys: set[str] = set(chain(*map(self.to_hf_key_list, self.state_dict()))) | ||
| expected_hf_keys: set[str] = set(chain(*map(self._to_hf_key_list, self.state_dict()))) |
There was a problem hiding this comment.
[A3，架构] 建议把最终 key 集合与 head 例外的严格校验留在 _load_params，并复用 load spec。
这个位置本来就负责严格预检，适合直接收敛规则。当前 expected_hf_keys 用 _to_hf_key_list 重新推导，而实际复制和缺头初始化用已应用 hf_key_mapping 的 LoadSpec；应使用同一事实来源。from_hf 保持“校验角色→加载→初始化缺失 head”的线性流程。
关键是同时调整复制后的 missing/unloaded 严格检查：只放行确定会被初始化的 value head，保留 backbone 检查，且不提前把未初始化参数记为 loaded。
def _load_params(self, checkpoint_loader, strict=True):
present = set(checkpoint_loader.weight_map)
expected = {
key for spec in self.load_spec_mapping.values()
for key in spec.global_hf_keys
}
allowed_missing, allowed_extra = self._head_load_exceptions(checkpoint_loader)
if strict:
missing = expected - present - allowed_missing
unexpected = present - expected - allowed_extra
if missing or unexpected:
raise RuntimeError(f"Missing: {missing}; unexpected: {unexpected}")
... # 保留当前递归复制、shape 校验，产生 loaded_keys/unloaded_keys/missing_keys。
if strict:
bad_missing = missing_keys - allowed_missing
bad_unloaded = {
name for name in unloaded_keys
if not set(self.load_spec_mapping[name].global_hf_keys) <= allowed_missing
}
if bad_missing or bad_unloaded:
raise RuntimeError(f"Missing: {bad_missing}; unloaded: {bad_unloaded}")
return loaded_keys, unloaded_keys, missing_keys # 保留原报告，尚未把缺头标为 loaded。
# from_hf 继续：
loaded, unloaded, missing = self._load_params(loader, strict=strict)
return self._init_missing_critic_head(loader, loaded, unloaded, missing)_head_load_exceptions 只对合法 actor→critic 的缺头情况提供例外：missing 限于会初始化的完整 value-head key 集合，extra 限于明确解析的 actor head；同类 checkpoint 返回空集合，critic→actor 的角色拒绝不应被 strict=False 绕过。模型族前缀/自定义映射要明确，不能忽略所有包含 lm_head 的 key，也不能豁免 backbone。所有 rank 根据全局 key map 一致决定初始化。
请通过公开 from_hf 测试自定义映射 strict 加载、actor→critic strict 缺头补齐、额外 backbone 缺失仍报错，以及成功加载的 value head 不被覆盖。该项是在既有行为上扩展有针对性的 strict 契约，不是已确认的新运行时故障。
| for i in range(self.num_hidden_layers) | ||
| ] | ||
|
|
||
| def build_head(self) -> nn.Linear: |
There was a problem hiding this comment.
改为 config.head_config.build() 和当前仓库使用惯例更一致
[F5 / P1，设计整改] 建议用 config.head_config.build(...) 表达 head 构建。
attention/router 已采用配置对象拥有 build 的模式。这里把 head_type 分支放进整个 TransformerConfig，后续构造规则仍会堆在模型配置中。建议将选择与构建收敛到 head 配置，提升 Locality，并让 head_config.type 替代独立 head_type，避免双重配置来源。
最小伪代码可用一个配置类，不必引入配置继承层级：
class HeadConfig(BaseModel):
type: Literal["lm_head", "value_head"] = "lm_head"
def build(self, hidden_size: int, vocab_size: int):
if self.type == "value_head":
return ValueHead(hidden_size, bias=False)
return LMHead(hidden_size, vocab_size, bias=False)
class TransformerConfig(XTunerBaseModelConfig):
head_config: HeadConfig = Field(default_factory=HeadConfig)
self.lm_head = config.head_config.build(config.hidden_size, config.vocab_size)共享尺寸由父配置传入，符合 MHAConfig.build(hidden_size=...) 惯例，无需重复保存尺寸来追求无参 build。请覆盖默认 actor、显式 critic、配置序列化以及 head 类型/形状；已有 head_type 配置如需兼容，在解析入口转换一次。这里是设计整改优先级。
| raise ValueError(f"Invalid head_type: {self.head_type}") | ||
|
|
||
| @property | ||
| def tie_head(self) -> bool: |
There was a problem hiding this comment.
没必要改名？
[F6 / P1，设计整改] 建议保留 tie_word_embeddings 为实际执行状态的唯一来源。
这里已比重复条件分支更集中，但 raw tie_word_embeddings=True 与 effective tie_head=False 同时存在，配置与执行多了一层语义。可以在统一 critic 配置构造/转换入口归一化为 False，构建、FSDP 重绑和 HF mapper 继续消费原字段，并删除 tie_head 派生属性。
# 统一 actor -> critic 配置转换入口；采用 F5 后改为设置 head_config。
critic_cfg = actor_cfg.model_copy(deep=True)
critic_cfg.head_type = "value_head"
critic_cfg.tie_word_embeddings = False
# 模型构建、FSDP 重绑、HF mapper 继续使用原字段。
if config.tie_word_embeddings:
self.lm_head.weight = self.embed_tokens.weight相同归一化规则也需覆盖直接构造和 HF 配置转换。当前测试创建配置后直接赋值 head_type，这些路径应统一；只加构造期 validator 覆盖不到后续赋值，也不要再在 key 查询时临时翻转字段。
请从 tie=True 的 actor 转换并验证 critic 导出配置为 False、head/embedding 不共享参数、保存重载一致，同时保持 actor 的原 tie 行为。
Stack 关系
本 PR 属于 Stack #2189，依赖顺序从底到顶为：
本层基于 packing #2185，PPO #2186 基于本层。 这是原 fork PR #2099 在主仓库中的 stack 版本，包含其最新三个提交。原 PR 保留用于查阅 review 讨论。
功能
ValueHead（Linear(hidden_size, 1)），通过TransformerConfig.head_type="value_head"选择；VL/compose 模型使用text_config.head_type。lm_head；HF 保存与加载使用顶层value_head.weight，不添加language_model.前缀。CriticLossConfig：不提供old_values时使用 MSE，提供时使用 clipped value loss。FSDPConfig.fp32_lm_head更名为fp32_head，覆盖 LM head 与 value head；RolloutConfig.fp32_lm_head保持不变。与 PPO 层的分工
本层提供 value-head 模型、HF 权重映射、critic loss 和对应配置适配。上层 #2186 提供 attached CriticWorker、共享样本缓存、value 收集、controller GAE 计算和 actor/critic 训练调度。
value-head 核心提交从原 #2186 的历史中拆出，因此上层不再重复包含该功能 diff。原 #2099 后续的测试注释修正和加载完成后的分布式 barrier 也包含在本层。
测试与验证
tests/engine/test_critic_sft.py覆盖 MTP 约束、HF key 映射、MSE/clipped loss、actor→critic 加载、保存与重新加载。