File size: 6,549 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 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 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 | # SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
#
# SPDX-License-Identifier: Apache-2.0
from types import SimpleNamespace
import pytest
from models.tt_transformers.tt.model_config import (
ModelArgs,
TensorGroup,
compute_galaxy_padded_vocab_size,
compute_galaxy_width_shard_cores,
compute_padded_vocab_size,
create_galaxy_ff1_out_reduce_scatter_memcfg,
should_pad_sampling_logits_to_power_of_2,
)
@pytest.mark.parametrize(
("vocab_size", "num_devices", "expected"),
[
(151936, 1, 151936),
(151936, 4, 151936),
(151936, 8, 152064),
(151936, 32, 152576),
(32001, 2, 32064),
],
)
def test_compute_padded_vocab_size(vocab_size, num_devices, expected):
padded_vocab_size = compute_padded_vocab_size(vocab_size, num_devices)
assert padded_vocab_size == expected
assert padded_vocab_size >= vocab_size
assert padded_vocab_size % (32 * num_devices) == 0
assert (padded_vocab_size // num_devices) % 32 == 0
def test_compute_padded_vocab_size_rejects_invalid_num_devices(expect_error):
with expect_error(ValueError, "num_devices must be >= 1"):
compute_padded_vocab_size(32000, 0)
@pytest.mark.parametrize(
("vocab_size", "num_devices", "expected"),
[
(128256, 32, 128 * 1024),
(151936, 32, 152576),
(152064, 32, 152576),
],
)
def test_compute_galaxy_padded_vocab_size(vocab_size, num_devices, expected):
padded_vocab_size = compute_galaxy_padded_vocab_size(vocab_size, num_devices)
assert padded_vocab_size == expected
assert padded_vocab_size >= vocab_size
assert padded_vocab_size % (32 * num_devices) == 0
@pytest.mark.parametrize(
("width", "expected"),
[
(768, 24),
(1024, 32),
(1280, 20),
(2048, 32),
(3584, 28),
(4096, 32),
],
)
def test_compute_galaxy_width_shard_cores(width, expected):
num_cores = compute_galaxy_width_shard_cores(width)
assert num_cores == expected
assert (width // num_cores) % 32 == 0
@pytest.mark.parametrize(
("hidden_dim", "expected_cores", "expected_shard_width"),
[
# Galaxy is 8x4. The reduce-scatter OUTPUT width is
# hidden_dim // mesh_rows // mesh_cols -- the op computes the shape itself as
# input_shape[dim] / ring_size, so this config describes the scattered width.
# Llama-70B: 28672 // 8 // 4 = 896 = 28 * 32, so the 7x4 grid is exact.
pytest.param(28672, 28, 32, id="llama-70b"),
# Qwen-72B padded: 32768 // 8 // 4 = 1024, which 28 cores cannot tile-align; 32 can.
pytest.param(32768, 32, 32, id="qwen-72b-padded"),
# 14336 // 8 // 4 = 448, which 28 cores cannot fill at tile granularity, so the
# shard rounds up to one tile and the grid over-covers (28 * 32 = 896 >= 448) --
# the same padding the demo applies when it covers 896 with 30 * 32 = 960.
pytest.param(14336, 28, 32, id="llama-8b-padded"),
],
)
def test_galaxy_ff1_out_reduce_scatter_memory_config(hidden_dim, expected_cores, expected_shard_width):
memory_config = create_galaxy_ff1_out_reduce_scatter_memcfg(hidden_dim, mesh_rows=8, mesh_cols=4)
assert memory_config.shard_spec.grid.num_cores() == expected_cores
assert memory_config.shard_spec.shape == [32, expected_shard_width]
@pytest.mark.parametrize("hidden_dim", [28672, 32768])
def test_galaxy_ff1_shard_grid_exactly_covers_the_scattered_width(hidden_dim):
"""The shard grid must cover the reduce-scatter output width exactly, no more."""
memory_config = create_galaxy_ff1_out_reduce_scatter_memcfg(hidden_dim, mesh_rows=8, mesh_cols=4)
scattered_width = hidden_dim // 8 // 4
covered = memory_config.shard_spec.grid.num_cores() * memory_config.shard_spec.shape[1]
assert covered == scattered_width
assert memory_config.shard_spec.shape[1] % 32 == 0
def test_galaxy_ff1_matches_the_validated_demo_ratio():
"""Cross-check against models/demos/llama3_70b_galaxy: 3840 in, 960 out, ratio 4."""
memory_config = create_galaxy_ff1_out_reduce_scatter_memcfg(28672, mesh_rows=8, mesh_cols=4)
covered = memory_config.shard_spec.grid.num_cores() * memory_config.shard_spec.shape[1]
per_device_width = 28672 // 8
assert covered == per_device_width // 4 == 896
assert covered <= 960 # the demo's 30-core layout pads 896 up to 960
def test_compute_galaxy_width_shard_cores_rejects_non_positive_width(expect_error):
with expect_error(ValueError, "positive multiple"):
compute_galaxy_width_shard_cores(0)
@pytest.mark.parametrize(
("base_model_name", "padded_vocab_size", "sampling_splits", "expected"),
[
("Llama-3.1-70B", 128256, 4, True),
("Llama-3.1-70B", 131072, 4, False),
("Llama-3.1-8B", 128256, 4, True),
],
)
def test_should_pad_sampling_logits_to_power_of_2(base_model_name, padded_vocab_size, sampling_splits, expected):
assert should_pad_sampling_logits_to_power_of_2(base_model_name, padded_vocab_size, sampling_splits) is expected
def test_should_pad_sampling_logits_to_power_of_2_rejects_invalid_sampling_splits(expect_error):
with expect_error(ValueError, "sampling_splits must be >= 1"):
should_pad_sampling_logits_to_power_of_2("Llama-3.1-70B", 128256, 0)
def _llama_model_args(device_name="P150"):
args = object.__new__(ModelArgs)
args.model_name = "Llama-3.1-8B-Instruct"
args.device_name = device_name
args.dram_grid_size = SimpleNamespace(x=8)
return args
@pytest.mark.parametrize(
("tensor_group", "n"),
[
(TensorGroup.WQKV, 6144),
(TensorGroup.WO, 4096),
(TensorGroup.FF2, 4096),
(TensorGroup.FF1_FF3, 14336),
],
)
def test_llama31_8b_p150_uses_projection_specific_dram_reader_counts(tensor_group, n):
args = _llama_model_args()
assert args.get_dram_sharded_matmul_num_workers(tensor_group, n) == 2
def test_llama31_8b_p150_uses_one_dram_reader_when_padded_shards_cannot_split_evenly():
args = _llama_model_args()
# PAD_MLP_CORES=9 pads the gate/up width to 14,400, or 57 tiles per P150 DRAM bank.
assert args.get_dram_sharded_matmul_num_workers(TensorGroup.FF1_FF3, 14400) == 1
def test_llama31_8b_dram_reader_counts_remain_default_on_unvalidated_devices():
args = _llama_model_args(device_name="P300")
for tensor_group in (TensorGroup.WQKV, TensorGroup.WO, TensorGroup.FF1_FF3, TensorGroup.FF2):
assert args.get_dram_sharded_matmul_num_workers(tensor_group, 14336) == 1
|