From 5a700b80fe6994684798cbbf440a273711fd0d26 Mon Sep 17 00:00:00 2001 From: Qingyang Wu Date: Sat, 22 Aug 2026 16:46:31 +0000 Subject: [PATCH 1/8] Add scalar LoRA value-head (critic) support for SAO-style RL The critic is a second LoRA session on the shared base model. The value head is a LoraLinear(hidden, 1) with a zero frozen base weight, so its factors ride the existing multi-adapter machinery (per-session copies, optimizer state, layouts, rank slicing) unchanged. Two new server losses, value_loss and value_prediction, consume the head's folded weight like lm_head; per-token returns/old_values flow through the packer target-aligned. Sampler exports drop the head; training saves keep it. Includes a pure-Python skip-observation GAE reference implementation (SAO, arXiv:2607.07508, Eq. 4-5). Part of #84 --- src/xorl/distributed/torch_parallelize.py | 10 +- src/xorl/lora/utils.py | 8 + src/xorl/ops/loss/__init__.py | 5 + src/xorl/ops/loss/value_loss.py | 165 ++++++++++++ src/xorl/rl/__init__.py | 6 + src/xorl/rl/advantages.py | 80 ++++++ src/xorl/server/orchestrator/packing.py | 5 + src/xorl/server/runner/checkpoint/manager.py | 16 +- src/xorl/server/runner/model_runner.py | 123 ++++++++- src/xorl/server/runner/utils/batch_utils.py | 1 + src/xorl/server/server_arguments.py | 24 ++ src/xorl/trainers/model_builder.py | 42 +++ tests/lora/test_value_head_export.py | 55 ++++ tests/ops/loss/test_value_loss.py | 244 ++++++++++++++++++ tests/rl/__init__.py | 0 tests/rl/test_advantages.py | 91 +++++++ .../orchestrator/test_packing_value_fields.py | 78 ++++++ tests/server/test_value_head_arguments.py | 54 ++++ 18 files changed, 1001 insertions(+), 6 deletions(-) create mode 100644 src/xorl/ops/loss/value_loss.py create mode 100644 src/xorl/rl/__init__.py create mode 100644 src/xorl/rl/advantages.py create mode 100644 tests/lora/test_value_head_export.py create mode 100644 tests/ops/loss/test_value_loss.py create mode 100644 tests/rl/__init__.py create mode 100644 tests/rl/test_advantages.py create mode 100644 tests/server/orchestrator/test_packing_value_fields.py create mode 100644 tests/server/test_value_head_arguments.py diff --git a/src/xorl/distributed/torch_parallelize.py b/src/xorl/distributed/torch_parallelize.py index 5211cd7c..dfd75366 100644 --- a/src/xorl/distributed/torch_parallelize.py +++ b/src/xorl/distributed/torch_parallelize.py @@ -567,6 +567,10 @@ def _experts_shard_placement_fn(param): if exact_dsv4_lm_head and getattr(parallel_state, "lm_head_tp_size", 1) != 8: raise RuntimeError("The exact DSV4-Flash lm head requires lm_head_tensor_parallel_size=8") if lm_head_mod is not None and (fsdp_sharded_lm_head_loss or exact_dsv4_lm_head): + if getattr(model, "value_head", None) is not None: + raise NotImplementedError( + "enable_value_head is not supported with fsdp_sharded_lm_head_loss or the exact lm-head lanes" + ) if parallel_state.tp_enabled: raise NotImplementedError("fsdp_sharded_lm_head_loss is not supported with tensor parallelism.") if not parallel_state.cp_enabled and not lm_head_tp: @@ -622,7 +626,11 @@ def _experts_shard_placement_fn(param): fully_shard(lm_head_mod, **fsdp_kwargs) logger.info_rank0("Using FSDP-sharded lm_head loss over the FSDP group.") elif not pp_enabled and fsdp_kwargs.get("reshard_after_forward", True) is not False: - last_modules = [m for m in [norm_mod, lm_head_mod] if m is not None] + # A scalar value head (critic) is consumed like lm_head: the loss reads + # its weight AFTER the model forward, so it must live in the same + # stay-gathered unit that norm.forward() unshards. + value_head_mod = getattr(model, "value_head", None) + last_modules = [m for m in [norm_mod, lm_head_mod, value_head_mod] if m is not None] if last_modules: last_fsdp_kwargs = dict(fsdp_kwargs) last_fsdp_kwargs["reshard_after_forward"] = False diff --git a/src/xorl/lora/utils.py b/src/xorl/lora/utils.py index 3307cbcb..eef2a0c8 100644 --- a/src/xorl/lora/utils.py +++ b/src/xorl/lora/utils.py @@ -1118,6 +1118,7 @@ def save_lora_checkpoint( transpose_moe_lora_to_peft: bool = True, lora_export_format: str = "peft", preserve_lora_dtype: bool = False, + exclude_value_head: bool = False, ) -> str: """ Save LoRA weights in PEFT-compatible format. @@ -1150,6 +1151,10 @@ def save_lora_checkpoint( preserve_lora_dtype: Keep LoRA tensor dtypes in the safetensors file instead of exporting bf16 weights. Use this for training-resume checkpoints; keep the default bf16 export for inference adapters. + exclude_value_head: Drop ``value_head`` factors from the export. Use + for inference/sampler adapters — serving engines have no scalar + value head module and would reject the unknown target. Training + checkpoints must keep the factors so critic sessions resume. Returns: Path to saved checkpoint directory @@ -1175,6 +1180,9 @@ def save_lora_checkpoint( else: lora_state_dict = slice_lora_state_dict_to_active_rank(model, lora_state_dict) + if exclude_value_head: + lora_state_dict = {key: value for key, value in lora_state_dict.items() if "value_head" not in key} + if lora_export_format == "dsv4_expert_banks": from xorl.models.transformers.deepseek_v4.exact_contract import ( # noqa: PLC0415 DSV4_FLASH_LOGICAL_FACTOR_COUNT, diff --git a/src/xorl/ops/loss/__init__.py b/src/xorl/ops/loss/__init__.py index e96383ee..8811fd1e 100644 --- a/src/xorl/ops/loss/__init__.py +++ b/src/xorl/ops/loss/__init__.py @@ -19,6 +19,7 @@ from xorl.ops.loss.opd_loss import OPDLossMetrics, opd_loss_function, opd_vocab_parallel_loss_function from xorl.ops.loss.policy_loss import policy_loss_function from xorl.ops.loss.reducers import Reducer, SequencePartial, TokenPartial +from xorl.ops.loss.value_loss import value_loss_function, value_prediction_function from xorl.ops.loss.vocab_parallel_cross_entropy import vocab_parallel_cross_entropy @@ -39,6 +40,8 @@ "policy_loss": policy_loss_function, "drgrpo": drgrpo_loss_function, "opd_loss": opd_loss_function, + "value_loss": value_loss_function, + "value_prediction": value_prediction_function, } @@ -72,5 +75,7 @@ def register_loss_function(name: str, fn: Callable) -> None: "opd_loss_function", "opd_vocab_parallel_loss_function", "policy_loss_function", + "value_loss_function", + "value_prediction_function", "vocab_parallel_cross_entropy", ] diff --git a/src/xorl/ops/loss/value_loss.py b/src/xorl/ops/loss/value_loss.py new file mode 100644 index 00000000..5654714b --- /dev/null +++ b/src/xorl/ops/loss/value_loss.py @@ -0,0 +1,165 @@ +"""Value-function (critic) losses for RL training with a scalar value head. + +The value head is a ``hidden_size -> 1`` projection consumed the same way the +lm_head is consumed by the CE-based losses: the loss receives the head's +effective (LoRA-folded) weight and applies it to the trunk's hidden states, +so gradients reach the head's adapter factors and the trunk in one backward. + +Both functions follow the server loss contract: per-token values are reduced +through the injected ``Reducer`` (the server passes raw-sum ``TokenPartial`` +reducers and defers normalization to ``optim_step``), and per-token outputs +ride the standard per-token channels of :class:`LossOutput`. +""" + +from __future__ import annotations + +from typing import Any, Dict, Optional + +import torch +import torch.nn.functional as F + +from xorl.ops.loss.loss_output import LossOutput +from xorl.ops.loss.reducers import Reducer, TokenPartial + + +def _project_values(hidden_states: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + """Project flattened hidden states through the scalar value head in fp32.""" + if weight.dim() != 2 or weight.size(0) != 1: + raise ValueError(f"Value head weight must have shape (1, hidden_size), got {tuple(weight.shape)}") + hidden_flat = hidden_states.reshape(-1, hidden_states.size(-1)) + return F.linear(hidden_flat.float(), weight.float()).squeeze(-1) + + +def value_loss_function( + hidden_states: torch.Tensor, + weight: torch.Tensor, + labels: torch.Tensor, + returns: torch.Tensor, + old_values: Optional[torch.Tensor] = None, + clip_range: float = 0.0, + vf_coef: float = 1.0, + ignore_index: int = -100, + loss_reducer: Optional[Reducer] = None, + metric_reducer: Optional[Reducer] = None, +) -> LossOutput: + """Masked squared-error value loss against per-token returns. + + Args: + hidden_states: Trunk hidden states, shape (batch, seq_len, hidden). + weight: Effective value-head weight, shape (1, hidden). For a LoRA + value head, pass the straight-through folded weight so gradients + reach the adapter factors. + labels: Target-aligned token ids, shape (batch, seq_len). Tokens equal + to ``ignore_index`` are masked out (the packer folds the client's + ``weights`` mask into labels). + returns: Per-token return targets R_t, shape (batch, seq_len), + target-aligned like ``advantages``. + old_values: Optional per-token values from before the update, enabling + PPO-style value clipping. + clip_range: If > 0 and ``old_values`` is provided, uses the clipped + value objective ``max((v-R)^2, (clip(v)-R)^2)``. + vf_coef: Multiplier on the final loss. + loss_reducer / metric_reducer: Reduction policies; ``None`` falls back + to the local token mean (does not compose across ranks). + + Returns: + LossOutput with ``loss = vf_coef * 0.5 * reduce(vf_error_sq)``, + per-token values in ``per_token_logprobs`` (the generic per-token + channel), and per-token squared errors in ``per_token_loss``. + """ + original_shape = labels.shape + labels_flat = labels.reshape(-1) + returns_flat = returns.reshape(-1).float() + + valid_mask = labels_flat != ignore_index + valid_mask_f = valid_mask.float() + valid_count = valid_mask.sum() + + if loss_reducer is None: + loss_reducer = TokenPartial(scale=valid_count.float()) + if metric_reducer is None: + metric_reducer = TokenPartial(scale=valid_count.float()) + + values = _project_values(hidden_states, weight) + + error_sq = (values - returns_flat).square() + clip_fraction = None + if clip_range > 0.0 and old_values is not None: + old_values_flat = old_values.reshape(-1).float() + values_clipped = old_values_flat + (values - old_values_flat).clamp(-clip_range, clip_range) + clipped_error_sq = (values_clipped - returns_flat).square() + clip_fraction = metric_reducer((clipped_error_sq > error_sq).float(), valid_mask_f).detach() + vf_error_sq = torch.maximum(error_sq, clipped_error_sq) + else: + vf_error_sq = error_sq + + vf_error_sq = vf_error_sq * valid_mask_f + loss = vf_coef * 0.5 * loss_reducer(vf_error_sq, valid_mask_f) + + with torch.no_grad(): + # Sum-composable moments: downstream normalization by the global valid + # token count turns these into global means, from which explained + # variance can be derived (EV = 1 - E[(R-V)^2] / (E[R^2] - E[R]^2)). + metrics: Dict[str, Any] = { + "value_mean": metric_reducer(values, valid_mask_f), + "value_sq_mean": metric_reducer(values.square(), valid_mask_f), + "return_mean": metric_reducer(returns_flat, valid_mask_f), + "return_sq_mean": metric_reducer(returns_flat.square(), valid_mask_f), + "value_error_sq_mean": metric_reducer(error_sq, valid_mask_f), + "valid_tokens": int(valid_count.item()), + } + if clip_fraction is not None: + metrics["value_clip_fraction"] = clip_fraction + + return LossOutput( + loss=loss, + per_token_logprobs=values.detach().view(original_shape), + per_token_loss=(vf_error_sq.detach()).view(original_shape), + metrics=metrics, + metric_ops={}, + ) + + +def value_prediction_function( + hidden_states: torch.Tensor, + weight: torch.Tensor, + labels: torch.Tensor, + ignore_index: int = -100, + loss_reducer: Optional[Reducer] = None, + metric_reducer: Optional[Reducer] = None, +) -> LossOutput: + """Forward-only per-token value prediction V(s_t). + + Intended for the no-grad ``forward`` op: the client reads the per-token + values from the per-token output channel and computes advantages (GAE) + itself. The loss is an exact zero that keeps a valid autograd graph, so an + accidental ``forward_backward`` produces zero gradients instead of failing. + + ``labels`` is used only for mask-aware metrics; values are returned for + every position so the client can apply its own action mask. + """ + original_shape = labels.shape + labels_flat = labels.reshape(-1) + valid_mask_f = (labels_flat != ignore_index).float() + valid_count = valid_mask_f.sum() + + if metric_reducer is None: + metric_reducer = TokenPartial(scale=valid_count.float()) + + values = _project_values(hidden_states, weight) + loss = (values * 0.0).sum() + + with torch.no_grad(): + metrics = { + "value_mean": metric_reducer(values, valid_mask_f), + "value_sq_mean": metric_reducer(values.square(), valid_mask_f), + "valid_tokens": int(valid_count.item()), + } + + return LossOutput( + loss=loss, + per_token_logprobs=values.detach().view(original_shape), + per_token_loss=None, + metrics=metrics, + metric_ops={}, + ) diff --git a/src/xorl/rl/__init__.py b/src/xorl/rl/__init__.py new file mode 100644 index 00000000..79f4c830 --- /dev/null +++ b/src/xorl/rl/__init__.py @@ -0,0 +1,6 @@ +"""RL utilities: advantage estimation for value-model (critic) training.""" + +from xorl.rl.advantages import compute_skip_observation_gae + + +__all__ = ["compute_skip_observation_gae"] diff --git a/src/xorl/rl/advantages.py b/src/xorl/rl/advantages.py new file mode 100644 index 00000000..b0c4cad9 --- /dev/null +++ b/src/xorl/rl/advantages.py @@ -0,0 +1,80 @@ +"""Skip-observation token-level GAE (SAO, arXiv:2607.07508, Eq. 4-5). + +Agentic trajectories interleave model actions with environment observations +the model did not generate. Standard GAE bootstraps across every adjacent +token, which propagates critic noise through observation tokens. The +skip-observation estimator chains the Bellman recursion across *action tokens +only*: the TD target of an action token bootstraps from the value of the next +action token, bridging any observation gap in between. + +Pure Python (no torch/numpy) so it can be lifted verbatim into xorl-client, +which has no torch dependency. Sequences follow the server's target-aligned +convention: index t describes the t-th target token, matching the per-token +values returned by the ``value_prediction`` loss and the ``advantages`` / +``returns`` fields of a training datum. + +Note: the server's packer masks tokens where ``advantages == 0.0``; that is +the desired behavior for non-action tokens (this module emits exactly 0.0 +there), but a true action-token advantage of exactly 0.0 would also be +masked. Nudge such values by a tiny epsilon if that matters for your reward +scale. +""" + +from __future__ import annotations + +from typing import List, Optional, Sequence, Tuple + + +def compute_skip_observation_gae( + rewards: Sequence[float], + values: Sequence[float], + action_mask: Optional[Sequence[int]] = None, + gamma: float = 1.0, + lam: float = 1.0, + bootstrap_value: float = 0.0, +) -> Tuple[List[float], List[float]]: + """Compute per-token advantages and value targets across action tokens. + + Args: + rewards: Per-token rewards, length T (typically zero everywhere except + the final action token of the trajectory). + values: Per-token value predictions V(s_t), length T (e.g. from a + ``forward`` call with ``loss_fn="value_prediction"``). + action_mask: Per-token 0/1 mask, length T; 1 marks model-generated + (action) tokens. ``None`` treats every token as an action token, + which reduces to standard token-level GAE. + gamma: Discount factor. + lam: GAE lambda. For length-adaptive GAE (VAPO), pass + ``1 - 1 / (alpha * num_action_tokens)``. + bootstrap_value: Value bootstrapped after the last action token + (0.0 for terminated trajectories). + + Returns: + ``(advantages, returns)`` lists of length T. Non-action tokens carry + 0.0 in both; action tokens carry the skip-observation GAE advantage + and the corresponding value target ``R_t = A_t + V(s_t)``. + """ + n = len(rewards) + if len(values) != n: + raise ValueError(f"rewards ({n}) and values ({len(values)}) must have the same length") + if action_mask is None: + action_indices = list(range(n)) + else: + if len(action_mask) != n: + raise ValueError(f"action_mask ({len(action_mask)}) must match rewards ({n})") + action_indices = [t for t in range(n) if action_mask[t]] + + advantages = [0.0] * n + returns = [0.0] * n + + next_advantage = 0.0 + next_value = float(bootstrap_value) + for t in reversed(action_indices): + delta = float(rewards[t]) + gamma * next_value - float(values[t]) + advantage = delta + gamma * lam * next_advantage + advantages[t] = advantage + returns[t] = advantage + float(values[t]) + next_advantage = advantage + next_value = float(values[t]) + + return advantages, returns diff --git a/src/xorl/server/orchestrator/packing.py b/src/xorl/server/orchestrator/packing.py index c9559a7d..ad72c975 100644 --- a/src/xorl/server/orchestrator/packing.py +++ b/src/xorl/server/orchestrator/packing.py @@ -76,6 +76,11 @@ "logprob_top_ks", "logprob_top_ps", "logprob_min_ps", + # Value-model (critic) per-token fields: aligned with labels[1:] like + # advantages/logprobs. Unlike advantages, a 0.0 entry does NOT mask the + # token — masking comes only from weights/target_tokens. + "returns", + "old_values", ) NORMALIZED_SAMPLING_METADATA_FIELDS = { diff --git a/src/xorl/server/runner/checkpoint/manager.py b/src/xorl/server/runner/checkpoint/manager.py index a50e6ac3..05c08cc6 100644 --- a/src/xorl/server/runner/checkpoint/manager.py +++ b/src/xorl/server/runner/checkpoint/manager.py @@ -364,7 +364,14 @@ def _slice_lora_state_dict_to_rank( sliced_state_dict[name] = tensor.narrow(rank_dim, 0, active_rank).contiguous() return sliced_state_dict - def _save_lora_weights(self, save_path: str, model_id: str, *, preserve_lora_dtype: bool = False) -> None: + def _save_lora_weights( + self, + save_path: str, + model_id: str, + *, + preserve_lora_dtype: bool = False, + exclude_value_head: bool = False, + ) -> None: """ Core LoRA saving logic: activate adapter, gather weights, write PEFT checkpoint. @@ -412,6 +419,7 @@ def _save_lora_weights(self, save_path: str, model_id: str, *, preserve_lora_dty lora_state_dict=lora_state_dict, lora_export_format=lora_export_format, preserve_lora_dtype=preserve_lora_dtype, + exclude_value_head=exclude_value_head, ) if adapter_session_spec is not None: write_session_spec(save_path, adapter_session_spec) @@ -1202,8 +1210,10 @@ def save_lora_only(self, lora_path: str, model_id: str = "default") -> Dict[str, local_error = None try: - # Save LoRA weights (collective operation) - self._save_lora_weights(lora_path, model_id) + # Save LoRA weights (collective operation). This is the sampler- + # facing export: serving engines have no value head, so its + # factors are dropped (training-resume saves keep them). + self._save_lora_weights(lora_path, model_id, exclude_value_head=True) except Exception as e: logger.error(f"Failed to save LoRA-only checkpoint for model_id={model_id}: {e}", exc_info=True) local_error = str(e) diff --git a/src/xorl/server/runner/model_runner.py b/src/xorl/server/runner/model_runner.py index d34654c5..86f74f89 100644 --- a/src/xorl/server/runner/model_runner.py +++ b/src/xorl/server/runner/model_runner.py @@ -87,6 +87,8 @@ opd_loss_function, opd_vocab_parallel_loss_function, policy_loss_function, + value_loss_function, + value_prediction_function, ) from xorl.optim import build_optimizer from xorl.server.runner.adapters import LoRAAdapterManager @@ -527,6 +529,36 @@ class ModelRunner: "logprob_top_ps", "logprob_min_ps", }, + "value_loss": { + "labels", + "target_tokens", + "weights", + "logprobs", + "advantages", + "returns", + "old_values", + "_original_position_ids", + "rollout_logprobs", + "logprob_temperatures", + "logprob_top_ks", + "logprob_top_ps", + "logprob_min_ps", + }, + "value_prediction": { + "labels", + "target_tokens", + "weights", + "logprobs", + "advantages", + "returns", + "old_values", + "_original_position_ids", + "rollout_logprobs", + "logprob_temperatures", + "logprob_top_ks", + "logprob_top_ps", + "logprob_min_ps", + }, "opd_loss": { "labels", "target_tokens", @@ -933,6 +965,14 @@ def _compile_registered_adapter_gradient_ownership( ), None, ) + value_head = next( + ( + candidate + for part in self._adapter_manager.model_parts + if (candidate := getattr(part, "value_head", None)) is not None + ), + None, + ) direct_parameter_ids: set[int] = set() managed_fsdp_parameter_ids: set[int] = set() producer_by_parameter_id: dict[int, ProducerFamily] = {} @@ -946,7 +986,9 @@ def _compile_registered_adapter_gradient_ownership( def _walk_module_tree(module: nn.Module, *, inherited_fsdp: bool = False, direct: bool = False) -> None: managed_fsdp = inherited_fsdp or isinstance(module, FSDPModule) - direct = direct or module is lm_head + # The value head's weight, like lm_head's, is consumed directly by + # the loss instead of through the module's forward. + direct = direct or module is lm_head or (value_head is not None and module is value_head) selector = getattr(module, "adapter_gradient_producer_family", None) selected = selector() if callable(selector) else selector producer = None @@ -1282,6 +1324,10 @@ def _walk_module_tree(module: nn.Module, *, inherited_fsdp: bool = False, direct guard_payload.update(expert_guard_by_parameter_id.get(id(parameter), {})) guard_payload.update(exact_lm_head_guard_by_parameter_id.get(id(parameter), {})) guard_payloads[name] = guard_payload + # Policy-loss steps never touch the value head, so its factors may + # legitimately have no gradient; every other adapter factor keeps + # the strict presence contract. + is_value_head_param = name.startswith("value_head.") or ".value_head." in name declarations[name] = ParameterOwnershipDeclaration( topology=topology, producer=producer, @@ -1289,7 +1335,9 @@ def _walk_module_tree(module: nn.Module, *, inherited_fsdp: bool = False, direct completed_domains=tuple(completed), capture_domains=tuple(capture), pending_domains=tuple(pending), - presence=GradientPresencePolicy.REQUIRED_IF_ACTIVE, + presence=GradientPresencePolicy.AUTHORIZED_ZERO + if is_value_head_param + else GradientPresencePolicy.REQUIRED_IF_ACTIVE, config_guard_fingerprint=self._adapter_gradient_hash(guard_payload), config_guard_fields=tuple(sorted(guard_payload.items())), managed_fsdp_shard=id(parameter) in managed_fsdp_parameter_ids, @@ -1551,6 +1599,17 @@ def _initialize_model(self): ) construction_target_modules = None if block_fp8_qlora_training else target_modules + enable_value_head = bool(self.lora_config.get("enable_value_head", False)) + if enable_value_head: + if not lora_enabled or enable_qlora: + raise ValueError("enable_value_head requires plain LoRA (enable_lora=true, enable_qlora=false)") + if target_modules is not None and "lm_head" in target_modules: + raise ValueError( + "enable_value_head requires lm_head excluded from the LoRA targets " + "(set train_unembed=false or pass explicit lora_target_modules): a value_loss backward " + "produces no lm-head adapter gradients, which the gradient-ownership plan would reject" + ) + model_dtype = resolve_training_model_dtype( enable_lora=lora_enabled, enable_qlora=enable_qlora, @@ -1587,6 +1646,7 @@ def _initialize_model(self): lora_target_manifest=self.lora_config.get("lora_target_manifest"), unfuse_for_lora=self.lora_config.get("unfuse_for_lora", False), moe_hybrid_shared_lora=self.lora_config.get("moe_hybrid_shared_lora", False), + enable_value_head=enable_value_head, enable_qlora=enable_qlora, block_fp8_qlora_training=block_fp8_qlora_training, quant_format=self.lora_config.get("quant_format", "nvfp4"), @@ -1673,6 +1733,13 @@ def _initialize_model(self): self.model = result.model self.model_config_obj = result.model_config + value_head = getattr(self.model, "value_head", None) + if value_head is not None: + # The value head's frozen base weight is not part of any base-model + # checkpoint; whatever the load path materialized, the contract is + # an exactly-zero base so V(s) lives entirely in the LoRA factors. + with torch.no_grad(): + value_head.weight.zero_() self._select_exact_dsv4_lora_export_format() if getattr(get_parallel_state(), "lm_head_tp_size", 1) > 1: sync_lm_head_tp_parameters( @@ -2064,6 +2131,21 @@ def _get_effective_lm_head_weight(self): """Get lm_head weight, merging LoRA delta on-the-fly if needed.""" return self._get_effective_lm_head_weight_for(self.model.lm_head) + def _get_effective_value_head_weight(self): + """Get the scalar value head's weight, merging its LoRA delta. + + The value head is a LoRA module with a zero frozen base weight, so the + effective weight IS the folded adapter delta; consuming it here (like + the lm_head weight) routes gradients through the same direct-output- + projection ownership lane. + """ + value_head = getattr(self.model, "value_head", None) + if value_head is None: + raise ValueError( + "value_loss/value_prediction require a value head; start the server with enable_value_head=true" + ) + return self._get_effective_lm_head_weight_for(value_head) + @staticmethod def _get_effective_lm_head_weight_for(lm_head): """Resolve the effective weight for an explicitly owned terminal head.""" @@ -6552,6 +6634,43 @@ def _profile_elapsed_ms(start: float) -> float: is_metric_ops, ) + elif loss_fn in {"value_loss", "value_prediction"}: + value_head_weight = self._get_effective_value_head_weight() + target_tokens = micro_batch.get("target_tokens", micro_batch.get("labels")) + if target_tokens is None: + raise ValueError(f"{loss_fn} requires target_tokens or labels for the valid-token mask") + + if loss_fn == "value_loss": + returns = micro_batch.get("returns") + if returns is None: + raise ValueError("value_loss requires per-token 'returns' in loss_fn_inputs") + _result = value_loss_function( + hidden_states=hidden_states, + weight=value_head_weight, + labels=target_tokens, + returns=returns, + old_values=micro_batch.get("old_values"), + clip_range=float(params.get("clip_range", 0.0)), + vf_coef=float(params.get("vf_coef", 1.0)), + loss_reducer=token_sum_reducer, + metric_reducer=token_sum_reducer, + ) + else: + _result = value_prediction_function( + hidden_states=hidden_states, + weight=value_head_weight, + labels=target_tokens, + metric_reducer=token_sum_reducer, + ) + local_loss_sum = _result.loss + # Per-token values ride the generic per-token channel (the client + # reads them from loss_fn_outputs[i].logprobs). + per_token_outputs["logprobs"] = _result.per_token_logprobs + if return_per_token and _result.per_token_loss is not None: + per_token_outputs["loss"] = _result.per_token_loss + is_metrics = _result.metrics + is_metric_ops = _result.metric_ops + elif loss_fn == "opd_loss": profile_loss_start = _profile_now() if profile_timings else 0.0 student_layer_hidden_states = None diff --git a/src/xorl/server/runner/utils/batch_utils.py b/src/xorl/server/runner/utils/batch_utils.py index 05dd0232..b68f7d32 100644 --- a/src/xorl/server/runner/utils/batch_utils.py +++ b/src/xorl/server/runner/utils/batch_utils.py @@ -372,6 +372,7 @@ def convert_batch_to_tensors(batch: Dict[str, Any], rank: int = 0) -> Dict[str, "logprob_min_ps", "values", "returns", + "old_values", "teacher_weights", "hidden_match_weights", "teacher_hidden_states", diff --git a/src/xorl/server/server_arguments.py b/src/xorl/server/server_arguments.py index 544b114d..0bdfdd3d 100644 --- a/src/xorl/server/server_arguments.py +++ b/src/xorl/server/server_arguments.py @@ -1112,6 +1112,15 @@ class ServerArguments: }, ) + enable_value_head: bool = field( + default=False, + metadata={ + "help": "Attach a scalar LoRA value head (critic) as model.value_head, enabling the 'value_loss' and " + "'value_prediction' loss functions for PPO/SAO-style RL. Requires plain LoRA, no pipeline parallelism, " + "and lm_head excluded from the LoRA targets (train_unembed=false or an explicit target list)." + }, + ) + # ======================================================================== # QLoRA Configuration # ======================================================================== @@ -1439,6 +1448,20 @@ def __post_init__(self): raise ValueError("pipeline_parallel_virtual_stages requires pipeline_parallel_size > 1.") if self.enable_lora and self.merge_lora_interval > 0: raise ValueError("merge_lora_interval is not supported with multi-adapter LoRA server training") + if self.enable_value_head: + if not self.enable_lora or self.enable_qlora: + raise ValueError("enable_value_head requires plain LoRA (enable_lora=true, enable_qlora=false)") + if self.pipeline_parallel_size > 1: + raise ValueError("enable_value_head is not supported with pipeline parallelism yet") + if self.fsdp_sharded_lm_head_loss or self.lm_head_tensor_parallel_size > 1: + raise ValueError( + "enable_value_head is not supported with fsdp_sharded_lm_head_loss or lm-head tensor parallelism" + ) + if self.lora_target_modules is not None and "lm_head" in self.lora_target_modules: + raise ValueError( + "enable_value_head requires lm_head excluded from the LoRA targets: a value_loss backward " + "produces no lm-head adapter gradients, which the gradient-ownership plan would reject" + ) if self.adapter_gradient_ownership_bucket_bytes <= 0: raise ValueError("adapter_gradient_ownership_bucket_bytes must be positive") if self.load_weights_mode not in {"grouped", "all_ranks", "skip"}: @@ -1631,6 +1654,7 @@ def to_config_dict(self) -> Dict[str, Any]: "unfuse_for_lora": self.unfuse_for_lora, "moe_hybrid_shared_lora": self.moe_hybrid_shared_lora, "lora_export_format": self.lora_export_format, + "enable_value_head": self.enable_value_head, "enable_qlora": self.enable_qlora, "block_fp8_qlora_training": self.block_fp8_qlora_training, "quant_format": self.quant_format, diff --git a/src/xorl/trainers/model_builder.py b/src/xorl/trainers/model_builder.py index cf224676..a9436b38 100644 --- a/src/xorl/trainers/model_builder.py +++ b/src/xorl/trainers/model_builder.py @@ -218,6 +218,7 @@ def build_training_model( lora_target_manifest: Optional[dict[str, Any] | str] = None, moe_hybrid_shared_lora: bool = False, unfuse_for_lora: bool = False, + enable_value_head: bool = False, # --- QLoRA --- enable_qlora: bool = False, block_fp8_qlora_training: bool = False, @@ -570,6 +571,11 @@ def build_training_model( moe_hybrid_shared_lora=moe_hybrid_shared_lora, ) + if enable_value_head: + if not enable_lora or enable_qlora: + raise ValueError("enable_value_head currently requires plain LoRA (enable_lora=True, enable_qlora=False)") + _attach_value_head(model, lora_rank=lora_rank, lora_alpha=lora_alpha) + # ------------------------------------------------------------------ # 4. LoRA + mixed precision: upcast trainable params to fp32 # ------------------------------------------------------------------ @@ -953,6 +959,42 @@ def _inject_lora( helper.print_device_mem_info("VRAM usage after LoRA injection") +def _attach_value_head(model: nn.Module, *, lora_rank: int, lora_alpha: int) -> None: + """Attach a scalar LoRA value head (critic) as ``model.value_head``. + + The head is a ``LoraLinear(hidden_size, 1)`` whose base weight is zero and + frozen: the value function lives entirely in the LoRA factors, so + ``V(s) = lora_B @ lora_A @ h * (alpha / r)``. A scalar head is rank-1, so + the factorization loses no expressivity, and because its parameters are + named ``value_head.lora_A``/``value_head.lora_B`` the multi-adapter + machinery (per-session copies, optimizer state, layouts, rank slicing, + deterministic init) manages them like any other adapter factor. + + The base weight is absent from base-model checkpoints; the server runner + re-zeroes it after weight loading. + """ + from xorl.lora.modules.linear import LoraLinear # noqa: PLC0415 + + lm_head_weight = model.lm_head.weight + value_head = LoraLinear( + in_features=model.config.hidden_size, + out_features=1, + r=lora_rank, + lora_alpha=lora_alpha, + bias=False, + device=lm_head_weight.device, + dtype=lm_head_weight.dtype, + ) + with torch.no_grad(): + if value_head.weight.device.type != "meta": + value_head.weight.zero_() + value_head.weight.requires_grad_(False) + model.value_head = value_head + logger.info_rank0( + f"Attached scalar LoRA value head: hidden_size={model.config.hidden_size}, r={lora_rank}, alpha={lora_alpha}" + ) + + def _deferred_qlora_quantize( model: nn.Module, weights_path: str, diff --git a/tests/lora/test_value_head_export.py b/tests/lora/test_value_head_export.py new file mode 100644 index 00000000..be3836e6 --- /dev/null +++ b/tests/lora/test_value_head_export.py @@ -0,0 +1,55 @@ +"""Value-head factors must stay out of sampler adapters but survive training saves.""" + +import json +import os + +import pytest +import torch +import torch.nn as nn +from safetensors.torch import load_file + +from xorl.lora.modules.linear import LoraLinear +from xorl.lora.utils import save_lora_checkpoint + + +pytestmark = pytest.mark.cpu + + +class _TinyLoraModel(nn.Module): + def __init__(self, hidden=8): + super().__init__() + self.q_proj = LoraLinear(hidden, hidden, r=2, lora_alpha=2, bias=False) + self.value_head = LoraLinear(hidden, 1, r=2, lora_alpha=2, bias=False) + with torch.no_grad(): + self.value_head.weight.zero_() + self.value_head.weight.requires_grad_(False) + self.config = None + + +def _saved_keys(save_path): + tensors = load_file(os.path.join(save_path, "adapter_model.safetensors")) + return set(tensors.keys()) + + +def test_sampler_export_drops_value_head(tmp_path): + model = _TinyLoraModel() + save_path = str(tmp_path / "sampler_adapter") + save_lora_checkpoint(model, save_path, base_model_name="tiny", exclude_value_head=True) + + keys = _saved_keys(save_path) + assert any("q_proj" in key for key in keys) + assert not any("value_head" in key for key in keys) + + with open(os.path.join(save_path, "adapter_config.json")) as f: + adapter_config = json.load(f) + assert "value_head" not in adapter_config["target_modules"] + + +def test_training_save_keeps_value_head(tmp_path): + model = _TinyLoraModel() + save_path = str(tmp_path / "training_adapter") + save_lora_checkpoint(model, save_path, base_model_name="tiny", preserve_lora_dtype=True) + + keys = _saved_keys(save_path) + assert any("value_head" in key and "lora_A" in key for key in keys) + assert any("value_head" in key and "lora_B" in key for key in keys) diff --git a/tests/ops/loss/test_value_loss.py b/tests/ops/loss/test_value_loss.py new file mode 100644 index 00000000..c6ce2244 --- /dev/null +++ b/tests/ops/loss/test_value_loss.py @@ -0,0 +1,244 @@ +import pytest +import torch + +from tests.ops.loss.conftest import assert_close +from xorl.lora.modules.linear import LoraLinear +from xorl.ops.loss import ( + TokenPartial, + get_loss_function, + value_loss_function, + value_prediction_function, +) + + +pytestmark = pytest.mark.cpu + +IGNORE_INDEX = -100 + + +@pytest.fixture +def inputs(): + torch.manual_seed(7) + batch, seq, hidden = 2, 6, 16 + hidden_states = torch.randn(batch, seq, hidden) / (hidden**0.5) + weight = torch.randn(1, hidden) + labels = torch.randint(0, 100, (batch, seq)) + mask = torch.tensor( + [[0, 0, 1, 1, 1, 1], [1, 1, 1, 0, 0, 1]], + dtype=torch.bool, + ) + labels[~mask] = IGNORE_INDEX + returns = torch.randn(batch, seq) + return { + "hidden_states": hidden_states, + "weight": weight, + "labels": labels, + "returns": returns, + "mask": mask, + } + + +def expected_values(data): + return (data["hidden_states"].reshape(-1, data["hidden_states"].size(-1)).float() @ data["weight"].float().T).view( + data["labels"].shape + ) + + +def test_registry_exposes_value_losses(): + assert get_loss_function("value_loss") is value_loss_function + assert get_loss_function("value_prediction") is value_prediction_function + + +def test_value_loss_matches_masked_mse(inputs): + unit = TokenPartial(scale=torch.tensor(1.0)) + output = value_loss_function( + hidden_states=inputs["hidden_states"], + weight=inputs["weight"], + labels=inputs["labels"], + returns=inputs["returns"], + loss_reducer=unit, + metric_reducer=unit, + ) + values = expected_values(inputs) + reference = 0.5 * ((values - inputs["returns"]) ** 2)[inputs["mask"]].sum() + assert_close(output.loss, reference) + # Per-token channel carries the values, not logprobs. + assert_close(output.per_token_logprobs, values) + + +def test_value_loss_raw_sum_reducer_contract(inputs): + """With scale=1 reducers, loss must be the raw masked sum (server contract).""" + unit = TokenPartial(scale=torch.tensor(1.0)) + full = value_loss_function( + hidden_states=inputs["hidden_states"], + weight=inputs["weight"], + labels=inputs["labels"], + returns=inputs["returns"], + loss_reducer=unit, + metric_reducer=unit, + ) + # Splitting the batch into two micro-batches must sum to the same loss. + parts = [] + for row in range(2): + parts.append( + value_loss_function( + hidden_states=inputs["hidden_states"][row : row + 1], + weight=inputs["weight"], + labels=inputs["labels"][row : row + 1], + returns=inputs["returns"][row : row + 1], + loss_reducer=unit, + metric_reducer=unit, + ).loss + ) + assert_close(full.loss, parts[0] + parts[1]) + + +def test_value_loss_default_reducer_is_token_mean(inputs): + output = value_loss_function( + hidden_states=inputs["hidden_states"], + weight=inputs["weight"], + labels=inputs["labels"], + returns=inputs["returns"], + ) + values = expected_values(inputs) + reference = 0.5 * ((values - inputs["returns"]) ** 2)[inputs["mask"]].mean() + assert_close(output.loss, reference) + + +def test_value_loss_clipping(inputs): + unit = TokenPartial(scale=torch.tensor(1.0)) + old_values = expected_values(inputs) + torch.randn_like(inputs["returns"]) + clip_range = 0.1 + output = value_loss_function( + hidden_states=inputs["hidden_states"], + weight=inputs["weight"], + labels=inputs["labels"], + returns=inputs["returns"], + old_values=old_values, + clip_range=clip_range, + loss_reducer=unit, + metric_reducer=unit, + ) + values = expected_values(inputs) + clipped = old_values + (values - old_values).clamp(-clip_range, clip_range) + reference = ( + 0.5 * torch.maximum((values - inputs["returns"]) ** 2, (clipped - inputs["returns"]) ** 2)[inputs["mask"]].sum() + ) + assert_close(output.loss, reference) + assert "value_clip_fraction" in output.metrics + + +def test_value_loss_vf_coef_scales_loss(inputs): + unit = TokenPartial(scale=torch.tensor(1.0)) + base = value_loss_function( + hidden_states=inputs["hidden_states"], + weight=inputs["weight"], + labels=inputs["labels"], + returns=inputs["returns"], + loss_reducer=unit, + metric_reducer=unit, + ).loss + scaled = value_loss_function( + hidden_states=inputs["hidden_states"], + weight=inputs["weight"], + labels=inputs["labels"], + returns=inputs["returns"], + vf_coef=0.25, + loss_reducer=unit, + metric_reducer=unit, + ).loss + assert_close(scaled, 0.25 * base) + + +def test_value_loss_metrics_are_sum_composable_means(inputs): + unit = TokenPartial(scale=torch.tensor(1.0)) + output = value_loss_function( + hidden_states=inputs["hidden_states"], + weight=inputs["weight"], + labels=inputs["labels"], + returns=inputs["returns"], + loss_reducer=unit, + metric_reducer=unit, + ) + values = expected_values(inputs) + mask = inputs["mask"] + assert_close(output.metrics["return_mean"], inputs["returns"][mask].sum()) + assert_close(output.metrics["value_mean"], values[mask].sum()) + assert_close(output.metrics["value_error_sq_mean"], ((values - inputs["returns"]) ** 2)[mask].sum()) + assert output.metrics["valid_tokens"] == int(mask.sum()) + + +def test_value_loss_gradients_flow_through_lora_value_head(inputs): + """The server consumes the head via its straight-through folded weight; + gradients must reach the LoRA factors and the hidden states, never the + frozen zero base weight.""" + value_head = LoraLinear(16, 1, r=4, lora_alpha=4, bias=False) + with torch.no_grad(): + value_head.weight.zero_() + torch.nn.init.normal_(value_head.lora_B, std=0.1) # nonzero so grads are nontrivial + value_head.weight.requires_grad_(False) + + hidden_states = inputs["hidden_states"].clone().requires_grad_(True) + effective_weight = value_head.weight + value_head.get_delta_weight().to(value_head.weight.dtype) + output = value_loss_function( + hidden_states=hidden_states, + weight=effective_weight, + labels=inputs["labels"], + returns=inputs["returns"], + loss_reducer=TokenPartial(scale=torch.tensor(1.0)), + metric_reducer=TokenPartial(scale=torch.tensor(1.0)), + ) + output.loss.backward() + assert value_head.lora_A.grad is not None and value_head.lora_A.grad.abs().sum() > 0 + assert value_head.lora_B.grad is not None and value_head.lora_B.grad.abs().sum() > 0 + assert hidden_states.grad is not None and hidden_states.grad.abs().sum() > 0 + assert value_head.weight.grad is None + + +def test_zero_base_lora_value_head_predicts_zero_at_init(inputs): + """Fresh critic (lora_B = 0, base = 0) must predict V(s) = 0 everywhere.""" + value_head = LoraLinear(16, 1, r=4, lora_alpha=4, bias=False) + with torch.no_grad(): + value_head.weight.zero_() + effective_weight = value_head.weight + value_head.get_delta_weight().to(value_head.weight.dtype) + output = value_prediction_function( + hidden_states=inputs["hidden_states"], + weight=effective_weight, + labels=inputs["labels"], + ) + assert torch.all(output.per_token_logprobs == 0.0) + + +def test_value_prediction_returns_values_everywhere(inputs): + output = value_prediction_function( + hidden_states=inputs["hidden_states"], + weight=inputs["weight"], + labels=inputs["labels"], + metric_reducer=TokenPartial(scale=torch.tensor(1.0)), + ) + # Values are returned for masked-out positions too (client applies its own + # action mask for GAE). + assert_close(output.per_token_logprobs, expected_values(inputs)) + assert output.per_token_loss is None + assert_close(output.loss, torch.tensor(0.0)) + + +def test_value_prediction_zero_loss_keeps_graph(inputs): + hidden_states = inputs["hidden_states"].clone().requires_grad_(True) + output = value_prediction_function( + hidden_states=hidden_states, + weight=inputs["weight"], + labels=inputs["labels"], + ) + output.loss.backward() # must not raise; gradients are exactly zero + assert_close(hidden_states.grad, torch.zeros_like(hidden_states)) + + +def test_value_loss_rejects_bad_weight_shape(inputs): + with pytest.raises(ValueError, match="1, hidden_size"): + value_loss_function( + hidden_states=inputs["hidden_states"], + weight=torch.randn(4, 16), + labels=inputs["labels"], + returns=inputs["returns"], + ) diff --git a/tests/rl/__init__.py b/tests/rl/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/rl/test_advantages.py b/tests/rl/test_advantages.py new file mode 100644 index 00000000..9267a6e4 --- /dev/null +++ b/tests/rl/test_advantages.py @@ -0,0 +1,91 @@ +import math + +import pytest + +from xorl.rl import compute_skip_observation_gae + + +pytestmark = pytest.mark.cpu + + +def reference_gae(rewards, values, gamma, lam, bootstrap): + """Textbook GAE over a dense action-only sequence.""" + n = len(rewards) + advantages = [0.0] * n + next_adv = 0.0 + for t in reversed(range(n)): + next_value = values[t + 1] if t + 1 < n else bootstrap + delta = rewards[t] + gamma * next_value - values[t] + next_adv = delta + gamma * lam * next_adv + advantages[t] = next_adv + return advantages + + +def test_matches_standard_gae_without_mask(): + rewards = [0.0, 0.0, 0.5, 0.0, 1.0] + values = [0.2, -0.1, 0.4, 0.3, 0.1] + gamma, lam = 0.99, 0.95 + adv, ret = compute_skip_observation_gae(rewards, values, gamma=gamma, lam=lam) + expected = reference_gae(rewards, values, gamma, lam, 0.0) + for a, e in zip(adv, expected): + assert math.isclose(a, e, rel_tol=1e-12, abs_tol=1e-12) + for r, a, v in zip(ret, adv, values): + assert math.isclose(r, a + v, rel_tol=1e-12, abs_tol=1e-12) + + +def test_skips_observation_tokens(): + # Layout: [action, obs, obs, action, obs, action] + rewards = [0.0, 0.0, 0.0, 0.0, 0.0, 1.0] + values = [0.5, 9.9, -9.9, 0.4, 9.9, 0.2] + action_mask = [1, 0, 0, 1, 0, 1] + gamma, lam = 1.0, 1.0 + adv, ret = compute_skip_observation_gae(rewards, values, action_mask, gamma=gamma, lam=lam) + + # Equivalent to dense GAE over the action-token subsequence only. + sub_adv = reference_gae([0.0, 0.0, 1.0], [0.5, 0.4, 0.2], gamma, lam, 0.0) + action_indices = [0, 3, 5] + for k, t in enumerate(action_indices): + assert math.isclose(adv[t], sub_adv[k], rel_tol=1e-12) + assert math.isclose(ret[t], sub_adv[k] + values[t], rel_tol=1e-12) + + # Observation tokens carry exactly 0.0 (the packer's mask convention) and + # their (garbage) values never influence any action's advantage. + for t in (1, 2, 4): + assert adv[t] == 0.0 + assert ret[t] == 0.0 + + +def test_observation_values_do_not_leak(): + rewards = [0.0, 0.0, 1.0] + values = [0.3, 123.0, 0.1] + base_adv, _ = compute_skip_observation_gae(rewards, [0.3, 0.0, 0.1], [1, 0, 1]) + leak_adv, _ = compute_skip_observation_gae(rewards, values, [1, 0, 1]) + assert base_adv == leak_adv + + +def test_bootstrap_value_for_truncated_trajectory(): + rewards = [0.0, 0.0] + values = [0.1, 0.2] + adv, ret = compute_skip_observation_gae(rewards, values, gamma=1.0, lam=1.0, bootstrap_value=0.7) + # Last action bootstraps from 0.7: delta = 0 + 0.7 - 0.2 = 0.5 + assert math.isclose(adv[1], 0.5, rel_tol=1e-12) + assert math.isclose(adv[0], (0.2 - 0.1) + 0.5, rel_tol=1e-12) + for r, a, v in zip(ret, adv, values): + assert math.isclose(r, a + v, rel_tol=1e-12) + + +def test_perfect_critic_terminal_reward_yields_zero_advantage(): + # gamma=1, lam=1, terminal reward 1.0; V(s_t)=1.0 at every action token is + # exactly correct, so all advantages vanish and returns equal 1.0. + rewards = [0.0, 0.0, 1.0] + values = [1.0, 1.0, 1.0] + adv, ret = compute_skip_observation_gae(rewards, values, gamma=1.0, lam=1.0) + assert all(math.isclose(a, 0.0, abs_tol=1e-12) for a in adv) + assert all(math.isclose(r, 1.0, rel_tol=1e-12) for r in ret) + + +def test_length_mismatch_raises(): + with pytest.raises(ValueError, match="same length"): + compute_skip_observation_gae([0.0], [0.0, 0.0]) + with pytest.raises(ValueError, match="action_mask"): + compute_skip_observation_gae([0.0, 0.0], [0.0, 0.0], [1]) diff --git a/tests/server/orchestrator/test_packing_value_fields.py b/tests/server/orchestrator/test_packing_value_fields.py new file mode 100644 index 00000000..c11d9ac7 --- /dev/null +++ b/tests/server/orchestrator/test_packing_value_fields.py @@ -0,0 +1,78 @@ +"""Packing of value-model (critic) per-token fields: returns / old_values.""" + +import pytest + +from xorl.data.constants import IGNORE_INDEX +from xorl.server.orchestrator.packing import SequentialPacker + + +pytestmark = [pytest.mark.cpu, pytest.mark.server] + + +def test_returns_pack_in_shifted_client_format(): + """xorl_client format (pre-shifted target_tokens): returns concatenate as-is.""" + data = [ + { + "input_ids": [1, 2, 3], + "target_tokens": [2, 3, 4], + "weights": [0.0, 1.0, 1.0], + "returns": [0.0, 0.7, 0.9], + "old_values": [0.0, 0.6, 0.8], + }, + { + "input_ids": [5, 6], + "target_tokens": [6, 7], + "weights": [1.0, 1.0], + "returns": [0.5, 0.4], + "old_values": [0.5, 0.3], + }, + ] + packer = SequentialPacker(enable_packing=True, log_stats=False, pad_to_multiple_of=1) + batches = packer.pack(data, max_seq_len=100, request_id="value") + + assert len(batches) == 1 + batch = batches[0] + assert batch["returns"] == [[0.0, 0.7, 0.9, 0.5, 0.4]] + assert batch["old_values"] == [[0.0, 0.6, 0.8, 0.5, 0.3]] + # weights=0 masks; returns=0.0 must NOT mask (unlike advantages). + assert batch["target_tokens"] == [[IGNORE_INDEX, 3, 4, 6, 7]] + + +def test_returns_shift_with_hf_style_labels(): + """HF format (unshifted labels): returns/old_values are target-aligned and + shift with labels[1:].""" + data = [ + { + "input_ids": [1, 2, 3, 4], + "labels": [10, 20, 30, 40], + "returns": [0.1, 0.2, 0.3, 0.4], + "old_values": [1.1, 1.2, 1.3, 1.4], + } + ] + packer = SequentialPacker(enable_packing=True, log_stats=False, pad_to_multiple_of=1) + batches = packer.pack(data, max_seq_len=100, request_id="value-hf") + + assert len(batches) == 1 + batch = batches[0] + assert batch["input_ids"] == [[1, 2, 3]] + assert batch["labels"] == [[20, 30, 40]] + assert batch["returns"] == [[0.2, 0.3, 0.4]] + assert batch["old_values"] == [[1.2, 1.3, 1.4]] + + +def test_zero_returns_do_not_mask_labels(): + """A legitimate return of exactly 0.0 keeps its token in the loss.""" + data = [ + { + "input_ids": [1, 2, 3], + "target_tokens": [2, 3, 4], + "weights": [1.0, 1.0, 1.0], + "returns": [0.0, 0.0, 0.0], + } + ] + packer = SequentialPacker(enable_packing=True, log_stats=False, pad_to_multiple_of=1) + batches = packer.pack(data, max_seq_len=100, request_id="value-zero") + + batch = batches[0] + assert batch["target_tokens"] == [[2, 3, 4]] + assert IGNORE_INDEX not in batch["target_tokens"][0] diff --git a/tests/server/test_value_head_arguments.py b/tests/server/test_value_head_arguments.py new file mode 100644 index 00000000..d70bf076 --- /dev/null +++ b/tests/server/test_value_head_arguments.py @@ -0,0 +1,54 @@ +"""Validation of the enable_value_head server-argument surface.""" + +import pytest + +from xorl.server.server_arguments import ServerArguments + + +pytestmark = [pytest.mark.cpu, pytest.mark.server] + + +def _args(**overrides): + base = {"model_path": "Qwen/Qwen3-8B"} + base.update(overrides) + return ServerArguments(**base) + + +def test_value_head_defaults_off(): + args = _args(enable_lora=True) + assert args.enable_value_head is False + assert args.to_config_dict()["lora"]["enable_value_head"] is False + + +def test_value_head_accepts_plain_lora(): + args = _args(enable_lora=True, enable_value_head=True) + assert args.to_config_dict()["lora"]["enable_value_head"] is True + + +def test_value_head_requires_lora(): + with pytest.raises(ValueError, match="plain LoRA"): + _args(enable_value_head=True) + + +def test_value_head_rejects_qlora(): + with pytest.raises(ValueError, match="plain LoRA"): + _args(enable_lora=True, enable_qlora=True, enable_value_head=True) + + +def test_value_head_rejects_pipeline_parallel(): + with pytest.raises(ValueError, match="pipeline parallelism"): + _args(enable_lora=True, enable_value_head=True, pipeline_parallel_size=2) + + +def test_value_head_rejects_fsdp_sharded_lm_head_loss(): + with pytest.raises(ValueError, match="fsdp_sharded_lm_head_loss"): + _args(enable_lora=True, enable_value_head=True, fsdp_sharded_lm_head_loss=True) + + +def test_value_head_rejects_explicit_lm_head_target(): + with pytest.raises(ValueError, match="lm_head excluded"): + _args( + enable_lora=True, + enable_value_head=True, + lora_target_modules=["q_proj", "lm_head"], + ) From 7f9989441956dd231a679cb8223f33f13c8336fa Mon Sep 17 00:00:00 2001 From: Qingyang Wu Date: Mon, 24 Aug 2026 05:59:57 +0000 Subject: [PATCH 2/8] GPU-validate the value head; add state_values, per-session frozen modules, docs GPU validation (2xH100 FSDP2 certification test + live Qwen3-8B E2E) surfaced and fixed two real integration bugs: - The value head now gets its OWN FSDP unit whose forward never runs, so its factors stay sharded DTensors at all times. The previous stay- gathered [norm, lm_head, value_head] grouping broke adapter layout validation after forward-only ops (a no-grad pass leaves the group materialized and FSDP2 reshard() cannot restore it). The loss consumes the folded delta via full_tensor() (direct DTensor lane). - Each rank can only differentiate its own factor shards, so the upstream weight gradient is now all-reduce-summed before reaching the factors (_SumGradAcrossRanks); staged numerators match analytic references exactly on 2 GPUs. Also: - Dedicated state_values field in LossFnOutput: value losses no longer reuse the logprobs channel on the wire. - Per-session frozen_module_patterns (SAO frozen-attention critic): matching factors are skipped at gradient staging, declared AUTHORIZED_ZERO, and provably stay at zero delta (verified on-disk in the E2E run and in the certification test). - Session normalization accepts the SDK's default dropout=0.0 as a no-op (previously every create_lora_training_client call was rejected). - xorl.rl.explained_variance (paper's critic diagnostic) derived from the value_loss moment metrics; value-model docs page; SAO example loop. Part of #84 --- docs/astro.config.mjs | 1 + .../docs/server-training/value-model.md | 76 +++ examples/server/sao_critic/README.md | 31 ++ examples/server/sao_critic/run_sao_loop.py | 115 +++++ src/xorl/distributed/torch_parallelize.py | 15 +- src/xorl/rl/__init__.py | 4 +- src/xorl/rl/advantages.py | 25 + src/xorl/server/api_server/api_types.py | 15 + src/xorl/server/api_server/server.py | 18 +- src/xorl/server/api_server/training_ops.py | 8 +- src/xorl/server/runner/adapters/manager.py | 38 ++ src/xorl/server/runner/model_runner.py | 64 ++- src/xorl/server/session_spec.py | 28 +- tests/rl/test_advantages.py | 37 +- .../runner/test_value_head_adapter_gpu.py | 453 ++++++++++++++++++ tests/server/test_session_spec_value_model.py | 48 ++ 16 files changed, 947 insertions(+), 29 deletions(-) create mode 100644 docs/src/content/docs/server-training/value-model.md create mode 100644 examples/server/sao_critic/README.md create mode 100644 examples/server/sao_critic/run_sao_loop.py create mode 100644 tests/server/runner/test_value_head_adapter_gpu.py create mode 100644 tests/server/test_session_spec_value_model.py diff --git a/docs/astro.config.mjs b/docs/astro.config.mjs index f0f3c21a..caf364f7 100644 --- a/docs/astro.config.mjs +++ b/docs/astro.config.mjs @@ -50,6 +50,7 @@ export default defineConfig({ ], }, { label: "Inference: xorl-sglang", slug: "server-training/sglang" }, + { label: "Value-Model (Critic) Training", slug: "server-training/value-model" }, { label: "Client SDK (xorl-client)", collapsed: true, diff --git a/docs/src/content/docs/server-training/value-model.md b/docs/src/content/docs/server-training/value-model.md new file mode 100644 index 00000000..a2bccf12 --- /dev/null +++ b/docs/src/content/docs/server-training/value-model.md @@ -0,0 +1,76 @@ +--- +title: "Value-Model (Critic) Training" +--- + +xorl supports training a **value model (critic)** alongside policy sessions for PPO/SAO-style RL — the recipe behind single-rollout asynchronous RL ([SAO, arXiv:2607.07508](https://arxiv.org/abs/2607.07508)), where one rollout per prompt replaces GRPO group sampling and a trained critic supplies the advantage baseline. + +## Design + +The critic is **just another LoRA session** on the shared base model, so it costs a LoRA adapter — not a second model. When the server is launched with `enable_value_head: true`, the model carries a scalar value head (`hidden_size → 1`) implemented as a LoRA module with a zero, frozen base weight: the value function lives entirely in per-session adapter factors, every session gets its own independent copy, and a fresh critic predicts exactly `V(s) = 0`. + +Two loss functions become available: + +| `loss_fn` | Op | Inputs (`loss_fn_inputs`) | Output | +|---|---|---|---| +| `value_prediction` | `forward` (no-grad) | `target_tokens`, `weights` | per-token `V(s_t)` in `LossFnOutput.state_values` | +| `value_loss` | `forward_backward` | `target_tokens`, `weights`, `returns`, optional `old_values` | masked squared error; `state_values` + per-token errors in `elementwise_loss` | + +`value_loss` params (via `loss_fn_params`): `vf_coef` (default 1.0) and `clip_range` (default 0.0 = off; with `old_values`, applies the PPO clipped-value objective). + +Like `advantages` and `logprobs`, the `returns` / `old_values` fields are **target-aligned** per-token vectors. Unlike `advantages`, a `returns` value of exactly `0.0` does **not** mask the token — masking comes only from `weights` / `target_tokens`. + +## Server configuration + +```yaml +enable_lora: true +enable_value_head: true +# lm_head must NOT be a LoRA target (a value_loss backward produces no +# lm-head adapter gradients): use an explicit list, or train_unembed: false. +lora_target_modules: ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"] +``` + +Current restrictions: plain LoRA only (no QLoRA), `pipeline_parallel_size: 1`, no `fsdp_sharded_lm_head_loss` / lm-head tensor parallelism. + +The value head never reaches inference: `save_weights_for_sampler` adapters exclude it (SGLang has no such module), while `save_state` training checkpoints keep it so critic sessions resume. + +## The SAO training loop + +```python +from xorl_client import ServiceClient, compute_skip_observation_gae, explained_variance + +svc = ServiceClient(base_url=SERVER) +policy = svc.create_lora_training_client(BASE_MODEL, model_id="policy") +critic = svc.create_lora_training_client(BASE_MODEL, model_id="critic") + +# Per completed rollout (single-rollout, asynchronous): +values_out = critic.forward(datums, loss_fn="value_prediction").result() +values = values_out.loss_fn_outputs[0].state_values.data # V(s_t) per token + +advantages, returns = compute_skip_observation_gae( + rewards, values, action_mask, gamma=1.0, lam=0.95, +) + +for _ in range(K): # faster value update (K=2 in the paper) + fb = critic.forward_backward(with_returns(datums, returns), loss_fn="value_loss").result() + critic.optim_step(critic_adam).result() + +policy.forward_backward(with_advantages(datums, advantages, rollout_logprobs), + loss_fn="policy_loss").result() # DIS: IcePop masking + rollout logprobs +policy.optim_step(policy_adam).result() +``` + +`compute_skip_observation_gae` implements the paper's skip-observation estimator (Eq. 4–5): the Bellman recursion chains across **action tokens only**, so critic noise never propagates through environment-feedback tokens the model did not generate. + +## Monitoring critic health + +Explained variance is the paper's key critic diagnostic (it should climb toward 1.0; near 0 the critic is no better than the mean return). Every `value_loss` step reports sum-composable moments that reduce to global means, from which: + +```python +ev = explained_variance( + value_error_sq_mean=metrics["is_value_error_sq_mean"], + return_mean=metrics["is_return_mean"], + return_sq_mean=metrics["is_return_sq_mean"], +) +``` + +The paper's frozen-attention critic corresponds to restricting the critic's trainable modules to MLP/MoE projections; per-session target masking is tracked as a follow-up — today `lora_target_modules` applies substrate-wide. diff --git a/examples/server/sao_critic/README.md b/examples/server/sao_critic/README.md new file mode 100644 index 00000000..59515441 --- /dev/null +++ b/examples/server/sao_critic/README.md @@ -0,0 +1,31 @@ +# SAO-style critic training (value model) + +Minimal single-rollout RL loop with a trained value model, following +[SAO (arXiv:2607.07508)](https://arxiv.org/abs/2607.07508): one rollout per +prompt, skip-observation GAE from a critic that shares the base model with the +policy as a second LoRA session, and the critic updated K× per policy step. + +## Server + +Launch a LoRA training server with the value head enabled (lm_head must not be +a LoRA target): + +```bash +python -m xorl.server.launcher --mode auto \ + --config examples/server/configs/lora/qwen3_8b_lora.yaml \ + --server.enable_value_head true --api-port 8300 +``` + +## Run + +```bash +python examples/server/sao_critic/run_sao_loop.py \ + --base-url http://127.0.0.1:8300 --model Qwen/Qwen3-8B --steps 20 +``` + +The script uses a toy reward (no environment needed) so the loop mechanics — +`value_prediction` → GAE → `value_loss` ×K → `policy_loss` — can be verified +end-to-end. Explained variance of the critic is printed each step; it should +climb toward 1.0. Swap in real rollouts (a SamplingClient against SGLang, +rollout logprobs, and a real reward) to make this a production loop; see +`docs/server-training/value-model`. diff --git a/examples/server/sao_critic/run_sao_loop.py b/examples/server/sao_critic/run_sao_loop.py new file mode 100644 index 00000000..60392f0e --- /dev/null +++ b/examples/server/sao_critic/run_sao_loop.py @@ -0,0 +1,115 @@ +"""Minimal SAO-style single-rollout RL loop with a trained value model. + +Demonstrates the loop mechanics against a live xorl training server started +with ``enable_value_head: true`` (see README.md). Rollouts are synthetic (a +toy terminal reward on fixed token sequences) so the critic/policy plumbing +can be verified without an environment: + + per rollout: + V(s_t) <- critic.forward(loss_fn="value_prediction") + A_t, R_t <- compute_skip_observation_gae(...) + critic <- K x forward_backward(loss_fn="value_loss") + optim_step + policy <- forward_backward(loss_fn="policy_loss") + optim_step + +Usage: + python run_sao_loop.py --base-url http://127.0.0.1:8300 --model Qwen/Qwen3-8B --steps 20 +""" + +import argparse +import random + +from xorl_client import ServiceClient, compute_skip_observation_gae, explained_variance +from xorl_client.types.adam_params import AdamParams +from xorl_client.types.datum import Datum +from xorl_client.types.model_input import ModelInput + + +def make_rollout(rng: random.Random, length: int = 12): + """A synthetic 'rollout': tokens, an action mask (first 3 tokens are + prompt), and a terminal reward correlated with the token pattern.""" + tokens = [rng.randrange(100, 5000) for _ in range(length)] + action_mask = [0] * 3 + [1] * (length - 4) # target-aligned, length-1 + reward = 1.0 if tokens[-1] % 2 == 0 else 0.0 + return tokens, action_mask, reward + + +def to_datum(tokens, action_mask, extra): + return Datum( + model_input=ModelInput.from_ints(tokens[:-1]), + loss_fn_inputs={ + "target_tokens": tokens[1:], + "weights": [float(m) for m in action_mask], + **extra, + }, + ) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--base-url", required=True) + parser.add_argument("--model", required=True) + parser.add_argument("--steps", type=int, default=20) + parser.add_argument("--critic-updates", type=int, default=2, help="K: critic steps per policy step") + parser.add_argument("--gamma", type=float, default=1.0) + parser.add_argument("--lam", type=float, default=0.95) + parser.add_argument("--policy-lr", type=float, default=1e-5) + parser.add_argument("--critic-lr", type=float, default=5e-5) + args = parser.parse_args() + + svc = ServiceClient(base_url=args.base_url) + policy = svc.create_lora_training_client(args.model, model_id="sao-policy") + critic = svc.create_lora_training_client(args.model, model_id="sao-critic") + policy_adam = AdamParams(learning_rate=args.policy_lr) + critic_adam = AdamParams(learning_rate=args.critic_lr) + rng = random.Random(0) + + for step in range(args.steps): + tokens, action_mask, reward = make_rollout(rng) + base = to_datum(tokens, action_mask, {}) + + # 1) Critic predicts V(s_t) for the rollout. + pred = critic.forward([base], loss_fn="value_prediction").result() + values = list(pred.loss_fn_outputs[0].state_values.data) + + # 2) Skip-observation GAE across action tokens (terminal reward). + rewards = [0.0] * (len(values) - 1) + [reward] + advantages, returns = compute_skip_observation_gae(rewards, values, action_mask, gamma=args.gamma, lam=args.lam) + + # 3) Faster value update: K critic steps per policy step. + ev = float("nan") + for _ in range(args.critic_updates): + fb = critic.forward_backward( + [to_datum(tokens, action_mask, {"returns": returns})], + loss_fn="value_loss", + ).result() + critic.optim_step(critic_adam).result() + metrics = fb.metrics + ev = explained_variance( + value_error_sq_mean=metrics.get("is_value_error_sq_mean", float("nan")), + return_mean=metrics.get("is_return_mean", float("nan")), + return_sq_mean=metrics.get("is_return_sq_mean", float("nan")), + ) + + # 4) Policy step. With real rollouts, ``logprobs`` are the sampler's + # behavior logprobs (DIS: the ratio is policy/rollout); the synthetic + # stand-in just exercises the wire format. + n = len(tokens) - 1 + policy.forward_backward( + [ + to_datum( + tokens, + action_mask, + {"advantages": advantages, "logprobs": [-2.0] * n}, + ) + ], + loss_fn="policy_loss", + ).result() + policy.optim_step(policy_adam).result() + + print(f"step {step:3d} reward {reward:.1f} critic EV {ev:+.3f}") + + print("done") + + +if __name__ == "__main__": + main() diff --git a/src/xorl/distributed/torch_parallelize.py b/src/xorl/distributed/torch_parallelize.py index dfd75366..1f1207e5 100644 --- a/src/xorl/distributed/torch_parallelize.py +++ b/src/xorl/distributed/torch_parallelize.py @@ -626,16 +626,21 @@ def _experts_shard_placement_fn(param): fully_shard(lm_head_mod, **fsdp_kwargs) logger.info_rank0("Using FSDP-sharded lm_head loss over the FSDP group.") elif not pp_enabled and fsdp_kwargs.get("reshard_after_forward", True) is not False: - # A scalar value head (critic) is consumed like lm_head: the loss reads - # its weight AFTER the model forward, so it must live in the same - # stay-gathered unit that norm.forward() unshards. - value_head_mod = getattr(model, "value_head", None) - last_modules = [m for m in [norm_mod, lm_head_mod, value_head_mod] if m is not None] + last_modules = [m for m in [norm_mod, lm_head_mod] if m is not None] if last_modules: last_fsdp_kwargs = dict(fsdp_kwargs) last_fsdp_kwargs["reshard_after_forward"] = False fully_shard(last_modules, **last_fsdp_kwargs) + # A scalar value head (critic) gets its own FSDP unit whose forward is + # never invoked: its parameters therefore remain sharded DTensors at all + # times (stable for adapter layout validation across forward-only ops), + # and the loss consumes the folded LoRA delta through the direct-DTensor + # lane (full_tensor of a Partial matmul), like direct lm-head factors. + value_head_mod = getattr(model, "value_head", None) + if value_head_mod is not None and not pp_enabled: + fully_shard(value_head_mod, **fsdp_kwargs) + # shard root model # Collect all _skip_fsdp experts params so they're also ignored by the # root-level fully_shard (layer-level already ignores them above, but the diff --git a/src/xorl/rl/__init__.py b/src/xorl/rl/__init__.py index 79f4c830..70bd2ec9 100644 --- a/src/xorl/rl/__init__.py +++ b/src/xorl/rl/__init__.py @@ -1,6 +1,6 @@ """RL utilities: advantage estimation for value-model (critic) training.""" -from xorl.rl.advantages import compute_skip_observation_gae +from xorl.rl.advantages import compute_skip_observation_gae, explained_variance -__all__ = ["compute_skip_observation_gae"] +__all__ = ["compute_skip_observation_gae", "explained_variance"] diff --git a/src/xorl/rl/advantages.py b/src/xorl/rl/advantages.py index b0c4cad9..a9ba956f 100644 --- a/src/xorl/rl/advantages.py +++ b/src/xorl/rl/advantages.py @@ -22,6 +22,7 @@ from __future__ import annotations +import math from typing import List, Optional, Sequence, Tuple @@ -78,3 +79,27 @@ def compute_skip_observation_gae( next_value = float(values[t]) return advantages, returns + + +def explained_variance( + value_error_sq_mean: float, + return_mean: float, + return_sq_mean: float, +) -> float: + """Critic explained variance from the ``value_loss`` moment metrics. + + ``EV = 1 - E[(R - V)^2] / Var(R)`` with ``Var(R) = E[R^2] - E[R]^2``. + The three inputs are exactly the (globally normalized) ``value_error_sq_mean``, + ``return_mean``, and ``return_sq_mean`` metrics a ``value_loss`` + forward_backward reports, so EV composes correctly across micro-batches + and ranks. EV is the paper's key critic-health diagnostic (SAO Fig. 4a): + it should climb toward 1.0 as the critic converges; near 0 the critic is + no better than predicting the mean return. + + Returns NaN when the return distribution is (numerically) constant, where + explained variance is undefined. + """ + return_variance = return_sq_mean - return_mean * return_mean + if not math.isfinite(return_variance) or return_variance <= 1e-12: + return float("nan") + return 1.0 - value_error_sq_mean / return_variance diff --git a/src/xorl/server/api_server/api_types.py b/src/xorl/server/api_server/api_types.py index f358540e..ab669fbb 100644 --- a/src/xorl/server/api_server/api_types.py +++ b/src/xorl/server/api_server/api_types.py @@ -175,6 +175,10 @@ class LossFnOutput(BaseModel): loss: Optional[float] = Field(default=None, description="Loss value (for backward compatibility)") logprobs: Optional[TensorData] = Field(default=None, description="Per-token log probabilities") elementwise_loss: Optional[TensorData] = Field(default=None, description="Per-token cross entropy loss") + state_values: Optional[TensorData] = Field( + default=None, + description="Per-token value predictions V(s_t) from the value_loss / value_prediction loss functions", + ) k3: Optional[float] = Field(default=None, description="Per-sample K3 KL divergence estimate") token_diagnostics: Optional[Dict[str, Any]] = Field( default=None, @@ -218,6 +222,14 @@ class LoRAConfigRequest(BaseModel): validation_alias=AliasChoices("lora_alpha", "alpha"), description="LoRA alpha override. Accepts Tinker's alpha alias.", ) + frozen_module_patterns: Optional[List[str]] = Field( + default=None, + description=( + "Substring patterns of adapter modules this session must NOT train (e.g. " + "['q_proj', 'k_proj', 'v_proj', 'o_proj'] for a frozen-attention critic). Matching adapter " + "factors keep their zero-delta initialization and never receive gradient updates." + ), + ) @model_validator(mode="before") @classmethod @@ -262,6 +274,9 @@ class LoRARuntimeConfig(BaseModel): lora_rank: int = Field(..., description="LoRA rank") lora_alpha: int = Field(..., description="LoRA alpha") + frozen_module_patterns: Optional[List[str]] = Field( + default=None, description="Adapter-module substring patterns this session does not train" + ) class OptimizerRuntimeConfig(BaseModel): diff --git a/src/xorl/server/api_server/server.py b/src/xorl/server/api_server/server.py index aa8a7099..5d0cbc0e 100644 --- a/src/xorl/server/api_server/server.py +++ b/src/xorl/server/api_server/server.py @@ -279,20 +279,28 @@ def _flatten_api_data(data_list) -> List[Dict[str, Any]]: return result @staticmethod - def _build_loss_fn_outputs(result: Dict[str, Any]): - """Build (loss_fn_outputs, loss_fn_output_type) from engine result.""" + def _build_loss_fn_outputs(result: Dict[str, Any], loss_fn: Optional[str] = None): + """Build (loss_fn_outputs, loss_fn_output_type) from engine result. + + The runner ships every per-token vector through one generic channel; + ``loss_fn`` names its semantics at the API boundary: for the value + losses the channel carries V(s_t) and is exposed as ``state_values``. + """ per_sample_outputs = result.get("per_sample_outputs", []) per_sample_k3 = result.get("per_sample_k3", []) + value_output = loss_fn in {"value_loss", "value_prediction"} if per_sample_outputs: outputs = [] for i, sample in enumerate(per_sample_outputs): - logprobs = sample.get("logprobs", []) + per_token = sample.get("logprobs", []) elementwise_loss = sample.get("elementwise_loss", []) k3_val = per_sample_k3[i] if i < len(per_sample_k3) else None + per_token_tensor = TensorData(data=per_token, dtype="float32", shape=[len(per_token)]) outputs.append( LossFnOutput( - logprobs=TensorData(data=logprobs, dtype="float32", shape=[len(logprobs)]), + state_values=per_token_tensor if value_output else None, + logprobs=None if value_output else per_token_tensor, elementwise_loss=TensorData( data=elementwise_loss, dtype="float32", shape=[len(elementwise_loss)] ), @@ -300,7 +308,7 @@ def _build_loss_fn_outputs(result: Dict[str, Any]): token_diagnostics=sample.get("token_diagnostics"), ) ) - return outputs, "CrossEntropyLossReturn" + return outputs, "ValueOutput" if value_output else "CrossEntropyLossReturn" # When no per-sample outputs, but we have per_sample_k3, create one output per sample if per_sample_k3: diff --git a/src/xorl/server/api_server/training_ops.py b/src/xorl/server/api_server/training_ops.py index 0b95b494..e39b373c 100644 --- a/src/xorl/server/api_server/training_ops.py +++ b/src/xorl/server/api_server/training_ops.py @@ -324,7 +324,9 @@ async def forward_backward(self, request: ForwardBackwardRequest) -> ForwardBack # Sanitize NaN/Inf values for JSON serialization result = _sanitize_nan_to_zero(result) - loss_fn_outputs, loss_fn_output_type = self._build_loss_fn_outputs(result) + loss_fn_outputs, loss_fn_output_type = self._build_loss_fn_outputs( + result, loss_fn=request.forward_backward_input.loss_fn + ) # Build metrics with tinker naming convention total_loss = result.get("loss", 0.0) @@ -429,7 +431,9 @@ async def forward(self, request: ForwardRequest) -> ForwardResponse: # Extract results (same format as forward_backward) result = _sanitize_nan_to_zero(output.outputs[0] if output.outputs else {}) - loss_fn_outputs, loss_fn_output_type = self._build_loss_fn_outputs(result) + loss_fn_outputs, loss_fn_output_type = self._build_loss_fn_outputs( + result, loss_fn=request.forward_input.loss_fn + ) total_loss = result.get("loss", 0.0) valid_tokens = result.get("valid_tokens", 1) diff --git a/src/xorl/server/runner/adapters/manager.py b/src/xorl/server/runner/adapters/manager.py index 2aea8674..be303514 100644 --- a/src/xorl/server/runner/adapters/manager.py +++ b/src/xorl/server/runner/adapters/manager.py @@ -1197,6 +1197,36 @@ def _current_layout_world_identity(self) -> Tuple[int, Tuple[int, ...]]: def _session_rank(session_spec: Dict[str, Any]) -> int: return int(session_spec["lora_config"]["lora_rank"]) + def frozen_parameter_fqns(self, session_spec: Dict[str, Any]) -> frozenset: + """Adapter-factor FQNs this session must never train. + + ``frozen_module_patterns`` are substring matches against the adapter + parameter names (e.g. ``"q_proj"`` freezes every attention query + factor for this session). Matching factors keep their zero-delta + initialization: they are skipped at gradient staging, declared + AUTHORIZED_ZERO in the ownership plan, and their optimizer moments + stay zero, so their per-session values never change. + """ + patterns = tuple(session_spec.get("lora_config", {}).get("frozen_module_patterns") or ()) + if not patterns: + return frozenset() + return frozenset(name for name in self._lora_param_names if any(pattern in name for pattern in patterns)) + + def _validate_frozen_module_patterns(self, session_spec: Dict[str, Any]) -> None: + patterns = tuple(session_spec.get("lora_config", {}).get("frozen_module_patterns") or ()) + if not patterns: + return + unmatched = [pattern for pattern in patterns if not any(pattern in name for name in self._lora_param_names)] + if unmatched: + raise ValueError( + f"frozen_module_patterns {unmatched!r} match no adapter parameters; " + f"adapter parameter names look like: {self._lora_param_names[:6]!r}" + ) + if len(self.frozen_parameter_fqns(session_spec)) == len(self._lora_param_names): + raise ValueError( + "frozen_module_patterns freeze every adapter parameter; the session would have nothing to train" + ) + @staticmethod def _session_alpha(session_spec: Dict[str, Any]) -> int: return int(session_spec["lora_config"]["lora_alpha"]) @@ -2218,6 +2248,7 @@ def register_adapter( self._validate_exact_glm_session_contract(session_spec) self._validate_session_rank_against_model_capacity(session_spec) + self._validate_frozen_module_patterns(session_spec) session_rank = self._session_rank(session_spec) session_alpha = self._session_alpha(session_spec) optimizer_config = session_spec["optimizer_config"] @@ -2591,10 +2622,17 @@ def stage_gradient_numerators( canonical_parameter_name(name): parameter for name, parameter in self.model.named_parameters() } layouts = {layout.fqn: layout for layout in state.tensor_layouts.values()} + frozen_fqns = self.frozen_parameter_fqns(state.session_spec) staged_fqns: list[str] = [] staged_numerators: dict[str, torch.Tensor] = {} for item in plan.parameters: parameter = named_parameters[item.fqn] + if item.fqn in frozen_fqns: + # Session-frozen factor: the backward may have produced a + # gradient (the module sits in the graph), but this session + # never trains it. Skipping the stage leaves its numerator at + # zero, so the optimizer step is an exact no-op for it. + continue if parameter.grad is None: if item.requires_local_gradient: raise AdapterGradientOwnershipError(f"Required adapter gradient is absent for {item.fqn!r}") diff --git a/src/xorl/server/runner/model_runner.py b/src/xorl/server/runner/model_runner.py index 86f74f89..b20459c8 100644 --- a/src/xorl/server/runner/model_runner.py +++ b/src/xorl/server/runner/model_runner.py @@ -446,6 +446,27 @@ def _sp_allreduce_kl_metrics( return metrics +class _SumGradAcrossRanks(torch.autograd.Function): + """Identity forward; all-reduce(SUM) the gradient across all ranks. + + Used for the directly-consumed value-head weight: every rank's local loss + share depends on the full replicated weight, but backward on a rank can + only produce that rank's contribution. Summing the upstream weight + gradient makes the sharded factor gradients complete, matching the + raw-sum loss contract (normalization happens once at optim_step). + """ + + @staticmethod + def forward(ctx, tensor: torch.Tensor) -> torch.Tensor: + return tensor + + @staticmethod + def backward(ctx, grad: torch.Tensor) -> torch.Tensor: + grad = grad.contiguous() + dist.all_reduce(grad, op=dist.ReduceOp.SUM) + return grad + + class ModelRunner: """ ModelRunner handles model operations on distributed GPUs using xorl infrastructure. @@ -1191,6 +1212,7 @@ def _walk_module_tree(module: nn.Module, *, inherited_fsdp: bool = False, direct output_group_size = len(local_group_memberships.get("output_projection_replica", (0,))) declarations: dict[str, ParameterOwnershipDeclaration] = {} guard_payloads: dict[str, dict[str, Any]] = {} + session_frozen_fqns = self._adapter_manager.frozen_parameter_fqns(state.session_spec) for name, layout in state.tensor_layouts.items(): parameter = named_parameters[name] if id(parameter) in direct_parameter_ids: @@ -1324,10 +1346,12 @@ def _walk_module_tree(module: nn.Module, *, inherited_fsdp: bool = False, direct guard_payload.update(expert_guard_by_parameter_id.get(id(parameter), {})) guard_payload.update(exact_lm_head_guard_by_parameter_id.get(id(parameter), {})) guard_payloads[name] = guard_payload - # Policy-loss steps never touch the value head, so its factors may - # legitimately have no gradient; every other adapter factor keeps - # the strict presence contract. + # Two classes of factors may legitimately lack a gradient: the + # value head during policy-loss steps, and factors the session + # explicitly froze via frozen_module_patterns. Everything else + # keeps the strict presence contract. is_value_head_param = name.startswith("value_head.") or ".value_head." in name + authorized_zero = is_value_head_param or name in session_frozen_fqns declarations[name] = ParameterOwnershipDeclaration( topology=topology, producer=producer, @@ -1336,7 +1360,7 @@ def _walk_module_tree(module: nn.Module, *, inherited_fsdp: bool = False, direct capture_domains=tuple(capture), pending_domains=tuple(pending), presence=GradientPresencePolicy.AUTHORIZED_ZERO - if is_value_head_param + if authorized_zero else GradientPresencePolicy.REQUIRED_IF_ACTIVE, config_guard_fingerprint=self._adapter_gradient_hash(guard_payload), config_guard_fields=tuple(sorted(guard_payload.items())), @@ -2132,19 +2156,37 @@ def _get_effective_lm_head_weight(self): return self._get_effective_lm_head_weight_for(self.model.lm_head) def _get_effective_value_head_weight(self): - """Get the scalar value head's weight, merging its LoRA delta. - - The value head is a LoRA module with a zero frozen base weight, so the - effective weight IS the folded adapter delta; consuming it here (like - the lm_head weight) routes gradients through the same direct-output- - projection ownership lane. + """Get the scalar value head's effective weight (its folded LoRA delta). + + The value head has a zero, frozen base weight by contract, so the + effective weight IS the adapter delta ``B[:, :r] @ A[:r] * (alpha/r)``. + Its FSDP unit never runs a forward, so the factors are sharded + DTensors: the delta matmul contracts over the sharded rank dim into a + Partial placement, and ``full_tensor()`` materializes the replicated + weight while keeping autograd routed back into the sharded factors — + the direct-output-projection ownership lane completes those Partial + gradients at capture time. + + Each rank's loss is a distinct share of the global raw-sum objective, + but both factor shards influence every rank's loss through the + replicated weight. A rank can only differentiate its own shard, so + the upstream weight gradient must be summed across ranks BEFORE it + reaches the factors (the FSDP post-backward hook does this for + module-managed adapters; direct consumption owns it here). """ + from torch.distributed.tensor import DTensor # noqa: PLC0415 + value_head = getattr(self.model, "value_head", None) if value_head is None: raise ValueError( "value_loss/value_prediction require a value head; start the server with enable_value_head=true" ) - return self._get_effective_lm_head_weight_for(value_head) + delta = value_head.get_delta_weight() + if isinstance(delta, DTensor): + delta = delta.full_tensor() + if dist.is_available() and dist.is_initialized() and dist.get_world_size() > 1: + delta = _SumGradAcrossRanks.apply(delta) + return delta @staticmethod def _get_effective_lm_head_weight_for(lm_head): diff --git a/src/xorl/server/session_spec.py b/src/xorl/server/session_spec.py index 7f313146..83a2f7a1 100644 --- a/src/xorl/server/session_spec.py +++ b/src/xorl/server/session_spec.py @@ -110,14 +110,31 @@ def normalize_lora_runtime_config( lora_config = _normalize_lora_config_keys(raw_lora_config) server_lora_config = dict(server_lora_config or {}) - for key in sorted(set(lora_config) - {"lora_rank", "lora_alpha"}): + # The tinker-style client SDK always includes ``dropout`` in its LoRA + # config payload. The server has no LoRA dropout, so 0.0 (the SDK + # default) is exactly the server behavior — accept it as a no-op and + # reject only a nonzero request. + dropout = lora_config.pop("dropout", None) + if dropout: + raise ValueError(f"LoRA dropout is not supported by the training server (requested dropout={dropout})") + + for key in sorted(set(lora_config) - {"lora_rank", "lora_alpha", "frozen_module_patterns"}): server_value = server_lora_config.get(key) if lora_config[key] != server_value: raise ValueError( - "Per-session LoRA config may only override rank and alpha. " + "Per-session LoRA config may only override rank, alpha, and frozen_module_patterns. " f"Unsupported override for {key!r}: {lora_config[key]!r} (server={server_value!r})." ) + frozen_module_patterns = lora_config.get("frozen_module_patterns") + if frozen_module_patterns is not None: + if not isinstance(frozen_module_patterns, (list, tuple)) or not all( + isinstance(pattern, str) and pattern for pattern in frozen_module_patterns + ): + raise ValueError("frozen_module_patterns must be a list of non-empty strings") + # Sorted + deduped so equivalent sessions hash to identical specs. + frozen_module_patterns = sorted(set(frozen_module_patterns)) + lora_rank = int(lora_config.get("lora_rank", default_rank)) lora_alpha = int(lora_config.get("lora_alpha", default_alpha)) @@ -131,10 +148,15 @@ def normalize_lora_runtime_config( "Increase server.max_lora_rank to support this session." ) - return { + normalized: Dict[str, Any] = { "lora_rank": lora_rank, "lora_alpha": lora_alpha, } + # Omitted (or empty) keeps legacy session specs byte-identical, so + # existing checkpoints and adapter-generation hashes are unaffected. + if frozen_module_patterns: + normalized["frozen_module_patterns"] = list(frozen_module_patterns) + return normalized def normalize_optimizer_config( diff --git a/tests/rl/test_advantages.py b/tests/rl/test_advantages.py index 9267a6e4..66ad09d7 100644 --- a/tests/rl/test_advantages.py +++ b/tests/rl/test_advantages.py @@ -1,8 +1,9 @@ import math +import random import pytest -from xorl.rl import compute_skip_observation_gae +from xorl.rl import compute_skip_observation_gae, explained_variance pytestmark = pytest.mark.cpu @@ -89,3 +90,37 @@ def test_length_mismatch_raises(): compute_skip_observation_gae([0.0], [0.0, 0.0]) with pytest.raises(ValueError, match="action_mask"): compute_skip_observation_gae([0.0, 0.0], [0.0, 0.0], [1]) + + +def test_explained_variance_matches_direct_computation(): + rng = random.Random(3) + returns = [rng.gauss(0.5, 1.3) for _ in range(500)] + values = [r + rng.gauss(0.0, 0.4) for r in returns] # decent critic + n = len(returns) + mean_r = sum(returns) / n + var_r = sum((r - mean_r) ** 2 for r in returns) / n + mse = sum((v - r) ** 2 for v, r in zip(values, returns)) / n + direct = 1.0 - mse / var_r + + ev = explained_variance( + value_error_sq_mean=mse, + return_mean=mean_r, + return_sq_mean=sum(r * r for r in returns) / n, + ) + assert math.isclose(ev, direct, rel_tol=1e-9) + assert 0.5 < ev < 1.0 + + +def test_explained_variance_perfect_and_mean_critic(): + returns = [0.0, 1.0, 2.0, 3.0] + mean_r = 1.5 + sq_mean = sum(r * r for r in returns) / 4 + # Perfect critic: zero error -> EV = 1. + assert explained_variance(0.0, mean_r, sq_mean) == pytest.approx(1.0) + # Mean-predicting critic: error variance == return variance -> EV = 0. + var_r = sq_mean - mean_r**2 + assert explained_variance(var_r, mean_r, sq_mean) == pytest.approx(0.0) + + +def test_explained_variance_undefined_for_constant_returns(): + assert math.isnan(explained_variance(0.1, 1.0, 1.0)) diff --git a/tests/server/runner/test_value_head_adapter_gpu.py b/tests/server/runner/test_value_head_adapter_gpu.py new file mode 100644 index 00000000..a591123d --- /dev/null +++ b/tests/server/runner/test_value_head_adapter_gpu.py @@ -0,0 +1,453 @@ +"""Real-FSDP certification of the scalar LoRA value head (critic) lifecycle. + +Validates, against analytically reconstructed references on 2 GPUs: + +- ownership compile: ``value_head.*`` factors classify as + DIRECT_OUTPUT_PROJECTION with AUTHORIZED_ZERO presence; trunk factors stay + DENSE_REPLICATED / REQUIRED_IF_ACTIVE; +- a ``value_loss``-style step (loss consumes the folded value-head weight the + way the server does) produces correct staged numerators and a correct Adam + step for BOTH trunk and value-head factors; +- a policy-style step on the SAME session (backward never touches the value + head) is accepted (AUTHORIZED_ZERO), transports exact-zero value-head + gradients (no stale-numerator leakage from the previous epoch), and matches + a zero-grad reference Adam step; +- a forward-only (no-grad) value_prediction pass between training steps + leaves the adapter layout contract intact (the value head's own FSDP unit + never unshards, so its factors stay sharded DTensors at all times); +- a second critic session with a smaller active rank slices the value head + correctly and stays isolated from the first session. +""" + +from __future__ import annotations + +import datetime +import os +import shutil +import tempfile +from pathlib import Path + +import pytest +import torch +import torch.distributed as dist +import torch.nn as nn +from torch.distributed.fsdp import fully_shard +from torch.distributed.tensor import Shard + +from xorl.data.constants import IGNORE_INDEX +from xorl.distributed.parallel_state import get_parallel_state, init_parallel_state +from xorl.lora.modules.linear import LoraLinear +from xorl.ops.loss import TokenPartial, value_loss_function, value_prediction_function +from xorl.server.runner.adapters.gradient_ownership import ( + GradientPresencePolicy, + GradientScaleState, + TopologyFamily, +) +from xorl.server.runner.adapters.manager import LoRAAdapterManager +from xorl.server.runner.model_runner import ModelRunner + + +pytestmark = [pytest.mark.server, pytest.mark.gpu] + +_HIDDEN = 4 + + +def _session_spec(rank: int, frozen_module_patterns: list[str] | None = None) -> dict: + lora_config = {"lora_rank": rank, "lora_alpha": rank, "seed": 37} + if frozen_module_patterns: + lora_config["frozen_module_patterns"] = frozen_module_patterns + return { + "base_model": "value-head-certification-fixture", + "is_lora": True, + "lora_config": lora_config, + "optimizer_config": { + "type": "adamw", + "learning_rate": 1e-2, + "weight_decay": 0.0, + "optimizer_dtype": "fp32", + "betas": [0.9, 0.95], + "eps": 1e-8, + "optimizer_kwargs": {}, + }, + } + + +def _logical_slice(value: torch.Tensor, layout) -> torch.Tensor: + slices = tuple( + slice(offset, offset + size) + for offset, size in zip(layout.active_global_offset, layout.active_storage_shape, strict=True) + ) + return value[slices].contiguous() + + +def _gather_logical_slots(state) -> dict[str, torch.Tensor]: + payload = { + name: { + "shape": layout.logical_shape, + "offset": layout.active_global_offset, + "storage_shape": layout.active_storage_shape, + "value": parameter.detach().cpu().contiguous(), + } + for name, parameter in state.local_params.items() + for layout in (state.tensor_layouts[name],) + } + gathered = [None] * dist.get_world_size() + dist.all_gather_object(gathered, payload) + result = {name: torch.zeros(item["shape"], dtype=item["value"].dtype) for name, item in payload.items()} + coverage = {name: torch.zeros_like(value, dtype=torch.int64) for name, value in result.items()} + for rank_payload in gathered: + for name, item in rank_payload.items(): + if item["value"].numel() == 0: + continue + slices = tuple( + slice(offset, offset + size) for offset, size in zip(item["offset"], item["storage_shape"], strict=True) + ) + result[name][slices].copy_(item["value"]) + coverage[name][slices].add_(1) + assert all(torch.all(mask == 1) for mask in coverage.values()) + return result + + +def _critic_objective( + reference: dict[str, nn.Parameter], + trunk_base: torch.Tensor, + inputs: list[torch.Tensor], + labels: list[torch.Tensor], + returns: list[torch.Tensor], + *, + with_value_head: bool, +) -> torch.Tensor: + """Unsharded mathematical objective mirroring the worker's loss exactly.""" + eff_trunk = trunk_base + reference["trunk.lora_B"] @ reference["trunk.lora_A"] + objective = torch.zeros((), device=trunk_base.device) + for x, y, r in zip(inputs, labels, returns, strict=True): + h = torch.nn.functional.linear(x, eff_trunk) + if with_value_head: + eff_vh = reference["value_head.lora_B"] @ reference["value_head.lora_A"] + values = torch.nn.functional.linear(h.float(), eff_vh.float()).squeeze(-1) + valid = (y != IGNORE_INDEX).float() + objective = objective + (0.5 * ((values - r).square() * valid)).sum() + else: + objective = objective + h.square().sum() + return objective + + +def _reference_gradients( + reference: dict[str, nn.Parameter], + trunk_base: torch.Tensor, + inputs, + labels, + returns, + *, + with_value_head: bool, +) -> dict[str, torch.Tensor]: + objective = _critic_objective(reference, trunk_base, inputs, labels, returns, with_value_head=with_value_head) + names = list(reference) + gradients = torch.autograd.grad(objective, [reference[name] for name in names], allow_unused=True) + return { + name: (gradient if gradient is not None else torch.zeros_like(reference[name])) + for name, gradient in zip(names, gradients, strict=True) + } + + +def _reference_step(optimizer, parameters, gradients, *, clip: float) -> float: + norm = float(torch.sqrt(sum(gradient.float().square().sum() for gradient in gradients.values()))) + coefficient = min(1.0, clip / (norm + 1e-6)) + for name, parameter in parameters.items(): + parameter.grad = gradients[name].mul(coefficient) + optimizer.step() + optimizer.zero_grad(set_to_none=True) + return norm + + +def _assert_local_state_matches_reference(state, parameters) -> None: + for name, local_parameter in state.local_params.items(): + layout = state.tensor_layouts[name] + torch.testing.assert_close( + local_parameter.float(), + _logical_slice(parameters[name].detach(), layout).float(), + rtol=5e-4, + atol=2e-4, + msg=name, + ) + + +def _run_worker() -> None: + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + device = torch.device("cuda", local_rank) + dist.init_process_group("nccl", device_id=device, timeout=datetime.timedelta(seconds=30)) + rank = dist.get_rank() + world = dist.get_world_size() + assert world == 2 + try: + init_parallel_state(dp_size=2, dp_shard_size=2, device_type="cuda") + + class _Model(nn.Module): + """Trunk adapter + scalar LoRA value head. Like production, the + value head has its own FSDP unit whose forward never runs: its + factors stay sharded DTensors and the loss consumes the folded + delta via full_tensor() (direct DTensor lane).""" + + def __init__(self) -> None: + super().__init__() + self.trunk = LoraLinear(_HIDDEN, _HIDDEN, r=4, lora_alpha=4, device=device, dtype=torch.float32) + self.value_head = LoraLinear(_HIDDEN, 1, r=4, lora_alpha=4, device=device, dtype=torch.float32) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.trunk(x) + + model = _Model() + with torch.no_grad(): + model.trunk.weight.copy_( + torch.linspace(-0.2, 0.3, model.trunk.weight.numel(), device=device).reshape_as(model.trunk.weight) + ) + model.value_head.weight.zero_() + model.trunk.weight.requires_grad_(False) + model.value_head.weight.requires_grad_(False) + trunk_base = model.trunk.weight.detach().clone() + + lora_b_ids = {id(model.trunk.lora_B), id(model.value_head.lora_B)} + + def _shard_placement(parameter: nn.Parameter) -> Shard: + return Shard(1) if id(parameter) in lora_b_ids else Shard(0) + + mesh = get_parallel_state().fsdp_mesh + fully_shard(model.trunk, mesh=mesh, shard_placement_fn=_shard_placement) + model.trunk.set_gradient_divide_factor(1.0) + # Production gives the value head its own unit whose forward never + # runs, so its factors stay sharded DTensors at all times. + fully_shard(model.value_head, mesh=mesh, shard_placement_fn=_shard_placement) + model.value_head.set_gradient_divide_factor(1.0) + fully_shard(model, mesh=mesh, shard_placement_fn=_shard_placement) + model.set_gradient_divide_factor(1.0) + + root = Path(tempfile.gettempdir()) / f"value-head-adapter-{os.environ['MASTER_PORT']}" + if rank == 0: + shutil.rmtree(root, ignore_errors=True) + dist.barrier() + manager = LoRAAdapterManager( + model, + device=device, + checkpoint_dir=str(root), + auto_save_on_eviction=False, + optimizer_fused=False, + optimizer_dtype="fp32", + weight_decay=0.0, + ) + assert set(manager._lora_param_names) == { + "trunk.lora_A", + "trunk.lora_B", + "value_head.lora_A", + "value_head.lora_B", + } + runner = ModelRunner.__new__(ModelRunner) + runner.model = model + runner._adapter_manager = manager + + session_ranks = {"critic": 4, "critic_r2": 2, "critic_frozen": 4} + session_frozen = {"critic_frozen": ["trunk"]} + references = {} + for model_id, active_rank in session_ranks.items(): + manager.register_adapter( + model_id, + session_spec=_session_spec(active_rank, session_frozen.get(model_id)), + initialize_fresh=True, + ) + runner._compile_registered_adapter_gradient_ownership(model_id) + state = manager.get_adapter_state(model_id) + + by_name = {item.fqn: item for item in state.gradient_ownership_plan.parameters} + assert by_name["trunk.lora_A"].topology is TopologyFamily.DENSE_REPLICATED + assert by_name["value_head.lora_A"].topology is TopologyFamily.DIRECT_OUTPUT_PROJECTION + assert by_name["value_head.lora_B"].topology is TopologyFamily.DIRECT_OUTPUT_PROJECTION + assert by_name["value_head.lora_A"].presence is GradientPresencePolicy.AUTHORIZED_ZERO + assert by_name["value_head.lora_B"].presence is GradientPresencePolicy.AUTHORIZED_ZERO + trunk_presence = ( + GradientPresencePolicy.AUTHORIZED_ZERO + if model_id in session_frozen + else GradientPresencePolicy.REQUIRED_IF_ACTIVE + ) + assert by_name["trunk.lora_A"].presence is trunk_presence + assert by_name["trunk.lora_B"].presence is trunk_presence + + logical = _gather_logical_slots(state) + reference_parameters = { + name: nn.Parameter(value.to(device), requires_grad=True) for name, value in logical.items() + } + references[model_id] = ( + reference_parameters, + torch.optim.AdamW( + reference_parameters.values(), lr=1e-2, betas=(0.9, 0.95), eps=1e-8, weight_decay=0.0 + ), + ) + + def _run_step(model_id: str, step_tag: float, *, with_value_head: bool) -> None: + state = manager.get_adapter_state(model_id) + reference_parameters, reference_optimizer = references[model_id] + other_ids = [other for other in session_ranks if other != model_id] + others_before = { + other: { + name: value.detach().clone() + for name, value in manager.get_adapter_state(other).local_params.items() + } + for other in other_ids + } + manager.prepare_forward(model_id) + assert model.value_head.active_r == session_ranks[model_id] + + local_input = ( + torch.arange(2 * _HIDDEN, device=device, dtype=torch.float32).reshape(2, _HIDDEN) / 17.0 + + rank * 0.13 + + step_tag + ) + local_labels = torch.tensor([[IGNORE_INDEX, 5]] if rank == 0 else [[7, 9]], device=device) + local_returns = torch.tensor([[0.3, -0.4]], device=device, dtype=torch.float32) + rank * 0.21 + step_tag + + all_inputs = [torch.empty_like(local_input) for _ in range(world)] + dist.all_gather(all_inputs, local_input) + all_labels = [torch.empty_like(local_labels) for _ in range(world)] + dist.all_gather(all_labels, local_labels) + all_returns = [torch.empty_like(local_returns) for _ in range(world)] + dist.all_gather(all_returns, local_returns) + + if with_value_head: + denominator = float(sum(int((y != IGNORE_INDEX).sum().item()) for y in all_labels)) + else: + denominator = float(sum(x.shape[0] for x in all_inputs)) + + raw_reference = _reference_gradients( + reference_parameters, + trunk_base, + all_inputs, + [y.squeeze(0) for y in all_labels], + [r.squeeze(0) for r in all_returns], + with_value_head=with_value_head, + ) + frozen_prefixes = tuple(session_frozen.get(model_id, ())) + for name in raw_reference: + if any(name.startswith(prefix) for prefix in frozen_prefixes): + raw_reference[name] = torch.zeros_like(raw_reference[name]) + normalized_reference = {name: gradient / denominator for name, gradient in raw_reference.items()} + + def _capture(_micro_batches, **_kwargs): + assert manager.begin_gradient_capture(model_id, scale_state=GradientScaleState.RAW_NUMERATOR) + hidden = model(local_input) + if with_value_head: + effective_weight = runner._get_effective_value_head_weight() + unit = TokenPartial(scale=torch.tensor(1.0, device=device)) + output = value_loss_function( + hidden_states=hidden.view(1, 2, _HIDDEN), + weight=effective_weight, + labels=local_labels, + returns=local_returns, + loss_reducer=unit, + metric_reducer=unit, + ) + output.loss.backward() + else: + hidden.square().sum().backward() + manager.stage_gradient_numerators(model_id, denominator=denominator, backward_completed=True) + return {"loss": 0.0} + + runner._forward_backward_impl = _capture + runner.forward_backward([], model_id=model_id) + scratch = manager.get_adapter_state(model_id).gradient_scratch + for fqn, staged_numerator in scratch.staged_numerators.items(): + layout = state.tensor_layouts[fqn] + if not layout.has_active_storage: + continue + torch.testing.assert_close( + staged_numerator, + _logical_slice(raw_reference[fqn], layout).float().reshape(staged_numerator.shape), + rtol=5e-4, + atol=2e-4, + msg=f"staged numerator mismatch: {fqn}", + ) + staged = set(scratch.staged_parameter_fqns) + if with_value_head: + assert {"value_head.lora_A", "value_head.lora_B"} <= staged + else: + assert not any(name.startswith("value_head.") for name in staged) + for prefix in session_frozen.get(model_id, ()): + assert not any(name.startswith(prefix) for name in staged), staged + runner.commit_forward_backward_completion(model_id) + + expected_norm = _reference_step(reference_optimizer, reference_parameters, normalized_reference, clip=0.05) + actual_norm = manager.optim_step(model_id, lr=1e-2, gradient_clip=0.05) + assert actual_norm == pytest.approx(expected_norm, rel=5e-4, abs=2e-4), ( + f"{model_id} step_tag={step_tag} with_value_head={with_value_head} " + f"actual={actual_norm} expected={expected_norm}" + ) + dist.barrier() + manager.commit_optimizer_publication(model_id) + _assert_local_state_matches_reference(state, reference_parameters) + + for other, before in others_before.items(): + for name, value in manager.get_adapter_state(other).local_params.items(): + torch.testing.assert_close(value, before[name], msg=f"{other}:{name}") + + def _run_forward_only(model_id: str) -> None: + """No-grad value_prediction pass (regression: a forward-only op + must not disturb the adapter layout contract for later steps).""" + manager.prepare_forward(model_id) + with torch.no_grad(): + hidden = model(torch.ones(2, _HIDDEN, device=device)) + effective_weight = runner._get_effective_value_head_weight() + output = value_prediction_function( + hidden_states=hidden.view(1, 2, _HIDDEN), + weight=effective_weight, + labels=torch.tensor([[5, 9]], device=device), + ) + assert output.per_token_logprobs.shape == (1, 2) + + # Step 1: value_loss on the full-rank critic (trunk + value-head grads). + _run_step("critic", 0.00, with_value_head=True) + # Forward-only op between training steps (this is what a client's + # value_prediction forward does) — the next step must still validate. + _run_forward_only("critic") + # Step 2: policy-style step on the SAME session — value head untouched. + # AUTHORIZED_ZERO must accept the absent gradients and transport exact + # zeros (no stale numerators from step 1). + _run_step("critic", 0.05, with_value_head=False) + # Step 3: value_loss again on the same session (post-hygiene epoch). + _run_step("critic", 0.10, with_value_head=True) + # Step 4: the rank-2 critic slices the value-head substrate and stays + # isolated from the rank-4 session. + _run_step("critic_r2", 0.15, with_value_head=True) + # Step 5: frozen-trunk critic (SAO frozen-attention analogue): trunk + # factor gradients are skipped at staging, its optimizer never moves + # them, and only the value head trains. + frozen_state = manager.get_adapter_state("critic_frozen") + frozen_trunk_before = { + name: value.detach().clone() + for name, value in frozen_state.local_params.items() + if name.startswith("trunk") + } + _run_step("critic_frozen", 0.20, with_value_head=True) + _run_step("critic_frozen", 0.25, with_value_head=True) + for name, before in frozen_trunk_before.items(): + torch.testing.assert_close(frozen_state.local_params[name], before, msg=f"frozen factor moved: {name}") + + if rank == 0: + print("VALUE_HEAD_ADAPTER_REAL_FSDP_CERTIFIED", flush=True) + finally: + dist.destroy_process_group() + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Requires two GPUs") +def test_value_head_adapter_lifecycle_real_fsdp() -> None: + from tests.distributed.distributed_utils import run_distributed_script + + result = run_distributed_script( + __file__, + num_gpus=2, + timeout=180, + extra_env={"XORL_VALUE_HEAD_ADAPTER_WORKER": "1"}, + ) + result.assert_success("value-head adapter real-FSDP lifecycle certification") + assert "VALUE_HEAD_ADAPTER_REAL_FSDP_CERTIFIED" in result.stdout + + +if os.environ.get("XORL_VALUE_HEAD_ADAPTER_WORKER") == "1": + _run_worker() diff --git a/tests/server/test_session_spec_value_model.py b/tests/server/test_session_spec_value_model.py new file mode 100644 index 00000000..2266cf82 --- /dev/null +++ b/tests/server/test_session_spec_value_model.py @@ -0,0 +1,48 @@ +"""Session-spec normalization for value-model sessions: frozen patterns, dropout.""" + +import pytest + +from xorl.server.session_spec import normalize_lora_runtime_config + + +pytestmark = [pytest.mark.cpu, pytest.mark.server] + + +def _normalize(raw): + return normalize_lora_runtime_config(raw, default_rank=32, default_alpha=32, max_lora_rank=64) + + +def test_frozen_module_patterns_normalize_sorted_and_deduped(): + config = _normalize({"rank": 8, "frozen_module_patterns": ["v_proj", "q_proj", "q_proj"]}) + assert config["frozen_module_patterns"] == ["q_proj", "v_proj"] + assert config["lora_rank"] == 8 + + +def test_frozen_module_patterns_omitted_keeps_legacy_spec_shape(): + """Absent (or empty) patterns must not appear in the spec at all, so + existing session hashes and checkpoints stay byte-identical.""" + assert set(_normalize({"rank": 8})) == {"lora_rank", "lora_alpha"} + assert set(_normalize({"rank": 8, "frozen_module_patterns": []})) == {"lora_rank", "lora_alpha"} + + +@pytest.mark.parametrize("bad", ["q_proj", [1, 2], [""], [None]]) +def test_frozen_module_patterns_reject_bad_types(bad): + with pytest.raises(ValueError, match="non-empty strings"): + _normalize({"frozen_module_patterns": bad}) + + +def test_sdk_default_dropout_is_accepted_as_noop(): + """The tinker-style SDK always sends dropout in its LoRA payload; 0.0 is + exactly the server behavior and must not be rejected.""" + config = _normalize({"rank": 8, "alpha": 16, "dropout": 0.0}) + assert config == {"lora_rank": 8, "lora_alpha": 16} + + +def test_nonzero_dropout_is_rejected(): + with pytest.raises(ValueError, match="dropout is not supported"): + _normalize({"dropout": 0.1}) + + +def test_other_overrides_still_rejected(): + with pytest.raises(ValueError, match="only override rank, alpha, and frozen_module_patterns"): + _normalize({"target_modules": ["q_proj"]}) From bdca80437d9e1d20d1da3911e51e4e9c3f216cba Mon Sep 17 00:00:00 2001 From: Qingyang Wu Date: Mon, 24 Aug 2026 06:27:28 +0000 Subject: [PATCH 3/8] Update LoRAConfigRequest schema-shape test for frozen_module_patterns The CI assertion pins the exact request-schema property set; the new per-session field is intentional (PR #85), so the pin moves with it. Part of #84 --- tests/server/api_server/test_api_types.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/tests/server/api_server/test_api_types.py b/tests/server/api_server/test_api_types.py index 548c24d7..c4323bcc 100644 --- a/tests/server/api_server/test_api_types.py +++ b/tests/server/api_server/test_api_types.py @@ -302,7 +302,13 @@ def test_optim_step_types(self): assert lora.lora_rank == 8 assert lora.lora_alpha == 16 assert lora.model_dump(exclude_none=True) == {"lora_rank": 8, "lora_alpha": 16} - assert set(LoRAConfigRequest.model_json_schema()["properties"]) == {"lora_rank", "lora_alpha"} + assert set(LoRAConfigRequest.model_json_schema()["properties"]) == { + "lora_rank", + "lora_alpha", + "frozen_module_patterns", + } + critic_lora = LoRAConfigRequest(rank=8, frozen_module_patterns=["q_proj", "k_proj"]) + assert critic_lora.frozen_module_patterns == ["q_proj", "k_proj"] create_request = CreateModelRequest( model_id="session-a", From 79be32fc2e2b9bcbc4cd66318b590be4707e8d00 Mon Sep 17 00:00:00 2001 From: Qingyang Wu Date: Mon, 24 Aug 2026 06:32:55 +0000 Subject: [PATCH 4/8] Use the exact wire metric keys (is_*:mean) in the EV docs and SAO example Part of #84 --- docs/src/content/docs/server-training/value-model.md | 6 +++--- examples/server/sao_critic/run_sao_loop.py | 6 +++--- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/docs/src/content/docs/server-training/value-model.md b/docs/src/content/docs/server-training/value-model.md index a2bccf12..c496656d 100644 --- a/docs/src/content/docs/server-training/value-model.md +++ b/docs/src/content/docs/server-training/value-model.md @@ -67,9 +67,9 @@ Explained variance is the paper's key critic diagnostic (it should climb toward ```python ev = explained_variance( - value_error_sq_mean=metrics["is_value_error_sq_mean"], - return_mean=metrics["is_return_mean"], - return_sq_mean=metrics["is_return_sq_mean"], + value_error_sq_mean=metrics["is_value_error_sq_mean:mean"], + return_mean=metrics["is_return_mean:mean"], + return_sq_mean=metrics["is_return_sq_mean:mean"], ) ``` diff --git a/examples/server/sao_critic/run_sao_loop.py b/examples/server/sao_critic/run_sao_loop.py index 60392f0e..b2fc7cca 100644 --- a/examples/server/sao_critic/run_sao_loop.py +++ b/examples/server/sao_critic/run_sao_loop.py @@ -85,9 +85,9 @@ def main(): critic.optim_step(critic_adam).result() metrics = fb.metrics ev = explained_variance( - value_error_sq_mean=metrics.get("is_value_error_sq_mean", float("nan")), - return_mean=metrics.get("is_return_mean", float("nan")), - return_sq_mean=metrics.get("is_return_sq_mean", float("nan")), + value_error_sq_mean=metrics.get("is_value_error_sq_mean:mean", float("nan")), + return_mean=metrics.get("is_return_mean:mean", float("nan")), + return_sq_mean=metrics.get("is_return_sq_mean:mean", float("nan")), ) # 4) Policy step. With real rollouts, ``logprobs`` are the sampler's From 15bae3a09b0ace5f29851b72254b13cb3925d815 Mon Sep 17 00:00:00 2001 From: Qingyang Wu Date: Mon, 24 Aug 2026 06:33:23 +0000 Subject: [PATCH 5/8] Document the per-session frozen-attention critic option Part of #84 --- .../content/docs/server-training/value-model.md | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/docs/src/content/docs/server-training/value-model.md b/docs/src/content/docs/server-training/value-model.md index c496656d..08ba6e70 100644 --- a/docs/src/content/docs/server-training/value-model.md +++ b/docs/src/content/docs/server-training/value-model.md @@ -73,4 +73,16 @@ ev = explained_variance( ) ``` -The paper's frozen-attention critic corresponds to restricting the critic's trainable modules to MLP/MoE projections; per-session target masking is tracked as a follow-up — today `lora_target_modules` applies substrate-wide. +## Frozen-attention critic + +The paper's frozen-attention critic (its strongest ablation) is a per-session option — it constrains only the critic, not the policy sessions sharing the substrate: + +```python +critic = svc.create_lora_training_client( + BASE_MODEL, + model_id="critic", + frozen_module_patterns=["q_proj", "k_proj", "v_proj", "o_proj"], +) +``` + +Patterns are substring matches against adapter parameter names. Matching factors keep their zero-delta initialization for this session: they are skipped at gradient staging and their optimizer state never moves, so the critic trains only its MLP factors and the value head. From b5b6e849547717bcff442ab7744e2afdb0cac0c5 Mon Sep 17 00:00:00 2001 From: Qingyang Wu Date: Mon, 24 Aug 2026 06:50:02 +0000 Subject: [PATCH 6/8] All-reduce the value-head weight gradient over its shard mesh, not WORLD Under expert parallelism with dp_shard=1, every rank runs the same batch: dense-module gradients are already complete per rank, and a WORLD sum would overcount by the EP factor. The correct group is exactly the mesh the head is sharded on (ranks holding shards AND seeing distinct data); for pure-DP setups that mesh IS the world, so behavior there is unchanged (re-certified on 2 GPUs). Part of #84 --- src/xorl/server/runner/model_runner.py | 22 +++++++++++++++------- 1 file changed, 15 insertions(+), 7 deletions(-) diff --git a/src/xorl/server/runner/model_runner.py b/src/xorl/server/runner/model_runner.py index b20459c8..b4545f79 100644 --- a/src/xorl/server/runner/model_runner.py +++ b/src/xorl/server/runner/model_runner.py @@ -447,24 +447,30 @@ def _sp_allreduce_kl_metrics( class _SumGradAcrossRanks(torch.autograd.Function): - """Identity forward; all-reduce(SUM) the gradient across all ranks. + """Identity forward; all-reduce(SUM) the gradient over the given group. Used for the directly-consumed value-head weight: every rank's local loss share depends on the full replicated weight, but backward on a rank can only produce that rank's contribution. Summing the upstream weight gradient makes the sharded factor gradients complete, matching the raw-sum loss contract (normalization happens once at optim_step). + + The group must be exactly the ranks that hold shards of the head AND see + distinct data — the mesh the head is sharded on. Ranks OUTSIDE that mesh + (e.g. expert-parallel replicas with dp=1) compute identical full-batch + gradients that must NOT be summed. """ @staticmethod - def forward(ctx, tensor: torch.Tensor) -> torch.Tensor: + def forward(ctx, tensor: torch.Tensor, group) -> torch.Tensor: + ctx.group = group return tensor @staticmethod - def backward(ctx, grad: torch.Tensor) -> torch.Tensor: + def backward(ctx, grad: torch.Tensor): grad = grad.contiguous() - dist.all_reduce(grad, op=dist.ReduceOp.SUM) - return grad + dist.all_reduce(grad, op=dist.ReduceOp.SUM, group=ctx.group) + return grad, None class ModelRunner: @@ -2183,9 +2189,11 @@ def _get_effective_value_head_weight(self): ) delta = value_head.get_delta_weight() if isinstance(delta, DTensor): + mesh = delta.device_mesh + group = mesh.get_group() if mesh.ndim == 1 else mesh._flatten().get_group() delta = delta.full_tensor() - if dist.is_available() and dist.is_initialized() and dist.get_world_size() > 1: - delta = _SumGradAcrossRanks.apply(delta) + if dist.is_available() and dist.is_initialized() and dist.get_world_size(group) > 1: + delta = _SumGradAcrossRanks.apply(delta, group) return delta @staticmethod From ebb4eb12c686eea87267858e578fd56a627398e6 Mon Sep 17 00:00:00 2001 From: Qingyang Wu Date: Thu, 27 Aug 2026 00:38:18 +0000 Subject: [PATCH 7/8] Raise register_session response timeout to 600s Rank-64 LoRA session registration on a 35B-A3B MoE (per-expert factor banks, EP=2) legitimately exceeds the hardcoded 60s; align with the other heavyweight operations' timeouts. --- src/xorl/server/backend/remote.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/xorl/server/backend/remote.py b/src/xorl/server/backend/remote.py index b436d53f..0eba28bd 100644 --- a/src/xorl/server/backend/remote.py +++ b/src/xorl/server/backend/remote.py @@ -414,7 +414,9 @@ async def register_session(self, model_id="default", session_spec=None, material materialize=materialize, ), request_id=request_id, - timeout=60.0, + # Rank-64 LoRA injection across a 35B-A3B MoE's expert banks takes + # well over a minute; match the other heavyweight operations. + timeout=600.0, ) async def register_adapter(self, model_id="default", lr=1e-5, request_id=None): From f9f3f2d87adafc4c909d717c56f418bdebe2c1af Mon Sep 17 00:00:00 2001 From: Qingyang Wu Date: Thu, 27 Aug 2026 02:01:15 +0000 Subject: [PATCH 8/8] Raise worker ack_timeout to 1800s The worker acknowledges new requests only between operations; a queued forward_backward behind a multi-minute MoE fb chunk exceeded the 300s ACK window and killed otherwise-healthy runs. --- src/xorl/server/launcher.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/xorl/server/launcher.py b/src/xorl/server/launcher.py index 0a1308ef..0fb6cbfb 100644 --- a/src/xorl/server/launcher.py +++ b/src/xorl/server/launcher.py @@ -230,7 +230,7 @@ def run_orchestrator( rank0_worker_address=rank0_worker_address, operation_timeout=operation_timeout, connection_timeout=3600.0, # 1 hour for loading large models (235B) + EP sharding + LoRA + Triton compilation - ack_timeout=300.0, # 5 min — weight sync can block workers for 40s+ on large MoE models + ack_timeout=1800.0, # workers ACK only between ops; large-MoE forward_backward chunks can run many minutes sample_packing_sequence_len=sample_packing_sequence_len, enable_packing=enable_packing, pad_to_multiple_of=pad_to_multiple_of,