tzhouamandClaude Fable 5.1
4e3e8a1ebc
Verbatim copy of the edge/omni-edge working tree (tracked diffs + untracked edge/local, tests/edge/local, docs/edge, pipeline files). Taken so that vLLM-Omni v0.29.0rc1 can be merged on a branch without touching the user's working tree. logs_edge/ and benchmark .npy outputs excluded. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
65 lines
2.3 KiB
Python
65 lines
2.3 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Shape arithmetic for the exported vision tower.
|
|
|
|
The tower itself needs a 29 GB checkpoint, so what is covered here is the part
|
|
that decides whether an export is even well-formed — and that is where a silent
|
|
mistake would land: a grid that does not divide by the merge size produces a
|
|
graph that exports, runs, and returns the wrong number of tokens to the LLM.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
|
|
import pytest
|
|
|
|
from vllm_omni.edge.local.artifacts.vision_tower import tower_shape
|
|
|
|
VISION_CONFIG = {
|
|
"depth": 27, "hidden_size": 1152, "in_channels": 3, "intermediate_size": 4304,
|
|
"num_heads": 16, "num_position_embeddings": 2304, "out_hidden_size": 5120,
|
|
"patch_size": 16, "spatial_merge_size": 2, "temporal_patch_size": 2,
|
|
}
|
|
|
|
|
|
@pytest.fixture
|
|
def checkpoint(tmp_path):
|
|
(tmp_path / "config.json").write_text(json.dumps({"vision_config": VISION_CONFIG}))
|
|
return tmp_path
|
|
|
|
|
|
def test_448_gives_784_patches_and_196_tokens(checkpoint):
|
|
shape = tower_shape(checkpoint, 448)
|
|
assert (shape.grid_h, shape.grid_w) == (28, 28)
|
|
assert shape.patches == 784
|
|
# 3 channels x 2 temporal x 16 x 16
|
|
assert shape.patch_dim == 1536
|
|
# 784 patches merged 2x2
|
|
assert shape.merged_tokens == 196
|
|
assert shape.out_hidden_size == 5120
|
|
|
|
|
|
def test_image_size_must_divide_by_the_patch_size(checkpoint):
|
|
with pytest.raises(ValueError, match="multiple of patch_size"):
|
|
tower_shape(checkpoint, 450)
|
|
|
|
|
|
def test_grid_must_divide_by_the_spatial_merge(checkpoint):
|
|
"""A 3x3 grid of patches cannot be merged 2x2. Exporting it anyway would
|
|
hand the language model a silently wrong number of image tokens."""
|
|
with pytest.raises(ValueError, match="spatial_merge_size"):
|
|
tower_shape(checkpoint, 48)
|
|
|
|
|
|
def test_a_checkpoint_without_a_tower_is_refused(tmp_path):
|
|
(tmp_path / "config.json").write_text(json.dumps({"model_type": "spark2_5"}))
|
|
with pytest.raises(ValueError, match="no vision_config"):
|
|
tower_shape(tmp_path, 448)
|
|
|
|
|
|
@pytest.mark.parametrize("size,patches,tokens", [(224, 196, 49), (448, 784, 196), (896, 3136, 784)])
|
|
def test_token_count_scales_with_area(checkpoint, size, patches, tokens):
|
|
shape = tower_shape(checkpoint, size)
|
|
assert (shape.patches, shape.merged_tokens) == (patches, tokens)
|