Llava Onevision2 Consistency

Purpose / 用途

Use this skill when validating whether a HuggingFace checkpoint and a Megatron/MCore checkpoint are behaviorally consistent in this repository.

在这个仓库里，需要验证 HuggingFace checkpoint 和 Megatron/MCore checkpoint 是否行为一致时，使用这个 skill。

There are two test systems in this repo:

本仓库有两套测试系统：

1. pytest test suite (recommended / 推荐)

• tests/consistency/conftest.py — session fixtures, HF→mcore conversion, Megatron initialization

• tests/consistency/test_model_consistency.py — 6 integration tests

• tests/consistency/test_consistency_utils.py — 10 utility functions + 11 unit tests

• tests/consistency/run_consistency_tests.sh — shell wrapper with auto-conversion + torchrun

• tests/consistency/conftest.py —— session 级 fixture、HF→mcore 转换、Megatron 初始化

• tests/consistency/test_model_consistency.py —— 6 个集成测试

• tests/consistency/test_consistency_utils.py —— 10 个工具函数 + 11 个单元测试

• tests/consistency/run_consistency_tests.sh —— shell 入口，自动转换 + torchrun

2. Legacy monolithic script (reference only / 仅供参考)

• examples/llava_onevision2/check_model_consistency.sh
• examples/llava_onevision2/check_model_consistency.py

仅作历史参考，新的工作请用 pytest 套件。

Architecture / 架构

Direction: HF → mcore

The pytest suite assumes only the HF checkpoint exists as input. The mcore checkpoint is generated automatically via conversion.

pytest 测试套件假设 只有 HF checkpoint 作为输入。mcore checkpoint 通过转换 自动生成。

HF auto-model (input)
  → convert_4b_hf_to_mcore.sh (auto-run by conftest.py or run_consistency_tests.sh)
    → mcore checkpoint (generated)
      → both models loaded → 6 tests run

Direction: mcore → HF (reverse / deploy / round-trip) / 反向：mcore → HF（部署 / 回环）

The pytest suite does not exercise the reverse path. For the p14m2 variant, two scripts ship for this:

# mcore → HF (auto-detects /release subdir; pass either form)
bash examples/llava_onevision2/convert/convert_4b_p14m2_mcore_to_hf.sh \
    /train_tmp/llava_onevision2_4b_p14m2_mcore_tp1pp1 \
    /train_tmp/llava_onevision2_4b_p14m2_hf_out \
    1 1

# Re-shard: mcore TP=1 PP=1 → mcore TP=2 PP=4 (round-trips through HF)
bash examples/llava_onevision2/convert/convert_4b_p14m2_mcore_to_release.sh \
    /src_mcore_tp1pp1 /dst_mcore_tp2pp4 2 4 0,12,12,12

Round-trip correctness (TP=1 PP=1, verified 2026-05-25): mcore → HF → mcore is bitwise identical to the original mcore checkpoint (588 non-empty tensors compared, max abs diff = 0.000e+00, 0 shape mismatches, 0 missing keys). This is the strongest correctness guarantee for the reverse path. Use this whenever changing TP/PP layout without retraining.

pytest 套件 不覆盖 反向路径。p14m2 variant 提供两个脚本（单次 mcore→HF 用于部署/推理 debug，mcore→release 用于通过 HF 中转改 TP/PP 切分）。 回环 mcore→HF→mcore 在 TP=1 PP=1 下与原始 mcore 逐位一致（588 个非空 tensor， max abs diff = 0.000e+00，0 形状不匹配，0 缺失键，2026-05-25 验证）。 在不重训的前提下改 TP/PP layout 时使用回环。

Path convention note: convert_4b_p14m2_mcore_to_hf.sh auto-detects <load>/release — pass either the parent dir or the explicit release path. Sibling scripts (4b, p14m3, p16m3, 8b, 30b) still require explicit /release.

路径约定说明：convert_4b_p14m2_mcore_to_hf.sh 会自动检测 <load>/release — 父目录或显式 release 路径都可以传。Sibling 脚本仍要求显式 /release。

