Implementation:Hiyouga LLaMA Factory Value Head

Knowledge Sources	Hiyouga_LLaMA_Factory
Domains	Reinforcement Learning, Reward Modeling
Last Updated	2026-02-06 19:00 GMT

Overview

Loads and prepares value head parameters for RLHF reward model training pipelines.

Description

The valuehead module provides two functions for value head management in RLHF workflows. load_valuehead_params loads value head weights (v_head.summary.weight and v_head.summary.bias) from either safetensors or PyTorch binary checkpoint files, supporting both local paths and HuggingFace Hub repository IDs. It first attempts safetensors loading and falls back to the standard PyTorch format. prepare_valuehead_model patches model-specific LM head references for non-standard architectures (LLaVA, ChatGLM, InternLM2) so that TRL's AutoModelForCausalLMWithValueHead wrapper can correctly locate the output embeddings layer.

Usage

Use load_valuehead_params when resuming training of a reward model or value head model to restore the value head weights from a checkpoint. Use prepare_valuehead_model when initializing a model for PPO or reward model training to ensure the LM head is correctly exposed for the value head wrapper.

Code Reference

Source Location

• Repository: Hiyouga_LLaMA_Factory
• File: src/llamafactory/model/model_utils/valuehead.py
• Lines: 1-72

Signature

def load_valuehead_params(
    path_or_repo_id: str,
    model_args: "ModelArguments",
) -> dict[str, torch.Tensor]

def prepare_valuehead_model(
    model: "PreTrainedModel",
) -> None

Import

from llamafactory.model.model_utils.valuehead import load_valuehead_params, prepare_valuehead_model

I/O Contract

Inputs

Name	Type	Required	Description
path_or_repo_id	str	Yes	Local path or HuggingFace Hub repository ID containing value head weights
model_args	ModelArguments	Yes	Contains cache_dir and hf_hub_token for downloading from HuggingFace Hub
model	PreTrainedModel	Yes	The pretrained model to prepare for value head attachment (for prepare_valuehead_model)

Outputs

Name	Type	Description
params	dict[str, torch.Tensor] or None	Dictionary with keys v_head.summary.weight and v_head.summary.bias, or None if no value head weights are found
None	None	prepare_valuehead_model modifies the model in-place, setting lm_head and _keys_to_ignore_on_save

Usage Examples

# Loading value head parameters from a checkpoint
from llamafactory.model.model_utils.valuehead import load_valuehead_params

vhead_params = load_valuehead_params("path/to/reward_model", model_args)
if vhead_params is not None:
    model.load_state_dict(vhead_params, strict=False)

Related Pages

• Implementation:Hiyouga_LLaMA_Factory_Model_Patcher - Calls prepare_valuehead_model via patch_model when add_valuehead is True