HRM-Text 技术解析：一个把“高阶推理”做进预训练框架的 1B 级文本模型仓库

如果只看 README，HRM-Text 最显眼的标签是“用更少算力和更少数据预训练 foundation model”。但从代码角度看，这个仓库真正有意思的地方，不只是参数规模或训练成本，而是它把几件通常分散在不同项目里的技术，压进了一条完整可跑通的链路里：

  • 分层递归的 HRM 模型结构
  • 面向指令/回答格式的 PrefixLM 数据组织
  • 基于 FlashAttention 3 的定制前缀注意力实现
  • 基于 PyTorch FSDP2 的大规模分布式训练
  • 从评测、推理到导出 Hugging Face 格式的一整套工具

这篇文章不复述 README，而是直接沿着仓库实现，拆解它的技术设计。

1. 先看这个仓库到底解决什么问题

这个仓库的主入口是 pretrain.py，默认配置在 config/cfg_pretrain.yaml。它并不是“给一个模型跑微调”的工程，而是一套从零开始预训练文本模型的框架。

默认配置里有几个值得注意的点：

  • 默认架构是 hrm，不是普通 Transformer
  • 默认 size 是 XL
  • global_batch_size 按 token 计数，默认是 196608
  • 默认训练 4 个 epoch
  • 默认启用 ema: 0.9999
  • 学习率策略是 warmup + cosine 的统一接口，但默认 lr_min_ratio: 1.0，等价于 warmup 后不衰减

这说明作者的重点不是做一个通用训练模板，而是围绕一套特定架构和训练设定，把工程路径尽量压短。

2. 模型核心：HRM 不是更深，而是“高层-低层”分工递归

仓库的核心模型实现在 models/baselines/hrm_nocarry_bp_warmup.py。默认配置在 config/arch/net/hrm.yaml：

name: baselines.hrm_nocarry_bp_warmup@HierarchicalReasoningModel
head: lm_head@LMHead

half_layers: True
H_cycles: 2
L_cycles: 3
bp_warmup_ratio: 0.2
bp_max_steps: 5

这套设计的关键不是“堆更多层”，而是把网络拆成两个递归层级：

  • H_level：高层推理模块
  • L_level：低层细化模块

二者本质上都还是 Transformer，只是运行方式不同。代码里 HierarchicalReasoningModelRecurrentBlock 内部包的就是 models/transformer.py 里的 Transformer。

默认 XL 规模在 config/arch/size/XL.yaml 中定义为：

  • n_layers: 32
  • hidden_size: 1536
  • num_heads: 12
  • expansion: 4

而 half_layers: True 会把层数均分给 H/L 两个模块，因此默认情况下：

  • H 模块 16 层
  • L 模块 16 层

然后在前向过程中做分层递归：

  1. 初始化 z_H = x
  2. 初始化 z_L = zL_init
  3. 每个 H cycle 内先连续跑多个 L cycle
  4. 再用 L 的结果更新一次 H

对应代码可以概括成：

for i in range(H_cycles):
    for k in range(L_cycles):
        z_L = L_level(z_L, z_H)
    z_H = H_level(z_H, z_L)

默认参数 H_cycles=2、L_cycles=3，意味着一次前向中，低层模块会被迭代 6 次，高层模块会被迭代 2 次。这个结构和“固定深度单次穿过”的普通 Transformer 差别很大，它更像是在同一段输入上做多轮内部计算。

3. HRM 的一个关键技巧：反向传播步数不是一开始就拉满

如果一个模型内部有多轮递归，训练最大的风险之一就是反向传播路径过长、显存和稳定性都变差。这个仓库没有简单粗暴地“全展开回传”，而是做了一个 warmup 机制。

在 HierarchicalReasoningModel.forward() 中，代码根据 bp_steps 控制哪些递归步骤打开梯度。更重要的是，compute_train_extra_args() 会随着训练进度，逐步把 bp_steps 从较小值抬到上限：

return dict(
    bp_steps=self.bp_min_steps + int(
        min(1, train_state.step / (train_state.total_steps * self.bp_warmup_ratio))
        * (self.bp_max_steps - self.bp_min_steps)
    )
)

默认配置下：

  • bp_min_steps = 2
  • bp_max_steps = 5
  • bp_warmup_ratio = 0.2

也就是说，训练前 20% 的进度里，反向传播会从较短路径逐渐增加到更长路径。这种做法很务实：

  • 前期先保证训练稳定和吞吐
  • 后期再把更深的递归信用分配打开

这不是论文式的“概念性递归”，而是工程上真的考虑了如何把递归训练跑起来。

