inference/test_vision.py

"""Focused regressions for the standalone Step-5 vision integration.

Run from inference/: python -m pytest -q test_vision.py
"""
import json
from unittest.mock import patch

import pytest
import torch
from safetensors import safe_open
from safetensors.torch import load_file, save_file

from convert import _shard_for, pre_shard
from generate import _load_rank_config, _load_rank_model, load_image_inputs, parse_args, validate_prompt_ids
from model import (
    STEP5_TP_LAYOUT_VERSION, Step5Config, Step5ForCausalLM,
    Step5ForConditionalGeneration, Step5VLConfig, config_from_dict,
)
from vision_encoder import EncoderRope2D, Step5VisionEncoder


def tiny_config(**kwargs):
    text = dict(
        hidden_size=32, intermediate_size=64, num_hidden_layers=1,
        num_nextn_predict_layers=3, vocab_size=64, num_attention_heads=4,
        num_attention_groups=2, head_dim=8, max_position_embeddings=128,
        layer_types=["sliding_attention"] * 4, sliding_window=16, yarn_only_types=["full_attention"],
        use_moe=False, sparse_config={"enabled": False, "num_provider_groups": 2},
        tp_layout_version=STEP5_TP_LAYOUT_VERSION,
    )
    text.update(kwargs.pop("text", {}))
    vision = dict(width=16, layers=2, heads=2, image_size=32, patch_size=4,
                  mlp_ratio=2, use_ln_pre=True, use_ln_post=False)
    vision.update(kwargs.pop("vision", {}))
    return Step5VLConfig(text_config=text, vision_config=vision,
                         image_token_id=70, image_token_len=4, patch_token_len=1, **kwargs)


def initialized_model(config):
    model = Step5ForConditionalGeneration(config)
    # TP linear layers allocate empty checkpoint storage; seed every parameter
    # explicitly because these tests exercise random fixtures, not trained weights.
    with torch.no_grad():
        for name, parameter in model.named_parameters():
            if "norm" in name and name.endswith("weight"):
                parameter.zero_()
            elif name.endswith("weight_scale_inv"):
                parameter.uniform_(0.5, 1.5)
            elif parameter.dtype == torch.float8_e4m3fn:
                parameter.copy_(torch.randn(parameter.shape) * 0.02)
            else:
                parameter.normal_(0, 0.02)
    return model


def test_configuration_and_quantization_roundtrip():
    quant = {"quant_method": "fp8", "modules_to_not_convert": [
        "vision_model.transformer.resblocks.3.attn.q_proj", "model.layers.4.moe.gate",
        "model.layers.88.moe.experts.w13_weight", "model.layers.89.moe.up_proj",
    ]}
    config = tiny_config(quantization_config=quant, tp_size=2)
    assert config.text_config.bf16_expert_layers == [88, 89]
    assert not hasattr(config, "quantization_config")
    assert not hasattr(config.text_config, "quantization_config")
    serialized = json.loads(json.dumps(config.to_dict()))
    restored = config_from_dict(serialized)
    assert restored.model_type == "step5_vl"
    assert restored.architectures == ["Step5ForConditionalGeneration"]
    assert restored.text_config.tp == 2
    assert restored.text_config.num_nextn_predict_layers == 3
    assert len(restored.text_config.to_dict()["layer_types"]) == 4
    assert restored.text_config.bf16_expert_layers == [88, 89]
    assert restored.to_dict()["quantization_config"] == quant
    assert isinstance(config_from_dict(restored.text_config.to_dict()), Step5Config)


def test_vision_shape_rectangular_cls_and_optional_layer_scale():
    config = tiny_config(vision={"use_cls_token": True, "ls_init_value": None}).vision_config
    encoder = Step5VisionEncoder(config).eval()
    output = encoder(torch.randn(2, 3, 16, 24))
    assert output.shape == (2, 24, 16)
    assert torch.isfinite(output).all()
    with pytest.raises(ValueError, match="positive multiples"):
        encoder(torch.randn(1, 3, 15, 16))
    with pytest.raises(ValueError, match="exceeds configured"):
        encoder(torch.randn(1, 3, 36, 32))


