Hugging Face's logo

Transformers documentation

HRM-Text

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

该模型于 2025-06-26 在 HF papers 上发布，并于 2026-05-18 贡献至 Hugging Face Transformers。

HRM-Text

概述

HRM-Text 是 Sapient AI 团队提出的分层推理模型（HRM，Hierarchical Reasoning Model）的改进型自回归语言建模变体。 它是一个基础模型，采用分层循环前向传播——两个 Transformer 堆栈（H 用于慢速、抽象的规划，L 用于快速、细致的计算）在嵌套循环中被复用：

for h in range(H_cycles):
    for l in range(L_cycles):
        z_L = L(z_L + z_H)
    z_H = H(z_H + z_L)

架构特性：

  • PrefixLM 注意力：指令 token 采用双向注意力，响应 token 采用因果注意力。由 config.prefix_lm 控制（默认值为 True）；标准形式参见 4D-masks 博客 / FlexAttention 博客。
  • 应用于 o_proj 之前的注意力输出的逐头 sigmoid 输出门控（Qwen3-Next 风格；参见 Qwen3NextAttention）。以单个融合的 gqkv_proj 形式存储的旧版检查点会在加载时通过已注册的 HRM-Text 检查点转换映射拆分为 gate_proj / q_proj / k_proj / v_proj。
  • 无参数 RMSNorm — 没有可学习缩放因子的 F.rms_norm。
  • L_bp_cycles — 来自 HRM 的 k-step grad trick。在训练时，仅 L_cycles 个底层迭代中末尾的 L_bp_cycles[i] 个会传播梯度； 更早的迭代在 torch.no_grad() 下运行，因此不会存储其激活值。对推理没有影响。

用法

HRM-Text-1B 是一个基础语言模型。它未附带chat_template，且本次发布有意不支持apply_chat_template——预训练期间使用的提示格式仍在演进中，带有稳定聊天模板的指令微调版本将在后续单独发布。请通过普通的AutoTokenizer + AutoModelForCausalLM.generate(...)来驱动该基础模型：

from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("sapientinc/HRM-Text-1B")
model = AutoModelForCausalLM.from_pretrained(
    "sapientinc/HRM-Text-1B", device_map="auto",
)

inputs = tokenizer("The quick brown fox", return_tensors="pt").to(model.device)
out = model.generate(**inputs, max_new_tokens=16, do_sample=False)
print(tokenizer.decode(out[0], skip_special_tokens=True))

注意力后端

"sdpa" 是默认选项，也是大多数工作负载的合适选择。"flex_attention" 受到支持，在长上下文场景下能带来收益——但它每次前向传播都会产生固定的 BlockMask 构建开销，这一开销无法被 HRM-Text 循环堆栈复用所带来的预期收益所摊薄。在单张 H100 上，使用已发布的 1.2B 基础检查点以及默认的 H_cycles=2、L_cycles=3 时，预填充延迟的参考值如下：

seq_lensdpaflex_attention推荐
6441 毫秒70 毫秒sdpa
25641 毫秒70 毫秒sdpa
102442 毫秒69 毫秒sdpa
204885 毫秒78 毫秒flex（约 1.1 倍）

因此，请根据工作负载选择后端：

# Default — short / medium context
model = AutoModelForCausalLM.from_pretrained("sapientinc/HRM-Text-1B", device_map="auto")

# Long context (≥ 2K tokens) — FlexAttention's per-block sparsity overtakes SDPA
model = AutoModelForCausalLM.from_pretrained(
    "sapientinc/HRM-Text-1B", device_map="auto", attn_implementation="flex_attention",
)

两种后端产生的 logits 是等价的（已验证端到端 top-1 与 torch 参考实现 100% 匹配）。"eager" 受支持且产生相同的 logits，但在现代硬件上通常不是最快的选项。它的主要用途是 output_attentions=True —— SDPA / FlexAttention 不返回每个头部的注意力权重，因此需要这些权重进行分析或可视化的 pass 应使用 attn_implementation="eager" 运行。

