Adding missing q_norm/k_norm weights to olmo2.json

[armsp]

I noticed an issue in merging step -

Loading weights:  83%|████████▎ | 301/363 [13:56<02:48,  2.72s/it]
Loading weights: 100%|██████████| 363/363 [13:56<00:00,  2.31s/it]
[transformers] �[1mOlmo2ForCausalLM LOAD REPORT�[0m from: /cluster/scratch/sharaj/merges/OLMo-2-1124-13B-TA7
Key                                           | Status  |
----------------------------------------------+---------+-
model.layers.{0...39}.self_attn.q_norm.weight | MISSING |
model.layers.{0...39}.self_attn.k_norm.weight | MISSING |

And the merged model failed to even load and generate -

============================================================

Generating for: /cluster/scratch/xxx/merges/OLMo-2-1124-13B-TA7

============================================================

INFO 07-20 21:44:31 [utils.py:233] non-default args: {'trust_remote_code': True, 'dtype': 'bfloat16', 'disable_log_stats': True, 'model': '/cluster/scratch/sharaj/merges/OLMo-2-1124-13B-TA7'}

INFO 07-20 21:44:31 [model.py:549] Resolved architecture: Olmo2ForCausalLM

INFO 07-20 21:44:31 [model.py:2013] Downcasting torch.float32 to torch.bfloat16.

INFO 07-20 21:44:31 [model.py:1678] Using max model len 4096

INFO 07-20 21:44:31 [scheduler.py:238] Chunked prefill is enabled with max_num_batched_tokens=16384.

INFO 07-20 21:44:31 [vllm.py:790] Asynchronous scheduling is enabled.

...
(EngineCore pid=1491941) INFO 07-20 21:44:39 [core.py:105] Initializing a V1 LLM engine (v0.19.1) with config: model='/cluster/scratch/sharaj/merges/OLMo-2-1124-13B-TA7', speculative_config=None, tokenizer='/cluster/scratch/sharaj/merges/OLMo-2-1124-13B-TA7', skip_tokenizer_init=False, tokenizer_mode=auto, revision=None, tokenizer_revision=None, trust_remote_code=True, dtype=torch.bfloat16, max_seq_len=4096, download_dir=None, load_format=auto, tensor_parallel_size=1, pipeline_parallel_size=1, data_parallel_size=1, decode_context_parallel_size=1, dcp_comm_backend=ag_rs, disable_custom_all_reduce=False, quantization=None, enforce_eager=False, enable_return_routed_experts=False, kv_cache_dtype=auto, device_config=cuda, structured_outputs_config=StructuredOutputsConfig(backend='auto', disable_any_whitespace=False, disable_additional_properties=False, reasoning_parser='', reasoning_parser_plugin='', enable_in_reasoning=False), observability_config=ObservabilityConfig(show_hidden_metrics_for_version=None, otlp_traces_endpoint=None, collect_detailed_traces=None, kv_cache_metrics=False, kv_cache_metrics_sample=0.01, cudagraph_metrics=False, enable_layerwise_nvtx_tracing=False, enable_mfu_metrics=False, enable_mm_processor_stats=False, enable_logging_iteration_details=False), seed=0, served_model_name=/cluster/scratch/sharaj/merges/OLMo-2-1124-13B-TA7, enable_prefix_caching=True, enable_chunked_prefill=True, pooler_config=None, compilation_config={'mode': <CompilationMode.VLLM_COMPILE: 3>, 'debug_dump_path': None, 'cache_dir': '', 'compile_cache_save_format': 'binary', 'backend': 'inductor', 'custom_ops': ['none'], 'splitting_ops': ['vllm::unified_attention', 'vllm::unified_attention_with_output', 'vllm::unified_mla_attention', 'vllm::unified_mla_attention_with_output', 'vllm::mamba_mixer2', 'vllm::mamba_mixer', 'vllm::short_conv', 'vllm::linear_attention', 'vllm::plamo2_mamba_mixer', 'vllm::gdn_attention_core', 'vllm::olmo_hybrid_gdn_full_forward', 'vllm::kda_attention', 'vllm::sparse_attn_indexer', 'vllm::rocm_aiter_sparse_attn_indexer', 'vllm::unified_kv_cache_update', 'vllm::unified_mla_kv_cache_update'], 'compile_mm_encoder': False, 'cudagraph_mm_encoder': False, 'encoder_cudagraph_token_budgets': [], 'encoder_cudagraph_max_images_per_batch': 0, 'compile_sizes': [], 'compile_ranges_endpoints': [16384], 'inductor_compile_config': {'enable_auto_functionalized_v2': False, 'size_asserts': False, 'alignment_asserts': False, 'scalar_asserts': False, 'combo_kernels': True, 'benchmark_combo_kernel': True}, 'inductor_passes': {}, 'cudagraph_mode': <CUDAGraphMode.FULL_AND_PIECEWISE: (2, 1)>, 'cudagraph_num_of_warmups': 1, 'cudagraph_capture_sizes': [1, 2, 4, 8, 16, 24, 32, 40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120, 128, 136, 144, 152, 160, 168, 176, 184, 192, 200, 208, 216, 224, 232, 240, 248, 256, 272, 288, 304, 320, 336, 352, 368, 384, 400, 416, 432, 448, 464, 480, 496, 512], 'cudagraph_copy_inputs': False, 'cudagraph_specialize_lora': True, 'use_inductor_graph_partition': False, 'pass_config': {'fuse_norm_quant': False, 'fuse_act_quant': False, 'fuse_attn_quant': False, 'enable_sp': False, 'fuse_gemm_comms': False, 'fuse_allreduce_rms': False}, 'max_cudagraph_capture_size': 512, 'dynamic_shapes_config': {'type': <DynamicShapesType.BACKED: 'backed'>, 'evaluate_guards': False, 'assume_32_bit_indexing': False}, 'local_cache_dir': None, 'fast_moe_cold_start': True, 'static_all_moe_layers': []}

(EngineCore pid=1491941) INFO 07-20 21:44:40 [parallel_state.py:1400] world_size=1 rank=0 local_rank=0 distributed_init_method=tcp://10.205.11.170:58379 backend=nccl

(EngineCore pid=1491941) INFO 07-20 21:44:40 [parallel_state.py:1716] rank 0 in world size 1 is assigned as DP rank 0, PP rank 0, PCP rank 0, TP rank 0, EP rank N/A, EPLB rank N/

...
(EngineCore pid=1491941) ERROR 07-20 21:46:04 [core.py:1108] raise ValueError(

(EngineCore pid=1491941) ERROR 07-20 21:46:04 [core.py:1108] ValueError: Following weights were not initialized from checkpoint: {'model.layers.7.self_attn.k_norm.weight', 'model.layers.8.self_attn.k_norm.weight', 'model.layers.38.self_attn.q_norm.weight', 'model.layers.6.self_attn.q_norm.weight', 'model.layers.4.self_attn.q_norm.weight', 'model.layers.39.self_attn.k_norm.weight', 'model.layers.38.self_attn.k_norm.weight', 'model.layers.0.self_attn.q_norm.weight', 'model.layers.25.self_attn.q_norm.weight', 'model.layers.14.self_attn.k_norm.weight', 'model.layers.24.self_attn.k_norm.weight', 'model.layers.24.self_attn.q_norm.weight', 'model.layers.26.self_attn.q_norm.weight', 'model.layers.1.self_attn.k_norm.weight', 'model.layers.7.self_attn.q_norm.weight', 'model.layers.18.self_attn.k_norm.weight', 'model.layers.26.self_attn.k_norm.weight', 'model.layers.29.self_attn.q_norm.weight', 'model.layers.9.self_attn.k_norm.weight', 'model.layers.33.self_attn.k_norm.weight', 'model.layers.27.self_attn.q_norm.weight', 'model.layers.35.self_attn.k_norm.weight', 'model.layers.17.self_attn.q_norm.weight', 'model.layers.35.self_attn.q_norm.weight', 'model.layers.11.self_attn.k_norm.weight', 'model.layers.34.self_attn.k_norm.weight', 'model.layers.8.self_attn.q_norm.weight', 'model.layers.3.self_attn.q_norm.weight', 'model.layers.9.self_attn.q_norm.weight', 'model.layers.20.self_attn.q_norm.weight', 'model.layers.11.self_attn.q_norm.weight', 'model.layers.17.self_attn.k_norm.weight', 'model.layers.33.self_attn.q_norm.weight', 'model.layers.30.self_attn.q_norm.weight', 'model.layers.32.self_attn.k_norm.weight', 'model.layers.32.self_attn.q_norm.weight', 'model.layers.25.self_attn.k_norm.weight', 'model.layers.20.self_attn.k_norm.weight', 'model.layers.12.self_attn.k_norm.weight', 'model.layers.1.self_attn.q_norm.weight', 'model.layers.2.self_attn.q_norm.weight', 'model.layers.30.self_attn.k_norm.weight', 'model.layers.37.self_attn.q_norm.weight', 'model.layers.4.self_attn.k_norm.weight', 'model.layers.5.self_attn.k_norm.weight', 'model.layers.36.self_attn.q_norm.weight', 'model.layers.3.self_attn.k_norm.weight', 'model.layers.37.self_attn.k_norm.weight', 'model.layers.36.self_attn.k_norm.weight', 'model.layers.21.self_attn.q_norm.weight', 'model.layers.23.self_attn.k_norm.weight', 'model.layers.14.self_attn.q_norm.weight', 'model.layers.0.self_attn.k_norm.weight', 'model.layers.15.self_attn.k_norm.weight', 'model.layers.13.self_attn.k_norm.weight', 'model.layers.13.self_attn.q_norm.weight', 'model.layers.21.self_attn.k_norm.weight', 'model.layers.16.self_attn.k_norm.weight', 'model.layers.12.self_attn.q_norm.weight', 'model.layers.10.self_attn.k_norm.weight', 'model.layers.29.self_attn.k_norm.weight', 'model.layers.5.self_attn.q_norm.weight', 'model.layers.22.self_attn.q_norm.weight', 'model.layers.31.self_attn.k_norm.weight', 'model.layers.34.self_attn.q_norm.weight', 'model.layers.22.self_attn.k_norm.weight', 'model.layers.23.self_attn.q_norm.weight', 'model.layers.19.self_attn.k_norm.weight', 'model.layers.16.self_attn.q_norm.weight', 'model.layers.10.self_attn.q_norm.weight', 'model.layers.2.self_attn.k_norm.weight', 'model.layers.31.self_attn.q_norm.weight', 'model.layers.19.self_attn.q_norm.weight', 'model.layers.6.self_attn.k_norm.weight', 'model.layers.28.self_attn.k_norm.weight', 'model.layers.39.self_attn.q_norm.weight', 'model.layers.28.self_attn.q_norm.weight', 'model.layers.15.self_attn.q_norm.weight', 'model.layers.27.self_attn.k_norm.weight', 'model.layers.18.self_attn.q_norm.weight'}

eu-g7-007:1491941:1492074 [0] NCCL INFO ENV/Plugin: Could not find: libnccl-env.so

eu-g7-007:1491941:1491941 [0] NCCL INFO ENV/Plugin: Closing env plugin ncclEnvDefault

❌ FAILED on OLMo-2-1124-13B-TA7: Engine core initialization failed. See root cause above. Failed core proc(s): {}

The Fix

mergekit's OLMo‑2 architecture definition is missing the QK‑norm weights, so every OLMo‑2 merge silently drops them.

The Olmo2ForCausalLM architecture definition omits the per‑layer QK‑norm weights self_attn.q_norm.weight and self_attn.k_norm.weight. These are important - Olmo2Attention applies them before RoPE (query_states = self.q_norm(self.q_proj(x)), key_states = self.k_norm(self.k_proj(x))) and they are trained. Because mergekit only processes tensors named in the architecture template, every OLMo‑2 merge silently drops all 2 × num_layers norm tensors regardless of merge method.

Downstream this can have either negative effects or straight up failures.

Fix: add the two weights to the layer template (mirrors qwen3.json). Same class of fix as #674 (cohere norm.weight).

Verified against allenai/OLMo-2-1124-13B: merged output goes from 363 to 443 tensors, matching the source models (40 q_norm + 40 k_norm).

I have also done further verification by generating with the models merged using this fix and they worked (the vLLM error from above never reappeared).

I have read the CLA Document and I hereby sign the CLA

---

Note

Low Risk
Single JSON template change with no runtime logic; it corrects incomplete weight coverage for OLMo-2 merges.

Overview
Fixes OLMo-2 merges omitting per-layer Q/K attention norm weights, which caused missing checkpoints and vLLM load failures.

The olmo2.json layer template now lists self_attn.q_norm.weight (after q_proj) and self_attn.k_norm.weight (after k_proj), so mergekit includes and merges these tensors like other architectures (e.g. Qwen3/Cohere).

^{Reviewed by Cursor Bugbot for commit 8f480e3. Bugbot is set up for automated code reviews on this repo. Configure here.}

[github-actions]

All contributors have signed the CLA ✍️ ✅
_{Posted by the CLA Assistant Lite bot.}

[armsp]

I have read the CLA Document and I hereby sign the CLA

[armsp]

recheck