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