Repository navigation
Feature/support olmo3 - #4445
learncat163 wants to merge 18 commits into
Conversation
|
Thanks for your contribution! |
Codecov Report❌ Patch coverage is Please upload reports for the commit 2c7ac27 to get more accurate results. ❌ Your patch status has failed because the patch coverage (71.10%) is below the target coverage (75.00%). You can increase the patch coverage or adjust the target coverage. Additional details and impacted files@@ Coverage Diff @@
## develop #4445 +/- ##
==========================================
Coverage ? 38.97%
==========================================
Files ? 481
Lines ? 90890
Branches ? 0
==========================================
Hits ? 35420
Misses ? 55470
Partials ? 0 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
|
/re-run all-failed |
|
/re-run all-failed |
|
/re-run all-failed |
|
稍等等待1天再重试吧，看起来现在其他人的PR也是一样的集成报错，看起来不是我得代码改动问题。 |
|
/re-run all-failed |
liuhao2638
left a comment
liuhao2638
left a comment
There was a problem hiding this comment.
已完成本轮初审。发现一个需要修复后再合入的并行加载/运行问题，另外还有配置校验和 tokenizer 导出遗漏，细节都已放在行级评论中。
| self.layer_types = [ | ||
| "sliding_attention" if (i + 1) % 4 != 0 else "full_attention" for i in range(self.num_hidden_layers) | ||
| ] | ||
| layer_type_validation(self.layer_types) |
There was a problem hiding this comment.
这里没有把 num_hidden_layers 传给 layer_type_validation，所以外部传入长度不匹配的 layer_types 时配置阶段不会报错，直到 Olmo3Attention.__init__ 访问 config.layer_types[layer_idx] 才会抛 IndexError，错误位置也不指向配置问题。其他带 layer_types 的配置（如 Qwen/Gemma）都会校验长度，建议在配置初始化时直接失败。
| layer_type_validation(self.layer_types) | |
| layer_type_validation(self.layer_types, self.num_hidden_layers) |
| self.q_proj = GeneralLinear.create( | ||
| config.hidden_size, | ||
| q_hidden_size, | ||
| has_bias=config.attention_bias, | ||
| config=config, | ||
| tp_plan="colwise", | ||
| ) | ||
| self.k_proj = GeneralLinear.create( | ||
| config.hidden_size, | ||
| kv_hidden_size, | ||
| has_bias=config.attention_bias, | ||
| config=config, | ||
| tp_plan="colwise", | ||
| ) | ||
| self.v_proj = GeneralLinear.create( | ||
| config.hidden_size, | ||
| kv_hidden_size, | ||
| has_bias=config.attention_bias, | ||
| config=config, | ||
| tp_plan="colwise", | ||
| ) | ||
| self.o_proj = GeneralLinear.create( | ||
| q_hidden_size, | ||
| config.hidden_size, | ||
| has_bias=config.attention_bias, | ||
| config=config, | ||
| tp_plan="rowwise", | ||
| ) | ||
|
|
||
| # QK-norm: RMSNorm applied to query and key after projection | ||
| self.q_norm = GeneralNorm.create( | ||
| config=config, | ||
| norm_type="rms_norm", | ||
| hidden_size=q_hidden_size, | ||
| has_bias=False, | ||
| norm_eps=config.rms_norm_eps, | ||
| ) | ||
| self.k_norm = GeneralNorm.create( | ||
| config=config, | ||
| norm_type="rms_norm", | ||
| hidden_size=kv_hidden_size, | ||
| has_bias=False, | ||
| norm_eps=config.rms_norm_eps, | ||
| ) |
There was a problem hiding this comment.
这里声明了 OLMo2/3 支持 tensor parallel 的 q_proj/k_proj/v_proj/o_proj，但后面 q_norm/k_norm 仍按完整 q_hidden_size/kv_hidden_size 创建权重，而且 GeneralLinear.create(..., tp_plan="colwise") 默认 gather_output=False，TP>1 时投影输出最后一维会被切分，和全量 norm 权重维度不一致。并且 Olmo2ForCausalLMPipe 还把 _get_tensor_parallel_mappings 绑定到没有实现该方法的 Olmo2Model，tensor_model_parallel_size > 1 从 safetensors/PD 权重加载时会走到基类 NotImplementedError。PR 描述里的训练配置包含 sharding/并行场景，这会直接影响多卡加载或前向。
请补齐 OLMo2/3 的 TP 方案：至少实现 _get_tensor_parallel_mappings 覆盖 q/k/v/o、MLP 和 lm_head/embedding 的 split/merge，并让 q/k norm 参数维度与 column-parallel 输出一致；如果暂时不支持 TP，应禁用这些 parallel Linear/pipe 暴露并在配置或加载时显式报错。建议加一个 tensor_model_parallel_size=2 的最小加载/前向冒烟测试，避免只覆盖单卡路径。
修复形态可以参考现有模型的映射结构，例如：
@classmethod
def _get_tensor_parallel_mappings(cls, config, is_split=True):
from functools import partial
from ..conversion_utils import split_or_merge_func
fn = split_or_merge_func(
is_split=is_split,
tensor_model_parallel_size=config.tensor_model_parallel_size,
tensor_parallel_rank=config.tensor_parallel_rank,
num_attention_heads=config.num_attention_heads,
)
actions = {"lm_head.weight": partial(fn, is_column=True)}
for layer_idx in range(config.num_hidden_layers):
prefix = f"{cls.base_model_prefix}.layers.{layer_idx}"
actions.update({
f"{prefix}.self_attn.q_proj.weight": partial(fn, is_column=True),
f"{prefix}.self_attn.k_proj.weight": partial(fn, is_column=True),
f"{prefix}.self_attn.v_proj.weight": partial(fn, is_column=True),
f"{prefix}.self_attn.o_proj.weight": partial(fn, is_column=False),
f"{prefix}.mlp.gate_proj.weight": partial(fn, is_column=True),
f"{prefix}.mlp.up_proj.weight": partial(fn, is_column=True),
f"{prefix}.mlp.down_proj.weight": partial(fn, is_column=False),
# 同时处理 q_norm/k_norm 在 TP 下的参数切分或改为 gather 后归一化
})
return actions|
|
||
| from ...utils.lazy_import import _LazyModule | ||
|
|
||
| import_structure = { |
There was a problem hiding this comment.
新增了 olmo3/tokenizer.py，但子包的 lazy import 没有把 tokenizer 加进 import_structure，TYPE_CHECKING 分支也没有导入它；顶层 paddleformers.transformers.__init__ 同样没有注册 "olmo3.tokenizer": ["Olmo3Tokenizer"]。这样 from paddleformers.transformers.olmo3 import Olmo3Tokenizer 和 tokenizer_class_from_name("Olmo3Tokenizer") 都找不到本地类，新增 tokenizer 文件实际上不可通过包 API 访问。
请同时补齐子包和顶层导出，例如：
# paddleformers/transformers/olmo3/__init__.py
import_structure = {
"tokenizer": ["Olmo3Tokenizer"],
"configuration": ["Olmo3Config"],
"modeling": [...],
}
if TYPE_CHECKING:
from .tokenizer import *# paddleformers/transformers/__init__.py
"olmo3.tokenizer": ["Olmo3Tokenizer"],
risemeup1111
left a comment
risemeup1111
left a comment
There was a problem hiding this comment.
已复查最新提交。配置校验和 tokenizer 导出这两处已经补上；TP 修复里还剩一个会影响多卡加载/前向的权重切分方向问题，细节已放在行级评论中。
|
|
||
| actions = { | ||
| "embed_tokens.weight": partial(fn, is_column=False), | ||
| "lm_head.weight": partial(fn, is_column=True), |
There was a problem hiding this comment.
这里把 lm_head.weight 按 column 维度切分会让 TP 权重形状不匹配。GeneralLMHead 创建的权重形状是 [vocab_size / tp, hidden_size]，并且分布式属性 split_axis = 0；但 partial(fn, is_column=True) 会切最后一维，得到 [vocab_size, hidden_size / tp]，从非 TP checkpoint 加载到 tensor_model_parallel_size > 1 时会写入错误形状，后续 lm head 前向也会失败。请按 vocab 维度切分，并补一个最小 TP=2 加载/前向冒烟测试覆盖这条路径。
| "lm_head.weight": partial(fn, is_column=True), | |
| "lm_head.weight": partial(fn, is_column=False), |
risemeup1111
left a comment
risemeup1111
left a comment
There was a problem hiding this comment.
已复查最新提交。之前指出的 lm_head.weight TP 切分方向已改为按 vocab 维度切分，并补充了对应 smoke test；当前未发现需要阻塞合入的问题。
PaddleFormers Log Analysis
日志分析报告
失败的测试case: 根本原因分析: 修复建议:
🔄 每次 Re-run 后自动更新 |
risemeup1111
left a comment
risemeup1111
left a comment
There was a problem hiding this comment.
已复查最新提交。此前指出的 TP 切分问题仍保持修复状态，这次预检查修复也没有发现新的代码阻塞点；Lint 已通过。当前 H20 集成任务是在容器启动阶段退出，模型测试步骤未执行，建议以后续 CI 重跑结果为准。
|
/re-run all-failed |
1 similar comment
|
/re-run all-failed |
|
This Pull Request is stale because it has been open for 60 days with no activity. 当前Pull Request 60天内无活动，被标记为stale。 |
…olmo3 # Conflicts: # paddleformers/cli/utils/llm_utils.py # paddleformers/transformers/__init__.py # paddleformers/transformers/auto/configuration.py # paddleformers/transformers/auto/modeling.py # paddleformers/transformers/olmo2/__init__.py # paddleformers/transformers/olmo2/configuration.py # paddleformers/transformers/olmo2/modeling.py # tests/transformers/olmo2/test_modeling.py
Paddle-Bot Review Board （review完成)Powered by Nyanpasu claude with Opus 5 默认推理级别, please check the suggestions carefully. |
| # QK-norm: RMSNorm applied to query and key after projection. | ||
| # Uses local head count to match column-parallel projection output dimension. | ||
| self.q_norm = GeneralNorm.create( | ||
| config=config, | ||
| norm_type="rms_norm", | ||
| hidden_size=self.head_dim * self.num_heads, | ||
| has_bias=False, | ||
| norm_eps=config.rms_norm_eps, | ||
| ) | ||
| self.k_norm = Olmo2RMSNorm( | ||
| kv_hidden_size, | ||
| eps=config.rms_norm_eps, | ||
| input_is_parallel=config.sequence_parallel, | ||
| self.k_norm = GeneralNorm.create( | ||
| config=config, | ||
| norm_type="rms_norm", | ||
| hidden_size=self.head_dim * self.num_key_value_heads, | ||
| has_bias=False, | ||
| norm_eps=config.rms_norm_eps, | ||
| ) |
There was a problem hiding this comment.
OLMo2/OLMo3 的 QK-norm 是对整段 q/k 投影输出（num_attention_heads * head_dim）做一次 RMSNorm，而不是逐 head 归一化。这里把 hidden_size 设为本 rank 的局部头数（self.num_heads / self.num_key_value_heads 在上方 tensor_model_parallel_size > 1 分支里已被整除），同时 _get_tensor_parallel_mappings 又按 is_column=True 切分 q_norm.weight / k_norm.weight，等价于让每张卡只用本地 head 子集计算 RMS 统计量。
rsqrt(mean(x^2)) 不是可按维度分解的逐元素量：TP>1 时每卡的 mean(x^2) 只覆盖自己那部分 head，与单卡/HF 参考实现的全维统计不同。后果是权重能正常加载、形状也对，但多卡前向与训练结果相对单卡静默不一致。
仓库里 Qwen3 的注释正好点出这个差别（paddleformers/transformers/qwen3/modeling.py 中 q_norm 处 "unlike olmo, only on the head dim!"）：Qwen3 的 q_norm 只作用在 head_dim 上，所以天然可切分；OLMo 的不是。
建议二选一：
- q/k 投影后先在 TP 组内把张量 all-gather 回全维再做 norm（或对
mean(x^2)在 TP 组内做一次 all-reduce），q_norm/k_norm权重保持全量复制、不进入 TP 切分映射； - 或者暂不支持 TP>1，在
__init__里显式抛错，避免给出错误结果。
另外建议补一个 tensor_model_parallel_size=2 与单卡输出数值一致的用例：当前 tests/transformers/olmo2/test_modeling.py 只断言了 lm_head.weight 的切分形状，覆盖不到这条路径。
| """Smoke tests for OLMo2 tensor-parallel weight splitting. | ||
|
|
||
| Covers the load path exercised when sharding a non-TP checkpoint onto | ||
| ``tensor_model_parallel_size > 1``. The lm_head is vocab-parallel: its weight | ||
| has shape ``[vocab_size / tp, hidden_size]`` with ``split_axis = 0``, so the | ||
| split action must slice the vocabulary dimension (``is_column=False``), not the | ||
| hidden dimension. This test guards against that regression. | ||
| """ |
There was a problem hiding this comment.
本次改动把 OLMo2 已有的功能测试整体删除，只保留两个 lm_head.weight TP 切分的形状断言，新文件里没有任何一处真正执行过模型前向。
被删掉的覆盖包括：Olmo2ModelTester / Olmo2ModelTest（前向输出 shape、causal-lm 带 labels 的 loss、position_ids 一致性、AutoModelForCausalLM.from_config），以及 Olmo2CompatibilityTest 与 HF torch 版本的数值对齐测试。
本 PR 的目标是新增 OLMo3，并不包含削减 OLMo2 的回归覆盖；而本 PR 恰好又改动了 OLMo2 的 norm 实现、apply_rotary_pos_emb 的 dtype 行为、投影层的 TP 配置与 Olmo2Config 签名。删除后，这些改动一旦引入回归，CI 不会再有任何信号，PR 描述中也没有说明删除原因。
建议保留原有的 Olmo2ModelTester / Olmo2ModelTest / Olmo2CompatibilityTest，把新增的 TP 切分用例作为追加的测试类；若原测试确实失败或已不适用，请在 PR 描述中说明具体原因与替代覆盖方式。
| self.head_dim = head_dim if head_dim is not None else self.hidden_size // self.num_attention_heads | ||
| self.rms_norm_eps = rms_norm_eps | ||
|
|
||
| self.diff_align_float32 = kwargs.pop("diff_align_float32", False) |
There was a problem hiding this comment.
diff_align_float32 是一个仅服务于对齐调试的开关，但它被加进了公共配置，且只实现了一半：
- 唯一消费点在
paddleformers/transformers/olmo2/modeling.py的Olmo2Model.forward，只把inputs_embedscast 成 float32，而后续各层 Linear / Norm 的权重仍是 bf16/fp16，matmul 会因 dtype 不一致报错，这个开关实际不可用； Olmo3Model.forward完全没有这段逻辑，同一个配置项在 OLMo3 上静默无效；- 全仓库没有任何测试或文档引用它，类 docstring 的 Args 中也未登记。
对齐实验直接用 dtype="float32" 加载即可达到同样目的。建议删除该字段以及 modeling.py 中对应分支，避免留下一个必然报错的公共配置项；若确有保留必要，请补上全链路 cast、docstring 说明，并让 OLMo3 行为一致。
| def _rope_scaling_validation(self): | ||
| """Validate the `rope_scaling` configuration.""" | ||
| if self.rope_scaling is None: | ||
| return | ||
|
|
||
| if not isinstance(self.rope_scaling, dict) or len(self.rope_scaling) != 2: | ||
| raise ValueError( | ||
| f"`rope_scaling` must be a dictionary with two fields, `type` and `factor`, got {self.rope_scaling}" | ||
| ) | ||
| rope_scaling_type = self.rope_scaling.get("type", None) | ||
| rope_scaling_factor = self.rope_scaling.get("factor", None) | ||
| if rope_scaling_type is None or rope_scaling_type not in ["linear", "dynamic"]: | ||
| raise ValueError( | ||
| f"`rope_scaling`'s type field must be one of ['linear', 'dynamic'], got {rope_scaling_type}" | ||
| ) | ||
| if rope_scaling_factor is None or not isinstance(rope_scaling_factor, float) or rope_scaling_factor <= 1.0: | ||
| raise ValueError(f"`rope_scaling`'s factor field must be a float > 1, got {rope_scaling_factor}") |
There was a problem hiding this comment.
_rope_scaling_validation 定义后没有任何调用点，而且与上方已经执行的 standardize_rope_params + rope_config_validation 规则冲突：它要求 rope_scaling 恰好包含两个字段、type 只能是 linear / dynamic，而 standardize_rope_params 会把 rope_theta 等键写入 rope 参数，并支持更多 rope_type。
保留这段死代码会让后续维护者误以为存在这层校验，建议删除。
| backend = self.config._attn_implementation | ||
| # Flashmask only consumes sparse row indices; preserve dense mask semantics. | ||
| if backend == "flashmask" and attention_mask is not None: | ||
| backend = "eager" | ||
| attention_interface: Callable = ALL_ATTENTION_FUNCTIONS[backend] |
There was a problem hiding this comment.
这里把 flashmask 静默降级为 eager，没有任何日志，用户显式配置的 _attn_implementation="flashmask" 在常见场景下被忽略：
- OLMo2：只要传入 padding
attention_mask，create_causal_mask_and_row_indices就会返回稠密 mask（masking_utils.py中is_fully_attended为 False），于是这里一律改走 eager。实测在无 flash 内核的环境下，不传 mask 时 flashmask 内核确实被调用并报错，传 padding mask 时却能完整跑通，确认发生了降级。 - OLMo3：
paddleformers/transformers/olmo3/modeling.py的 mask 构造在_attn_implementation == "flashmask"时无条件生成稠密 full mask，而 sliding 分支本来也总是返回稠密 mask，因此 OLMo3 在任何输入下都不会真正用到 flashmask（同样的环境可全程跑通，未触发内核）。
影响是长序列训练/推理会退回到 O(L^2) 稠密 mask + eager attention 的显存与耗时，而调用方看不到任何提示。
建议：
- 发生降级时用
logger.warning_once明确提示一次，便于定位； - OLMo3 侧的稠密 full mask 只在确实需要的「带缓存多 token」场景（
kv_seq_len and seq_length > 1）构造，不要以 backend 作为触发条件； - 更彻底的方案是把 padding mask 转换为
attn_mask_startend_row_indices后交给 flashmask，保留融合算子收益。
| norm.weight.set_value(weights) | ||
| expected = values / np.sqrt(np.mean(values**2, axis=-1, keepdims=True) + config.rms_norm_eps) | ||
| expected *= weights | ||
| np.testing.assert_allclose(norm(paddle.to_tensor(values)).numpy(), expected, rtol=1e-6, atol=1e-6) |
There was a problem hiding this comment.
这一行长 122 字符，超过仓库 black 配置的 119 列上限，当前 Lint 流水线因此失败（black ... reformatted tests/transformers/olmo2/test_modeling.py，make lint 返回非零），会卡住合入。
执行 pre-commit run --files tests/transformers/olmo2/test_modeling.py 即可自动修复为：
np.testing.assert_allclose(
norm(paddle.to_tensor(values)).numpy(), expected, rtol=1e-6, atol=1e-6
)
support Olmo3
1. 模型概述
Olmo2/Olmo3 是 AI2 (Allen Institute for AI) 发布的开源语言模型系列。
本次 PR 将 Olmo2 和 Olmo3 模型从 HuggingFace transformers 迁移到 PaddleFormers。
2. 权重信息
Aistudio存储的权重信息
支持直接加载huggingface 或者 aistudio平台上的权重信息，PT和PD2种格式都支持。
3. 对齐测试
3.1 对齐背景
Olmo3依赖于特定的Transformers的版本（4.57.0），太新的Transformers（5.3.0）推理会有点问题。
但是PaddleFormers降级为Transformers==4.57.0的情况下，PaddleFormers运行时会报告Tokenizer相关的组件缺失的报错。所以这里的对齐测试都是分离的：
3.2 前向对齐前10个tokens对齐
执行单元测试
tests/transformers/olmo3/test_modeling.py的test_diff_generated_tokens测试案例，assert会断言前10个token完全一致。3.3 前向对齐logits diff 1e-2
执行单元测试
tests/transformers/olmo3/test_modeling.py的test_diff_first_token_logits，这里断言稍微需要放宽为3e-2。经过多次测试发现，其diff的差异，因为是bfloat16的原因，在多层积累误差之后，最后一层输出diff一直是2e-2，所以前向对齐标准这里从1e-2放宽为3e-2。受限于最小模型尺寸就是8B，直接float32加载在一般工作站上比较困难。这里进行了多种方案对比
可以看到在float32的CPU模式下，是可以满足精度对齐要求。
4. ms-swift对比验证
因为基础模型就是8B起步，全参微调会显存OOM，使用lora微调对比验证。截取前30步的数据(max_length=2000,max_step=200)：
其中，ms-swift的注册模版如下：
swift训练参数
swift sft \ --model "$MODEL_PATH" \ --model_type olmo3 \ --tuner_type lora \ --lora_rank 8 \ --lora_alpha 16 \ --lora_dropout 0.05 \ --template olmo3 \ --custom_register_path "$REGISTER_SCRIPT" \ --dataset "$TRAIN_DATA" \ --max_length 2048 \ --per_device_train_batch_size 1 \ --learning_rate 2e-4 \ --optim adamw_torch \ --weight_decay 0 \ --max_steps 200 \ --warmup_steps 2 \ --gradient_accumulation_steps 1 \ --num_train_epochs 1 \ --bf16 true \ --seed 42 \ --max_grad_norm -1 \ --output_dir "$OUTPUT_DIR" \ --logging_steps 1 \ --save_strategy no \paddleformers-cli 训练参数
5. CE验证
目前没有多卡环境，只能运行单卡测试。
5.1 PT lora训练
5.2 热启动测试
热启动配置
热启动日志：
可以正常之前的训练并输出
5.3 2轮loss对比
5.4 VLLM推理
可以正常使用前面导出的ckpt，并正确进行推理出正确的结果。