def test_rope_coordinates_and_cls_subset():
    rope = EncoderRope2D(8, 8, 8, use_cls_token=True)
    q = torch.ones(1, 1, 7, 8)
    result, _ = rope(q, q, (2, 3))
    assert torch.equal(result[:, :, 0], q[:, :, 0])
    inv = 10000 ** (-torch.arange(0, 4, 2).float() / 4)
    coords = [(0, 0)] + [(x + 1, y + 1) for y in range(2) for x in range(3)]
    angles = torch.stack([torch.cat([x * inv, y * inv]).repeat_interleave(2) for x, y in coords])
    sign = torch.tensor([-1., 1.] * 4)
    expected = angles.cos() + sign * angles.sin()
    torch.testing.assert_close(result[0, 0], expected)


def test_crop_order_and_precomputed_embeddings():
    model = Step5ForConditionalGeneration(tiny_config()).eval()
    overview = torch.randn(2, 4, 32)
    crops = torch.randn(3, 1, 32)
    with patch.object(model, "_encode_images", side_effect=[overview, crops]):
        actual = model.get_multimodal_embeddings(
            pixel_values=torch.zeros(2, 3, 32, 32),
            patch_pixel_values=torch.zeros(3, 3, 16, 16), num_patches=[2, 1],
        )
    expected = torch.cat([crops[0], crops[1], overview[0], crops[2], overview[1]])
    torch.testing.assert_close(actual, expected)
    assert model.get_multimodal_embeddings(image_embeds=overview).shape == (8, 32)
    with pytest.raises(ValueError, match="mutually exclusive"):
        model.get_multimodal_embeddings(pixel_values=torch.empty(1), image_embeds=overview)
    with pytest.raises(ValueError, match="sum\\(num_patches\\)"):
        model.get_multimodal_embeddings(pixel_values=torch.zeros(2, 3, 32, 32), num_patches=[1, 0])
    with pytest.raises(ValueError, match="nonnegative integer"):
        model.get_multimodal_embeddings(pixel_values=torch.zeros(2, 3, 32, 32), num_patches=[0.5, 0])
    with pytest.raises(ValueError, match="overview images"):
        model.get_multimodal_embeddings(num_patches=[0])


def make_checkpoint(path, config):
    path.mkdir()
    model = initialized_model(config)
    state = model.state_dict()
    for i in range(config.text_config.num_hidden_layers,
                   config.text_config.num_hidden_layers + config.text_config.num_nextn_predict_layers):
        state[f"model.layers.{i}.mtp_block.unused.weight"] = torch.ones(1)
    save_file(state, str(path / "model.safetensors"))
    (path / "config.json").write_text(json.dumps(config.to_dict()))
    (path / "model.safetensors.index.json").write_text(json.dumps({
        "weight_map": {key: "model.safetensors" for key in state},
    }))
    return model


def test_conversion_and_strict_loader(tmp_path):
    original = make_checkpoint(tmp_path / "source", tiny_config())
    assert pre_shard(str(tmp_path / "source"), str(tmp_path / "unwritten"), 1, check_only=True) == 0
    assert not (tmp_path / "unwritten").exists()
    assert pre_shard(str(tmp_path / "source"), str(tmp_path / "converted"), 1) == 0
    directory = str(tmp_path / "converted" / "tp1")
    config = _load_rank_config(directory, 0)
    loaded = _load_rank_model(directory, 0, config, torch.device("cpu"))
    assert isinstance(loaded, Step5ForConditionalGeneration)
    for key, tensor in original.state_dict().items():
        torch.testing.assert_close(loaded.state_dict()[key], tensor)
    with torch.inference_mode():
        pixels = torch.randn(1, 3, 32, 32)
        # Nonpersistent 2D RoPE buffers must be rebuilt after HF meta loading.
        expected_features = original.get_multimodal_embeddings(pixel_values=pixels)
        actual_features = loaded.get_multimodal_embeddings(pixel_values=pixels)
        assert torch.isfinite(actual_features).all()
        torch.testing.assert_close(actual_features, expected_features, rtol=0, atol=0)
    weight = original.vision_model.transformer.resblocks[0].attn.out_proj.weight
    assert _shard_for("vision_model.transformer.resblocks.0.attn.out_proj.weight")(weight, None, 2, 1, config).shape == weight.shape


