EvolvingLMMs-Lab/LLaVA-OneVision-2

Purpose

Use this skill when writing git commit messages. It enforces a formal Conventional Commits style suitable for open-source collaboration and future changelog generation, while still keeping messages aligned with this repository's tone.

Format

<type>: <subject>

• Default style: formal Conventional Commits (feat:, fix:, docs:, test:, chore:, etc.)
• No scope parentheses by default — this repo usually does not use type(scope): style; add a scope only if the user explicitly asks
• Subject: imperative mood, lowercase first word, no trailing period, max ~72 chars
• Body (optional): blank line after subject, wrap at 72 chars, explain why not what

Types

Notes:

• feat and fix are strongly preferred for code changes
• Use docs for skill files under .opencode/skills/
• Avoid legacy bare messages like updated, missing _extra_state, or add OV2 SP unless the user explicitly asks to preserve old repo style
• Revert and Merge are generated by git — don't manually write these types

Rules

1. Subject Line

• Imperative mood: "add", "fix", "refactor" — NOT "added", "fixes", "refactoring"
• Lowercase after type prefix: feat: add ... not feat: Add ...
• No period at the end
• Max 72 characters (hard limit for git log formatting)
• What changed, not how: fix: correct rotary_base override in 8B config not fix: changed value from 1000000 to 8000000

2. Body (when needed)

Use a body when:

• The change is non-obvious (e.g., subtle bug fix, architectural decision)
• There are constraints or caveats others should know
• The PR has multiple commits and each needs context

feat: add YaRN RoPE scaling to Megatron training path

Wire YarnRotaryEmbedding into QwenModel when --rope-type=yarn is set.
All default parameters match the official Qwen3-8B config.

Affected files:
- qwen_model.py: construct YarnRotaryEmbedding, return (emb, mscale) tuple
- attention.py: extract mscale before RoPE application
- arguments.py: add 7 CLI args for YaRN configuration

3. PR-linked Commits

When a commit will be part of a PR (most cases in this repo):

