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

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

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

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

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.

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!"
'

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
'

# 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

# 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

# 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")

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}"

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}"