实习19-HRM
HRM-Text Method 阅读笔记
目标：解释论文中 H/L 两个模块、
z_H/z_L的初始化和前向数据流，以及梯度与 loss 如何设计。下面的说明主要对照本仓库实现：models/baselines/hrm_nocarry_bp_warmup.py、models/transformer.py、models/layers.py、models/lm_head.py、pretrain.py。
0. 一句话总览
HRM-Text 不是把一个普通 Transformer 从第 1 层跑到第 N 层，而是维护两个隐藏状态：
z_H：高层状态，初始化为 token embedding，最后也用它接 LM head 预测词。z_L：低层状态，初始化为固定向量zL_init，在循环里被低层模块反复更新。
然后用两个 Transformer 模块按嵌套循环交替更新：
z_H=token_embedding(input_ids)z_L=zL_initforiinrange(H_cycles):forjinrange(L_cycles):z_L=L_level(z_L,z_H)# 等价于 L_Transformer(z_L + z_H)z_H=H_level(z_H,z_L)# 等价于 H_Transformer(z_H + z_L)logits=lm_head(z_H)loss=CE(logits,labels)默认配置在config/arch/net/hrm.yaml中：
half_layers:TrueH_cycles:2L_cycles:3bp_warmup_ratio:0.2bp_max_steps:5也就是说每次高层更新之前，低层会先更新 3 次；整个前向里低层更新2 * 3 = 6次，高层更新 2 次。
1. 示例：L 模块和 H 模块的数据流动
1.1 概念区别：快/慢不是指单层算子速度，而是指状态更新频率
论文里说：
- L模块是 fast low-level module：快速低层模块。
- H模块是 slow high-level module：低速高层模块。
这里的“快”和“慢”最重要的含义是循环频率：
L_level在每个H_cycle内会连续运行L_cycles次，所以：更新更频繁（低速高层）。H_level只在每组低层循环结束后运行一次，所以：更新更慢（高层低速）。
以默认H_cycles=2,L_cycles=3为例，顺序是：
初始化: z_H = embedding, z_L = zL_init 第 1 个 H cycle: L 更新 1: z_L <- L(z_L + z_H) L 更新 2: z_L <- L(z_L + z_H) L 更新 3: z_L <- L(z_L + z_H) H 更新 1: z_H <- H(z_H + z_L) 第 2 个 H cycle: L 更新 4: z_L <- L(z_L + z_H) L 更新 5: z_L <- L(z_L + z_H) L 更新 6: z_L <- L(z_L + z_H) H 更新 2: z_H <- H(z_H + z_L) 输出: final z_H所以 L 的角色更像“在当前高层语义/计划z_H的指导下做多步局部细化”；H 的角色更像“吸收低层多步计算后的结果，再更新全局/高层表示”。
1.2 代码实现：H 和 L 都是 Transformer stack，但参数不同、调用频率不同
核心定义在models/baselines/hrm_nocarry_bp_warmup.py：
classHierarchicalReasoningModel(nn.Module):def__init__(self,config_dict:dict)->None:super().__init__()config=HierarchicalReasoningModelConfig(**config_dict)ifconfig.half_layers:assertconfig.n_layers%2==0,"n_layers must be divisible by 2."config.n_layers//=2# Reasoning Layers# TODO: Asymmetric.self.H_level=HierarchicalReasoningModelRecurrentBlock(TransformerConfig(**(config.model_dump()|config.H_override)))self.L_level=HierarchicalReasoningModelRecurrentBlock(config)- H_level 和 L_level 的结构是一样的，16 层，参数可以不一致；
classHierarchicalReasoningModelRecurrentBlock(nn.Module):def__init__(self,config:TransformerConfig)->None:super().__init__()self.core=Transformer(config)defforward(self,hidden_states:Tensor,input_injection:Tensor,**kwargs)->Tensor:returnself.core(hidden_states+input_injection,**kwargs)这说明每个 H/L 模块本质上都是一个Transformer(config)，区别不在于 block 内部算子，而在于外层如何循环调用它们。
H/L 创建代码：
self.H_level=HierarchicalReasoningModelRecurrentBlock(TransformerConfig(**(config.model_dump()|config.H_override)))self.L_level=HierarchicalReasoningModelRecurrentBlock(config)要点：
H_level和L_level是两个独立模块，参数不共享。- 默认
H_override: {}，所以 H 和 L 的网络结构一样。 - 代码注释里也写了 TODO：未来可以让 H/L 非对称，比如不同层数、隐藏维度、注意力类型等。
- 当前仓库默认还设置了
half_layers: True：如果总层数是 32，则先把n_layers //= 2，因此 H 有 16 层，L 也有 16 层。不是 H+L 各 32 层，而是把配置层数拆成两半。
1.3 H 和 L 模块内部细节：TransformerBlock = Attention + SwiGLU MLP
H/L 的core是models/transformer.py里的Transformer，代码如下：
classHierarchicalReasoningModelRecurrentBlock(nn.Module):def__init__(self,config:TransformerConfig)->None:super().__init__()self.core=Transformer(config)# Create cache functionself.create_cache=self.core.create_cachedefforward(self,hidden_states:Tensor,input_injection:Tensor,**kwargs)->Tensor:# Input injection (add)# TODO: Try better alternatives, such as GRU / gating in the following papers# Alternatively, "fixed" gating that does not depend on hidden state is also worth trying# E.g. only depends on position and index of hidden_states dimension# https://arxiv.org/pdf/1910.06764# https://arxiv.org/pdf/2202.10447# TODO: Asymmetric fusion is also worth trying. assign different number of tokens to H and L.returnself.core(hidden_states+input_injection,**kwargs)- hidden_states 和 input_injection 做了个简单的融合，代表低层高速和高层低速的 state；
Transformer模块代码如下：
- 非常传统标准的 model，一共 16 层 layer，rope + layers ；
classTransformer(nn.Module):def__init__(self,config:TransformerConfig)->None:super().__init__()self