任何 FlashAttention 变体——包括 FA 2/3/4 以及可能未遵循 flash_attention_* 命名规范的 HF Hub 内核实现——在 config.prefix_lm=True（默认值）时，都会在初始化阶段被 HrmTextModel 拒绝。FA 后端仅支持因果掩码与非因果掩码，无法表示 PrefixLM 的 4-D 叠加掩码。对于 PrefixLM，请使用 "sdpa"（默认）或 "flex_attention"。将 config.prefix_lm=False 会使掩码变为纯因果掩码，并重新启用 FA——这对于仅使用因果掩码的微调或推理路径非常有用，在这些场景中 FA 是最快的选项。

PrefixLM 训练

对于尊重指令/回复边界的有监督微调，数据整理器（data collator）应在输出 input_ids 的同时输出 token_type_ids——指令内的位置取 1，回复和填充位置取 0。模型将 token_type_ids == 1 的所有位置视为同一个双向块的一部分；其余位置保持因果性：

import torch

def collate_prefixlm(batch, pad_token_id=0, ignore_label_id=-100):
    """`batch[i] = {"instruction_ids": [...], "response_ids": [...]}`."""
    full_ids = [b["instruction_ids"] + b["response_ids"] for b in batch]
    prefix_lens = [len(b["instruction_ids"]) for b in batch]
    max_len = max(len(ids) for ids in full_ids)

    input_ids = torch.full((len(batch), max_len), pad_token_id, dtype=torch.long)
    token_type_ids = torch.zeros_like(input_ids)
    labels = torch.full_like(input_ids, ignore_label_id)
    attention_mask = torch.zeros_like(input_ids)

    for i, (ids, plen) in enumerate(zip(full_ids, prefix_lens)):
        input_ids[i, : len(ids)] = torch.tensor(ids)
        token_type_ids[i, :plen] = 1                      # bidirectional prefix
        labels[i, plen : len(ids)] = input_ids[i, plen : len(ids)]  # loss on response only
        attention_mask[i, : len(ids)] = 1
    return {
        "input_ids": input_ids,
        "token_type_ids": token_type_ids,
        "attention_mask": attention_mask,
        "labels": labels,
    }

有关接受的形状，请参阅 HrmTextModel.forward()。

HrmTextConfig

class transformers.HrmTextConfig

< >

( transformers_version: str | None = Nonearchitectures: list[str] | None = Noneoutput_hidden_states: bool | None = Falsereturn_dict: bool | None = Truedtype: str | torch.dtype | None = Nonechunk_size_feed_forward: int = 0is_encoder_decoder: bool = Falseid2label: dict[int, str] | dict[str, str] | None = Nonelabel2id: dict[str, int] | dict[str, str] | None = Noneproblem_type: Literal['regression', 'single_label_classification', 'multi_label_classification'] | None = Nonevocab_size: int = 151808hidden_size: int = 1536intermediate_size: int = 4096num_hidden_layers: int = 16num_attention_heads: int = 12hidden_act: str = 'silu'max_position_embeddings: int = 2048initializer_range: float = 0.02rms_norm_eps: float = 1e-06use_cache: bool = Truepad_token_id: int | None = Nonebos_token_id: int | None = Noneeos_token_id: int | list[int] | None = Nonetie_word_embeddings: bool = Falserope_parameters: transformers.modeling_rope_utils.RopeParameters | dict | None = Noneattention_bias: bool = Falseattention_dropout: int | float | None = 0.0mlp_bias: bool = Falsehead_dim: int = 128H_cycles: int = 2L_cycles: int = 3L_bp_cycles: list[int] | None = Noneembedding_scale: float | None = Noneprefix_lm: bool = Truenum_layers_per_stack: int | None = None )

