File size: 5,586 Bytes
e34d4af | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 | # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Weight-free compatibility tests for Transformers 4 and 5.
Load standalone remote-code modules without model weights. V5-specific
conversion tests are skipped when the conversion API is unavailable.
"""
import importlib.util
import sys
from pathlib import Path
import pytest
import torch
ROOT = Path(__file__).resolve().parent.parent
def _load(name: str):
spec = importlib.util.spec_from_file_location(name, ROOT / f"{name}.py")
module = importlib.util.module_from_spec(spec)
sys.modules[name] = module
spec.loader.exec_module(module)
return module
modeling_vit = _load("modeling_vit")
modeling_qformer = _load("modeling_qformer")
def test_qformer_initialization_and_attention_masks() -> None:
from transformers import BertConfig
config = BertConfig(
hidden_size=16, num_hidden_layers=1, num_attention_heads=2,
intermediate_size=32, vocab_size=64,
)
model = modeling_qformer.BertModel(config).eval()
if modeling_vit._HAS_TRANSFORMERS5_CONVERSION_API:
assert hasattr(model, "all_tied_weights_keys")
assert model.get_head_mask(None, 1) == [None]
mask = model.invert_attention_mask(torch.tensor([[1, 0]]))
assert mask.shape == (1, 1, 1, 2)
assert mask[0, 0, 0, 0] == 0
assert mask[0, 0, 0, 1] < 0
def test_pruning_helper_is_vendored() -> None:
heads, index = modeling_qformer.find_pruneable_heads_and_indices([0, 1], 4, 2, set())
assert heads == {0, 1}
assert index.tolist() == [4, 5, 6, 7]
def test_bicubic_interpolation_matches_reference() -> None:
pos_embed = torch.arange(1 * 257 * 8, dtype=torch.float32).reshape(1, 257, 8)
out = modeling_vit._bicubic_interpolate_pos_embed(pos_embed, 1, 24, 24)
assert out.shape == (1, 577, 8)
# Independent reference: same math as the historical interpolate_pos_embed.
extra = pos_embed[:, :1]
tokens = pos_embed[:, 1:].reshape(-1, 16, 16, 8).permute(0, 3, 1, 2)
tokens = torch.nn.functional.interpolate(tokens, size=(24, 24), mode="bicubic", align_corners=False)
tokens = tokens.permute(0, 2, 3, 1).flatten(1, 2)
reference = torch.cat((extra, tokens), dim=1)
assert torch.equal(out, reference)
assert torch.equal(out[:, :1], pos_embed[:, :1]) # cls token untouched
def test_bicubic_interpolation_is_noop_on_matching_grid() -> None:
pos_embed = torch.rand(1, 577, 8)
out = modeling_vit._bicubic_interpolate_pos_embed(pos_embed, 1, 24, 24)
assert torch.equal(out, pos_embed)
def test_interpolation_rejects_non_square_source_and_resizes_equal_area_rectangle() -> None:
with pytest.raises(ValueError, match="square patch grid"):
modeling_vit._bicubic_interpolate_pos_embed(torch.rand(1, 258, 8), 1, 24, 24)
source = torch.arange(1 * 257 * 8, dtype=torch.float32).reshape(1, 257, 8)
resized = modeling_vit._bicubic_interpolate_pos_embed(source, 1, 8, 32)
assert resized.shape == source.shape
assert not torch.equal(resized[:, 1:], source[:, 1:])
state_dict = {"visual_encoder.pos_embed": source.clone()}
modeling_vit.interpolate_pos_embed(
"visual_encoder.pos_embed", 256, 257, state_dict, target_h=8, target_w=32
)
assert torch.equal(state_dict["visual_encoder.pos_embed"], resized)
def test_interpolate_pos_embed_writes_only_on_mismatch() -> None:
state_dict = {"visual_encoder.pos_embed": torch.rand(1, 257, 8)}
modeling_vit.interpolate_pos_embed(
"visual_encoder.pos_embed", 256, 257, state_dict, target_h=16, target_w=16
)
assert state_dict["visual_encoder.pos_embed"].shape == (1, 257, 8) # no-op
@pytest.mark.skipif(
not modeling_vit._HAS_TRANSFORMERS5_CONVERSION_API,
reason="Transformers 5 conversion API not available",
)
def test_v5_conversion_mapping_registered() -> None:
from transformers.conversion_mapping import get_checkpoint_conversion_mapping
modeling_vit.register_transformers5_pos_embed_conversion()
converters = get_checkpoint_conversion_mapping("cosmos-embed1")
assert converters and len(converters) == 1
assert converters[0].source_patterns == ["visual_encoder.pos_embed"]
@pytest.mark.skipif(
not modeling_vit._HAS_TRANSFORMERS5_CONVERSION_API,
reason="Transformers 5 conversion API not available",
)
def test_v5_operation_interpolates_257_to_577() -> None:
class _PatchEmbed:
patch_shape = (24, 24)
class _Visual:
patch_embed = _PatchEmbed()
class _Model:
visual_encoder = _Visual()
op = modeling_vit.InterpolatePosEmbed()
tensor = torch.arange(1 * 257 * 4, dtype=torch.bfloat16).reshape(1, 257, 4)
out = op.convert(
{"visual_encoder.pos_embed": [tensor]},
["visual_encoder.pos_embed"],
["visual_encoder.pos_embed"],
model=_Model(),
)
result = out["visual_encoder.pos_embed"][0]
assert result.shape == (1, 577, 4)
assert result.dtype == torch.bfloat16
|