| |
| |
| |
|
|
| from pathlib import Path |
| from types import SimpleNamespace |
|
|
| import torch |
|
|
| from models.tt_transformers.tt import attention as attention_module |
|
|
|
|
| class _StopAfterBiasConstruction(Exception): |
| pass |
|
|
|
|
| class _FakeDecoderOptimizations: |
| def get_tensor_dtype(self, **_): |
| return None |
|
|
| def get_math_fidelity(self, **_): |
| return None |
|
|
|
|
| def test_galaxy_qkv_bias_uses_weight_matching_2d_sharding_and_survives_column_reduce(monkeypatch, expect_error): |
| mesh_device = object() |
| mesh_mapper_calls = [] |
| as_tensor_calls = [] |
|
|
| monkeypatch.setattr(attention_module.ttnn, "ShardTensorToMesh", lambda *args, **kwargs: object()) |
| monkeypatch.setattr(attention_module.ttnn, "ReplicateTensorToMesh", lambda *args, **kwargs: object()) |
| monkeypatch.setattr(attention_module.ttnn, "from_torch", lambda *args, **kwargs: object()) |
|
|
| def shard_tensor_2d_mesh(device, *, dims, mesh_shape): |
| mapper = object() |
| mesh_mapper_calls.append((device, dims, mesh_shape, mapper)) |
| return mapper |
|
|
| def as_tensor(tensor, **kwargs): |
| as_tensor_calls.append((tensor, kwargs)) |
| return SimpleNamespace(shape=(1, tensor.shape[-1])) |
|
|
| monkeypatch.setattr(attention_module.ttnn, "ShardTensor2dMesh", shard_tensor_2d_mesh) |
| monkeypatch.setattr(attention_module.ttnn, "as_tensor", as_tensor) |
| monkeypatch.setattr( |
| attention_module.ttnn, |
| "reshape", |
| lambda *args, **kwargs: (_ for _ in ()).throw(_StopAfterBiasConstruction("bias captured")), |
| ) |
|
|
| configuration = SimpleNamespace( |
| num_devices=32, |
| dim=8192, |
| n_heads=64, |
| head_dim=128, |
| max_seq_len=32768, |
| max_batch_size=32, |
| n_kv_heads=8, |
| min_kv_prefill_shard_seqlen=0, |
| ccl_dtype=None, |
| MAX_QKV_MM_SEQ_LEN=2048, |
| tile_size=32, |
| tile_padded_batch_rows=32, |
| rms_norm_add_unit_offset=False, |
| use_qk_fused=False, |
| use_hf_rope=False, |
| arch_name="wormhole_b0", |
| max_grid_size=None, |
| compute_kernel_config_hifi2=None, |
| compute_kernel_config_hifi2_fp16=None, |
| compute_kernel_config_hifi4=None, |
| layer_types=None, |
| cluster_shape=[8, 4], |
| is_multichip=True, |
| dummy_weights=False, |
| qkv_size=10240, |
| get_model_config=lambda: {}, |
| ccl_topology=lambda: None, |
| get_state_dict_prefix=lambda *_: "layers.0.attention", |
| ) |
| args = SimpleNamespace(decoders_optimizations=_FakeDecoderOptimizations()) |
| state_dict = { |
| "layers.0.attention.wq.bias": torch.arange(16), |
| "layers.0.attention.wk.bias": torch.arange(100, 108), |
| "layers.0.attention.wv.bias": torch.arange(200, 208), |
| } |
|
|
| with expect_error(_StopAfterBiasConstruction, "bias captured"): |
| attention_module.Attention( |
| mesh_device=mesh_device, |
| tt_ccl=None, |
| args=args, |
| state_dict=state_dict, |
| weight_cache_path=Path("cache"), |
| layer_num=0, |
| dtype=None, |
| transformation_mats=None, |
| configuration=configuration, |
| ) |
|
|
| assert len(mesh_mapper_calls) == 1 |
| assert mesh_mapper_calls[0][:3] == (mesh_device, (-1, None), [8, 4]) |
| qkv_bias, tensor_kwargs = as_tensor_calls[0] |
| assert tensor_kwargs["mesh_mapper"] is mesh_mapper_calls[0][3] |
| assert tensor_kwargs["cache_file_name"].name.endswith("wqkv_bias_prefill_sharded_2d_col_reduce_4") |
|
|
| |
| |
| |
| expected_device_shards = [ |
| torch.tensor([2 * i, 2 * i + 1, 100 + i, 200 + i], dtype=qkv_bias.dtype) for i in range(8) |
| ] |
| reduced_device_shards = [shard * configuration.cluster_shape[1] for shard in torch.chunk(qkv_bias, 8)] |
| assert all(torch.equal(actual, expected) for actual, expected in zip(reduced_device_shards, expected_device_shards)) |
|
|