参数

  • vocab_size (int, 可选, 默认值为 151808) — 模型的词汇表大小。定义了 input_ids 可以表示的不同 token 数量。
  • hidden_size (int，可选，默认为 1536) — 隐藏表示的维度。
  • intermediate_size (int，可选，默认为 4096) — MLP 表示的维度。
  • num_hidden_layers (int，可选，默认为 16) — Transformer 解码器中的隐藏层数量。
  • num_attention_heads (int，可选，默认为 12) — Transformer 解码器中每个注意力层的注意力头数量。
  • hidden_act (str，可选，默认为 silu) — 解码器中的非线性激活函数（函数或字符串）。例如，"gelu"、 "relu"、"silu" 等。
  • max_position_embeddings (int，可选，默认为 2048) — 该模型可能使用的最大序列长度。
  • initializer_range (float，可选，默认为 0.02) — 用于初始化所有权重矩阵的 truncated_normal_initializer 的标准差。
  • rms_norm_eps (float，可选，默认为 1e-06) — RMS 归一化层使用的 epsilon 值。
  • use_cache (bool，可选，默认为 True) — 模型是否应返回最后一次键/值注意力（并非所有模型都使用）。仅当 config.is_decoder=True 或模型为仅解码器生成式模型时相关。
  • pad_token_id (int, 可选) — 词表中用于填充的 token id。
  • bos_token_id (int, 可选) — 词汇表中用于表示流起始的 Token id。
  • eos_token_id (Union[int, list[int]], 可选) — 词汇表中用于表示流结束（end-of-stream）的 token id。
  • tie_word_embeddings (bool，可选，默认为 False) — 是否根据模型的 tied_weights_keys 映射来绑定权重嵌入。
  • rope_parameters (Union[~modeling_rope_utils.RopeParameters, dict], 可选) — 包含 RoPE 嵌入配置参数的字典。该字典应包含 rope_theta 的值，并且如果希望使用 RoPE 处理更长的 max_position_embeddings，则可选地包含用于缩放的参数。
  • attention_bias (bool，可选，默认为 False) — 在自注意力机制中，是否在 query、key、value 和 output 投影层中使用偏置。
  • attention_dropout (Union[int, float]，可选，默认为 0.0) — 注意力概率的 dropout 比率。
  • mlp_bias (bool，可选，默认为 False) — 是否在 MLP 层的 up_proj、down_proj 和 gate_proj 层中使用偏置。
  • head_dim (int，可选，默认为 128) — 注意力头维度。如果为 None，则默认为 hidden_size // num_attention_heads
  • H_cycles (int，可选，默认为 2) — 高层循环的次数。
  • L_cycles（int，可选，默认值为 3）— 每个 H-cycle 的低层循环次数。
  • L_bp_cycles（list[int]，可选，默认为 [2]）— 训练时的梯度路由列表；在模型内部会用 1 进行左填充，直至长度为 L_cycles。 在推理时不起作用。
  • embedding_scale (float, 可选) — Token 嵌入的乘数。如果为 None，则默认为 1 / initializer_range。
  • prefix_lm（bool，可选，默认为 True） — 指令 token 采用双向注意力，响应 token 采用因果注意力。
  • num_layers_per_stack (int, 可选) — 每个 H / L 堆栈中 Transformer 块的实际数量。在首次构建时自动设置：作为 num_hidden_layers 传入的值会在此处被记住，随后 num_hidden_layers 会被重写为 num_layers_per_stack * H_cycles * (L_cycles + 1)，以便 DynamicCache(config=...) 在循环前向过程中为每次唯一的注意力调用预分配一个槽位。首次构建时请勿直接设置此参数——请将每个堆栈的实际数量作为 num_hidden_layers 传入，并让 __post_init__ 对其进行拆分。

这是用于存储 HrmTextModel 配置的类。它根据指定的参数实例化 Hrm Text 模型，并定义模型架构。使用默认值实例化配置将得到与 sapientinc/HRM-Text-1B 类似的配置。

配置对象继承自 PreTrainedConfig，可用于控制模型输出。有关更多信息，请参阅 PreTrainedConfig 的文档。

HrmTextModel

class transformers.HrmTextModel

< >

( config: HrmTextConfig )

参数

  • config (HrmTextConfig) — 模型的配置类，包含模型的所有参数。使用配置文件进行初始化时，不会加载与模型关联的权重，仅加载配置。请查看 from_pretrained() 方法以加载模型权重。

裸 Hrm Text Text Model 输出原始隐藏状态，顶部不附加任何特定头。

该模型继承自 PreTrainedModel。请查阅父类文档，了解该库为所有模型实现的通用方法（例如下载或保存、调整输入嵌入大小、剪枝头等）。