@pytest.mark.parametrize("multimodal", [False, True])
def test_checkpoint_identity_normalizes_to_step5(tmp_path, multimodal):
    config = tiny_config()
    payload = config.to_dict() if multimodal else config.text_config.to_dict()
    payload["model_type"] = "step4_vl" if multimodal else "step4"
    payload["architectures"] = ["Step4ForConditionalGeneration" if multimodal else "Step4ForCausalLM"]
    payload["auto_map"] = {"AutoConfig": "old_config.Step4Config"}
    if multimodal:
        payload["text_config"]["model_type"] = "step4"
        payload["text_config"]["architectures"] = ["Step4ForCausalLM"]
        payload["text_config"]["auto_map"] = {"AutoConfig": "old_config.Step4Config"}
    restored = config_from_dict(payload)
    assert restored.model_type == ("step5_vl" if multimodal else "step5")
    encoded = json.dumps(restored.to_dict())
    assert "step4" not in encoded.lower()
    restored.save_pretrained(tmp_path)
    reloaded = config_from_dict(json.loads((tmp_path / "config.json").read_text()))
    assert reloaded.to_dict()["auto_map"] == restored.to_dict()["auto_map"]
    assert reloaded.to_dict()["architectures"] == restored.to_dict()["architectures"]
    assert reloaded.get_text_config().num_nextn_predict_layers == 3
    with pytest.raises(ValueError, match="unsupported model_type"):
        config_from_dict(dict(payload, model_type="unrelated_model"))


def test_legacy_checkpoint_conversion_and_shard_loading(tmp_path):
    original = make_checkpoint(tmp_path / "source", tiny_config())
    source_config = tmp_path / "source" / "config.json"
    source_config.write_text(source_config.read_text().replace("step5", "step4").replace("Step5", "Step4"))
    pre_shard(str(tmp_path / "source"), str(tmp_path / "converted"), 1)
    directory = tmp_path / "converted" / "tp1"
    emitted = (directory / "config-r0.json").read_text()
    assert "step4" not in emitted.lower()
    shard = directory / "model-r0.safetensors"
    state = load_file(str(shard))
    with safe_open(str(shard), framework="pt") as handle:
        metadata = handle.metadata()
    assert metadata["step5_tp_layout"] == STEP5_TP_LAYOUT_VERSION
    assert "step4_tp_layout" not in metadata
    metadata["step4_tp_layout"] = metadata.pop("step5_tp_layout")
    save_file(state, str(shard), metadata=metadata)
    # Existing shards and their per-rank configs must remain loadable together.
    (directory / "config-r0.json").write_text(emitted.replace("step5", "step4").replace("Step5", "Step4"))
    config = _load_rank_config(str(directory), 0)
    loaded = _load_rank_model(str(directory), 0, config, torch.device("cpu"))
    for name, tensor in original.state_dict().items():
        torch.testing.assert_close(loaded.state_dict()[name], tensor, rtol=0, atol=0)
    metadata["step5_tp_layout"] = "incompatible-layout"
    save_file(state, str(shard), metadata=metadata)
    with pytest.raises(ValueError, match="step5_tp_layout metadata"):
        _load_rank_model(str(directory), 0, config, torch.device("cpu"))
    metadata["step5_tp_layout"] = STEP5_TP_LAYOUT_VERSION
    metadata["step4_tp_layout"] = "incompatible-layout"
    save_file(state, str(shard), metadata=metadata)
    with pytest.raises(ValueError, match="step4_tp_layout metadata"):
        _load_rank_model(str(directory), 0, config, torch.device("cpu"))
    metadata.pop("step5_tp_layout")
    metadata.pop("step4_tp_layout")
    save_file(state, str(shard), metadata=metadata)
    with pytest.raises(ValueError, match="missing step5_tp_layout"):
        _load_rank_model(str(directory), 0, config, torch.device("cpu"))


def tiny_moe_config(*, quantized=True):
    return tiny_config(
        text=dict(hidden_size=256, intermediate_size=256, num_hidden_layers=2,
                  layer_types=["sliding_attention"] * 5, num_attention_heads=8,
                  head_dim=32, vocab_size=128, eos_token_id=[127], use_moe=True,
                  moe_layer_list=[0, 1], moe_num_experts=8, moe_top_k=2,
                  moe_intermediate_size=256, share_expert_dim=256),
        quantization_config={"quant_method": "fp8", "modules_to_not_convert": [
            "model.layers.1.moe.experts.w13_weight",
            "model.layers.1.moe.experts.w2_weight",
            "vision_model.transformer.resblocks.0.attn.q_proj",
        ]} if quantized else None,
    )