Test file structure / 测试文件结构

tests/consistency/
├── __init__.py                    # empty package init / 空包初始化
├── conftest.py                    # 9 session fixtures (209 lines) / 9 个 session fixture（209 行）
├── test_consistency_utils.py      # 10 utilities + 11 unit tests (373 lines, DO NOT MODIFY) / 10 个工具 + 11 个单元测试（373 行，不要改）
├── test_model_consistency.py      # 6 integration tests (402 lines) / 6 个集成测试（402 行）
└── run_consistency_tests.sh       # shell wrapper (60 lines) / shell 入口（60 行）

Fixtures in conftest.py / conftest.py 中的 fixtures

下表列出 conftest.py 暴露的 session 级 fixture，及其用途和默认来源：

What the 6 tests check / 6 个测试检查什么

test_weight_consistency (fast)

Compares all mapped weights between HF and mcore vision models:

比较 HF 和 mcore 视觉模型之间所有映射权重：

• Patch embedding (conv weight + bias) / patch embedding（卷积 weight + bias）
• Class embedding / class embedding
• Pre/post layer norms / 前/后 layer norm
• Per-layer (24 layers): QKV weight/bias, projection, MLP fc1/fc2, layer norms / 每层（24 层）：QKV weight/bias、projection、MLP fc1/fc2、layer norms
• QKV layout conversion via convert_hf_qkv_to_mcore_layout (interleaved Q/K/V per head) / 通过 convert_hf_qkv_to_mcore_layout 做 QKV 布局转换（每 head 交织 Q/K/V）
• TP-aware gathering via _maybe_gather_tp_weight / 通过 _maybe_gather_tp_weight 做 TP-aware gather
• Threshold: cosine > 0.9999 / 阈值：cosine > 0.9999

test_vision_encoder_consistency_336px (fast)

Compares forward_debug outputs at 4 strategic points:

在 4 个关键点比较 forward_debug 输出：

• after_patch_embed — patch embedding output / patch embedding 输出
• rotary_pos_emb — rotary position embedding (aligned via align_rotary_debug_tensors) / 旋转位置编码（通过 align_rotary_debug_tensors 对齐）
• after_pre_layernorm — after pre-layernorm / 经过 pre-layernorm 之后
• before_adapter — final encoder output before adapter / 进入 adapter 之前的最终 encoder 输出
• Threshold: cosine > 0.99 / 阈值：cosine > 0.99

test_mllm_after_merger_336px (fast)

Compares vision + adapter pipeline output:

比较视觉 + adapter pipeline 输出：

• HF: forward_debug['after_merger'] / HF：forward_debug['after_merger']
• mcore: vision_model() → adapter() / mcore：vision_model() → adapter()
• Threshold: cosine > 0.99 / 阈值：cosine > 0.99

test_encoder_layer_wise_consistency (slow)

Layer-by-layer comparison of all 24 encoder layers:

逐层比较所有 24 个 encoder 层：

• layer_{i}_input and layer_{i}_output for each layer / 每层的 layer_{i}_input 和 layer_{i}_output
• input_hidden_states — initial encoder input / 初始 encoder 输入
• final_output — final encoder output / 最终 encoder 输出
• Uses align_encoder_debug_tensors for shape alignment / 用 align_encoder_debug_tensors 做形状对齐
• Threshold: cosine > 0.99 / 阈值：cosine > 0.99

test_llm_output_consistency (slow)

End-to-end LLM logits comparison:

端到端 LLM logits 比较：

• Loads LlavaOnevision2ForConditionalGeneration (HF) and full mcore model / 加载 HF 的 LlavaOnevision2ForConditionalGeneration 和完整 mcore 模型
• Tokenizes prompt with image, runs forward pass on both / 用图像 tokenize prompt，两边都跑 forward
• Compares output logits / 比较输出 logits
• Threshold: cosine > 0.99 / 阈值：cosine > 0.99