该模型也是 PyTorch torch.nn.Module 的子类。 请将其作为常规的 PyTorch Module 使用，有关通用用法和行为的所有事项，请参阅 PyTorch 文档。

forward

< >

( input_ids: typing.Optional[torch.LongTensor] = Noneattention_mask: typing.Optional[torch.Tensor] = Noneposition_ids: typing.Optional[torch.LongTensor] = Nonepast_key_values: transformers.cache_utils.Cache | None = Nonetoken_type_ids: typing.Optional[torch.LongTensor] = Noneinputs_embeds: typing.Optional[torch.Tensor] = Noneuse_cache: bool | None = None**kwargs: *<class 'transformers.utils.generic.TransformersKwargs'> ) → BaseModelOutputWithPast or tuple(torch.Tensor)

参数

  • input_ids (torch.LongTensor 形状为 (batch_size, sequence_length), ，可选) — 输入序列中各 token 在词汇表中的索引。默认情况下，填充（padding）将被忽略。

    可以使用 AutoTokenizer 获取索引。有关详细信息，请参阅 PreTrainedTokenizer.encode() 和 PreTrainedTokenizer.call()。

    什么是输入 ID？

  • attention_mask (torch.Tensor 形状为 (batch_size, sequence_length), ，可选) — 用于避免对填充 token 索引执行注意力机制的掩码。掩码值取自 [0, 1]:
    • 1 表示未被屏蔽的 token，
    • 对于被掩码的 token，值为 0。

    什么是注意力掩码？

  • position_ids (torch.LongTensor 形状为 (batch_size, sequence_length), 可选) — 输入序列中每个 token 在位置嵌入中的位置索引。取值范围为 [0, config.n_positions - 1].

    什么是位置 ID？

  • past_key_values (~cache_utils.Cache, 可选) — 预计算的隐藏状态（自注意力块和交叉注意力块中的键和值），可用于加速顺序解码。这通常由 past_key_values 在解码的先前阶段返回，当 use_cache=True 或 config.use_cache=True.

    仅允许将 Cache 实例作为输入，请参阅我们的 kv cache 指南。 如果未传入 past_key_values，则默认会初始化 DynamicCache。

    模型将输出与输入相同的缓存格式。

    如果使用了 past_key_values，用户应仅输入未处理的 input_ids（即未向该模型提供其过去键值状态的那些），其形状为 (batch_size, unprocessed_length)，而不是形状为 (batch_size, sequence_length) 的全部 input_ids。

  • token_type_ids (torch.LongTensor，形状为 (batch, seq_len)，可选) — 每个位置的双向/因果指示器。token_type_ids == 1 的 token 构成一个双向块；其余所有位置均为因果。
  • inputs_embeds（形状为 (batch_size, sequence_length, hidden_size) 的 torch.Tensor，可选） — 可选地，您可以选择直接传入嵌入表示，而不是传入 input_ids。如果您希望对 input_ids 索引如何转换为相应向量拥有比模型内部嵌入查找矩阵更多的控制权，这将非常有用。
  • use_cache (bool，可选) — 如果设置为 True，则会返回 past_key_values 键值状态，可用于加速解码（参见 past_key_values）。

返回值

BaseModelOutputWithPast 或 tuple(torch.Tensor)

一个 BaseModelOutputWithPast 或 torch.Tensor 元组（如果传入了 return_dict=False 或 config.return_dict=False），其包含的元素取决于配置（HrmTextConfig）和输入。

HrmTextModel 的 forward 方法重写了 __call__ 特殊方法。