@pytest.mark.parametrize("quantized", [False, True], ids=["bf16", "mixed_fp8_bf16"])
def test_experts_conversion_and_hf_loading(tmp_path, quantized):
    original = make_checkpoint(tmp_path / "source", tiny_moe_config(quantized=quantized))
    pre_shard(str(tmp_path / "source"), str(tmp_path / "converted"), 2)
    directory = str(tmp_path / "converted" / "tp2")
    config = _load_rank_config(directory, 0)
    assert config.text_config.bf16_expert_layers == ([1] if quantized else [0, 1])
    loaded = _load_rank_model(directory, 0, config, torch.device("cpu"))
    stored = load_file(str(tmp_path / "converted" / "tp2" / "model-r0.safetensors"))
    for key, actual in loaded.state_dict().items():
        expected = stored[key]
        if key == "lm_head.weight":
            expected = expected[:64]
        assert actual.dtype == expected.dtype, key
        torch.testing.assert_close(actual.float(), expected.float(), rtol=0, atol=0, msg=key)
    if quantized:
        assert loaded.model.layers[0].moe.gate_proj.weight.dtype == torch.float8_e4m3fn
        assert loaded.model.layers[0].moe.gate_proj.weight_scale_inv.dtype == torch.float32
    else:
        assert loaded.model.layers[0].moe.gate_proj.weight.dtype == torch.bfloat16
        assert not any("weight_scale" in key for key in stored)
        assert loaded.model.layers[0].moe.gate_proj.weight_scale_inv is None
    assert loaded.model.layers[1].moe.gate_proj.weight.dtype == torch.bfloat16
    assert loaded.model.layers[1].moe.gate_proj.weight_scale_inv is None
    # Independently verify that only the expert axis was partitioned.
    torch.testing.assert_close(loaded.model.layers[0].moe.gate_proj.weight.float(),
                               original.model.layers[0].moe.gate_proj.weight[:4].float(), rtol=0, atol=0)
    with torch.inference_mode():
        pixels = torch.randn(1, 3, 32, 32)
        torch.testing.assert_close(loaded.get_multimodal_embeddings(pixel_values=pixels),
                                   original.get_multimodal_embeddings(pixel_values=pixels), rtol=0, atol=0)


@pytest.mark.parametrize("quantization", [None, {}])
def test_unquantized_moe_configuration_roundtrip(quantization):
    raw = tiny_moe_config(quantized=False).to_dict()
    raw["text_config"]["quantization_config"] = quantization
    # Derived fields from another checkpoint must not override its quantization metadata.
    raw["text_config"]["bf16_expert_layers"] = []
    config = config_from_dict(raw)
    for restored in (config, config_from_dict(config.to_dict()),
                     config_from_dict(config.text_config.to_dict())):
        text = restored.get_text_config()
        assert text.bf16_expert_layers == [0, 1]
        assert not any(text.is_fp8_expert_layer(i) for i in range(2))
        assert not hasattr(text, "quantization_config")


def test_unsupported_expert_quantization_is_rejected():
    with pytest.raises(ValueError, match="quant_method='fp8'"):
        tiny_config(quantization_config={"quant_method": "int8"})


@pytest.mark.parametrize("corruption", ["missing", "unexpected", "shape", "rank"])
def test_loader_rejects_incomplete_or_wrong_rank_shards(tmp_path, corruption):
    make_checkpoint(tmp_path / "source", tiny_config())
    pre_shard(str(tmp_path / "source"), str(tmp_path / "converted"), 1)
    directory = tmp_path / "converted" / "tp1"
    shard = directory / "model-r0.safetensors"
    with safe_open(str(shard), framework="pt") as handle:
        metadata = handle.metadata()
    state = load_file(str(shard))
    key = "vision_model.conv1.weight"
    if corruption == "missing":
        del state[key]
    elif corruption == "unexpected":
        state["vision_model.unexpected.weight"] = torch.ones(1)
    elif corruption == "shape":
        state[key] = state[key][:1].clone()
    else:
        metadata["rank"] = "1"
    save_file(state, str(shard), metadata=metadata)
    config = _load_rank_config(str(directory), 0)
    with patch.object(Step5ForConditionalGeneration, "from_pretrained") as load:
        with pytest.raises(ValueError, match="key mismatch|shape|rank metadata"):
            _load_rank_model(str(directory), 0, config, torch.device("cpu"))
        load.assert_not_called()