4. 底座并不花哨：本质仍然是一个干净的 Transformer 实现

虽然主角是 HRM，但底层积木很标准，主要在 models/transformer.py 和 models/layers.py：

  • RMSNorm
  • RoPE
  • SwiGLU
  • 多头注意力
  • KV cache

TransformerBlock 支持 pre-norm 和 post-norm，默认配置选择 pre。初始化也做成了可切换策略：

  • fixed_normal
  • lecun_normal
  • megatron

默认 XL 配置使用的是 lecun_normal。这说明仓库并没有把实验性的复杂性塞进每一层，而是把“新意”主要放在整体结构和训练流程上。

这其实是一个优点：底层 Transformer 足够常规，HRM 的收益和问题都更容易定位。

5. 数据管线的重点不是“读取样本”，而是把样本压成 PrefixLM 训练格式

数据加载逻辑在 dataset_new.py。这个文件最关键的不是 mmap 读取，而是它如何把 instruction/response 数据拼成 PrefixLM。

仓库假定上游 data_io 已经生成了这些文件：

  • tokens.npy
  • metadata.json
  • epoch_x/inst_start.npy
  • epoch_x/inst_len.npy
  • epoch_x/resp_start.npy
  • epoch_x/resp_len.npy

加载时，每条样本会被拆成：

  • inst：前缀，通常是指令或条件
  • resp：目标回答

然后组织成一种适合 PrefixLM 的形式：

  • inputs = inst + resp[:-1]
  • labels = ignore(inst 部分) + resp
  • position_ids 对前缀和回答连续编号

其中一个很关键的细节是：

batch["labels"].append(
    np.full(len(i) - 1, fill_value=IGNORE_LABEL_ID)
    if self.config.target_only else i[1:]
)

默认 target_only=True，意味着训练时只监督 Answer，不监督 Instruction。这个选择非常符合指令式生成任务的直觉：前缀用于条件化，回答才是主要学习目标。

6. 这里的 PrefixLM 不是“打个 mask”就完了，而是单独做了注意力内核路径

这个仓库最有技术含量的部分之一，在 models/flash_attention_prefixlm_v2.py。

普通 Causal LM 的注意力很简单：所有 token 只能看左边。PrefixLM 不一样：

  • 前缀区内部可以双向注意
  • 回答区必须因果注意
  • 回答区可以看到前缀区

仓库的实现不是在 Python 里拼一个巨大 mask，而是直接围绕 FlashAttention 3 做了两段式计算：

  1. 前缀区做一次 bidirectional attention
  2. 回答区做一次 causal attention

代码里很直白：

  • Fwd pass 1 (bidirectional)
  • Fwd pass 2 (causal)

这背后有两个工程收益：

  • 避免大尺寸显式 attention mask 带来的额外开销
  • 能更自然地复用变长序列的 FlashAttention 路径

此外，这个文件还自己包了一层 _custom_flash_attn_forward()，注释里明确写了这是为了绕开上游 FA3 的一个 issue。这种处理方式很典型：作者不是等依赖修复，而是先在仓库内把训练链路打通。

7. 为了喂饱 GPU，batch 不是按样本数，而是按 token 预算做动态打包

训练吞吐的另一个关键点在 multipack_sampler.py。

这个采样器的目标不是“每卡分到同样多样本”，而是“每卡分到长度尽可能均衡的一批 token”。它做了两件事：

  • 用 batch_max_length 把 batch 预算定义成 token 上限
  • 用 LPT，也就是 Longest Processing Time first，做分布式装箱

作者甚至把核心分配逻辑用 numba JIT 了，包括：

  • lpt_check
  • lpt_with_result
  • allocate

为什么这样设计？因为对于注意力模型来说，算力消耗和序列长度不是线性关系，而更接近二次复杂度。只按样本数均分，常常会导致：

  • 某些 rank 的 token 很短，GPU 空转
  • 某些 rank 的 token 很长，成为全局瓶颈

这个采样器用 LPT 近似求解多机装箱问题，目标是提高 token slot 利用率，并平衡二次 attention 开销。对大模型训练来说，这比“普通 DistributedSampler + padding”更接近真正有意义的优化。

8. 训练系统的骨架：Hydra 负责配置，Pydantic 负责落地约束

入口文件 pretrain.py 很值得一看，因为它没有陷入“配置到处传 dict”的常见混乱，而是做了比较干净的结构化处理。

几个核心对象：

  • ArchConfig
  • DataConfig
  • PretrainConfig
  • TrainState

