File size: 1,146 Bytes
babe29f | 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 | # SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
# SPDX-License-Identifier: Apache-2.0
from pathlib import Path
from models.tt_transformers.tt.lm_head import _lm_head_cache_file_name
def test_galaxy_2d_lm_head_uses_distinct_cache_key():
cache_root = Path("weights")
common = {
"dummy_weights": False,
"num_splits": 1,
"shard": 0,
"width": 131072,
"mode": 0,
"mesh_shape": (8, 4),
}
flat = _lm_head_cache_file_name(cache_root, galaxy_2d=False, **common)
galaxy_2d = _lm_head_cache_file_name(cache_root, galaxy_2d=True, **common)
assert flat.name == "output_lm_head_1_split_shard_0_131072_mode_0"
assert galaxy_2d.name == "output_lm_head_1_split_shard_0_131072_mode_0_galaxy_2d_8x4_v1"
assert flat != galaxy_2d
def test_dummy_lm_head_does_not_use_weight_cache():
assert (
_lm_head_cache_file_name(
Path("weights"),
dummy_weights=True,
num_splits=1,
shard=0,
width=131072,
mode=0,
galaxy_2d=True,
mesh_shape=(8, 4),
)
is None
)
|