test_hf_loading_consistency (slow)

Validates HF model loading methods are equivalent:

验证 HF 模型加载方式等价：

• from_pretrained() vs manual load_file() from safetensors / from_pretrained() 对比从 safetensors 手动 load_file()
• Compares all vision weights (exact match via np.allclose) / 比较所有 vision 权重（用 np.allclose 做精确匹配）
• Compares forward_debug outputs (cosine > 0.9999) / 比较 forward_debug 输出（cosine > 0.9999）

Environment variables / 环境变量

How to run / 怎么跑

All Python must run inside the container llava_megatron_container_ax.

所有 Python 必须在容器 llava_megatron_container_ax 内运行。

Quick: run non-slow tests with auto-conversion / 快速：跑非 slow 测试 + 自动转换

# Inside container, from repo root:
# 在容器内、仓库根目录执行：
bash tests/consistency/run_consistency_tests.sh

Run all tests including slow / 跑全部测试（含 slow）

bash tests/consistency/run_consistency_tests.sh -m ""

Custom TP/PP / 自定义 TP/PP

TP=2 PP=1 MASTER_PORT=29501 bash tests/consistency/run_consistency_tests.sh

Skip conversion (pre-existing mcore checkpoint) / 跳过转换（已有 mcore checkpoint）

MCORE_CHECKPOINT_PATH=/path/to/existing bash tests/consistency/run_consistency_tests.sh

Run only unit tests (no GPU needed, works on host) / 只跑单元测试（不需要 GPU，host 上也能跑）

pytest tests/consistency/test_consistency_utils.py -v

Run specific integration test / 跑指定的集成测试

bash tests/consistency/run_consistency_tests.sh -k test_weight_consistency

What run_consistency_tests.sh does / run_consistency_tests.sh 做了什么

1. Validates HF_MODEL_PATH and TEST_IMAGE_PATH exist / 校验 HF_MODEL_PATH 和 TEST_IMAGE_PATH 存在
2. If MCORE_CHECKPOINT_PATH is empty, runs convert_4b_hf_to_mcore.sh to generate it / 如果 MCORE_CHECKPOINT_PATH 为空，跑 convert_4b_hf_to_mcore.sh 生成
3. Exports all env vars for conftest.py / 为 conftest.py 导出所有环境变量
4. Sets PYTHONPATH to include transformers_impl/llavaonevision2, aiak_megatron, repo root / 把 transformers_impl/llavaonevision2、aiak_megatron、仓库根目录加入 PYTHONPATH
5. Launches torchrun --nproc_per_node=$((TP*PP)) with pytest / 用 torchrun --nproc_per_node=$((TP*PP)) 启动 pytest

What conftest.py does for Megatron init / conftest.py 如何初始化 Megatron

Since pytest has its own arg parsing, Megatron CLI args can't be passed via command line. The solution:

由于 pytest 有自己的参数解析，Megatron CLI 参数不能通过命令行传递。解决方案：

1. Shell script exports env vars (HF_MODEL_PATH, MCORE_CHECKPOINT_PATH, CONSISTENCY_TEST_TP/PP, etc.) / shell 脚本导出环境变量（HF_MODEL_PATH、MCORE_CHECKPOINT_PATH、CONSISTENCY_TEST_TP/PP 等）
2. conftest.py reads env vars, temporarily overrides sys.argv with constructed Megatron CLI args / conftest.py 读取环境变量，临时把 sys.argv 替换成构造好的 Megatron CLI 参数
3. Calls parse_arguments() + initialize_aiak_megatron() inside the override / 在替换期内调用 parse_arguments() + initialize_aiak_megatron()
4. Restores sys.argv afterward / 完事后恢复 sys.argv

How to interpret failures / 如何解读失败

Priority order for diagnosis / 诊断优先顺序