配置来源是 Hydra，但真正进入训练逻辑前，会先转成 Pydantic 模型。这么做的好处是：

  • CLI override 依旧保留 Hydra 的灵活性
  • 一旦进入 Python 逻辑，字段结构更稳定
  • 默认值、可选项和校验关系更清晰

模型类的加载则通过 utils/functions.py 里的 load_model_class() 动态完成，配置中用 module@class 的字符串形式指向实现。这让同一套训练框架可以挂不同结构：

  • hrm
  • transformer
  • trm
  • rins
  • ut

也就是说，这个仓库不是只能跑 HRM，它其实还承担了一个 baseline 平台的角色。

9. FSDP2 的用法很“工程化”：既包块，也包整体

分布式训练部分也是这个仓库的重点。pretrain.py 里定义了 apply_fsdp()，然后做了两层包装：

  1. 递归遍历模块，把每个 TransformerBlock 都 fully_shard
  2. 最外层整个模型再 fully_shard

同时它还设置了：

  • MixedPrecisionPolicy
  • reshard_after_forward=False
  • set_gradient_divide_factor(1.0)
  • set_force_sum_reduction_for_comms(True)

这里能看出作者对训练行为是有明确判断的：

  • reshard_after_forward=False 是用显存换通信，优先减少通信开销
  • 梯度不做默认的除法，因为注释里明确写了 Adam 类优化器对 scale invariant
  • reduction 逻辑也显式调整，以便更可控

这类代码和“能跑 FSDP”不一样，它明显是在为了大规模训练吞吐做定制。

10. 优化器、学习率和 EMA：都是为稳定预训练服务

默认优化器不是 PyTorch 内置 AdamW，而是仓库自带的 models/adam_atan2.py 中的 AdamATan2。虽然本文不展开这个文件，但从训练脚本的接入方式能看出它被当成一等公民：

  • 训练时直接创建
  • checkpoint 时同步保存 optimizer state
  • 推理加载 checkpoint 时也会临时构建 optimizer，只为支持状态恢复和 EMA 权重切换

学习率更新在 update_lr() 中完成，接口上是：

  • 前期线性 warmup
  • 后期 cosine schedule

但默认配置 lr_min_ratio: 1.0，所以实际行为是：

  • warmup 到目标学习率
  • 后续基本维持常数学习率

这个设计说明仓库保留了调度器弹性，但参考实验更偏向稳定、简单的训练曲线。

11. 推理实现说明了作者很在意“真实可用”，不是只会离线评测

推理相关代码在 simple_inference_engine.py，实现上很有几个意思。

第一，它不是单纯 model.generate() 风格，而是自己把生成过程拆成两个阶段：

  • prefill
  • decode

并且都做了 torch.compile：

  • _prefill()
  • _batched_decode()

第二，它显式维护了 GPU KV cache、cache length、last tokens，并在批内做流水式调度。代码注释里反复出现“Overlap”，说明作者在优化：

  • CPU 侧 prompt 准备
  • GPU 侧 prefill/decode
  • 生成结果回收

第三，checkpoint 的加载逻辑也比较完整：

  • 自动读取 all_config.yaml
  • 自动读取 train_metadata.yaml
  • 自动探测最新 epoch
  • 支持切换到 EMA 权重
  • 自动从训练 metadata 中恢复 tokenizer

这意味着仓库的训练产物不是“只能靠训练脚本自己解释”的半成品，而是被当成可部署、可评测、可导出的模型资产来管理的。

12. 评测层做了一个很实用的优化：按 generation config 分组

评测入口在 evaluation/main.py。它不是一股脑遍历 benchmark 然后逐个调用模型，而是先把任务按 generation 参数分组。

为什么这么做？因为不同 benchmark 可能共享相同的：

  • temperature
  • max_tokens
  • stop
  • prompt_template

如果每个 benchmark 单独跑，会制造大量推理气泡。这个仓库先把 generation config 相同的 benchmark 合并成一个批次，再一次性生成，最后再切片回填到各 benchmark 的 metrics。

这是一个非常典型、也非常实用的系统优化：不改变模型，只优化评测编排方式。

同时，评测引擎做成了双后端：

  • SimpleEngine：直接加载本仓库 checkpoint 推理
  • VLLMEngine：针对导出后的兼容模型走 vLLM

因此它既照顾了研究期实验，也为后续服务化留了接口。

13. 导出到 Hugging Face 格式，不只是 rename 权重这么简单

模型导出在 conversion/convert_to_hf.py。