• The commit message should be self-contained — don't rely on PR description
• Reference issue numbers if applicable: fix: loading dataloader (#121)
• The (#N) suffix is added by GitHub on squash-merge — don't add it manually in local commits

4. Squash / Reword / Multi-author Rules

When cleaning up commit history before a PR:

• Prefer squashing small fixup commits by day: one coherent commit per day of work, unless a day contains clearly unrelated changes that should stay separate.
• If a squashed commit includes work from multiple authors, the final commit message must preserve every contributor with Co-authored-by trailers.
• Use each contributor's canonical Git name and email from the original commits.
• Rewording a commit message normally preserves the original Author; do not use --reset-author unless the user explicitly requests it.
• If the branch was already pushed, any squash/reword rewrites history and must be pushed with --force-with-lease, never plain --force.

Example multi-author squash commit:

docs: explain distributed offline packing workflow

Group one day of packing notes into a single guide so the PR history stays
reviewable while preserving each contributor from the original commits.

Co-authored-by: Alice Zhang <alice@example.com>
Co-authored-by: Bob Chen <bob@example.com>

5. Multi-file Changes

For commits touching many files, the subject should describe the intent, not list files:

# GOOD
feat: add YaRN RoPE support to QwenModel

# BAD
update qwen_model.py, attention.py, arguments.py, provider.py, config.py, model.py, mid_training.sh

6. What NOT to Write

• No generic messages: "update code", "fix bug", "misc changes"
• No WIP commits to main: "wip", "temp", "test" (use branches)
• No file lists as subject: "update foo.py and bar.py"
• No past tense: "fixed", "added", "updated" — use imperative
• No emoji unless the repo uses them (this repo does not)

Decision Tree

Is it a new capability?
  YES → feat: <what it enables>
  NO ↓

Is it a bug fix?
  YES → fix: <what was broken>
  NO ↓

Is it restructuring without behavior change?
  YES → refactor: <what was reorganized>
  NO ↓

Is it adding a new file/script/example?
  YES → docs: if documentation/skill, test: if test fixture, chore: if repo maintenance, feat: if user-visible capability
  NO ↓

Is it a config/dependency/version update?
  YES → build: for dependencies/build/Docker, ci: for CI, chore: for general maintenance
  NO ↓

Is it trivial (typo, comment, formatting)?
  YES → docs: for documentation typo, style: for formatting-only code changes
  NO → pick the closest type

Examples from This Repository

These are real commits — use them as reference for tone and length:

fix loading dataloader (#121)
missing _extra_state (#117)
Refactor model consistency checks and enhance encoder loading (#120)
add OV2 SP (#112)
update 30ba3b (#97)
feat: add MoE merge support for Qwen3-30B-A3B (#95)
refactor: move patch position block layout into task encoder (#89)
dev assert (#81)

Prefer the formal Conventional Commits shape for new local commits, even though older history contains mixed styles.

Better current-style examples:

feat: add packed SFT dataset assembly workflow
fix: prevent cross-sample attention leakage in packed runs
docs: explain cu_lengths attention boundaries
test: add HF/Megatron consistency checks for PP=2
chore: remove internal-only skill references
refactor: simplify Megatron checkpoint layout detection
build: update flash attention to v2.7.0

Quick Template

For most changes, just fill in:

<type>: <imperative verb> <what> [in/for/to <where>]

Examples:

• feat: add custom pipeline layer splitting for PP>2
• fix: correct rotary_base override for 8B config
• docs: document distributed offline packing workflow
• test: add OneVision2 consistency fixture for PP=2
• chore: remove internal-only skill references
• refactor: extract ViT encoder into standalone module
• build: update flash attention to v2.7.0

Purpose / 用途

Use this skill when reasoning about attention boundaries in the LLaVA-OneVision2 forward pass — specifically how cu_lengths and patch_positions control attention at different stages of the model.

在分析 LLaVA-OneVision2 前向传播中的 attention 边界时使用这个 skill——具体来说，cu_lengths 和 patch_positions 如何在模型的不同阶段控制 attention。

This skill is specifically for:

• Understanding the difference between ViT-level and LLM-level attention control
• Debugging packed vs non-packed attention behavior
• Reasoning about cross-sample isolation in packed sequences
• Understanding why patch_positions grouping does NOT carry into the LLM

这个 skill 专门用于：

• 理解 ViT 层和 LLM 层 attention 控制的区别
• 调试 packed 和 non-packed 的 attention 行为
• 分析 packed 序列中跨样本隔离机制
• 理解为什么 patch_positions 的分组不会延续到 LLM 中

Key Files / 关键文件

Core Concept: Two Independent Attention Control Mechanisms / 核心概念：两套独立的 Attention 控制机制

Overview Diagram / 概览图

┌─────────────────────────────────────────────────┐
│  ViT Encoder                                    │
│                                                 │
│  Control: patch_positions (temporal dimension)  │
│  Effect:  Local/shared attention                │
│           e.g. 4 images share one attention     │
│           window via same temporal index        │
│                                                 │
│  Output:  visual embeddings                     │
└──────────────────┬──────────────────────────────┘
                   │  (embeddings replace image
                   │   placeholder tokens)
                   ▼
┌─────────────────────────────────────────────────┐
│  LLM Decoder                                    │
│                                                 │
│  Control: cu_lengths (cumulative sub-seq lens)  │
│  Effect:  Determines attention domain           │
│                                                 │
│  NON-PACKED: cu_lengths == [[0]]                │
│    → full causal attention                      │
│    → ALL tokens see ALL previous tokens         │
│    → patch_positions grouping is GONE           │
│                                                 │
│  PACKED: cu_lengths = [0, a, a+b, ...]          │
│    → block-diagonal causal attention            │
│    → sub-sequences isolated from each other     │
│    → within each sub-seq: full causal           │
└─────────────────────────────────────────────────┘

Mechanism Details / 机制详解

1. cu_lengths Generation / cu_lengths 的产生

Source A: Offline Packed Data (PackedCaptioningSample)

# task_encoder.py → pack_selected_samples()
cu_lengths = [0]
for sample in samples:
    current_length += sample.total_len
    cu_lengths.append(current_length)
# Result: [0, 512, 1024, 1389] — 3 sub-samples packed together

Each sub-sample was independently encoded (tokenized + image processed) then concatenated into one long sequence. cu_lengths records the boundaries.

每个子样本独立编码（tokenize + 图像处理），然后拼接成一个长序列。cu_lengths 记录边界。

Source B: Non-Packed Single Samples

# task_encoder.py → batch()
if self.is_packing_enabled or int(os.environ.get("OFFLINE_PACKED_DATA", 0)) == 1:
    cu_lengths = torch.stack([s.cu_lengths for s in samples])
else:
    cu_lengths = torch.tensor([[0]], dtype=torch.int32)  # dummy value

Non-packed samples get cu_lengths = [[0]] with shape [1, 1].

非 packed 样本得到 cu_lengths = [[0]]，shape 为 [1, 1]。

2. Forward Function Branching / 前向函数分支

In pretrain_llava_onevision2.py:

if cu_lengths.shape == torch.Size([1, 1]):
    # ===== NON-PACKED PATH =====
    # Uses attn_mask for padding_causal attention
    # Every token attends to all previous tokens (full causal)
    # packed_seq_params = None
    for i in range(attn_mask.shape[0]):
        loss_mask[i, (attn_mask[i] == False).sum() - 1] = 0
else:
    # ===== PACKED PATH =====
    # micro-batch-size must be 1 for packing
    assert cu_lengths.shape[0] == 1
    attn_mask = None  # not needed — cu_seqlens defines boundaries
    packed_seq_params = PackedSeqParams(
        qkv_format="thd",
        cu_seqlens_q=cu_lengths[0],      # → Flash Attention kernel
        cu_seqlens_kv=cu_lengths[0],
        max_seqlen_q=max_lengths[0].item(),
        max_seqlen_kv=max_lengths[0].item(),
    )

3. LLM Attention Behavior / LLM 层的 Attention 行为

Non-packed (cu_lengths == [[0]]) → Full Causal Attention：

• The entire sequence is ONE attention domain
• Token at position i can attend to positions 0..i
• ALL visual tokens (from any image) + ALL text tokens are mutually visible
• Compute cost: O(n²) where n = total sequence length
• packed_seq_params = None → standard causal mask

整个序列是一个 attention 域，所有 visual token（来自任何图片）和所有 text token 互通可见。计算量 O(n²)。

Packed (cu_lengths = [0, a, a+b, ...]) → Block-Diagonal Causal Attention：

• Each sub-sequence [cu_lengths[i], cu_lengths[i+1]) is an independent attention domain
• Within each sub-sequence: full causal attention
• Between sub-sequences: ZERO attention (completely isolated)
• Implemented via Flash Attention's cu_seqlens parameter
• Compute cost: O(a² + b² + c² + ...) << O((a+b+c+...)²)

每个子序列 [cu_lengths[i], cu_lengths[i+1]) 是独立的 attention 域，子序列之间完全隔离。通过 Flash Attention 的 cu_seqlens 参数实现。

4. Sequence Parallelism Padding / 序列并行填充

When args.sequence_parallel is enabled, the sequence must be divisible by TP size (and TP×CP×2 if CP > 1). For packed sequences, padding tokens are appended as a dummy extra sub-sequence:

当启用 sequence_parallel 时，序列长度必须被 TP size 整除。对 packed 序列，padding token 作为一个额外的 dummy 子序列追加：

if packed_seq_params is not None:
    new_end_q = packed_seq_params.cu_seqlens_q[-1:] + pad_size
    packed_seq_params = PackedSeqParams(
        cu_seqlens_q=torch.cat([packed_seq_params.cu_seqlens_q, new_end_q]),
        cu_seqlens_kv=torch.cat([packed_seq_params.cu_seqlens_kv, new_end_kv]),
        max_seqlen_q=max(packed_seq_params.max_seqlen_q, pad_size),
        max_seqlen_kv=max(packed_seq_params.max_seqlen_kv, pad_size),
    )

Critical Insight: patch_positions Scope / 关键洞察：patch_positions 的作用域

ViT Layer: patch_positions Controls Attention Grouping

In the ViT encoder, patch_positions has a temporal dimension (t, h, w). Images sharing the same temporal index share one attention window. For example, 4 images treated as "video frames" share attention via their temporal coordinates.

在 ViT encoder 中，patch_positions 有时间维度 (t, h, w)。共享相同时间索引的图片共享一个 attention window。例如，4 张图片被当作"视频帧"通过时间坐标共享 attention。

LLM Layer: patch_positions Has NO Effect on Attention

patch_positions is NOT used to control LLM attention. It is only passed through the data pipeline for potential use in position embeddings or other purposes, but the LLM's attention boundaries are controlled EXCLUSIVELY by cu_lengths.

patch_positions 不控制 LLM 的 attention。 它只是在数据 pipeline 中传递，可能用于 position embedding 等目的，但 LLM 的 attention 边界完全由 cu_lengths 控制。

This means:

这意味着：

Why This Matters / 为什么这很重要

If you assume ViT-level grouping persists into the LLM, you will misunderstand the compute profile:

如果你假设 ViT 层的分组延续到 LLM 中，会误解计算特征：

• Non-packed: LLM always does full causal attention over the entire sequence. A sample with 4 high-res images has O((4×img_tokens + text_tokens)²) attention cost — there is NO per-image isolation in the LLM.

• Packed: The isolation is between SAMPLES (sub-sequences), not between images within a sample. A packed sample containing samples A (2 images) and B (1 image) isolates A from B, but within A, both images + text are fully visible to each other.

• Non-packed: LLM 总是对整个序列做 full causal attention。一个包含 4 张高分辨率图片的样本，attention 代价为 O((4×img_tokens + text_tokens)²)——LLM 中没有按图片隔离。

• Packed: 隔离是在样本（子序列）之间，不是在同一样本内的图片之间。一个包含样本 A（2 张图）和样本 B（1 张图）的 packed 样本，A 和 B 互相隔离，但 A 内部的两张图 + 文本完全互通。

SFT Path: Attention Mask Based cu_seqlens / SFT 路径：基于 attention_mask 的 cu_seqlens

In the SFT training path (sft/utils.py), cu_seqlens can also be derived from the attention mask using sample-ID encoding:

在 SFT 训练路径中，cu_seqlens 也可以从 attention mask 推导：

# sft/utils.py → _get_packed_sequence_params()
# attention_mask encodes sample IDs: [[1,1,2,2,2,3,3,4,5,5,5,0,0]]
# → cu_seqlens = [0, 2, 5, 7, 8, 11, 13]
reduced_mask = torch.bincount(attention_mask.view(-1), minlength=max_num + 1)
cu_seqlens = reduced_mask[1:].cumsum(dim=0).to(torch.int32)
cu_seqlens[-1] = attention_mask.shape[1]  # include padding
cu_seqlens = torch.cat((zero, cu_seqlens))

This achieves the same block-diagonal attention as the pretrain path's cu_lengths mechanism.

这与 pretrain 路径的 cu_lengths 机制实现相同的 block-diagonal attention。

Quick Reference Table / 快速参考表

Common Pitfalls / 常见误区

1. Assuming ViT attention grouping carries into LLM — It does NOT. patch_positions only affects ViT; LLM uses cu_lengths.

   假设 ViT 的 attention 分组延续到 LLM — 不会。patch_positions 只影响 ViT；LLM 使用 cu_lengths。

2. Confusing packed sample isolation with image-level isolation — cu_lengths boundaries separate SAMPLES, not images within a sample.

   混淆 packed 样本隔离和图片级隔离 — cu_lengths 边界分隔的是样本，不是样本内的图片。

3. Forgetting micro-batch-size=1 constraint for packing — The code asserts cu_lengths.shape[0] == 1 in the packed path.

   忘记 packing 要求 micro-batch-size=1 — 代码在 packed 路径中断言 cu_lengths.shape[0] == 1。

4. Ignoring SP padding for packed sequences — When sequence parallelism is enabled, padding tokens are added as a dummy sub-sequence in cu_seqlens, not ignored.

   忽略 packed 序列的 SP 填充 — 启用序列并行时，padding token 作为 dummy 子序列加入 cu_seqlens，不是被忽略。

Purpose / 用途

Use this skill when packing a large SFT JSONL (hundreds of thousands to millions of samples) into Energon WebDataset shards at a fixed sequence length, using offline_packing/auto_pipe.sh parallelized across multiple nodes.

当需要把大规模 SFT JSONL（几十万到几百万样本）按固定序列长度打包成 Energon WebDataset shards，并通过 offline_packing/auto_pipe.sh 在多台机器上并行处理时，使用这个 skill。

Prerequisites / 前置条件

• All nodes share the same NFS mount (data + repo + output)
• Same docker image on every node (must contain transformers, energon, project repo)
• offline_packing/auto_pipe.sh and stage scripts s1_split_json_to_samples.py … s4_bins_to_webdataset.py
• Tokenizer + image processor available locally (HF format model dir)
• Source JSONL where each line has images, prompts, captions (multi-turn list-of-lists) and image paths are usable as-is

所有节点共享同一个 NFS（数据+代码+输出）。每台机器使用同一个 docker 镜像，里面要有 transformers、energon、项目代码。需要 offline_packing/auto_pipe.sh 和 s1–s4 四个 stage 脚本。Tokenizer 和 image processor 是本地 HF 格式目录。源 JSONL 每行包含 images / prompts / captions（多轮是 list-of-list），图片路径可直接使用。

Architecture / 架构

Pipeline Stages / 流水线阶段

JSONL (N samples)
  ├─ s1_split_json_to_samples.py   # validate + drop bad/missing-image samples
  │                                # output: per-sample serialized records
  ├─ s2_compute_token_lengths.py   # tokenize prompts/captions, compute image-patch tokens
  │                                # output: length array per sample
  ├─ s3_bin_packing.py             # BFD (Best-Fit-Decreasing) into bins of capacity L
  │                                # output: bin assignment
  └─ s4_bins_to_webdataset.py      # write tar shards + idx + .nv-meta/{dataset.yaml,split.yaml,sample_loader.py}

auto_pipe.sh runs all four stages sequentially on one node for one input file. To use N nodes, split the JSONL into N parts and run auto_pipe.sh independently on each — s3 BFD does NOT shard across nodes, so each node packs its own slice.

auto_pipe.sh 在一台机器上对一个输入文件顺序跑完四个 stage。要用 N 台机器就把 JSONL 切成 N 份，每台独立跑一次 auto_pipe.sh——s3 BFD 不能跨节点共享，所以每台只 pack 自己的那一份。

Key Design Decisions / 关键设计

• Per-node shard prefix: pass distinct --shard-prefix <name>_<a|b|...> so tar files don't collide on shared NFS
• Sample class: choose based on data shape — PackedCaptioningSample for image+text packed turns, MultiMixQASample for QA-style. The class controls the auto-generated sample_loader.py
• --no-npy flag: only set when JSONL has no precomputed patch_positions field. Without --no-npy s2 expects per-sample .npy files; missing files cause noisy warnings AND fall back to slow real-image tokenization
• Sequence length L: scan token lengths first to confirm drop rate is acceptable
• Output layout per node: <output_root>/node_<x>/webdataset/{*.tar, *.idx, .nv-meta/}
• Top-level Metadataset yaml combines all node outputs into one logical dataset

Step-by-step Workflow / 操作步骤

0. Pre-flight: choose L / 预检：选 L

Scan token lengths once on the full JSONL using s2_compute_token_lengths.py (or a quick standalone script). Pick L so that drop rate is acceptable (typically <0.1%).

# Inside container, on any single node
python offline_packing/s2_compute_token_lengths.py \
    --jsonl <path/to/full.jsonl> \
    --tokenizer <path/to/tokenizer> \
    --image-processor Qwen2_5_VLProcessor \
    --factor 48 --min-pixels 3136 --max-pixels 4000000 \
    --output <path/to/token_lens.txt>

# Then quickly inspect distribution (max, p99, count > L) before committing to L

1. Split JSONL across nodes / 切分 JSONL

TOTAL=$(wc -l < full.jsonl)
HALF=$(( (TOTAL + 1) / 2 ))
split -l $HALF -d --additional-suffix=.jsonl full.jsonl part_
# produces part_00.jsonl, part_01.jsonl

For >2 nodes, adjust -l accordingly.

2. Mount NFS on every node / 挂载 NFS

Make sure the same NFS is mounted on every node at the same path so paths in JSONL and outputs match.

3. Launch container on every node / 每台启动容器

Use the project's standard docker image with the repo bind-mounted. Working directory should be the repo root.

4. Run auto_pipe.sh in parallel on each node / 并行启动

On node A:

cd <repo_root>
bash offline_packing/auto_pipe.sh \
    --jsonl <data_root>/part_00.jsonl \
    --tokenizer <tokenizer_path> \
    --image-processor Qwen2_5_VLProcessor \
    --factor 48 --min-pixels 3136 --max-pixels 4000000 \
    --image-root / \
    --sample-class PackedCaptioningSample \
    --shard-prefix <dataset_name>_a \
    --output-dir <output_root>/node_a \
    --seq-len 4096 \
    --no-npy \
    2>&1 | tee <log_dir>/node_a.log

On node B (in parallel):

bash offline_packing/auto_pipe.sh \
    --jsonl <data_root>/part_01.jsonl \
    ... \
    --shard-prefix <dataset_name>_b \
    --output-dir <output_root>/node_b \
    ... \
    2>&1 | tee <log_dir>/node_b.log

[!IMPORTANT]

• --shard-prefix must differ between nodes so tar filenames don't collide
• --output-dir must differ between nodes
• --image-root / if JSONL paths are absolute; otherwise set it to the image root prefix
• Add --no-npy if the JSONL has no patch_positions field

5. Verify outputs / 验证产物

# Tar count per node should match s4 log
ls <output_root>/node_a/webdataset/*.tar | wc -l
ls <output_root>/node_b/webdataset/*.tar | wc -l

# Bin count + capacity utilization printed by s3
grep -E "(bins|util|efficiency)" <log_dir>/node_a.log

# Inspect one tar to confirm sample schema
mkdir -p /tmp/tar_inspect && cd /tmp/tar_inspect
tar -xf <output_root>/node_a/webdataset/<prefix>-000000.tar
ls | head
python -c "import json; d=json.load(open(open(__import__('glob').glob('*.json')[0]).name)); print(list(d.keys()))"

Expected JSON top-level keys for PackedCaptioningSample: images, prompts, captions, sample_count, patch_positions, timestamp_decimal.

patch_positions=[[""]] is normal under --no-npy — the auto-generated sample_loader.py handles it via sample.get(..., None).

6. Write top-level Metadataset yaml / 写顶层 Metadataset yaml

__module__: megatron.energon
__class__: Metadataset
splits:
  train:
    datasets:
      - weight: <num_samples_in_part_00>
        path: <output_root>/node_a/webdataset
        subflavors:
          augmentation: false
      - weight: <num_samples_in_part_01>
        path: <output_root>/node_b/webdataset
        subflavors:
          augmentation: false

Notes / 注意:

• weight is sample-count proportional (use the original input line count of each part, not the bin count)
• Use absolute paths — Energon resolves them as-is
• Add val split only if you actually need one; for SFT-only pipelines omit it
• subflavors is optional; here we mark augmentation: false since data is already packed

Common Pitfalls / 常见坑

Pitfall 1: forgetting --no-npy

Without it s2 hunts for per-sample .npy files and either spams warnings OR falls back to the slow real-image tokenization path. Always check the source JSONL for a patch_positions field first.

Pitfall 2: same --shard-prefix on multiple nodes

Tar filenames collide on shared NFS, second node overwrites the first. Use distinct suffixes (_a, _b, _n0, _n1, …).

Pitfall 3: starting fresh while old rm -rf still running on NFS

Removing millions of small files from NFS can take 10+ minutes. Don't wait — use a fresh _v2 output dir, and mv (atomic rename on same FS) when done if you want the original name back.

Pitfall 4: du -sh / rm -rf exceeding bash 120s timeout

Run them as nohup ... & and poll the PID, or run them in tmux.

Pitfall 5: container vs host timezone skew

Container time may differ from host time but wall clock is the same. Don't be confused by log timestamps when comparing across docker exec sessions.

Pitfall 6: choosing the wrong --sample-class

The class determines the auto-generated sample_loader.py and how downstream code unpacks tars. Confirm by reading the dataclass file (e.g. aiak_training_llm/data/multimodal/flavors/packed_captioning.py) and matching its fields to your JSONL shape.

Pitfall 7: weight confusion in Metadataset yaml

Use sample counts (or proportional integers), not bin counts. Energon samples each dataset proportional to weight.

Performance Reference / 性能参考

For ~390k samples per node at L=4096 on a multi-core machine with NFS storage:

• s1 (split + validate): ~30 min (NFS small-file IO bound; gets worse if concurrent rm is running)
• s2 (token lengths, with --no-npy): ~5 min
• s3 (BFD bin packing): <1 min
• s4 (tar writing): ~5 min
• Total wall time per node: ~40–45 min

Two nodes in parallel ≈ same wall time as one node, so 2× throughput.

Typical packing efficiency at L=4096 with avg ~10 samples/bin: >99% capacity utilization.

Quick Sanity Checklist / 快速自检清单

Before declaring done:

• Tar count per node matches s4 log
• s3 log shows >95% capacity utilization
• At least one tar inspected; JSON keys match the chosen sample class
• sample_count in tar JSON > 0 (not all 1; means packing actually worked)
• Top-level dataset.yaml exists with absolute paths and correct weights
• (Optional) Energon load smoke test passes

Related Files / 相关文件

• offline_packing/auto_pipe.sh — pipeline driver
• offline_packing/s1_split_json_to_samples.py — JSONL → per-sample records
• offline_packing/s2_compute_token_lengths.py — token length computation
• offline_packing/s3_bin_packing.py — BFD bin packing
• offline_packing/s4_bins_to_webdataset.py — tar + idx + .nv-meta writer
• aiak_training_llm/data/multimodal/flavors/packed_captioning.py — PackedCaptioningSample dataclass
• aiak_training_llm/data/multimodal/flavors/multi_mix_qa.py — MultiMixQASample dataclass
• aiak_megatron/examples/multimodal/sft_dataset.yaml — Metadataset yaml template

Purpose / 用途

Use this skill when analyzing, debugging, or tuning LengthPoolSortDataset — the cross-rank length synchronization mechanism used in this repository's training pipeline.

在分析、调试或调优本仓库训练 pipeline 中的跨 rank 长度同步机制 LengthPoolSortDataset 时，使用这个 skill。

This skill is specifically for:

• aiak_training_llm/data/multimodal/length_sort_dataset.py
• Understanding why training speed improves when length_sort_pool_size > 0
• Tuning pool_size for optimal multi-GPU efficiency
• Diagnosing rank synchronization bottlenecks

这个 skill 专门用于：

• aiak_training_llm/data/multimodal/length_sort_dataset.py
• 理解为什么 length_sort_pool_size > 0 时训练速度提升
• 调优 pool_size 以获得最佳多卡效率
• 排查 rank 间同步瓶颈

Core mechanism / 核心机制

Three-step pipeline / 三步流水线

上游 dataset → 累积 pool_size 个 sample → 按序列长度排序 → 用确定性 seed shuffle → 逐个 yield

for batch_idx, sample in enumerate(self.dataset):
    pool.append(sample)
    if len(pool) >= self.pool_size:
        pool.sort(key=self.key_fn)                     # 1. 按长度排序
        shuffle_seed = 42 + batch_idx                   # 2. 确定性 seed
        random.Random(shuffle_seed).shuffle(pool)       # 3. 同 seed shuffle
        for s in pool:
            yield s
        pool.clear()

Pipeline position / 在 pipeline 中的位置

CrudeWebdataset → ShuffleBuffer → cook_crude_sample → encode_sample
    → LengthPoolSortDataset → BatchDataset → EpochizeDataset → LogSampleDataset

Inserted after encode_sample (where total_len / tokens are available) and before BatchDataset.

插在 encode_sample 之后（此时已有 total_len / tokens）、BatchDataset 之前。

Activated by: --length-sort-pool-size N (where N > 0).

通过 --length-sort-pool-size N（N > 0）激活。

Why it accelerates training / 为什么能加速训练

The problem / 问题

In multi-GPU data-parallel training, all ranks must synchronize at each step (gradient all-reduce). If different ranks process samples of very different lengths, fast ranks idle waiting for slow ranks.

多卡数据并行训练中，所有 rank 每步都要同步（梯度 all-reduce）。如果不同 rank 处理的 sample 长度差异很大，快的 rank 空等慢的 rank。

无 pool sort:
  Rank 0 step 100: 长度 200 → 0.5s    ┐
  Rank 1 step 100: 长度 5000 → 3.0s   ├→ 所有 rank 等 3.0s
  Rank 2 step 100: 长度 300 → 0.6s    ┘

The solution / 解决方案

Sort + same-seed shuffle ensures all ranks yield samples of approximately the same length at the same time.

排序 + 同 seed shuffle 保证所有 rank 在同一时刻输出近似相同长度的 sample。

有 pool sort:
  Rank 0 step 100: 长度 ~200 → 0.5s   ┐
  Rank 1 step 100: 长度 ~200 → 0.5s   ├→ 所有 rank 等 0.5s
  Rank 2 step 100: 长度 ~200 → 0.5s   ┘

Why it works — step by step / 为什么有效——逐步分析

Step 1: Sort aligns the i-th position across ranks / 排序对齐各 rank 第 i 个位置

Each rank's data comes from different shards of the same dataset, so length distributions are approximately identical. After sorting, the i-th position in each rank's pool corresponds to the i-th quantile of its length distribution. Similar distributions → similar quantiles → similar lengths.

各 rank 的数据来自同一数据集的不同 shard，长度分布近似相同。排序后，各 rank pool 中第 i 个位置对应各自长度分布的第 i 个分位数。分布相似 → 分位数相似 → 长度相似。

Rank 0 排序后: [100, 102, 105, 108, ..., 4998, 5000]
Rank 1 排序后: [101, 103, 106, 109, ..., 4999, 5001]
Rank 2 排序后: [ 99, 104, 107, 110, ..., 4997, 5002]

Step 2: Same seed preserves alignment after shuffle / 同 seed 保持 shuffle 后的对齐

shuffle_seed = 42 + batch_idx is identical across all ranks → same permutation applied to all pools. Since the i-th position had similar lengths, the shuffled output at the same position still has similar lengths.

shuffle_seed = 42 + batch_idx 对所有 rank 相同 → 相同排列应用于所有 pool。由于第 i 个位置的长度相似，shuffle 后同一位置的输出长度仍然相似。

permutation = [3, 0, 4, 1, 2]:
  Rank 0 输出: [108, 100, 5000, 102, 105]
  Rank 1 输出: [109, 101, 5001, 103, 106]
  Rank 2 输出: [110,  99, 5002, 104, 107]
  → 同一位置长度近似一致

Why both sort AND shuffle are needed / 为什么排序和 shuffle 缺一不可

Sort solves cross-rank synchronization. Shuffle solves temporal uniformity. Both are required.

排序解决跨 rank 同步问题，shuffle 解决时序均匀性问题。二者缺一不可。

pool_size tuning / pool_size 调优

Larger pool_size → more accurate quantile estimation across ranks → better synchronization → less idle time.

pool_size 越大 → 各 rank 分位数估计越准 → 同步越好 → 空等越少。

Rule of thumb: pool_size should be significantly larger than batch_size, ideally 10x-100x.

经验法则：pool_size 应远大于 batch_size，理想情况下 10x-100x。

Multi-worker behavior (num_workers > 1) / 多 worker 行为

When num_workers > 1, each worker runs an independent LengthPoolSortDataset instance with its own pool.

当 num_workers > 1 时，每个 worker 运行独立的 LengthPoolSortDataset 实例，各自维护独立 pool。

• Correctness: No issue. Each worker sorts its own shard subset independently.

• Synchronization effect: Diluted. DataLoader interleaves outputs from multiple workers, partially disrupting the within-pool length ordering.

• Recommendation: num_workers=1 gives the strongest synchronization effect. If num_workers > 1 is needed for I/O throughput, increase pool_size proportionally.

• 正确性：无问题。每个 worker 独立排序自己的 shard 子集。

• 同步效果：被稀释。DataLoader 交替取多个 worker 的输出，部分打乱 pool 内的长度顺序。

• 建议：num_workers=1 同步效果最强。如果需要 num_workers > 1 提升 I/O 吞吐，可按比例增大 pool_size。

Checkpoint resume caveat / checkpoint 恢复注意事项

save_state / restore_state delegate directly to the upstream dataset. The pool's internal state (accumulated but not yet yielded samples) is NOT saved.

save_state / restore_state 直接委托给上游 dataset。pool 内部状态（已累积但未 yield 的 sample）不会被保存。

On resume, up to pool_size - 1 samples may be re-ordered differently or skipped. This is generally negligible for large datasets.

恢复时，最多 pool_size - 1 个 sample 可能会被不同地排序或跳过。对于大数据集来说通常可以忽略。

key_fn / 排序键

key_fn=lambda s: getattr(s, "total_len", len(getattr(s, "tokens")))

Prefers total_len attribute (set by encode_sample), falls back to len(tokens).

优先使用 total_len 属性（由 encode_sample 设置），回退到 len(tokens)。

What to check during debugging / 调试时要检查什么

• Is --length-sort-pool-size actually set and > 0?

• What is the actual length distribution of the dataset? (High variance → more benefit from pool sort)

• Are all ranks receiving data from the same underlying dataset with similar length distributions?

• Is num_workers > 1 diluting the synchronization effect?

• After enabling, did step time variance across steps decrease?

• Is memory usage acceptable for the chosen pool_size?

• --length-sort-pool-size 是否真正设置且 > 0？

• 数据集的实际长度分布如何？（方差越大 → pool sort 收益越大）

• 所有 rank 是否从同一底层数据集接收长度分布相似的数据？

• num_workers > 1 是否稀释了同步效果？

• 启用后，各 step 的 step time 方差是否降低了？

• 所选 pool_size 的内存占用是否可接受？

Expected outputs when using this skill / 使用本 skill 时的期望输出

When asked to analyze a training speed issue related to length sorting, return:

当被要求分析与长度排序相关的训练速度问题时，应返回：

1. Whether LengthPoolSortDataset is active and with what pool_size

2. The dataset's length distribution characteristics (high/low variance)

3. The relationship between pool_size, batch_size, and num_workers

4. Whether cross-rank synchronization is the bottleneck

5. Recommended pool_size adjustment if needed

6. LengthPoolSortDataset 是否激活，pool_size 是多少

7. 数据集的长度分布特征（高/低方差）

8. pool_size、batch_size、num_workers 之间的关系

9. 跨 rank 同步是否是瓶颈

10. 如需要，推荐的 pool_size 调整

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 测试。

Purpose / 用途

Use this skill when diagnosing, designing, or converting Megatron/Megatron-Core checkpoints that may use TP, PP, and EP.

在排查、设计或转换使用 TP、PP、EP 的 Megatron / Megatron-Core checkpoint 时，使用这个 skill。

Core rule / 核心规则

• TP only: mp_rank_{tp}

• TP + PP: mp_rank_{tp}_{pp}

• TP + PP + EP: mp_rank_{tp}_{pp}_{ep}

• 只有 TP：mp_rank_{tp}

• TP + PP：mp_rank_{tp}_{pp}

• TP + PP + EP：mp_rank_{tp}_{pp}_{ep}

The key discriminator is whether expert parallelism participates in checkpoint sharding.

真正的分界点是：expert parallelism 是否参与了 checkpoint 切分。

• If EP is present, treat the checkpoint layout as 3D.

• If EP is absent, treat the checkpoint layout as non-EP and use 1D or 2D.

• 如果存在 EP，就按 3D 布局处理。

• 如果不存在 EP，就按非 EP 布局处理，即 1D 或 2D。

Mental model / 心智模型

Megatron does not treat pp > 1 as meaning 3D by itself.

Megatron 不会因为 pp > 1 就自动把 checkpoint 视为 3D。

• PP adds a pipeline index.

• EP adds an expert index.

• The third coordinate exists because EP exists, not because PP exists.

• PP 只是在目录里增加 pipeline 这一维。

• EP 才会增加 expert 这一维。

• 第三维存在的原因是 EP 存在，而不是因为 PP 存在。

So even if tp=1 and pp=1, once EP is enabled the checkpoint naming is still conceptually 3D because ranks are addressed by (tp, pp, ep).

所以即使 tp=1 且 pp=1，只要启用了 EP，checkpoint 在语义上仍然是 3D，因为 rank 仍然由 (tp, pp, ep) 共同定位。

Practical interpretation / 实际使用解释

When reading or converting checkpoints:

在读取或转换 checkpoint 时：

1. First decide whether EP exists in the checkpoint contract.

2. If EP exists, require mp_rank_{tp}_{pp}_{ep}.

3. If EP does not exist, read as mp_rank_{tp} or mp_rank_{tp}_{pp}.

4. Do not infer 3D solely from pipeline_model_parallel_size > 1.

5. 先判断这个 checkpoint 契约里是否存在 EP。

6. 如果存在 EP，就要求目录是 mp_rank_{tp}_{pp}_{ep}。

7. 如果不存在 EP，就按 mp_rank_{tp} 或 mp_rank_{tp}_{pp} 去读。

8. 不要仅凭 pipeline_model_parallel_size > 1 就推断它一定是 3D。

Typical failure pattern / 典型错误模式

Bad assumption:

错误假设：

• pp > 1 so loader chooses a 3D reader.

• 只要 pp > 1，loader 就应该走 3D reader。

Why it fails:

为什么会失败：

• Dense non-MoE checkpoints with TP+PP usually use mp_rank_{tp}_{pp} only.

• A 3D loader then looks for an EP coordinate that is not present.

• 普通 dense、非 MoE 的 TP+PP checkpoint，通常只有 mp_rank_{tp}_{pp}。

• 这时 3D loader 会去找并不存在的 EP 坐标，最终报错。

Recommended repo-local contract / 当前仓库建议契约

For this repository, follow this rule:

这个仓库建议遵循以下规则：

• if expert_parallel_size is passed, require 3D

• if expert_parallel_size is not passed, use non-EP loading

• 如果传了 expert_parallel_size，就强制按 3D 处理

• 如果没有传 expert_parallel_size，就走非 EP 的加载逻辑

This matches Megatron's path-building logic better than using pp > 1 as the branch condition.

这个规则比“用 pp > 1 作为分支条件”更贴近 Megatron 自己的路径生成逻辑。

What to check during debugging / 调试时要检查什么

• What are the actual shard directory names under the checkpoint root?

• Was expert_parallel_size provided by the caller?

• Is the model dense or MoE?

• Is the loader branching on EP or incorrectly branching on PP?

• If conversion failed, which exact mp_rank_* pattern was expected and which one exists on disk?

• checkpoint 根目录下，实际 shard 目录名是什么？

• 调用方是否传入了 expert_parallel_size？

• 当前模型是 dense 还是 MoE？

• loader 是按 EP 分支，还是错误地按 PP 分支？

• 如果转换失败，程序期望的 mp_rank_* 模式是什么，磁盘上实际又是什么？

Expected outputs when using this skill / 使用本 skill 时的期望输出

When asked to analyze a checkpoint issue, return:

当你被要求分析 checkpoint 问题时，应该返回：

1. the inferred layout class: 1D, 2D, or 3D

2. the reason for that classification

3. the expected directory naming pattern

4. whether the caller should use a non-EP loader or an EP-aware loader

5. any mismatch between runtime arguments and on-disk shard layout

6. 推断出的布局类别：1D、2D 或 3D

7. 这样分类的原因

8. 期望的目录命名模式

9. 调用方应使用非 EP loader 还是 EP-aware loader

10. 运行时参数与磁盘上 shard 布局之间是否存在不匹配

Purpose / 用途

Use this skill when merging a standalone ViT encoder and LLM into a unified LlavaOnevision2 HuggingFace checkpoint, and when validating that the merged weights and inference outputs are consistent with the originals.

当需要将独立的 ViT encoder 和 LLM 合并成统一的 LlavaOnevision2 HuggingFace checkpoint，并验证合并后权重和推理输出与原始模型一致时，使用这个 skill。

Prerequisites / 前置条件

• Container llava_megatron_container_ax running with GPU access
• All paths below assume execution inside the container at /workspace/LLaVA-OneVision-2
• PYTHONPATH=transformers_impl:. must be set for all Python commands
• For large models, use tmpfs (/train_tmp) for I/O performance

容器 llava_megatron_container_ax 需启动并有 GPU 访问权限。以下所有路径假设在容器内 /workspace/LLaVA-OneVision-2 执行。所有 Python 命令需设置 PYTHONPATH=transformers_impl:.。大模型建议用内存盘 /train_tmp。

Architecture / 架构

What merge_ov2 does / merge_ov2 做了什么

ViT encoder (e.g. onevision_encoder_patch16_0424)
  + LLM (e.g. Qwen3-4B-Instruct-2507)
  + Processor (e.g. lmms-lab-encoder/LLaVA-OneVision-2-8B-Instruct)
  → Unified LlavaOnevision2ForConditionalGeneration checkpoint

Key transformations during merge:

合并时的关键转换：

1. ViT weights are prefixed with visual. (e.g. encoder.layers.0.self_attn.q_proj.weight → visual.encoder.layers.0.self_attn.qkv.weight)
2. QKV fusion: separate q_proj / k_proj / v_proj are concatenated into fused self_attn.qkv (introduces ~1e-7 bf16 divergence)
3. LLM weights are prefixed with language_model. (e.g. model.layers.0.self_attn.q_proj.weight → language_model.model.layers.0.self_attn.q_proj.weight)
4. Adapter (multi_modal_projector) is randomly initialized if no adapter checkpoint is provided
5. layernorm_post from ViT is dropped (not used in LlavaOnevision2)
6. class_embedding may not exist in some ViT encoders (e.g. patch16 variant)

Source code layout / 源码结构

transformers_impl/merge_ov2/
├── __main__.py          # CLI entry point
├── cli.py               # Argument parsing for merge / validate / dry-run
├── remap.py             # Weight key remapping logic
├── loader.py            # Weight loading from source checkpoints
├── save.py              # Save merged checkpoint
├── io.py                # I/O utilities
├── utils.py             # Shared utilities
├── variants/
│   ├── dense.py         # Dense model variant
│   └── moe.py           # MoE model variant
└── validators/
    ├── vit_layerwise.py   # ViT layer-wise weight validator
    ├── vit_blockorder.py  # ViT block-order validator (patch14+sms=2 only)
    ├── llm_parallel.py    # LLM parallel validator
    ├── llm_sequential.py  # LLM sequential validator
    └── e2e.py             # End-to-end validator

CLI Reference / CLI 参考

Subcommand: merge

Remap + load + (optional validate) + save.

PYTHONPATH=transformers_impl:. python -m merge_ov2 merge \
  --variant dense \
  --vit /path/to/vit_encoder \
  --llm /path/to/llm \
  --processor /path/to/processor \
  --out /path/to/output \
  --spatial-merge-size 2 \
  --target-dtype bf16 \
  --vit-validator-strategy layerwise

Subcommand: validate

Validate an already-merged checkpoint against original sources.

验证已合并的 checkpoint 与原始模型的一致性。

PYTHONPATH=transformers_impl:. python -m merge_ov2 validate \
  --variant dense \
  --ckpt /path/to/merged_checkpoint \
  --vit /path/to/vit_encoder \
  --llm /path/to/llm \
  --processor /path/to/processor \
  --vit-validator-strategy layerwise

Subcommand: dry-run

Remap only; report load coverage; no save.

仅做 remap，报告加载覆盖率，不保存。

PYTHONPATH=transformers_impl:. python -m merge_ov2 dry-run \
  --variant dense \
  --vit /path/to/vit_encoder \
  --llm /path/to/llm \
  --processor /path/to/processor
# --spatial-merge-size defaults to 2 (the current production p14m2 variant);
# pass --spatial-merge-size 3 explicitly for the legacy p14m33 / p16m3 layouts.

Concrete Example / 具体示例

Merging Qwen3-4B + onevision-encoder-large-lang-tf57 (patch14, sms=2, current production)

docker exec llava_megatron_container_ax bash -c '
cd /workspace/LLaVA-OneVision-2 && \
PYTHONPATH=transformers_impl:. python -u -m merge_ov2 merge \
  --variant dense \
  --vit /train_tmp/onevision-encoder-large-lang-tf57 \
  --llm /train_tmp/Qwen3-4B-Instruct-2507 \
  --processor lmms-lab-encoder/LLaVA-OneVision-2-8B-Instruct \
  --out /train_tmp/llava_onevision2_4b_p14m2 \
  --target-dtype bf16 \
  --vit-validator-strategy layerwise \
  --img /train_tmp/sample.jpg \
  --sample-text "Hello, world!"
'

--spatial-merge-size is omitted because the CLI defaults to 2. Output checkpoint config: patch_size=14, spatial_merge_size=2, image_size=28*N, hidden_size=1024, 24 ViT layers, 36 LLM layers. The sample image must be a multiple of patch_size * spatial_merge_size = 28 in both dimensions (e.g. 504x504 — still works since 504 = 28 × 18).

--spatial-merge-size 省略，CLI 缺省为 2。输出 checkpoint 的配置: patch_size=14, spatial_merge_size=2, image_size=28*N, hidden_size=1024, 24 ViT 层, 36 LLM 层。样图尺寸必须是 patch_size * spatial_merge_size = 28 的整数倍（如 504x504，因 504 = 28 × 18 仍合法）。

Legacy 4b_p14m33 (sms=3) merge: same command, append --spatial-merge-size 3 and change --out to /train_tmp/llava_onevision2_4b_p14m33. Image size must then be a multiple of 42.

旧的 4b_p14m33（sms=3）合并：同样命令，加 --spatial-merge-size 3，--out 改为 /train_tmp/llava_onevision2_4b_p14m33。样图尺寸需为 42 的整数倍。

Variant Cheat Sheet / Variant 参数对照表

OV2 4B has two real-world variants. The --patch-size (set by the ViT checkpoint), --spatial-merge-size, and --vit-validator-strategy must all match — using the wrong validator silently crashes inside reshape ops.

OV2 4B 有两套真实使用的 variant。--patch-size（由 ViT checkpoint 决定）、 --spatial-merge-size 和 --vit-validator-strategy 必须配齐 —— 用错 validator 会在 reshape 里直接崩。

Why layerwise is required for non-(patch14+sms=2) variants / 为什么非 patch14+sms=2 的 variant 必须用 layerwise: vit_blockorder.py's reshape hard-codes patch_size=14, spatial_merge_size=2. Any other combo (sms=3, or different patch_size with sms≠2) breaks the reshape: RuntimeError: shape '[...]' is invalid for input of size N. 4b and 4b_p14m2 are the only variants that satisfy the hardcoded assumption, so they may use either blockorder or layerwise; everything else must use layerwise.

vit_blockorder.py 的 reshape 写死了 patch_size=14, spatial_merge_size=2 的尺寸假设。任何其他组合（sms=3，或非 2 的 sms）都会让维度对不上， 抛 RuntimeError: shape '[...]' is invalid for input of size N。 只有 4b 和 4b_p14m2 满足硬编码假设，可以用 blockorder 或 layerwise； 其余 variant 必须用 layerwise。

Post-Merge Validation / 合并后验证

Always prefer the validate subcommand over hand-written scripts. It runs the same three validators (vit, llm, e2e) used during merge, against an already-saved checkpoint. Use this when:

• you skipped validation during merge (--validate-skip)
• you manually patched a saved checkpoint and want to re-verify
• you downloaded a checkpoint and want to confirm parity

优先使用 validate 子命令而不是手写脚本。它会对一份已保存的 checkpoint 跑和 merge 时同样的三个 validator（vit、llm、e2e）。适用场景：

• merge 时跳过了验证（--validate-skip）
• 手动改过 checkpoint 后重验证
• 下载了 checkpoint 想确认一致性

docker exec llava_megatron_container_ax bash -c '
cd /workspace/LLaVA-OneVision-2 && \
PYTHONPATH=transformers_impl:. python -m merge_ov2 validate \
  --variant dense \
  --ckpt /train_tmp/llava_onevision2_4b_p14m2 \
  --vit /train_tmp/onevision-encoder-large-lang-tf57 \
  --llm /train_tmp/Qwen3-4B-Instruct-2507 \
  --processor lmms-lab-encoder/LLaVA-OneVision-2-8B-Instruct \
  --img /train_tmp/sample.jpg \
  --sample-text "Hello, world!" \
  --vit-validator-strategy layerwise
'

Skipping validators / 跳过 validator

--validate-skip is append-style; pass it once per validator to skip. Three validators exist: vit, llm, e2e. Run only the ones you need to save GPU time.

--validate-skip 是 append 模式，每跳一个就传一次。一共三个 validator： vit、llm、e2e。只跑你需要的，省 GPU 时间。

# Skip e2e only (run vit + llm) / 只跳 e2e
... --validate-skip e2e

# Validate only LLM (skip vit + e2e) / 只验证 LLM
... --validate-skip vit --validate-skip e2e

# Skip everything (no parity check) / 全跳过
... --validate-skip vit --validate-skip llm --validate-skip e2e

When you skip vit, --qwen-processor and --img may be omitted; when you skip llm, --sample-text may be omitted. The CLI enforces only the flags required by the validators you actually run.

跳了 vit 时可以省略 --qwen-processor 和 --img；跳了 llm 时可以 省略 --sample-text。CLI 只对实际要跑的 validator 强制要求对应的参数。

Next step: Megatron conversion + consistency / 下一步：转 Megatron + 一致性测试

After merge + validate succeeds, the standard next step is HF → mcore conversion plus the 6-check end-to-end HF↔mcore consistency suite. That workflow lives in the llava-onevision2-consistency skill — load it with skill(name="llava-onevision2-consistency").

merge + validate 通过后，标准下一步是 HF → mcore 转换 + 6 项 HF↔mcore 端到端 一致性测试。流程在 llava-onevision2-consistency skill 里 —— skill(name="llava-onevision2-consistency") 加载。

Reverse direction: mcore → HF (deploy / round-trip) / 反向：mcore → HF（部署 / 回环验证）

For the p14m2 variant, the reverse conversion ships as two scripts:

# 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

Round-trip mcore → HF → mcore (TP=1 PP=1) is bitwise identical to the original mcore checkpoint (588 non-empty tensors, max abs diff = 0.000e+00). This is the strongest correctness guarantee for the reverse path. Use round-trip when changing TP/PP layout without retraining.

对于 p14m2 variant，反向转换提供两个脚本（mcore→HF 用于部署/debug/推理， mcore→release 用于通过 HF 中转改 TP/PP 切分）。回环 mcore→HF→mcore 在 TP=1 PP=1 下与原始 mcore checkpoint 逐位一致（588 个非空 tensor，max abs diff = 0.000e+00）， 这是反向路径正确性的最强保证。在不重训的前提下改 TP/PP layout 时使用回环。

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

convert_4b_p14m2_mcore_to_hf.sh 会自动检测 <load>/release — 父目录 (/path/to/mcore_ckpt) 或显式 release 路径 (/path/to/mcore_ckpt/release) 都可以传。Sibling 脚本（4b、p14m3、p16m3、8b、30b） 仍然要求显式 /release 路径。

Manual validation scripts (OOM fallback) / 手写脚本（OOM 后备方案）

Only use these when the built-in validators OOM (e.g. 30B+ MoE on a single 80 GB GPU). They load only the visual or language sub-component to halve memory pressure.

只在内置 validator OOM 时用（如 30B+ MoE 单卡 80 GB）。手写脚本只加载 visual 或 language 子组件来缓解显存压力。

Test 1: ViT Weight Consistency / ViT 权重一致性

Compare every ViT weight tensor between the original encoder and the merged checkpoint.

逐 tensor 比较原始 encoder 和合并 checkpoint 中的 ViT 权重。

# Inside container, PYTHONPATH=transformers_impl:.
import torch
from safetensors.torch import load_file
import torch.nn.functional as F

merged_dir = "/train_tmp/llava_onevision2_4b_p14m2"
vit_dir = "/train_tmp/onevision-encoder-large-lang-tf57"

# Load weights
merged_w = {}
for f in ["model-00001-of-00002.safetensors", "model-00002-of-00002.safetensors"]:
    merged_w.update(load_file(f"{merged_dir}/{f}"))
vit_w = load_file(f"{vit_dir}/model.safetensors")

# Compare non-QKV weights (direct mapping with "visual." prefix)
for vit_key, vit_val in vit_w.items():
    if "q_proj" in vit_key or "k_proj" in vit_key or "v_proj" in vit_key:
        continue
    if "layernorm_post" in vit_key:  # dropped in merge
        continue
    merged_key = f"visual.{vit_key}"
    merged_val = merged_w[merged_key]
    cos = F.cosine_similarity(vit_val.flatten().float(), merged_val.flatten().float(), dim=0)
    assert cos > 0.9999, f"FAIL {vit_key}: cos={cos}"

# Compare QKV weights (fused: q+k+v → qkv)
for i in range(24):  # 24 encoder layers
    for suffix in ["weight", "bias"]:
        q = vit_w[f"encoder.layers.{i}.self_attn.q_proj.{suffix}"]
        k = vit_w[f"encoder.layers.{i}.self_attn.k_proj.{suffix}"]
        v = vit_w[f"encoder.layers.{i}.self_attn.v_proj.{suffix}"]
        fused = torch.cat([q, k, v], dim=0)
        merged = merged_w[f"visual.encoder.layers.{i}.self_attn.qkv.{suffix}"]
        cos = F.cosine_similarity(fused.flatten().float(), merged.flatten().float(), dim=0)
        assert cos > 0.9999, f"FAIL QKV layer {i} {suffix}: cos={cos}"

print("ViT weight consistency OK")

Expected: All cosine similarities > 0.9999. QKV may have ~1e-7 divergence due to bf16 cat.

预期: 所有余弦相似度 > 0.9999。QKV 因 bf16 拼接可能有 ~1e-7 的微小差异。

Test 2: ViT Inference Consistency / ViT 推理一致性

Full forward pass through independent patch_embed → layernorm_pre → 24 encoder layers.

独立的 patch_embed → layernorm_pre → 24 encoder 层的完整前向传播。

import torch, sys
import torch.nn.functional as F
from PIL import Image
from transformers import AutoModel, AutoModelForCausalLM, CLIPImageProcessor

DEVICE = torch.device("cuda:0")
DTYPE = torch.bfloat16
merged_dir = "/train_tmp/llava_onevision2_4b_p14m2"
vit_dir = "/train_tmp/onevision-encoder-large-lang-tf57"
sms, patch_size = 2, 14
pixel_unit = patch_size * sms  # 28

# Use a small synthetic image (must be multiple of pixel_unit)
image = Image.new("RGB", (504, 504), color="red")
h, w = 504, 504

# Load merged model's visual component
model = AutoModelForCausalLM.from_pretrained(
    merged_dir, torch_dtype=DTYPE,
    low_cpu_mem_usage=True, trust_remote_code=True,
    attn_implementation="flash_attention_2",
)
merged_visual = model.model.visual.to(DEVICE).eval()
del model.model.language_model
import gc; gc.collect()

# Load original ViT — IMPORTANT: use AutoModel + trust_remote_code, do NOT
# import from `transformers_impl/onevision_encoder` (local copy can drift
# from the modeling_*.py shipped inside the checkpoint, producing sim≈-0.02)
orig_vit = AutoModel.from_pretrained(
    vit_dir, torch_dtype=DTYPE, trust_remote_code=True,
    attn_implementation="flash_attention_2",
).to(DEVICE).eval()

# Prepare pixel values
clip_proc = CLIPImageProcessor.from_pretrained(vit_dir)
clip_px = clip_proc(images=image, return_tensors="pt",
                    do_resize=False, do_center_crop=False)["pixel_values"]
clip_px = clip_px.to(dtype=DTYPE, device=DEVICE)
grid_h, grid_w = h // patch_size, w // patch_size

# Use orig_vit's full forward (row-major output)
with torch.no_grad(), torch.amp.autocast("cuda", dtype=DTYPE):
    orig_out = orig_vit(clip_px).last_hidden_state  # (1, N, D)

    # Merged visual: block layout forward
    # NOTE: canonical RoPE helpers live in merge_ov2/utils.py — do NOT copy
    # them inline. The Megatron-side canonical implementation is in
    # aiak_training_llm/models/llava_onevision2/onevision_encoder_model.py
    # but cannot be imported from transformers_impl (would create a reverse
    # dep on the training framework).
    from merge_ov2.utils import (
        convert_rope_to_block_layout_by_positions,
        rowmajor_to_block,
    )

    def extract_block_patches(img_tensor, ps, s):
        b, c, ph, pw = img_tensor.shape
        h2, w2 = ph // ps, pw // ps
        patches = img_tensor.reshape(b, c, h2, ps, w2, ps).permute(0, 2, 4, 1, 3, 5).reshape(h2, w2, c, ps, ps)
        h_m, w_m = h2 // s, w2 // s
        patches = patches.reshape(h_m, s, w_m, s, c, ps, ps).permute(0, 2, 1, 3, 4, 5, 6).contiguous()
        return patches.reshape(-1, c, ps, ps)

    block_patches = extract_block_patches(clip_px, ps=patch_size, s=sms)
    merged_pre = merged_visual.layernorm_pre(merged_visual.embeddings(block_patches).unsqueeze(0))

    # Build RoPE
    grid_thw = torch.tensor([[1, grid_h, grid_w]], device=DEVICE)
    t_idx = torch.arange(1, device=DEVICE, dtype=torch.float32)
    h_idx = torch.arange(grid_h, device=DEVICE, dtype=torch.float32)
    w_idx = torch.arange(grid_w, device=DEVICE, dtype=torch.float32)
    mt, mh, mw = torch.meshgrid(t_idx, h_idx, w_idx, indexing="ij")
    patch_positions = torch.stack([mt, mh, mw], dim=-1).reshape(-1, 3)
    merged_freqs = merged_visual.video_rope.forward_from_positions(patch_positions)
    merged_freqs = convert_rope_to_block_layout_by_positions(
        merged_freqs, patch_positions, spatial_merge_size=sms, grid_thw=grid_thw)
    block_rope = torch.cat([merged_freqs, merged_freqs], dim=-1).unsqueeze(0)

    # Run encoder layers
    merged_h = merged_pre
    for i in range(len(merged_visual.encoder.layers)):
        merged_h = merged_visual.encoder.layers[i](
            merged_h, attention_mask=None, rotary_pos_emb=block_rope,
            output_attentions=False, cu_seqlens=None, max_seqlen=None)[0]

    # Convert orig row-major output to block layout for comparison
    # (rowmajor_to_block already imported from merge_ov2.utils above)
    orig_block = rowmajor_to_block(orig_out[0], 1, grid_h, grid_w, sms)
    cos = F.cosine_similarity(merged_h[0].flatten().float(), orig_block.flatten().float(), dim=0)
    diff = (merged_h[0] - orig_block).abs().mean().item()
    print(f"ViT inference: cos={cos:.8f}, diff={diff:.8e}")
    # bf16 24-layer accumulation: realistic min cos ≈ 0.98, not 0.999.
    # See "bf16 numerical thresholds" in Known Issues below.
    assert cos > 0.98, f"ViT inference mismatch: cos={cos}"

Expected: cos ≥ 0.98 (bf16 24-layer accumulation). Use fp32 for cos ≥ 0.999.

预期: cos ≥ 0.98（bf16 24 层累积）。要 cos ≥ 0.999 请用 fp32。

Note on image size: Use small images (e.g. 480x480) to avoid GPU OOM. The image dimensions must be multiples of patch_size * spatial_merge_size.

关于图像大小: 用小图（如 480x480）避免 GPU OOM。图像尺寸必须是 patch_size * spatial_merge_size 的整数倍。

Test 3: LLM Inference Consistency / LLM 推理一致性

Pure text forward pass comparing logits from the original LLM vs the merged model's language_model.

纯文本前向传播，比较原始 LLM 和合并模型的 language_model 的 logits。

import torch
import torch.nn.functional as F
from transformers import AutoModelForCausalLM, AutoTokenizer

DEVICE = torch.device("cuda:0")
DTYPE = torch.bfloat16
merged_dir = "/train_tmp/llava_onevision2_4b_p14m2"
llm_dir = "/train_tmp/Qwen3-4B-Instruct-2507"

tokenizer = AutoTokenizer.from_pretrained(llm_dir, trust_remote_code=True)
input_ids = tokenizer("Hello, world!", return_tensors="pt")["input_ids"].to(DEVICE)

# Load original LLM
orig_llm = AutoModelForCausalLM.from_pretrained(llm_dir, torch_dtype=DTYPE,
                                                 trust_remote_code=True).to(DEVICE).eval()
with torch.no_grad():
    orig_logits = orig_llm(input_ids).logits
del orig_llm
import gc; gc.collect(); torch.cuda.empty_cache()

# Load merged model's language_model
merged = AutoModelForCausalLM.from_pretrained(merged_dir, torch_dtype=DTYPE,
                                               low_cpu_mem_usage=True, trust_remote_code=True)
merged_lm = merged.model.language_model.to(DEVICE).eval()
del merged.model.visual
gc.collect()

with torch.no_grad():
    merged_logits = merged_lm(input_ids).logits

cos = F.cosine_similarity(orig_logits.flatten().float(), merged_logits.flatten().float(), dim=0)
diff = (orig_logits - merged_logits).abs().max().item()
print(f"LLM logits: cos={cos:.8f}, max_diff={diff:.8e}")
# bf16 logits: cos ≈ 0.9999, max_diff < 5e-2 is healthy. fp32 gives diff = 0.
assert cos > 0.999, f"LLM logits mismatch: cos={cos}"
assert diff < 5e-2, f"LLM logits diff too large: {diff}"

Expected: bf16 → cos ≈ 0.9999, max_diff < 5e-2 (LLM weights copied verbatim, only RMSNorm/MLP bf16 noise). fp32 → cos = 1.0, diff = 0.0.

预期: bf16 → cos ≈ 0.9999, max_diff < 5e-2（LLM 权重直接复制，只有 RMSNorm/MLP 的 bf16 噪声）。fp32 → cos = 1.0, diff = 0.0。

What is NOT tested / 未覆盖的部分

Known Issues & Workarounds / 已知问题和解决方案

1. blockorder ViT validator only works for patch14+sms=2

vit_blockorder.py does a reshape that hard-codes patch14+sms=2 dimensions. Any other combination (patch16+sms=3, patch14+sms=3, etc.) crashes with RuntimeError: shape '[...]' is invalid for input of size N. Use --vit-validator-strategy layerwise for everything except the legacy 4b (patch14+sms=2) variant.

vit_blockorder.py 的 reshape 写死了 patch14+sms=2 的尺寸。其他任何组合 （patch16+sms=3、patch14+sms=3 等）都会崩。除了旧的 4b（patch14+sms=2） 之外，全部用 --vit-validator-strategy layerwise。

2. GPU OOM during validation

The built-in validator loads the full merged model + original ViT simultaneously. For 4B+ models on a single 80GB GPU, this may OOM. Workaround: use the standalone scripts above (they load only the visual component, deleting language_model first).

内置 validator 同时加载完整合并模型和原始 ViT。4B+ 模型在单张 80GB GPU 上可能 OOM。解决方案：用上面的独立脚本（只加载 visual 部分，先删 language_model）。

3. Original ViT embeddings output shape mismatch

Some ViT encoders output (N, 1, D) from embeddings() while merged visual outputs (N, D). The full model() forward handles this internally, so use orig_vit(pixel_values).last_hidden_state instead of calling embeddings() + layers manually for the original ViT.

有些 ViT encoder 的 embeddings() 输出 (N, 1, D) 而合并后的 visual 输出 (N, D)。用 orig_vit(pixel_values).last_hidden_state 调用完整 forward 而非手动逐层调用。

4. Block layout conversion for comparison

Original ViT outputs features in row-major order; merged visual uses block layout (grouped by spatial_merge_size). Use rowmajor_to_block() to align before comparison.

原始 ViT 输出 row-major 顺序的特征；合并后的 visual 使用 block layout（按 spatial_merge_size 分组）。比较前用 rowmajor_to_block() 对齐。

5. CRITICAL — Load original ViT via AutoModel + trust_remote_code, NOT a local import

When validating against the original ViT, use:

from transformers import AutoModel
orig_vit = AutoModel.from_pretrained(
    vit_dir, torch_dtype=DTYPE, trust_remote_code=True,
    attn_implementation="flash_attention_2",
)

Do NOT do from onevision_encoder import OneVisionEncoderModel from transformers_impl/onevision_encoder/. The local copy of modeling_onevision_encoder.py can drift from the modeling_*.py shipped inside the checkpoint directory (e.g. RoPE construction, attention impl, embedding signature). When they disagree, layerwise sim collapses to ~−0.024 even though the weights are byte-identical, and the failure mode looks like "wrong weights" but isn't. This bit us during the 4b_p14m33 merge against onevision-encoder-large-lang-tf57.

AutoModel + trust_remote_code always loads the modeling code that ships with the checkpoint, guaranteeing parity with whoever produced the weights.

验证 orig ViT 时必须用 AutoModel.from_pretrained(..., trust_remote_code=True)， 不要 from onevision_encoder import OneVisionEncoderModel。本地的 transformers_impl/onevision_encoder/ 与 checkpoint 自带的 modeling 文件 可能漂移（RoPE、attention 实现、embedding 接口），导致权重一致但 sim ≈ −0.024， 错觉是"权重错了"，实际是 modeling 代码不匹配。AutoModel + trust_remote_code 保证加载 checkpoint 自带的 modeling，与产 checkpoint 的环境完全一致。

6. bf16 numerical thresholds (validators tuned for bf16, not fp32)

The built-in validators are tuned for --target-dtype bf16. Realistic thresholds:

Why so loose for ViT? 24 transformer layers in bf16 accumulate ~2% relative error end-to-end. A characteristic healthy bf16 signature is "mid-layer cos dips to 0.98 then climbs back to 0.99 by the last layer" — this is bf16 RoPE accumulation noise, not a weight bug. If you're seeing cos < 0.95 at every layer (not just middle), suspect modeling-code drift (Issue #5), not weight error.

内置 validator 的阈值是按 bf16 调的。ViT layerwise 中间层 cos 跌到 0.98、 末层回升到 0.99 是 bf16 RoPE 累积噪声的健康特征，不是权重错位。 如果是"每一层"都 < 0.95（不只是中间层），怀疑 modeling 代码漂移 （见 Issue #5），不是权重问题。

7. Canonical helpers live in merge_ov2/utils.py — do NOT copy

convert_rope_to_block_layout, convert_rope_to_block_layout_by_positions, _infer_hw_from_positions, rowmajor_to_block are all canonical in merge_ov2/utils.py. The Megatron-side definition in aiak_training_llm/models/llava_onevision2/onevision_encoder_model.py:604 is the upstream reference, but transformers_impl/ cannot import from aiak_training_llm/ (would create a reverse dep on the training framework), hence the controlled re-implementation in merge_ov2/utils.py.

When writing manual debug scripts, always import from merge_ov2.utils:

from merge_ov2.utils import (
    convert_rope_to_block_layout_by_positions,
    rowmajor_to_block,
    cosine_similarity,
    load_image,
)

Do not copy these functions inline (we accumulated a 100-line drift in vit_layerwise.py this way before consolidating). Old broken imports like from llavaonevision2.modeling_llava_onevision2_moe import convert_rope_to_block_layout_by_positions never worked — that function never existed in that module.

convert_rope_to_block_layout*、rowmajor_to_block 等 helper 在 merge_ov2/utils.py 是 canonical 定义。不要 inline 复制（会和 utils 版本漂移）。Megatron 侧 aiak_training_llm/.../onevision_encoder_model.py:604 是上游参考，但 transformers_impl/ 不能反向 import 训练框架。

8. cli.py does not call logging.set_verbosity_info() — validators use print(flush=True)

The CLI does not raise transformers.logging verbosity, so any logger.info(...) inside validators (which run as part of merge/validate) is swallowed at the default WARNING level. To work around this, vit_layerwise.py and peers emit progress via print(..., flush=True) instead of logger.info.

If you want logger-style output instead, the proper fix is to add logging.set_verbosity_info() in cli.py near argument parsing — but that changes behavior for all subcommands, so it's been left as tech debt for now.

cli.py 没调 logging.set_verbosity_info()，validator 里的 logger.info 会被默认 WARNING 等级吞掉。所以 vit_layerwise.py 等用 print(..., flush=True) 输出进度，是绕过这个症状的权宜之计。要根治 就在 cli.py 加 logging.set_verbosity_info()，但会影响所有子命令的行为， 暂作 tech debt。

Quick Reference: Key Weight Mappings / 快速参考：关键权重映射

Purpose / 用途

Use this skill when you set up or debug training-side sample packing for LLaVA-OneVision2 — i.e. when you need to decide which env vars to export in a training shell script (Stage-1 / Stage-1.5 / Stage-2) and want to understand why both OFFLINE_PACKING_BMR and OFFLINE_PACKED_DATA must be 1 to actually get padding-free attention.

在配置或调试 LLaVA-OneVision2 训练侧的样本 packing 时使用——比如要决定在训练 shell 脚本（Stage-1 / Stage-1.5 / Stage-2）中导出哪些环境变量，以及为什么必须 OFFLINE_PACKING_BMR=1 和 OFFLINE_PACKED_DATA=1 同时打开才能真正获得 padding-free 的 attention。

This skill is specifically for:

• Choosing the correct env var combination in training scripts
• Diagnosing cross-sample attention leakage in packed runs
• Understanding why cu_lengths is a dummy [[0]] in some runs and a real [B, P+1] tensor in others
• Avoiding the well-known OFFLINE_PACKING_VQA red herring (it is dead code)

Companion skill: cu-lengths-attention-flow covers the consumer side (how cu_lengths is fed into ViT/LLM attention). This skill covers the producer + gate side.

姊妹 skill：cu-lengths-attention-flow 讲消费端（cu_lengths 如何送入 ViT/LLM attention）。本 skill 讲生产端 + 开关。

TL;DR / 一句话总结

For packed training to work end-to-end, both env vars must be 1:

export OFFLINE_PACKING_BMR='1'   # data-layer gate: build real cu_lengths
export OFFLINE_PACKED_DATA='1'   # batch-layer gate: forward real cu_lengths to model

Setting only one is a silent bug. OFFLINE_PACKING_VQA is dead code; do not rely on it.

The Three Env Vars / 三个环境变量真相表

💡 The OFFLINE_PACKING_VQA red herring is the #1 source of confusion. Newcomers see it in shell scripts and assume it controls VQA packing. It does not. There is no third packing branch in task_encoder.py — only the BMR branch and the legacy captioning fallback.

💡 OFFLINE_PACKING_VQA 这个红鲱鱼是头号困惑源。新人在 shell 脚本里看到它，以为它控制 VQA packing。并不。task_encoder.py 里没有第三个 packing 分支——只有 BMR 分支和老的 captioning fallback。

Two-Stage Gate Architecture / 两段式 Gate 架构

Packing in this codebase is split into two orthogonal gates that must both fire. Understanding this is the whole point of the skill.

本仓库的 packing 拆成两个正交的 gate，必须都触发。理解这一点就是本 skill 的核心。

Gate 1 — Data Layer (OFFLINE_PACKING_BMR)

Where: aiak_training_llm/data/multimodal/task_encoder.py, inside the PackedCaptioningSample branch of the encoder dispatch (encode_sample ~line 186).

What it does:

• For each entry inside the packed sample (for idx in range(n_orig_sample):), if OFFLINE_PACKING_BMR == 1, it builds a MultiMixQASample carrying the full chat-format messages ({role: user, content: prompt}, {role: assistant, content: caption}) and routes it through encode_multi_mix_qa().
• If OFFLINE_PACKING_BMR != 1, it falls back to a plain CaptioningSample and encode_captioning() — losing the multi-turn / multi-image structure required for SFT.
• After the per-entry loop, regardless of the BMR flag, it calls self.pack_selected_samples(l_Qwen2VLImageTaskSample) (line 277), which constructs the per-sub-sample cumulative lengths cu_lengths = [0, len₁, len₁+len₂, ...] and attaches them to the resulting ImageTaskSamplePacked (line 473).

Net effect: enables the correct per-sub-sample encoding and produces real s.cu_lengths on each sample.

作用：启用正确的逐子样本编码，并在每个样本上产出真正的 s.cu_lengths。

⚠️ Even with BMR off, pack_selected_samples still attaches a cu_lengths tensor to the sample. But the sub-samples were encoded via the wrong path (legacy captioning), so the resulting boundaries don't match what the LLM actually sees. BMR off + PACKED_DATA on is a hidden corruption, not just a missing-feature.

⚠️ 即使 BMR 关掉，pack_selected_samples 仍然会给样本挂上 cu_lengths 张量。但子样本走的是错误的编码路径（老 captioning），结果 boundary 和 LLM 实际看到的 token 序列对不上。BMR 关 + PACKED_DATA 开是隐性数据损坏，不只是缺特性。

Gate 2 — Batch Layer (OFFLINE_PACKED_DATA)

Where: aiak_training_llm/data/multimodal/task_encoder.py:359-365, inside batch() (the collate function).

What it does:

# Cumulative sample lengths are needed for packing, otherwise use dummy values.
cu_lengths = torch.tensor([[0]], dtype=torch.int32)
max_lengths = torch.tensor([[0]], dtype=torch.int32)

if self.is_packing_enabled or int(os.environ.get("OFFLINE_PACKED_DATA", 0)) == 1:
    cu_lengths = torch.stack([s.cu_lengths for s in samples])
    max_lengths = torch.tensor([s.max_length for s in samples], dtype=torch.int32)

• Default: emit a dummy cu_lengths of shape [1, 1] containing only [[0]].
• When OFFLINE_PACKED_DATA == 1 (or the energon online-packing flag is set): stack the real per-sample cu_lengths produced by Gate 1 into shape [B, P+1].

Net effect: decides whether the consumer (model forward) sees real packing offsets or a dummy that says "no packing".

作用：决定消费端（模型 forward）看到的是真实的 packing 偏移，还是一个表示"没有 packing"的 dummy。

Why both gates must fire / 为什么必须两个都开

The consumer side at aiak_training_llm/train/pretrain/pretrain_llava_onevision2.py:153-168:

packed_seq_params = None
...
if cu_lengths.shape == torch.Size([1, 1]):
    pass                        # treat as not packed
else:
    assert cu_lengths.shape[0] == 1, "micro-batch-size must be 1 for packing"
    packed_seq_params = PackedSeqParams(
        qkv_format="thd",
        cu_seqlens_q=cu_lengths[0],
        cu_seqlens_kv=cu_lengths[0],
        ...
    )

So:

🔥 The "BMR=1, PACKED_DATA=0" footgun is the most dangerous combination. Training does not crash. Loss curves look reasonable. But every sub-sample in a packed sequence can attend to every other sub-sample's prefix. Use this skill's TL;DR snippet to avoid it.

🔥 "BMR=1, PACKED_DATA=0" 这个组合最危险。训练不会挂，loss 曲线看着也正常。但 packed 序列里每个子样本都能 attend 到别的子样本的 prefix。用本 skill 顶部的 TL;DR 片段避开它。

MBS=1 Hard Requirement / MBS=1 硬性要求

pretrain_llava_onevision2.py:157:

assert cu_lengths.shape[0] == 1, "micro-batch-size must be 1 for packing"

When packing is on, cu_lengths has shape [B, P+1] where B = micro_batch_size and P = number of sub-samples in a packed sequence. The current PackedSeqParams construction only handles B=1 (it indexes cu_lengths[0]). Therefore:

• --micro-batch-size 1 is mandatory for any packed training run.
• Increase throughput via --global-batch-size (gradient accumulation), pipeline parallelism, or longer --seq-length, not via MBS.
• If you forget, the assert fires immediately on the first batch.

打开 packing 时，cu_lengths 形状是 [B, P+1]，B = micro batch size，P = 一个 packed 序列里的子样本数。当前 PackedSeqParams 构造只处理 B=1（取 cu_lengths[0]）。所以：

• 打包训练必须 --micro-batch-size 1。
• 想提吞吐就调 --global-batch-size（梯度累积）、PP 并行、或更长的 --seq-length，不要调 MBS。
• 忘了的话第一个 batch 就 assert 挂掉。

End-to-End Flow / 端到端流程图

┌─────────────────────────────────────────────────────────────────┐
│ Offline preprocessing (auto_pipe.sh, separate skill)            │
│   Produces WebDataset shards with PackedCaptioningSample format │
└──────────────────────────────┬──────────────────────────────────┘
                               │
                    Energon dataloader yields PackedCaptioningSample
                               │
                               ▼
┌─────────────────────────────────────────────────────────────────┐
│ task_encoder.encode_sample()                                    │
│   if OFFLINE_PACKING_BMR == 1:                  ◄── GATE 1     │
│     for each sub-sample → MultiMixQASample → encode_multi_mix_qa │
│   else:                                                         │
│     for each sub-sample → CaptioningSample → encode_captioning  │
│   pack_selected_samples(l_samples)                              │
│     → ImageTaskSamplePacked with cu_lengths=[0,L1,L1+L2,...]   │
└──────────────────────────────┬──────────────────────────────────┘
                               │
                               ▼
┌─────────────────────────────────────────────────────────────────┐
│ task_encoder.batch()                                            │
│   if is_packing_enabled or OFFLINE_PACKED_DATA==1:  ◄── GATE 2 │
│     cu_lengths = stack([s.cu_lengths for s in samples])         │
│   else:                                                         │
│     cu_lengths = [[0]]    # dummy, signals "not packed"         │
└──────────────────────────────┬──────────────────────────────────┘
                               │ batch dict broadcast via tensor_parallel
                               ▼
┌─────────────────────────────────────────────────────────────────┐
│ pretrain_llava_onevision2.get_batch_on_this_tp_rank()           │
│   if cu_lengths.shape == [1,1]: packed_seq_params = None        │
│   else:                                                         │
│     assert cu_lengths.shape[0] == 1   # MBS=1 required          │
│     packed_seq_params = PackedSeqParams(                        │
│       qkv_format="thd",                                         │
│       cu_seqlens_q=cu_lengths[0],                               │
│       cu_seqlens_kv=cu_lengths[0], ...)                         │
└──────────────────────────────┬──────────────────────────────────┘
                               │
                               ▼
                Model forward → flash-attn varlen
                (see cu-lengths-attention-flow skill)

Recipe: Correct Stage-N Script Snippet / 正确的训练脚本片段

# ───────────────────────────────────────────────────────────
# Packing env vars — both REQUIRED for padding-free training
# Set both to '1' when DATA_PATH points to offline-packed shards
#   (PackedCaptioningSample format, e.g. produced by auto_pipe.sh)
# Leave both as '0' (or unset) for unpacked datasets.
# Mixed states are silent bugs — see offline-packing-env-vars skill.
# ───────────────────────────────────────────────────────────
export OFFLINE_PACKING_BMR='1'
export OFFLINE_PACKED_DATA='1'

# Hard requirement when packing is on
MBS=1
# Throughput knobs: GBS via grad-accum, longer SEQ_LEN, more PP — not MBS

For an A/B control run that uses the same packed dataset but disables packing semantics (to measure the leakage cost), set both to '0'. Setting only BMR=1 or only PACKED_DATA=1 is not a valid configuration — it is a bug.

如果想做 A/B 对照，用同一份 packed 数据但关闭 packing 语义（为了量化 leakage 损失），两个都设 '0'。只开一个不是合法配置，是 bug。

Concrete Stage-1 A/B Pair (this repo) / 本仓库的 Stage-1 A/B 对照

• examples/llava_onevision2/quick_start_4b/stage_1_alignment_p16m3_packed.sh — production: BMR=1, PACKED_DATA=1.
• examples/llava_onevision2/quick_start_4b/stage_1_alignment_p16m3_packed_bmr_only.sh — A/B control: BMR=1, PACKED_DATA=0. Note: this is the dangerous combo described above; it is named bmr_only deliberately to study the leakage effect, not as a recommended setting.

If you copy _bmr_only.sh for a real production run, you will get cross-sub-sample attention leakage. Always confirm intent.

如果你把 _bmr_only.sh 拷去做正式训练，就会得到跨子样本 attention leakage。务必确认是有意为之。

Diagnostics / 排查清单

If your packed training looks "off" (loss too low / too smooth / model overfits prefixes):

1. grep -n 'OFFLINE_PACKING_BMR\|OFFLINE_PACKED_DATA' your_script.sh — both should be '1'.
2. Add a one-shot print in task_encoder.batch() after line 365: print('cu_lengths.shape:', cu_lengths.shape). Expect [1, P+1] with P >= 2. If you see [1, 1], Gate 2 is closed.
3. Add a print in pretrain_llava_onevision2.py after line 168: print('packed_seq_params:', packed_seq_params). Should be a real PackedSeqParams, not None.
4. Confirm MBS=1 in the shell (--micro-batch-size 1). Otherwise the assert at line 157 fires and you wouldn't be reading this.
5. Confirm dataset is actually packed: cat $DATA_PATH/.../webdataset/.nv-meta/.info.yaml — look for shard structure produced by auto_pipe.sh (PackedCaptioningSample).
6. Do not add OFFLINE_PACKING_VQA=1 thinking it helps. It does nothing in this codebase.

Cross-References / 交叉引用

• Producer pipeline (how the packed shards are built): distributed-offline-packing skill.
• Consumer attention semantics (how cu_lengths is interpreted by ViT and LLM): cu-lengths-attention-flow skill.
• Dataloader length-balancing across ranks: length-pool-sort-dataset skill.

Source File Index / 源文件索引