尽管前向传播的实现需要在此函数中定义，但之后应调用 Module 实例，而不是直接调用此函数，因为前者会负责执行预处理和后处理步骤，而后者则会静默地忽略这些步骤。

  • last_hidden_state（形状为 (batch_size, sequence_length, hidden_size) 的 torch.Tensor）— 模型最后一层输出的隐藏状态序列。

    如果使用了 past_key_values，则仅输出形状为 (batch_size, 1, hidden_size) 的序列的最后一个隐藏状态。

  • past_key_values（Cache，可选，在传入 use_cache=True 或 config.use_cache=True 时返回）—— 它是一个 Cache 实例。更多详情，请参阅我们的 kv cache 指南。

    包含预计算的隐藏状态（自注意力块中的键和值，以及如果 config.is_encoder_decoder=True 则还包括交叉注意力块中的键和值），可用于（参见 past_key_values 输入）加速顺序解码。

  • hidden_states (tuple(torch.Tensor)，可选，当传入 output_hidden_states=True 或 config.output_hidden_states=True 时返回) — 一个 torch.Tensor 元组（如果模型具有嵌入层，则包含嵌入层的输出，以及每一层的输出），形状为 (batch_size, sequence_length, hidden_size)。

    模型在每一层输出处的隐藏状态，以及可选的初始嵌入输出。

  • attentions (tuple(torch.Tensor)，可选，当传入 output_attentions=True 或 config.output_attentions=True 时返回) — 一个 torch.Tensor 元组（每层一个），形状为 (batch_size, num_heads, sequence_length, sequence_length)。

    注意力权重，即经过注意力 softmax 之后的权重，用于在自注意力头中计算加权平均值。

HrmTextForCausalLM

class transformers.HrmTextForCausalLM

< >

( config: HrmTextConfig )

参数

  • config (HrmTextConfig) — 模型的配置类，包含模型的所有参数。使用配置文件进行初始化时，不会加载与模型关联的权重，仅加载配置。请查看 from_pretrained() 方法以加载模型权重。

用于因果语言建模的 Hrm 文本模型。

该模型继承自 PreTrainedModel。请查阅父类文档，了解该库为所有模型实现的通用方法（例如下载或保存、调整输入嵌入大小、剪枝头等）。

该模型也是 PyTorch torch.nn.Module 的子类。 请将其作为常规的 PyTorch Module 使用，有关通用用法和行为的所有事项，请参阅 PyTorch 文档。

forward

< >

( input_ids: typing.Optional[torch.LongTensor] = Noneattention_mask: typing.Optional[torch.Tensor] = Noneposition_ids: typing.Optional[torch.LongTensor] = Nonepast_key_values: transformers.cache_utils.Cache | None = Nonetoken_type_ids: typing.Optional[torch.LongTensor] = Noneinputs_embeds: typing.Optional[torch.Tensor] = Nonelabels: typing.Optional[torch.LongTensor] = Noneuse_cache: bool | None = Nonelogits_to_keep: typing.Union[int, torch.Tensor] = 0**kwargs: *<class 'transformers.utils.generic.TransformersKwargs'> ) → CausalLMOutputWithPast or tuple(torch.Tensor)