1. test_weight_consistency — If this fails, all other tests are unreliable / 这个挂了，其他测试都不可信
2. test_vision_encoder_consistency_336px — Strategic checkpoint comparison / 关键 checkpoint 点比较
3. test_mllm_after_merger_336px — Vision + adapter pipeline health / 视觉 + adapter pipeline 健康度
4. test_encoder_layer_wise_consistency — May fail due to debug alignment, not real bugs / 可能因 debug 对齐问题失败，未必是真 bug
5. test_llm_output_consistency — Full end-to-end, most sensitive to any discrepancy / 完整端到端，对任何偏差最敏感
6. test_hf_loading_consistency — HF-only test, independent of mcore / 仅 HF 的测试，与 mcore 无关

Common failure causes / 常见失败原因

Key weight mapping / 关键权重映射

QKV weights need layout conversion: HF stores [Q_all, K_all, V_all], mcore stores interleaved [Q_h0, K_h0, V_h0, Q_h1, K_h1, V_h1, ...].

QKV 权重需要布局转换：HF 存储 [Q_all, K_all, V_all]，mcore 存储交织的 [Q_h0, K_h0, V_h0, Q_h1, K_h1, V_h1, ...]。

Known repo-local lessons / 当前仓库已知经验

1. Rotary debug representation must be aligned

HF and Megatron expose different rotary_pos_emb debug shapes:

HF 和 Megatron 暴露不同形状的 rotary_pos_emb debug 张量：

• HF: (1, S, 64)
• Megatron: (S, 32)

The align_rotary_debug_tensors function handles this by squeezing batch dim and concatenating mcore's half-dim.

align_rotary_debug_tensors 函数通过去掉 batch 维度并拼接 mcore 的半维度来处理。

2. PP-aware testing is necessary

When PP > 1, not every pipeline stage owns vision_model, adapter, or decoder post-process outputs. Tests must skip non-owner stages.

当 PP > 1 时，不是每个 pipeline stage 都拥有 vision_model、adapter 或 decoder 后处理输出。测试必须跳过非 owner stage。

3. TP-aware weight comparison is necessary

When TP > 1, use _maybe_gather_tp_weight to gather shards before comparison. It gathers along first dim for QKV/FC1, last dim for proj/FC2.

当 TP > 1 时，用 _maybe_gather_tp_weight 在比较前 gather shards。QKV/FC1 沿第一维 gather，proj/FC2 沿最后一维。

4. HF and mcore use the same pixel value 2x2 memory layout

No pixel value conversion is needed between HF and mcore models.

HF 和 mcore 模型使用相同的 2x2 内存布局，无需转换 pixel values。

5. Encoder-layer-wise failures may be debug-layout issues

If weight_consistency + merger pass but encoder_layer_wise fails in late layers, suspect debug capture semantics rather than real model bugs.

如果 weight_consistency + merger 通过但 encoder_layer_wise 在后面层失败，优先怀疑 debug 捕获语义而非模型真错。

Minimal troubleshooting checklist / 最小排查清单

If the run fails, check in this order:

如果运行失败，按以下顺序排查：

1. Is the container running? docker exec -it llava_megatron_container_ax bash

2. Does HF_MODEL_PATH exist and contain safetensors files?

3. Did the HF→mcore conversion succeed? Check stderr output.

4. Does the container have enough GPUs for TP * PP?

5. Is MASTER_PORT already in use? Try a different port.

6. Did test_weight_consistency fail? → Fix this first before investigating other tests.

7. Is the failure in a @pytest.mark.slow test? → Run fast tests first with default marker filter.

8. 容器是否在运行？docker exec -it llava_megatron_container_ax bash

9. HF_MODEL_PATH 是否存在且包含 safetensors 文件？

10. HF→mcore 转换是否成功？检查 stderr 输出。

11. 容器 GPU 数量是否满足 TP * PP？

12. MASTER_PORT 是否被占用？换一个端口试试。

13. test_weight_consistency 是否失败？→ 先修这个再看其他测试。

14. 失败的是否是 @pytest.mark.slow 测试？→ 先用默认 marker 跑 fast 测试。