diff --git a/tests/engine/test_critic_sft.py b/tests/engine/test_critic_sft.py new file mode 100644 index 0000000000..df48a56cea --- /dev/null +++ b/tests/engine/test_critic_sft.py @@ -0,0 +1,418 @@ +"""Critic model unit tests on Qwen3.5-VL (compose) models. + +Coverage: +- value_head models reject MTP +- Qwen3.5-VL value-head HF keys are top-level (value_head.weight, no language_model. prefix) +- SFT value loss (returns only): forward, loss value, and backward +- PPO critic value loss (returns + old_values): forward, loss value, and backward +- load a Qwen3.5-VL actor HF checkpoint into a critic, save_hf, and reload +""" + +from __future__ import annotations + +import json +import os +import tempfile +import unittest +from pathlib import Path + +import torch +import torch.distributed as dist +from safetensors.torch import save_file +from torch.distributed.tensor import DTensor + +from xtuner._testing import DeterministicDDPTestCase +from xtuner.v1.config import AdamWConfig, FSDPConfig +from xtuner.v1.data_proto import SequenceContext +from xtuner.v1.engine.train_engine import TrainEngine +from xtuner.v1.loss.critic_loss import CriticLossConfig +from xtuner.v1.model.base import BaseModel, XTunerBaseModelConfig +from xtuner.v1.model.compose.qwen3_5.qwen3_5_config import Qwen3_5_VLMoE35BA3Config +from xtuner.v1.model.moe.qwen3_5_text import Qwen3_5_VLTextMoE35BA3BConfig +from xtuner.v1.module.head import ValueHead +from xtuner.v1.module.mtp.config import MTPConfig +from xtuner.v1.utils.device import get_device + + +DEVICE = get_device() +CLIPRANGE_VALUE = 0.2 +QWEN3_5_MOE_PATH = os.environ.get("QWEN3_5_MOE_PATH") + + +class _MappedCriticConfig(XTunerBaseModelConfig): + hidden_size: int = 4 + head_type: str = "value_head" + hf_key_mapping: dict[str, str] | None = {r"^value_head\.weight$": "critic.weight"} + + +class _MappedCriticModel(BaseModel): + def __init__(self, config: _MappedCriticConfig): + super().__init__(config) + self.lm_head = torch.nn.Linear(config.hidden_size, 1, bias=False) + self._init_load_spec() + + def to_hf_key_list(self, key: str) -> list[str]: + return [key] + + +def _critic_model_cfg() -> Qwen3_5_VLMoE35BA3Config: + model_cfg = Qwen3_5_VLMoE35BA3Config( + freeze_vision=True, + freeze_projector=True, + only_llm_forward=True, + compile_cfg=False, + ) + model_cfg.text_config.head_type = "value_head" + return model_cfg + + +class TestCriticCustomHFKeyMapping(unittest.TestCase): + def test_loaded_mapped_critic_head_is_not_reinitialized(self): + """A value head loaded through hf_key_mapping must retain checkpoint values.""" + model = _MappedCriticModel(_MappedCriticConfig()) + expected = torch.tensor([[1.0, 2.0, 3.0, 4.0]], dtype=torch.float32) + + with tempfile.TemporaryDirectory() as tmpdir: + save_file({"critic.weight": expected}, str(Path(tmpdir) / "model.safetensors")) + (Path(tmpdir) / "model.safetensors.index.json").write_text( + json.dumps( + { + "metadata": {"head_type": "value_head"}, + "weight_map": {"critic.weight": "model.safetensors"}, + } + ) + ) + loaded_keys, unloaded_keys, missing_keys = model.from_hf(tmpdir, strict=False) + + self.assertIn("lm_head.weight", loaded_keys) + self.assertNotIn("lm_head.weight", unloaded_keys) + self.assertNotIn("critic.weight", missing_keys) + torch.testing.assert_close(model.lm_head.weight, expected) + + +def _mse_value_loss( + vpred: torch.Tensor, + returns: torch.Tensor, + loss_weight: torch.Tensor, +) -> torch.Tensor: + """Hand-computed SFT critic loss used to document / check CriticLossContext. + + Args: + vpred (torch.Tensor): Current value predictions. + returns (torch.Tensor): Per-token regression targets. + loss_weight (torch.Tensor): Normalized per-token weights. + + Returns: + torch.Tensor: Scalar sum_t w_t * 0.5 * (v - R)^2. + """ + vpred = vpred.squeeze(-1) if vpred.dim() > 0 and vpred.size(-1) == 1 else vpred + returns = returns.squeeze(-1) if returns.dim() > 0 and returns.size(-1) == 1 else returns + loss_weight = loss_weight.squeeze(-1) if loss_weight.dim() > 0 and loss_weight.size(-1) == 1 else loss_weight + per_token = 0.5 * (vpred - returns) ** 2 + return (per_token * loss_weight).sum() + + +def _clipped_value_loss( + vpred: torch.Tensor, + old_values: torch.Tensor, + returns: torch.Tensor, + loss_weight: torch.Tensor, + cliprange_value: float, +) -> torch.Tensor: + """Hand-computed PPO critic loss used to document / check CriticLossContext. + + Args: + vpred (torch.Tensor): Current value predictions. + old_values (torch.Tensor): Frozen values used for clipping. + returns (torch.Tensor): Per-token regression targets. + loss_weight (torch.Tensor): Normalized per-token weights. + cliprange_value (float): Clip range eps. + + Returns: + torch.Tensor: Scalar sum_t w_t * 0.5 * max((v-R)^2, (v_clip-R)^2). + """ + vpred = vpred.squeeze(-1) if vpred.dim() > 0 and vpred.size(-1) == 1 else vpred + old_values = old_values.squeeze(-1) if old_values.dim() > 0 and old_values.size(-1) == 1 else old_values + returns = returns.squeeze(-1) if returns.dim() > 0 and returns.size(-1) == 1 else returns + loss_weight = loss_weight.squeeze(-1) if loss_weight.dim() > 0 and loss_weight.size(-1) == 1 else loss_weight + v_clip = old_values + (vpred - old_values).clamp(-cliprange_value, cliprange_value) + per_token = 0.5 * torch.maximum((vpred - returns) ** 2, (v_clip - returns) ** 2) + return (per_token * loss_weight).sum() + + +def _full_tensor(tensor: torch.Tensor) -> torch.Tensor: + if isinstance(tensor, DTensor): + return tensor.full_tensor() + return tensor + + +def _local_bf16(tensor: torch.Tensor) -> torch.Tensor: + if isinstance(tensor, DTensor): + tensor = tensor.to_local() + return tensor.detach().to(torch.bfloat16).contiguous().cpu() + + +def _text_module(model: torch.nn.Module) -> torch.nn.Module: + return model.language_model + + +def _critic_engine(model_cfg: Qwen3_5_VLMoE35BA3Config) -> TrainEngine: + ep_size = dist.get_world_size() + model_cfg.text_config.ep_size = ep_size + return TrainEngine( + model_cfg=model_cfg, + optim_cfg=AdamWConfig(lr=1e-4, foreach=False), + fsdp_cfg=FSDPConfig( + cpu_offload=False, + torch_compile=False, + fp32_value_head=True, + recompute_ratio=0.0, + ep_size=ep_size, + ), + ) + + +class TestCriticSFT(DeterministicDDPTestCase): + def test_critic_rejects_mtp(self): + pg = self.create_pg(str(DEVICE)) + with self.assertRaisesRegex(ValueError, "does not support MTP"): + Qwen3_5_VLTextMoE35BA3BConfig(head_type="value_head", mtp_config=MTPConfig(num_layers=1)) + try: + dist.destroy_process_group(pg) + except (AssertionError, ValueError): + pass + + def test_critic_qwen35_value_head_hf_keys(self): + """Qwen3.5-VL maps runtime lm_head.weight to top-level value_head.weight.""" + pg = self.create_pg(str(DEVICE)) + with torch.device("meta"): + model = _critic_model_cfg().build() + self.assertEqual( + model.language_model._to_hf_key_list("lm_head.weight"), + ["value_head.weight"], + ) + try: + dist.destroy_process_group(pg) + except (AssertionError, ValueError): + pass + + def test_critic_accepts_xtuner_actor_checkpoint_head_metadata(self): + """An XTuner actor checkpoint may initialize a critic with a new value head.""" + pg = self.create_pg(str(DEVICE)) + with torch.device("meta"): + model = _critic_model_cfg().build() + + with tempfile.TemporaryDirectory() as tmpdir: + index_path = Path(tmpdir) / "model.safetensors.index.json" + index_path.write_text(json.dumps({"metadata": {"head_type": "lm_head"}})) + model.language_model._validate_head_checkpoint(Path(tmpdir)) + + # A critic checkpoint cannot initialize an actor because it has no vocabulary head. + index_path.write_text(json.dumps({"metadata": {"head_type": "value_head"}})) + model.language_model.config.head_type = "lm_head" + with self.assertRaisesRegex(RuntimeError, "does not match config head_type"): + model.language_model._validate_head_checkpoint(Path(tmpdir)) + + try: + dist.destroy_process_group(pg) + except (AssertionError, ValueError): + pass + + def _make_batch(self, model_cfg: Qwen3_5_VLMoE35BA3Config, loss_cfg: CriticLossConfig): + seq_len = 32 + prompt_len = 8 + vocab_size = model_cfg.text_config.vocab_size + input_ids = torch.randint(0, vocab_size, (1, seq_len), device=DEVICE) + shifted_labels = input_ids.clone() + shifted_labels[:, :prompt_len] = loss_cfg.ignore_idx + returns = torch.zeros(1, seq_len, device=DEVICE, dtype=torch.float32) + returns[:, prompt_len:] = 1.0 + seq_ctx = SequenceContext.from_input_ids((input_ids,), device=str(DEVICE)) + return seq_ctx, shifted_labels, returns + + def _assert_loss_matches(self, actual: torch.Tensor, expected_local: torch.Tensor, name: str): + expected = expected_local.clone() + dist.all_reduce(expected, op=dist.ReduceOp.SUM) + self.assertTrue( + abs(actual.item() - expected.item()) < 1e-4, + f"{name} mismatch: got {actual.item()}, expected {expected.item()}", + ) + + def _assert_head_backward(self, engine: TrainEngine, loss: torch.Tensor, *, require_nonzero: bool = True): + engine.optimizer.zero_grad() + loss.backward() + head_grad = _text_module(engine.model).lm_head.weight.grad + self.assertIsNotNone(head_grad) + assert head_grad is not None + head_grad = _full_tensor(head_grad) + self.assertTrue(torch.isfinite(head_grad).all()) + if require_nonzero: + self.assertTrue(head_grad.abs().sum().item() > 0) + + @unittest.skipUnless(QWEN3_5_MOE_PATH, "QWEN3_5_MOE_PATH is not set") + def test_critic_sft_loss(self): + """SFT path: only returns is provided, loss is unclipped MSE.""" + pg = self.create_pg(str(DEVICE)) + torch.manual_seed(0) + + model_cfg = _critic_model_cfg() + loss_cfg = CriticLossConfig() + engine = _critic_engine(model_cfg) + engine.from_hf(QWEN3_5_MOE_PATH, strict=False) + + lm_head = _text_module(engine.model).lm_head + head_weight = _full_tensor(lm_head.weight) + self.assertEqual(tuple(head_weight.shape), (1, model_cfg.text_config.hidden_size)) + self.assertEqual(lm_head.out_features, 1) + self.assertIsInstance(lm_head, ValueHead) + + seq_ctx, shifted_labels, returns = self._make_batch(model_cfg, loss_cfg) + engine.model.train() + loss_ctx = loss_cfg.build( + data={"shifted_labels": shifted_labels, "returns": returns}, + sp_mesh=None, + ) + assert loss_ctx is not None + self.assertIsNone(loss_ctx.loss_kwargs.old_values) + loss_ctx = loss_cfg.loss_ctx_cls.build_batches([loss_ctx])[0] + out = engine.model(seq_ctx=seq_ctx, loss_ctx={"lm": loss_ctx}) + assert out.loss is not None + assert out.logits is not None + assert loss_ctx.loss_kwargs.loss_weight is not None + expected_local = _mse_value_loss( + out.logits.detach().float(), + returns, + loss_ctx.loss_kwargs.loss_weight, + ) + self._assert_loss_matches(out.loss, expected_local, "sft mse loss") + self._assert_head_backward(engine, out.loss) + + torch.cuda.empty_cache() + try: + dist.destroy_process_group(pg) + except (AssertionError, ValueError): + pass + + @unittest.skipUnless(QWEN3_5_MOE_PATH, "QWEN3_5_MOE_PATH is not set") + def test_critic_clipped_value_loss(self): + """PPO critic path: old_values is provided, loss is clipped value loss.""" + pg = self.create_pg(str(DEVICE)) + torch.manual_seed(0) + + model_cfg = _critic_model_cfg() + loss_cfg = CriticLossConfig(cliprange_value=CLIPRANGE_VALUE) + engine = _critic_engine(model_cfg) + engine.from_hf(QWEN3_5_MOE_PATH, strict=False) + + seq_ctx, shifted_labels, returns = self._make_batch(model_cfg, loss_cfg) + engine.model.train() + with torch.no_grad(): + freeze_out = engine.model(seq_ctx=seq_ctx, loss_ctx=None) + assert freeze_out.logits is not None + # v_old is shifted away from R=1 so the clipped term is the larger square. + old_values = freeze_out.logits.float().squeeze(-1).detach() - 1.0 + + loss_ctx = loss_cfg.build( + data={ + "shifted_labels": shifted_labels, + "returns": returns, + "old_values": old_values, + }, + sp_mesh=None, + ) + assert loss_ctx is not None + self.assertIsNotNone(loss_ctx.loss_kwargs.old_values) + loss_ctx = loss_cfg.loss_ctx_cls.build_batches([loss_ctx])[0] + out = engine.model(seq_ctx=seq_ctx, loss_ctx={"lm": loss_ctx}) + assert out.loss is not None + assert out.logits is not None + assert loss_ctx.loss_kwargs.loss_weight is not None + + mse_local = _mse_value_loss( + out.logits.detach().float(), + returns, + loss_ctx.loss_kwargs.loss_weight, + ) + expected_local = _clipped_value_loss( + out.logits.detach().float(), + old_values, + returns, + loss_ctx.loss_kwargs.loss_weight, + CLIPRANGE_VALUE, + ) + self.assertGreater( + abs(expected_local.item() - mse_local.item()), + 1e-6, + "clipped value loss should differ from SFT MSE when clip is active", + ) + self._assert_loss_matches(out.loss, expected_local, "clipped value loss") + # Saturated clip can zero dL/dv; only require a finite backward. + self._assert_head_backward(engine, out.loss, require_nonzero=False) + + torch.cuda.empty_cache() + try: + dist.destroy_process_group(pg) + except (AssertionError, ValueError): + pass + + @unittest.skipUnless(QWEN3_5_MOE_PATH, "QWEN3_5_MOE_PATH is not set") + def test_critic_load_actor_and_save_hf(self): + """Load Qwen3.5-VL actor HF weights into a critic, save to a temp dir, then reload.""" + pg = self.create_pg(str(DEVICE)) + + model_cfg = _critic_model_cfg() + engine = _critic_engine(model_cfg) + engine.from_hf(QWEN3_5_MOE_PATH, strict=False) + + text = _text_module(engine.model) + self.assertIsInstance(text.lm_head, ValueHead) + self.assertEqual(text.lm_head.out_features, 1) + self.assertEqual(tuple(_full_tensor(text.lm_head.weight).shape), (1, model_cfg.text_config.hidden_size)) + + saved_head = _local_bf16(text.lm_head.weight) + saved_embed = _local_bf16(text.embed_tokens.weight) + + with tempfile.TemporaryDirectory() as tmpdir: + save_dir = [tmpdir] + dist.broadcast_object_list(save_dir, src=0) + save_path = Path(save_dir[0]) + engine.save_hf(str(save_path)) + dist.barrier() + + if dist.get_rank() == 0: + index = json.loads((save_path / "model.safetensors.index.json").read_text()) + self.assertEqual(index["metadata"]["head_type"], "value_head") + self.assertIn("value_head.weight", index["weight_map"]) + self.assertNotIn("lm_head.weight", index["weight_map"]) + self.assertNotIn("language_model.value_head.weight", index["weight_map"]) + self.assertNotIn("language_model.lm_head.weight", index["weight_map"]) + + dist.barrier() + del engine + torch.cuda.empty_cache() + + restored = _critic_engine(_critic_model_cfg()) + restored.from_hf(save_path, strict=True) + + restored_text = _text_module(restored.model) + restored_head = _local_bf16(restored_text.lm_head.weight) + restored_embed = _local_bf16(restored_text.embed_tokens.weight) + torch.testing.assert_close(restored_head, saved_head, atol=0, rtol=0) + torch.testing.assert_close(restored_embed, saved_embed, atol=0, rtol=0) + + del restored + torch.cuda.empty_cache() + dist.barrier() # keep tmpdir until all ranks finish from_hf + + try: + dist.destroy_process_group(pg) + except (AssertionError, ValueError): + pass + + @property + def world_size(self) -> int: + return int(os.getenv("XTUNER_TEST_WORLD_SIZE", "8")) + + @property + def destroy_pg_upon_exit(self) -> bool: + return False diff --git a/xtuner/v1/config/fsdp.py b/xtuner/v1/config/fsdp.py index 7335d6d7a5..89bc5dff6b 100644 --- a/xtuner/v1/config/fsdp.py +++ b/xtuner/v1/config/fsdp.py @@ -31,7 +31,8 @@ class FSDPConfig(BaseModel): # TODO: (caoweihan) Convert `torch.dtype` to `Annotated` for compatibility with cyclopts param_dtype: Annotated[torch.dtype, Parameter(help="Data type for model parameters")] = torch.bfloat16 reduce_dtype: Annotated[torch.dtype, Parameter(help="Data type for reduction operations")] = torch.bfloat16 - fp32_lm_head: Annotated[bool, Parameter(help="Use float32 for language model head")] = False + fp32_lm_head: Annotated[bool, Parameter(help="Use float32 for the language model head")] = False + fp32_value_head: Annotated[bool, Parameter(help="Use float32 for the critic value head")] = False # TODO: deprecate `torch_compile` in favor of `compile_cfg` in XTunerBaseModelConfig torch_compile: Annotated[bool, Parameter(help="Enable model compilation for faster inference")] = True mesh_prefix: Annotated[str, Parameter(help="Prefix for device mesh configuration in distributed training")] = ( diff --git a/xtuner/v1/loss/__init__.py b/xtuner/v1/loss/__init__.py index 227dffa93f..9ef560e3af 100644 --- a/xtuner/v1/loss/__init__.py +++ b/xtuner/v1/loss/__init__.py @@ -2,6 +2,7 @@ from .base_loss_ctx import BaseLossConfig, BaseLossContext, BaseLossKwargs from .ce_loss import CELossConfig, CELossContext, LMHeadLossContext from .chunk_loss import ChunkLoss +from .critic_loss import CriticLossConfig, CriticLossContext, CriticLossKwargs from .moe_loss import ( BalancingLossConfig, BalancingLossContext, @@ -26,6 +27,9 @@ "CELossContext", "CELossConfig", "ChunkLoss", + "CriticLossConfig", + "CriticLossContext", + "CriticLossKwargs", "BaseLossConfig", "BaseLossContext", "BaseLossKwargs", diff --git a/xtuner/v1/loss/critic_loss.py b/xtuner/v1/loss/critic_loss.py new file mode 100644 index 0000000000..52d1dea7fa --- /dev/null +++ b/xtuner/v1/loss/critic_loss.py @@ -0,0 +1,208 @@ +# Copyright (c) OpenMMLab. All rights reserved. +from typing import Any, cast + +import torch +import torch.distributed as dist +import torch.nn.functional as F +from pydantic import model_validator +from torch.distributed.device_mesh import DeviceMesh +from torch.distributed.nn.functional import all_reduce + +from xtuner.v1.loss.base_loss_ctx import BaseLossConfig, BaseLossContext, BaseLossKwargs +from xtuner.v1.loss.utils import sp_split +from xtuner.v1.utils.device import get_device + + +DEVICE = get_device() + + +class CriticLossConfig(BaseLossConfig): + """Configuration for PPO critic clipped value-loss. + + Unlike GRPO / PPO policy loss, this config does not inherit + BaseRLLossConfig. It is used by ValueHead, not the language-model + vocabulary head. + + Only mode="eager" is supported: the value head output is 1-dim and + needs no chunking. Setting mode="chunk" raises a validation error. + + Args: + cliprange_value (float): Clip range applied to the difference between + the current value and old_values. Defaults to 0.2. + ignore_idx (int): Token index excluded from the loss. Defaults to -100. + """ + + cliprange_value: float = 0.2 + + @model_validator(mode="after") + def _reject_chunk_mode(self): + if self.mode != "eager": + raise ValueError( + "CriticLossConfig only supports mode='eager'; the value head output is 1-dim and needs no chunking" + ) + return self + + @property + def loss_ctx_cls(self) -> type["CriticLossContext"]: + return CriticLossContext + + @property + def _loss_kwargs_cls(self) -> type["CriticLossKwargs"]: + return CriticLossKwargs + + def build( + self, + data: dict, + sp_mesh: DeviceMesh | None = None, + ) -> "CriticLossContext | None": + """Build critic loss context from data dict. + + Args: + data (dict): Data dictionary containing critic-specific fields: + - shifted_labels (torch.Tensor): The shifted labels + - returns (torch.Tensor): Per-token regression targets + - old_values (torch.Tensor | None): Frozen values for clipping + (optional, can be omitted for SFT MSE) + sp_mesh (DeviceMesh | None): Sequence parallel device mesh + + Returns: + CriticLossContext | None: The built loss context, or None if required + fields are missing. + """ + if "shifted_labels" not in data or "returns" not in data: + return None + + shifted_labels = data["shifted_labels"] + returns = data["returns"] + old_values = data.get("old_values", None) + + loss_kwargs = CriticLossKwargs( + shifted_labels=shifted_labels, + returns=returns, + old_values=old_values, + ).to(DEVICE) + if sp_mesh is not None and sp_mesh.size() > 1: + loss_kwargs = loss_kwargs.sp_split(sp_mesh) + + return CriticLossContext(self, loss_kwargs) + + +class CriticLossKwargs(BaseLossKwargs): + """Keyword arguments for critic loss computation. + + Args: + shifted_labels (torch.Tensor): The shifted labels for the input sequences. + returns (torch.Tensor): Per-token regression targets (GAE return or SFT target). + old_values (torch.Tensor | None): Frozen values used for clipping. If None, + the loss is unclipped MSE. + loss_weight (torch.Tensor | None): Weights for each token in the loss. + """ + + shifted_labels: torch.Tensor + returns: torch.Tensor + old_values: torch.Tensor | None = None + loss_weight: torch.Tensor | None = None + + def sp_split(self, sp_mesh: DeviceMesh) -> "CriticLossKwargs": + self.shifted_labels = sp_split(self.shifted_labels, sp_mesh=sp_mesh, split_dim=1, padding_value=-100) + self.returns = sp_split(self.returns, sp_mesh=sp_mesh, split_dim=1, padding_value=0.0) + if self.old_values is not None: + self.old_values = sp_split(self.old_values, sp_mesh=sp_mesh, split_dim=1, padding_value=0.0) + return self + + def to(self, device: torch.device | str) -> "CriticLossKwargs": + self.shifted_labels = self.shifted_labels.to(device) + self.returns = self.returns.to(device) + if self.old_values is not None: + self.old_values = self.old_values.to(device) + return self + + +class CriticLossContext(BaseLossContext): + """Critic loss context for PPO value-head training. + + Args: + loss_cfg (CriticLossConfig): Configuration for critic loss computation. + loss_kwargs (CriticLossKwargs): Keyword arguments required for loss calculation. + """ + + loss_cfg: CriticLossConfig + loss_kwargs: CriticLossKwargs + + def __init__(self, loss_cfg: CriticLossConfig, loss_kwargs: CriticLossKwargs): + super().__init__(loss_cfg, loss_kwargs) + + @staticmethod + def build_batches(loss_ctx_list: list["CriticLossContext"]) -> list["CriticLossContext"]: # type: ignore[override] + assert len(loss_ctx_list) > 0, "loss_ctx_list can not be empty" + + loss_cfg = loss_ctx_list[0].loss_cfg + loss_weight_list: list[torch.Tensor] = [] + for loss_ctx in loss_ctx_list: + shifted_labels = loss_ctx.loss_kwargs.shifted_labels + loss_weight = torch.ones_like(shifted_labels, dtype=torch.float32) + loss_weight[shifted_labels == loss_cfg.ignore_idx] = 0.0 + loss_ctx.loss_kwargs.loss_weight = loss_weight + loss_weight_list.append(loss_weight) + + rank_denominator = sum(loss_weight.sum() for loss_weight in loss_weight_list) + rank_denominator = cast(torch.Tensor, rank_denominator) + global_denominator = rank_denominator + if dist.is_initialized(): + dist.all_reduce(global_denominator, op=dist.ReduceOp.SUM) + + if global_denominator == 0: + global_denominator.add_(1) + + for loss_ctx in loss_ctx_list: + loss_ctx._batch_size = len(loss_ctx_list) + assert loss_ctx.loss_kwargs.loss_weight is not None + loss_ctx.loss_kwargs.loss_weight = loss_ctx.loss_kwargs.loss_weight / global_denominator + return loss_ctx_list + + def loss_fn( + self, + hidden_states: torch.Tensor, + head_weight: torch.Tensor, + head_bias: torch.Tensor | None, + loss_kwargs: CriticLossKwargs, + ) -> tuple[torch.Tensor, tuple[torch.Tensor | None, dict[str, Any]]]: + values = F.linear(hidden_states, head_weight, head_bias).float() + values_squeezed = values.squeeze(-1) + + shifted_labels = loss_kwargs.shifted_labels + returns = loss_kwargs.returns + old_values = loss_kwargs.old_values + loss_weight = loss_kwargs.loss_weight + assert loss_weight is not None, "loss_weight can not be None" + + returns = returns.squeeze(-1) + loss_weight = loss_weight.squeeze(-1) + if old_values is not None: + old_values = old_values.squeeze(-1) + + rank_grad_tokens = (shifted_labels != self.loss_cfg.ignore_idx).sum() + if rank_grad_tokens == 0: + loss = values.sum() * 0 + elif old_values is None: + per_token = 0.5 * (values_squeezed - returns) ** 2 + loss = (per_token * loss_weight).sum() + else: + cliprange = self.loss_cfg.cliprange_value + v_clip = old_values + (values_squeezed - old_values).clamp(-cliprange, cliprange) + per_token = 0.5 * torch.maximum((values_squeezed - returns) ** 2, (v_clip - returns) ** 2) + loss = (per_token * loss_weight).sum() + + return loss, (values, {}) + + def forward( + self, + hidden_states: torch.Tensor, + head_weight: torch.Tensor, + head_bias: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, tuple[torch.Tensor | None, dict[str, Any]]]: + assert self.loss_kwargs is not None, "loss_kwargs must be set before calling forward" + loss, (values, extra_info) = self.loss_fn(hidden_states, head_weight, head_bias, self.loss_kwargs) + if dist.is_initialized(): + loss = all_reduce(loss, op=dist.ReduceOp.SUM, group=dist.group.WORLD) + return loss, (values, extra_info) diff --git a/xtuner/v1/model/base.py b/xtuner/v1/model/base.py index fa3712374e..62d3659c20 100644 --- a/xtuner/v1/model/base.py +++ b/xtuner/v1/model/base.py @@ -6,7 +6,7 @@ import re from concurrent.futures import Future, ThreadPoolExecutor, wait from dataclasses import dataclass -from functools import reduce +from functools import partial, reduce from importlib import import_module from itertools import chain from math import prod @@ -223,6 +223,7 @@ class TransformerConfig(XTunerBaseModelConfig): linear_attention: Annotated[GatedDeltaNetConfig | KDAConfig | None, Parameter(group="model")] = None mlp_bias: Annotated[bool, Parameter(group="model")] = False tie_word_embeddings: Annotated[bool, Parameter(group="model")] = False + head_type: Annotated[Literal["lm_head", "value_head"], Parameter(group="model")] = "lm_head" model_type: Annotated[str | None, Parameter(group="model")] = None # TODO: yehaochen maybe should be removed generate_config: GenerateConfig | None = None return_hidden_states: Annotated[bool, Parameter(group="model")] = False @@ -399,6 +400,39 @@ def layers_type(self) -> list[Literal["full_attention", "sliding_attention", "li for i in range(self.num_hidden_layers) ] + def build_head(self) -> nn.Linear: + """Build the output head. + + Returns: + LMHead when head_type="lm_head" + ValueHead (out_features=1) when head_type="value_head" + """ + from xtuner.v1.module.head import LMHead, ValueHead + + if self.head_type == "value_head": + return ValueHead(self.hidden_size, bias=False) + elif self.head_type == "lm_head": + return LMHead(self.hidden_size, self.vocab_size, bias=False) + else: + raise ValueError(f"Invalid head_type: {self.head_type}") + + @property + def tie_head(self) -> bool: + """Whether the output head shares weights with the input embeddings. + + A value head never ties with embeddings, regardless of + ``tie_word_embeddings``. + """ + return self.tie_word_embeddings and self.head_type != "value_head" + + def head_fp32(self, fsdp_cfg: FSDPConfig) -> bool: + """Whether the output head runs in float32 under FSDP. + + A value head reads the dedicated ``fp32_value_head`` switch instead of + ``fp32_lm_head``. + """ + return fsdp_cfg.fp32_value_head if self.head_type == "value_head" else fsdp_cfg.fp32_lm_head + class ModelOutputs(PydanticBaseModel): model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") @@ -598,9 +632,10 @@ def from_hf( if isinstance(hf_path, Path): hf_path = str(hf_path) + 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 loaded_keys, unloaded_keys, missing_keys + return self._init_missing_critic_head(hf_loader, loaded_keys, unloaded_keys, missing_keys) def scale_and_reduce_grad(self): return @@ -613,6 +648,145 @@ def cal_grad_norm(self, grads: list[DTensor], dtype=torch.float32): def to_hf_key_list(self, key: str) -> list[str]: raise NotImplementedError() + def _remap_critic_hf_keys(self, hf_keys: list[str]) -> list[str]: + """Map internal lm_head.* names to HF value_head.*. + + 内部参数名 lm_head.weight → HF 名 value_head.weight。 只在落盘时使用 value_head,运行时仍使用 lm_head。 + """ + remapped: list[str] = [] + for key in hf_keys: + remapped.append(re.sub(r"(^|\.)lm_head(?=\.|$)", r"\1value_head", key)) + return remapped + + def _to_hf_key_list(self, name: str) -> list[str]: + """Convert an internal parameter name to HuggingFace key names. + + Save/load 都走这里,而不是直接调 to_hf_key_list。 + + * 普通参数:走子类 to_hf_key_list(加 model. 前缀、layer 重命名、 + VL 的 language_model. 前缀、tie 时 lm_head → embed_tokens)。 + * value_head 的 lm_head.*:子类通过 ``config.tie_head`` 判断 tie, + value head 天然不 tie,不会被映射成 embed_tokens;这里再做 + lm_head → value_head 改名。 + """ + 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)) + else: + return self.to_hf_key_list(name) + + def _head_checkpoint_meta(self) -> dict[str, Any]: + """Checkpoint meta written into model.safetensors.index.json + metadata.""" + return { + "head_type": getattr(self.config, "head_type", "lm_head"), + } + + def _validate_head_checkpoint(self, hf_path: Path) -> None: + """Validate that the model's head type matches the checkpoint head + type. + + A checkpoint exported by an actor (head_type="lm_head") may initialize a + critic model (head_type="value_head"). Its value head is absent and + is initialized later by ``_init_missing_critic_head``. The reverse direction + is invalid because a critic checkpoint has no vocabulary head. + + Official HF checkpoints without this metadata are skipped. Vision and projector modules have no head_type and + are skipped. + """ + if not hasattr(self.config, "head_type"): + return + index_path = hf_path / "model.safetensors.index.json" + if not index_path.is_file(): + return + meta = json.loads(index_path.read_text()).get("metadata") or {} + if "head_type" not in meta: + return + cfg_head_type = self.config.head_type + ckpt_head_type = meta.get("head_type", "lm_head") + if cfg_head_type == "value_head" and ckpt_head_type == "lm_head": + log_rank0.info( + "Initializing a value_head critic from an lm_head actor checkpoint; " + "the missing value head will be randomly initialized." + ) + return + if ckpt_head_type != cfg_head_type: + raise RuntimeError( + f"Checkpoint head_type={ckpt_head_type!r} does not match config head_type={cfg_head_type!r}." + ) + + @torch.no_grad() + def _init_missing_critic_head( + self, + hf_loader: HFCheckpointLoader, + loaded_keys: set[str], + unloaded_keys: set[str], + missing_keys: set[str], + ) -> tuple[set[str], set[str], set[str]]: + """Randomly initialize a critic ValueHead missing from the checkpoint, + and reconcile the key sets accordingly. + + The decision is based on the global checkpoint key map, not on the + rank-local ``loaded_keys`` returned by ``_load_params``. For a small + value head, FSDP can leave all non-zero local shards on one rank while + other ranks own empty shards and report the runtime parameter as + loaded. A rank-local decision would therefore make only one rank enter + the DTensor collective used by random initialization. + + ``hf_key_mapping`` is respected because the load spec stores the mapped + checkpoint keys. The initialized weight is drawn from ``Normal(0, + 1 / (hidden_size + 1))``. + + When the ValueHead is randomly initialized, the key sets are updated: + ``loaded_keys`` / ``unloaded_keys`` hold runtime names (the XTuner + runtime critic module is still named ``lm_head.weight``), while + ``missing_keys`` holds HF names. + + Returns: + The possibly-updated (loaded_keys, unloaded_keys, missing_keys). + """ + + # 非critic模型直接返回 + if getattr(self.config, "head_type", "lm_head") != "value_head": + return loaded_keys, unloaded_keys, missing_keys + + # critic 必须有可初始化的线性 head。 + lm_head = getattr(self, "lm_head", None) + if not isinstance(lm_head, nn.Linear): + raise TypeError( + "Cannot initialize critic value head: config has head_type='value_head' " + "but model.lm_head is not nn.Linear." + ) + + # runtime 参数名 lm_head.weight 对应的 HF checkpoint key。 + # global_hf_keys 中已经应用过 hf_key_mapping,因此这里不需要再查一次 mapping。 + # 映射后的 HF key 都存在,说明 checkpoint 提供了完整 value head。 + # 此时不需要随机初始化;任意 key 缺失则走随机初始化。 + head_load_spec = self.load_spec_mapping.get("lm_head.weight") + if head_load_spec is None: + raise ValueError( + "Critic model has head_type='value_head' but no load spec for " + "'lm_head.weight'. Check model initialization and _init_load_spec()." + ) + if all(hf_loader.is_key_exist(key) for key in head_load_spec.global_hf_keys): + return loaded_keys, unloaded_keys, missing_keys + + from xtuner.v1.utils import init_params + + value_head_std = 1.0 / (self.config.hidden_size + 1) + init_params(lm_head.weight, partial(nn.init.normal_, mean=0.0, std=value_head_std)) + log_rank0.info( + "Value head (runtime module lm_head) is initialized with " + f"Normal(mean=0, std={value_head_std:.6g}) because the checkpoint has no value_head weights; " + "the rest of the model is loaded from the checkpoint." + ) + + # 随机初始化了 ValueHead,对 loaded_keys、unloaded_keys、missing_keys 进行对账 + loaded_keys.add("lm_head.weight") + unloaded_keys.discard("lm_head.weight") + missing_keys -= set(head_load_spec.global_hf_keys) + return loaded_keys, unloaded_keys, missing_keys + def trainable_parameters(self): params = [(name, param) for name, param in self.named_parameters() if param.requires_grad] return params @@ -1073,7 +1247,7 @@ def _init_load_spec(self) -> None: for name, param in self.state_dict().items(): name = self._clean_param_name(name) - _hf_keys = self.to_hf_key_list(name) + _hf_keys = self._to_hf_key_list(name) if not self.config.hf_key_mapping: hf_keys = _hf_keys @@ -1680,7 +1854,7 @@ def _write_hf_index_and_config(self, hf_dir: Path | str, weight_map: Mapping[str raise RuntimeError("Internal Error, both self.config.hf_config and self._hf_path are None") with open(hf_dir / "model.safetensors.index.json", "w") as f: - index = {"weight_map": dict(weight_map), "metadata": {}} + index = {"weight_map": dict(weight_map), "metadata": self._head_checkpoint_meta()} json.dump(index, f, indent=2, ensure_ascii=False) def _save_hf( @@ -1785,7 +1959,7 @@ def _save_hf( # write or overwrite `model.safetensors.index.json` with open(hf_dir / "model.safetensors.index.json", "w") as f: - index = {"weight_map": weight_map, "metadata": {}} + index = {"weight_map": weight_map, "metadata": self._head_checkpoint_meta()} json.dump(index, f, indent=2, ensure_ascii=False) if dist.is_initialized(): @@ -1793,7 +1967,7 @@ def _save_hf( 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()))) expected_keys = set(self.state_dict()) if strict and matched_hf_keys != expected_hf_keys: diff --git a/xtuner/v1/model/compose/base.py b/xtuner/v1/model/compose/base.py index 05f13d9ca7..a5cff67633 100644 --- a/xtuner/v1/model/compose/base.py +++ b/xtuner/v1/model/compose/base.py @@ -157,6 +157,10 @@ def from_hf(self, hf_path: str | Path, strict=True): if missing: raise RuntimeError(f"Missing parameters from {hf_path}: {list(missing)}. ") + def _head_checkpoint_meta(self) -> dict: + """Forward text-module head_type into the compose safetensors index.""" + return self.language_model._head_checkpoint_meta() + def save_hf(self, hf_dir: Path | str, save_dtype: torch.dtype = torch.bfloat16, safetensors_prefix: str = "model"): hf_dir = Path(hf_dir) self.language_model.save_hf(hf_dir, save_dtype, "model-language") @@ -172,7 +176,11 @@ def save_hf(self, hf_dir: Path | str, save_dtype: torch.dtype = torch.bfloat16, if dist.get_rank() == 0: with open(hf_dir / "model.safetensors.index.json", "w") as f: - json.dump({"weight_map": weight_map_dict, "metadata": {}}, f, indent=4) + json.dump( + {"weight_map": weight_map_dict, "metadata": self._head_checkpoint_meta()}, + f, + indent=4, + ) dist.barrier() def async_save_hf( diff --git a/xtuner/v1/model/dense/dense.py b/xtuner/v1/model/dense/dense.py index d9de4dc087..1aaf7eb96b 100644 --- a/xtuner/v1/model/dense/dense.py +++ b/xtuner/v1/model/dense/dense.py @@ -30,7 +30,6 @@ from xtuner.v1.module import ( GatedDeltaNetConfig, KDAConfig, - LMHead, MHAConfig, MLAConfig, RMSNorm, @@ -60,13 +59,13 @@ def __init__(self, config: TransformerConfig): super().__init__(config) self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps, type=config.rms_norm_type) - self.lm_head = LMHead(config.hidden_size, config.vocab_size, bias=False) + self.lm_head = config.build_head() self.layers = self.build_layers(config) self.rotary_emb = self.build_rotary_embedding(config) self.embed_tokens = self.build_embeddings(config) # Make sure it works properly when not using fsdp - if config.tie_word_embeddings: + if config.tie_head: self.lm_head.weight = self.embed_tokens.weight # TODO(@yehaochen): 把这两行移除 _maybe_compile_layers 要把 compile 相关的 setting 放到 fsdp_config 之外 @@ -222,10 +221,10 @@ def fully_shard( mp_policy = MixedPrecisionPolicy( param_dtype=self.fsdp_config.param_dtype, reduce_dtype=fsdp_config.reduce_dtype ) - if self.fsdp_config.fp32_lm_head: - lm_head_mp_policy = MixedPrecisionPolicy(param_dtype=torch.float32, reduce_dtype=torch.float32) + if self.config.head_fp32(self.fsdp_config): + head_mp_policy = MixedPrecisionPolicy(param_dtype=torch.float32, reduce_dtype=torch.float32) else: - lm_head_mp_policy = mp_policy + head_mp_policy = mp_policy num_recompute_layers = int(self.config.num_hidden_layers * self.fsdp_config.recompute_ratio) generator = torch.Generator() @@ -264,7 +263,7 @@ def fully_shard( self._fully_shard( mesh=self.fsdp_mesh if self.hsdp_mesh is None else self.hsdp_mesh, - mp_policy=lm_head_mp_policy if self.config.tie_word_embeddings else mp_policy, + mp_policy=(head_mp_policy if self.config.tie_head else mp_policy), reshard_after_forward=self.fsdp_config.reshard_after_forward, offload_policy=CPUOffloadPolicy() if self.fsdp_config.cpu_offload else None, module=self.embed_tokens, @@ -280,7 +279,7 @@ def fully_shard( self._fully_shard( mesh=self.fsdp_mesh if self.hsdp_mesh is None else self.hsdp_mesh, - mp_policy=lm_head_mp_policy, + mp_policy=head_mp_policy, reshard_after_forward=self.fsdp_config.reshard_after_forward, offload_policy=CPUOffloadPolicy() if self.fsdp_config.cpu_offload else None, module=self.lm_head, @@ -298,7 +297,7 @@ def fully_shard( self._to_empty_meta() # Make sure it works properly when using fsdp - if self.config.tie_word_embeddings: + if self.config.tie_head: self.lm_head.weight = self.embed_tokens.weight return self diff --git a/xtuner/v1/model/dense/qwen2.py b/xtuner/v1/model/dense/qwen2.py index f4c6c53bcf..efcc9047ab 100644 --- a/xtuner/v1/model/dense/qwen2.py +++ b/xtuner/v1/model/dense/qwen2.py @@ -15,7 +15,7 @@ class Qwen2Dense(Dense): def to_hf_key_list(self, key: str) -> list[str]: - if self.config.tie_word_embeddings and "lm_head" in key: + if self.config.tie_head and "lm_head" in key: key = key.replace("lm_head", "embed_tokens") if "layers" in key or "embed_tokens" in key: diff --git a/xtuner/v1/model/dense/qwen3.py b/xtuner/v1/model/dense/qwen3.py index 3f75dc3f26..432733f383 100644 --- a/xtuner/v1/model/dense/qwen3.py +++ b/xtuner/v1/model/dense/qwen3.py @@ -15,7 +15,7 @@ class Qwen3Dense(Dense): def to_hf_key_list(self, key: str) -> list[str]: - if self.config.tie_word_embeddings and "lm_head" in key: + if self.config.tie_head and "lm_head" in key: key = key.replace("lm_head", "embed_tokens") if "layers" in key or "embed_tokens" in key: diff --git a/xtuner/v1/model/dense/qwen3_5_text.py b/xtuner/v1/model/dense/qwen3_5_text.py index e56d2aced2..02e7c330ab 100644 --- a/xtuner/v1/model/dense/qwen3_5_text.py +++ b/xtuner/v1/model/dense/qwen3_5_text.py @@ -15,7 +15,7 @@ def to_hf_key_list(self, key: str) -> list[str]: # Emit the standalone language-model layout (``model.<...>``). The VLM nesting # under ``model.language_model.`` is applied by the config's ``hf_key_mapping`` # so this tower stays unaware of how it is composed into the VLM. - if self.config.tie_word_embeddings and "lm_head" in key: + if self.config.tie_head and "lm_head" in key: key = key.replace("lm_head", "embed_tokens") if "layers" in key: diff --git a/xtuner/v1/model/dense/qwen3vl_text.py b/xtuner/v1/model/dense/qwen3vl_text.py index 1294fe4401..8d380b6d8a 100644 --- a/xtuner/v1/model/dense/qwen3vl_text.py +++ b/xtuner/v1/model/dense/qwen3vl_text.py @@ -13,7 +13,7 @@ class Qwen3VLTextDense(Qwen3Dense): def to_hf_key_list(self, key: str) -> list[str]: - if self.config.tie_word_embeddings and "lm_head" in key: + if self.config.tie_head and "lm_head" in key: key = key.replace("lm_head", "embed_tokens") if "layers" in key or "embed_tokens" in key: diff --git a/xtuner/v1/model/moe/glm52/glm52.py b/xtuner/v1/model/moe/glm52/glm52.py index fe1cf80a8e..a4e1fe82b4 100644 --- a/xtuner/v1/model/moe/glm52/glm52.py +++ b/xtuner/v1/model/moe/glm52/glm52.py @@ -195,7 +195,7 @@ def _call_decoder_layer( return cast(GLM52MoEDecoderLayerOutput, layer_results) def to_hf_key_list(self, key: str) -> list[str]: - if self.config.tie_word_embeddings and "lm_head" in key: + if self.config.tie_head and "lm_head" in key: key = key.replace("lm_head", "embed_tokens") if key.startswith("mtp_block."): diff --git a/xtuner/v1/model/moe/glm53/glm53.py b/xtuner/v1/model/moe/glm53/glm53.py index b0c814bf65..1a8e84307c 100644 --- a/xtuner/v1/model/moe/glm53/glm53.py +++ b/xtuner/v1/model/moe/glm53/glm53.py @@ -240,7 +240,7 @@ def build_mtp_block(self, config: "Glm53TextMoEConfig"): # type: ignore[overrid return MTPBlock(mtp_config=mtp_config, mtp_layers=[mtp_layer]) def to_hf_key_list(self, key: str) -> list[str]: - if self.config.tie_word_embeddings and "lm_head" in key: + if self.config.tie_head and "lm_head" in key: key = key.replace("lm_head", "embed_tokens") if key.startswith("mtp_block."): diff --git a/xtuner/v1/model/moe/moe.py b/xtuner/v1/model/moe/moe.py index 2dcfe91175..6388e95ca5 100644 --- a/xtuner/v1/model/moe/moe.py +++ b/xtuner/v1/model/moe/moe.py @@ -9,7 +9,7 @@ import torch.distributed as dist import torch.nn.functional as F from cyclopts import Parameter -from pydantic import ConfigDict +from pydantic import ConfigDict, model_validator from torch import nn from torch.distributed._functional_collectives import all_reduce from torch.distributed.device_mesh import DeviceMesh, init_device_mesh @@ -58,7 +58,6 @@ GatedDeltaNetConfig, GreedyRouterConfig, KDAConfig, - LMHead, MHAConfig, MLAConfig, NoAuxRouter, @@ -184,6 +183,13 @@ class MoEConfig(TransformerConfig): # keeping it unsharded after forward to avoid repeated all-gathers. embed_reshard_after_forward: bool = True + @model_validator(mode="after") + def _validate_critic_head(self) -> Self: + """Reject MTP on value-head (critic) models.""" + if self.head_type == "value_head" and self.mtp_config is not None: + raise ValueError("head_type='value_head' does not support MTP; set mtp_config=None") + return self + def build(self) -> "MoE": from xtuner.v1.model.moe.moe import MoE @@ -261,7 +267,7 @@ def __init__(self, config: MoEConfig): self.ep_tp_mesh = None self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps, type=config.rms_norm_type) - self.lm_head = LMHead(config.hidden_size, config.vocab_size, bias=False) + self.lm_head = config.build_head() self.layers = self.build_layers(config) self.rotary_emb = self.build_rotary_embedding(config) @@ -1398,10 +1404,10 @@ def fully_shard( self.mp_policy = MixedPrecisionPolicy( param_dtype=self.fsdp_config.param_dtype, reduce_dtype=fsdp_config.reduce_dtype ) - if self.fsdp_config.fp32_lm_head: - lm_head_mp_policy = MixedPrecisionPolicy(param_dtype=torch.float32, reduce_dtype=torch.float32) + if self.config.head_fp32(self.fsdp_config): + head_mp_policy = MixedPrecisionPolicy(param_dtype=torch.float32, reduce_dtype=torch.float32) else: - lm_head_mp_policy = self.mp_policy + head_mp_policy = self.mp_policy self._init_device_mesh(fsdp_config) if self.config.float8_cfg is not None: @@ -1489,7 +1495,7 @@ def fully_shard( self._fully_shard( mesh=self.fsdp_mesh if self.hsdp_mesh is None else self.hsdp_mesh, - mp_policy=lm_head_mp_policy, + mp_policy=head_mp_policy, reshard_after_forward=False, offload_policy=CPUOffloadPolicy() if self.fsdp_config.cpu_offload else None, module=self.lm_head, diff --git a/xtuner/v1/module/__init__.py b/xtuner/v1/module/__init__.py index c183f4f30b..474683bc6e 100644 --- a/xtuner/v1/module/__init__.py +++ b/xtuner/v1/module/__init__.py @@ -10,7 +10,7 @@ MultiLatentAttention, ) from .dispatcher import NaiveDispatcher, TorchAll2AllDispatcher, build_dispatcher -from .lm_head import LMHead +from .head import LMHead, ValueHead from .rms_norm import RMSNorm from .rope import ( Qwen3VLTextRotaryEmbedding, @@ -46,6 +46,7 @@ "GreedyRouterConfig", "RouterResults", "LMHead", + "ValueHead", "Qwen3VLTextRotaryEmbedding", "get_rope_embedding", "RotaryEmbeddingProtocol", diff --git a/xtuner/v1/module/head/__init__.py b/xtuner/v1/module/head/__init__.py new file mode 100644 index 0000000000..decd23b6bd --- /dev/null +++ b/xtuner/v1/module/head/__init__.py @@ -0,0 +1,2 @@ +from .lm_head import LMHead +from .value_head import ValueHead diff --git a/xtuner/v1/module/lm_head/lm_head.py b/xtuner/v1/module/head/lm_head.py similarity index 100% rename from xtuner/v1/module/lm_head/lm_head.py rename to xtuner/v1/module/head/lm_head.py diff --git a/xtuner/v1/module/head/value_head.py b/xtuner/v1/module/head/value_head.py new file mode 100644 index 0000000000..259598cc4c --- /dev/null +++ b/xtuner/v1/module/head/value_head.py @@ -0,0 +1,71 @@ +from typing import Any, TypeAlias + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.distributed.tensor import DTensor +from typing_extensions import overload + +from xtuner.v1.loss.critic_loss import CriticLossContext + + +Loss: TypeAlias = torch.Tensor +Values: TypeAlias = torch.Tensor +Weight: TypeAlias = torch.Tensor | DTensor +Bias: TypeAlias = torch.Tensor | DTensor | None +HiddenStates: TypeAlias = torch.Tensor + + +class ValueHead(nn.Linear): + """Scalar value head for critic models. + + ValueHead is an nn.Linear(hidden_size, 1) aligned with LMHead: loss_ctx is None returns fp32 values (forward-only / + GAE), otherwise the clipped value loss from CriticLossContext. + + The runtime module name is still lm_head; the HuggingFace disk name is value_head. + """ + + def __init__(self, hidden_size: int, bias: bool = False) -> None: + super().__init__(hidden_size, 1, bias=bias) + + @overload # type: ignore[override] + def forward( + self, hidden_states: HiddenStates, loss_ctx: None = None + ) -> tuple[None, tuple[Values | None, dict[str, Any]]]: ... + + @overload # type: ignore[override] + def forward( + self, hidden_states: HiddenStates, loss_ctx: CriticLossContext + ) -> tuple[Loss, tuple[Values | None, dict[str, Any]]]: ... + + def forward( # type: ignore[override] + self, hidden_states: torch.Tensor, loss_ctx: CriticLossContext | None = None + ) -> tuple[Loss | None, tuple[Values | None, dict[str, Any]]]: + """Forward pass of the scalar value head.""" + if isinstance(self.weight, DTensor): + w = self.weight.to_local() + if self.bias is not None: + assert isinstance(self.bias, DTensor), "Bias should be a DTensor if weight is a DTensor" + b = self.bias.to_local() + else: + b = None + else: + w = self.weight + b = self.bias + if loss_ctx is None: + values = F.linear(hidden_states, w, b) + return None, (values.float(), {}) + else: + return loss_ctx.forward(hidden_states, w, b) + + @overload # type: ignore + def __call__( + self, hidden_states: HiddenStates, loss_ctx: None = None + ) -> tuple[None, tuple[Values | None, dict[str, Any]]]: ... + + @overload # type: ignore + def __call__( + self, hidden_states: HiddenStates, loss_ctx: CriticLossContext + ) -> tuple[Loss, tuple[Values | None, dict[str, Any]]]: ... + + __call__ = nn.Module.__call__ diff --git a/xtuner/v1/module/lm_head/__init__.py b/xtuner/v1/module/lm_head/__init__.py deleted file mode 100644 index 943e3e2e44..0000000000 --- a/xtuner/v1/module/lm_head/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from .lm_head import LMHead diff --git a/xtuner/v1/profiler/prober_utils.py b/xtuner/v1/profiler/prober_utils.py index 93a0070202..25c8fe80a1 100644 --- a/xtuner/v1/profiler/prober_utils.py +++ b/xtuner/v1/profiler/prober_utils.py @@ -5,11 +5,12 @@ import torch import torch.nn as nn -from xtuner.v1.model.moe.moe import DenseDecoderLayer, LMHead, MoEBlock, MoEDecoderLayer +from xtuner.v1.model.moe.moe import DenseDecoderLayer, MoEBlock, MoEDecoderLayer from xtuner.v1.module.attention.gated_deltanet import FusedRMSNormGated, GatedDeltaNet, has_fused_rms_norm_gated from xtuner.v1.module.attention.mha import MultiHeadAttention from xtuner.v1.module.attention.mla import MultiLatentAttention from xtuner.v1.module.decoder_layer.moe_decoder_layer import MoEGate, MoEMLP +from xtuner.v1.module.head import LMHead, ValueHead from xtuner.v1.module.linear.linear import _Linear from xtuner.v1.module.rms_norm.rms_norm import RMSNorm from xtuner.v1.module.rope.rope import FourierEmbedding, Qwen3VLTextRotaryEmbedding, RotaryEmbedding @@ -53,7 +54,7 @@ def register_prober_list(model: nn.Module): elif isinstance(module, MoEBlock): wrapped = ProberList.wrap_moe_block_forward(module.forward, name) module.forward = types.MethodType(wrapped, module) # type: ignore - elif isinstance(module, LMHead): + elif isinstance(module, (LMHead, ValueHead)): wrapped = ProberList.wrap_lm_head_forward(module.forward, name) module.forward = types.MethodType(wrapped, module) # type: ignore diff --git a/xtuner/v1/rl/loss/__init__.py b/xtuner/v1/rl/loss/__init__.py index 45d3e1d069..217582ecea 100644 --- a/xtuner/v1/rl/loss/__init__.py +++ b/xtuner/v1/rl/loss/__init__.py @@ -1,3 +1,5 @@ +from xtuner.v1.loss.critic_loss import CriticLossConfig, CriticLossContext, CriticLossKwargs + from .base_loss import ( BaseRLLossConfig, BaseRLLossContext, diff --git a/xtuner/v1/rl/loss/distillation_loss.py b/xtuner/v1/rl/loss/distillation_loss.py index 084d671dbb..d633c54f62 100644 --- a/xtuner/v1/rl/loss/distillation_loss.py +++ b/xtuner/v1/rl/loss/distillation_loss.py @@ -117,11 +117,12 @@ def build( self, data: dict, sp_mesh: DeviceMesh | None = None, - device: str | torch.device = DEVICE, + device: str | torch.device | None = None, ) -> "DistillationLossContext | None": if "shifted_labels" not in data or "advantages" not in data: return None - + if device is None: + device = DEVICE loss_kwargs = DistillationLossKwargs( shifted_labels=data["shifted_labels"], advantages=data["advantages"], diff --git a/xtuner/v1/rl/trainer/worker.py b/xtuner/v1/rl/trainer/worker.py index f753225d00..d38a486400 100644 --- a/xtuner/v1/rl/trainer/worker.py +++ b/xtuner/v1/rl/trainer/worker.py @@ -56,6 +56,7 @@ from xtuner.v1.rl.loss import ( BaseRLLossConfig, BaseRLLossContext, + CriticLossConfig, kl_penalty, ) from xtuner.v1.rl.model_utils import build_frozen_model @@ -191,7 +192,9 @@ class WorkerConfig(BaseModel): Args: model_cfg (TransformerConfig): Model architecture configuration. optim_cfg (OptimConfig): Optimizer configuration for training. - loss_cfg (BaseRLLossConfig): Loss function configuration for RL training. + loss_cfg (BaseRLLossConfig | CriticLossConfig): Policy RL loss or critic + value-loss configuration. TrainingWorker.fit still follows the actor + (advantages / logprobs) path; critic training is not wired yet. lr_cfg (LRConfig): Learning rate scheduler configuration. fsdp_cfg (FSDPConfig): Fully Sharded Data Parallel configuration. load_from (str | Path): Path to load the main model from. @@ -242,7 +245,7 @@ class WorkerConfig(BaseModel): model_config = ConfigDict(title="Worker config", extra="forbid", arbitrary_types_allowed=True) model_cfg: TransformerConfig | BaseComposeConfig optim_cfg: OptimConfig - loss_cfg: BaseRLLossConfig + loss_cfg: BaseRLLossConfig | CriticLossConfig lr_cfg: LRConfig fsdp_cfg: FSDPConfig load_from: str | Path # TODO: 把 actor 和 ref 配置分离 @@ -396,7 +399,7 @@ def __init__( self._engine = self._build_engine(worker_cfg) self._has_ref = False - if worker_cfg.loss_cfg.use_kl_loss: + if isinstance(worker_cfg.loss_cfg, BaseRLLossConfig) and worker_cfg.loss_cfg.use_kl_loss: self._has_ref = True if worker_cfg.ref_load_from is None: worker_cfg.ref_load_from = worker_cfg.load_from @@ -1333,7 +1336,8 @@ def _fit( for k, v in extra_info_dict.items() if isinstance(v, (torch.Tensor, int, float)) } - extra_info_dict = loss_cfg.finalize_metrics(extra_info_dict, DEVICE) + if isinstance(loss_cfg, BaseRLLossConfig): + extra_info_dict = loss_cfg.finalize_metrics(extra_info_dict, DEVICE) train_step_info.pop("total_loss") # type: ignore[misc] max_memory = DEVICE_MODULE.max_memory_allocated() / (1024**3) # type: ignore[attr-defined] reserved_memory = DEVICE_MODULE.max_memory_reserved() / (1024**3) # type: ignore[attr-defined]