参数

  • input_ids (torch.LongTensor 形状为 (batch_size, sequence_length), ，可选) — 输入序列中各 token 在词汇表中的索引。默认情况下，填充（padding）将被忽略。

    可以使用 AutoTokenizer 获取索引。有关详细信息，请参阅 PreTrainedTokenizer.encode() 和 PreTrainedTokenizer.call()。

    什么是输入 ID？

  • attention_mask (torch.Tensor 形状为 (batch_size, sequence_length), ，可选) — 用于避免对填充 token 索引执行注意力机制的掩码。掩码值取自 [0, 1]:
    • 1 表示未被屏蔽的 token，
    • 对于被掩码的 token，值为 0。

    什么是注意力掩码？

  • position_ids (torch.LongTensor 形状为 (batch_size, sequence_length), 可选) — 输入序列中每个 token 在位置嵌入中的位置索引。取值范围为 [0, config.n_positions - 1].

    什么是位置 ID？

  • past_key_values (~cache_utils.Cache, 可选) — 预计算的隐藏状态（自注意力块和交叉注意力块中的键和值），可用于加速顺序解码。这通常由 past_key_values 在解码的先前阶段返回，当 use_cache=True 或 config.use_cache=True.

    仅允许将 Cache 实例作为输入，请参阅我们的 kv cache 指南。 如果未传入 past_key_values，则默认会初始化 DynamicCache。

    模型将输出与输入相同的缓存格式。

    如果使用了 past_key_values，用户应仅输入未处理的 input_ids（即未向该模型提供其过去键值状态的那些），其形状为 (batch_size, unprocessed_length)，而不是形状为 (batch_size, sequence_length) 的全部 input_ids。

  • token_type_ids (torch.LongTensor，形状为 (batch, seq_len)，可选) — 每个位置的双向/因果指示器。token_type_ids == 1 的 token 构成一个双向块；其余所有位置均为因果。
  • inputs_embeds（形状为 (batch_size, sequence_length, hidden_size) 的 torch.Tensor，可选） — 可选地，您可以选择直接传入嵌入表示，而不是传入 input_ids。如果您希望对 input_ids 索引如何转换为相应向量拥有比模型内部嵌入查找矩阵更多的控制权，这将非常有用。
  • labels (torch.LongTensor 形状为 (batch_size, sequence_length), 可选) — 用于计算因果语言建模损失的标签。索引应在 [-100, 0, ..., config.vocab_size] (参见 input_ids 文档字符串）。索引设置为 -100 的 token 将被忽略（掩码），损失仅针对标签在 [0, ..., config.vocab_size].

    标签在内部会进行移位，因此使用labels[..., 1:]来预测logits[..., :-1, :]。 请传入与input_ids匹配的未移位标签（即labels = input_ids）。在传入之前，请勿自行对标签进行移位。

    如果标签已经移位（例如在使用序列并行时），请将其作为 shift_labels 关键字参数传入，系统将直接使用该值。为了计算损失，仍然必须传入 labels。

  • use_cache (bool，可选) — 如果设置为 True，则会返回 past_key_values 键值状态，可用于加速解码（参见 past_key_values）。
  • logits_to_keep (Union[int, torch.Tensor], 可选，默认为 0) — 如果为 int，则计算最后 logits_to_keep 个 token 的 logits。如果为 0，则计算所有 input_ids 的 logits（特殊情况）。生成过程只需要最后一个 token 的 logits，仅针对该 token 进行计算可以节省内存，对于长序列或大词汇量而言，这种节省相当显著。 如果为 torch.Tensor，则必须是一维张量，对应于序列长度维度中需要保留的索引。 这在使用打包张量格式（将批次和序列长度合并为单一维度）时非常有用。

返回值

CausalLMOutputWithPast 或 tuple(torch.Tensor)

一个 CausalLMOutputWithPast 或 torch.Tensor 元组（如果传入了 return_dict=False 或 config.return_dict=False），其包含的元素取决于配置（HrmTextConfig）和输入。

HrmTextForCausalLM 的 forward 方法重写了 __call__ 特殊方法。

尽管前向传播的实现需要在此函数中定义，但之后应调用 Module 实例，而不是直接调用此函数，因为前者会负责执行预处理和后处理步骤，而后者则会静默地忽略这些步骤。

  • loss (torch.Tensor，形状为 (1,)，可选，当提供 labels 时返回) — 语言建模损失（用于下一个 token 预测）。

  • logits（形状为 (batch_size, sequence_length, config.vocab_size) 的 torch.Tensor）— 语言建模头的预测分数（SoftMax 之前每个词汇表 token 的分数）。

  • past_key_values（Cache，可选，在传入 use_cache=True 或 config.use_cache=True 时返回）—— 它是一个 Cache 实例。更多详情，请参阅我们的 kv cache 指南。

    包含预计算的隐藏状态（自注意力块中的键和值），可用于（参见 past_key_values 输入）加速顺序解码。

  • hidden_states (tuple(torch.Tensor)，可选，当传入 output_hidden_states=True 或 config.output_hidden_states=True 时返回) — 一个 torch.Tensor 元组（如果模型具有嵌入层，则包含嵌入层的输出，以及每一层的输出），形状为 (batch_size, sequence_length, hidden_size)。

    模型在每一层输出处的隐藏状态，以及可选的初始嵌入输出。

  • attentions (tuple(torch.Tensor)，可选，当传入 output_attentions=True 或 config.output_attentions=True 时返回) — 一个 torch.Tensor 元组（每层一个），形状为 (batch_size, num_heads, sequence_length, sequence_length)。

    注意力权重，即经过注意力 softmax 之后的权重，用于在自注意力头中计算加权平均值。

示例：

Update on GitHub

←HerBERT HunYuanDenseV1→