def test_resume_rejects_old_text_only_shard(tmp_path):
    make_checkpoint(tmp_path / "source", tiny_config())
    pre_shard(str(tmp_path / "source"), str(tmp_path / "converted"), 2)
    # A matching shard must resume successfully when converting from a TP=1 config.
    assert pre_shard(str(tmp_path / "source"), str(tmp_path / "converted"), 2, start_rank=1) == 0
    shard = tmp_path / "converted" / "tp2" / "model-r0.safetensors"
    with safe_open(str(shard), framework="pt") as handle:
        metadata = handle.metadata()
    state = {key: tensor for key, tensor in load_file(str(shard)).items()
             if not key.startswith(("vision_model.", "vit_large_projector."))}
    save_file(state, str(shard), metadata=metadata)
    with patch("convert._build_rank_state") as build:
        with pytest.raises(ValueError, match="shard key mismatch"):
            pre_shard(str(tmp_path / "source"), str(tmp_path / "converted"), 2, start_rank=1)
        build.assert_not_called()


def test_cli_inputs(tmp_path):
    path = tmp_path / "images.safetensors"
    save_file({"image_embeds": torch.randn(4, 32)}, str(path))
    result = load_image_inputs(str(path), tiny_config())
    assert result["image_embeds"].shape == (4, 32)
    with pytest.raises(ValueError, match="step5_vl"):
        load_image_inputs(str(path), tiny_config().text_config)
    with pytest.raises(SystemExit):
        parse_args(["--tp-dir", ".", "--prompt", "test", "--image-inputs", str(path)])
    limits = dict(vocab_size=64, max_position_embeddings=128, max_new_tokens=1)
    validate_prompt_ids([5, 70, 6], image_token_id=70, **limits)
    with pytest.raises(ValueError, match="outside vocabulary"):
        validate_prompt_ids([5, 70, 6], **limits)
    with pytest.raises(ValueError, match="outside vocabulary"):
        validate_prompt_ids([5, 71, 6], image_token_id=70, **limits)


@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for packed text inference")
def test_cuda_prefill_decode_text_parity_and_single_vision_call():
    torch.manual_seed(11)
    device = torch.device("cuda")
    config = tiny_config()
    model = initialized_model(config).to(device).eval()
    text = Step5ForCausalLM(config.text_config).to(device).eval()
    text.load_state_dict({key: value for key, value in model.state_dict().items()
                          if key.startswith(("model.", "lm_head."))}, strict=True)
    pixels = torch.randn(1, 3, 32, 32, device=device)
    ids = torch.tensor([5, 70, 70, 70, 70, 6, 7], device=device)
    lengths = torch.tensor([7], dtype=torch.int32, device=device)
    with torch.inference_mode():
        cache_a = model.allocate_cache(batch=1, max_tokens=16, device=device)
        cache_b = text.allocate_cache(batch=1, max_tokens=16, device=device)
        features = model.get_multimodal_embeddings(pixel_values=pixels)
        embeds = model.model.embed_tokens(ids.masked_fill(ids == 70, 0))
        embeds[ids == 70] = features
        actual = model(ids, cache_a, lengths, pixel_values=pixels).logits
        assert torch.isfinite(actual).all()
        expected = text(ids, cache_b, lengths, inputs_embeds=embeds).logits
        torch.testing.assert_close(actual, expected, rtol=0, atol=0)
        token = actual.argmax(-1)
        one = torch.ones(1, dtype=torch.int32, device=device)
        torch.testing.assert_close(model(token, cache_a, one, is_decode=True).logits,
                                   text(token, cache_b, one, is_decode=True).logits, rtol=0, atol=0)
        before = cache_a.lengths.clone()
        with pytest.raises(ValueError, match="placeholders"):
            model(ids, cache_a, lengths, image_embeds=features[:1])
        assert torch.equal(cache_a.lengths, before)
        with pytest.raises(ValueError, match="only accepted during prefill"):
            model(token, cache_a, one, is_decode=True, pixel_values=pixels)
        with pytest.raises(ValueError, match="require image inputs"):
            model(ids, cache_a, lengths)
        with patch.object(model.vision_model, "forward", wraps=model.vision_model.forward) as encode:
            output = model.generate_greedy([ids.tolist(), [8, 9]], max_new_tokens=3,
                                          device=device, eos_token_ids={-1}, pixel_values=pixels)
            assert [len(row) for row in output] == [3, 3]
            assert encode.call_count == 1
        assert text.generate_greedy([[5, 6]], max_new_tokens=2, device=device, eos_token_ids={-1}) == model.generate_greedy([[5, 6]], max_new_tokens=2, device=device, eos_token_ids={-1})