这个脚本做了三件事：

  1. 加载原始训练 checkpoint
  2. 重映射参数名
  3. 构建 Hugging Face 风格的 config.json 和 model.safetensors

值得注意的是，它不仅改名，还把 HRM 结构里的关键信息编码进 HF config，比如：

  • H_cycles
  • L_cycles
  • L_bp_steps
  • prefix_lm
  • embedding_scale

这说明作者在设计导出时考虑的是“语义兼容”，而不是只把 tensor 文件转存一下。

14. 从这个仓库能看出什么技术判断

如果把 HRM-Text 当成一个工程样本来看，它反映出几条很鲜明的技术判断。

14.1 架构创新要和训练系统一起设计

很多项目提出新结构，但训练链路还是套旧模板，结果是论文里能讲、仓库里跑不动。HRM-Text 明显不是这样：

  • 递归结构配了 BP warmup
  • PrefixLM 配了定制 FlashAttention 路径
  • 变长数据配了 multipack sampler
  • 大规模训练配了 FSDP2

也就是说，模型结构、数据格式、注意力内核和训练系统是联动设计的。

14.2 这里的优化更关注“单位算力产出”，而不是单点极限

仓库大量实现细节都在服务一个目标：让有限的 GPU 小时尽可能产生更强的模型。

比如：

  • 用 token budget 做 batch，而不是固定样本数
  • 减少 padding 和长度不均衡
  • 用 PrefixLM 让 instruction/response 数据更高效地进入模型
  • 用 FSDP2 和 mixed precision 控制训练成本

这类优化不会像“参数翻倍”那样直观，但往往更接近真实世界里团队最关心的问题。

14.3 它是“研究代码”和“产品化代码”的中间态

这个仓库保留了很强的研究味道：

  • 多个 baseline 可切换
  • 很多模块留有实验注释
  • 配置暴露充分

但它又不像很多 research repo 那样只提供最短复现路径，因为它同时补上了：

  • checkpoint 管理
  • inference engine
  • benchmark framework
  • HF export

所以它更像一个“小而完整”的预训练系统。

15. 如果你想读这个仓库，推荐按什么顺序看

我建议按下面顺序阅读：

  1. README.md
    先理解作者希望这个项目解决什么问题。
  2. config/cfg_pretrain.yaml
    先看默认训练设定，知道它默认跑的是哪条路径。
  3. pretrain.py
    理解训练主流程、配置解析、FSDP 和优化器接入。
  4. dataset_new.py
    看 PrefixLM batch 是怎么组织出来的。
  5. multipack_sampler.py
    看变长样本怎么动态打包到多卡上。
  6. models/baselines/hrm_nocarry_bp_warmup.py
    看 HRM 递归结构本体。
  7. models/flash_attention_prefixlm_v2.py
    看 PrefixLM 注意力如何映射到高效 kernel。
  8. simple_inference_engine.py
    看训练产物如何被真正拿来生成。
  9. evaluation/main.py 和 conversion/convert_to_hf.py
    看评测与导出闭环。

16. 总结

HRM-Text 的价值，不只是“又一个 1B 模型仓库”。它更像是在回答一个更具体的问题：

如果不走“无限堆数据、无限堆算力”的路线，能不能通过结构设计和系统优化，把预训练这件事做得更便宜、更高效？

从代码实现看，这个仓库给出的答案是相当明确的：

  • 用分层递归结构提高单位参数的计算深度
  • 用 PrefixLM 贴合 instruction/response 数据形态
  • 用 FlashAttention 3 和 multipack sampler 保住吞吐
  • 用 FSDP2、EMA、评测与导出工具补齐完整工程闭环

不论你是否认同 HRM 这条架构路线，这个仓库都很值得看。因为它展示了一个重要事实：真正有竞争力的训练系统，往往不是某一个点子特别新，而是模型、数据、内核、分布式和工具链同时被认真设计过。

Logo
AtomGit开源社区

AtomGit 是由开放原子开源基金会联合 CSDN 等生态伙伴共同推出的新一代开源与人工智能协作平台。平台坚持“开放、中立、公益”的理念，把代码托管、模型共享、数据集托管、智能体开发体验和算力服务整合在一起，为开发者提供从开发、训练到部署的一站式体验。

更多推荐

cover

Java程序员速看！转行AI大模型，高薪风口轻松入局_程序员转行AI大模型教程（非常详细）

avatar AtomGit开源社区
cover

Anthropic收购Stainless：AI Agent时代的连接革命

avatar AtomGit开源社区
cover

Java并发编程：内存可见性与synchronized同步机制

avatar AtomGit开源社区