tenstorrent / tt-metal

:metal: TT-NN operator library, and TT-Metalium low level kernel programming model.
Apache License 2.0
440 stars 63 forks source link

TTNN `as_tensor` fails when creating multi-device tensors in tile layout #13824

Open esmalTT opened 3 days ago

esmalTT commented 3 days ago

Summary

TTNN as_tensor fails when creating multi-device tensors in tile layout that are in bfloat8 format. The following test demonstrates this:

@pytest.mark.parametrize("width", [16032 * 2, 16032])
@pytest.mark.parametrize("dtype", [ttnn.bfloat8_b, ttnn.bfloat16])
def test_multi_device_tensor(mesh_device, width, dtype, use_program_cache):
    a = torch.rand([1, 1, 32, width])
    a = ttnn.as_tensor(a, mesh_mapper=ttnn.ShardTensorToMesh(mesh_device, -1), layout=ttnn.TILE_LAYOUT, dtype=dtype)

Running will show a failure on the following test case: test_multi_device_tensor[wormhole_b0-True-dtype0-16032] with the following error:

E       RuntimeError: TT_THROW @ ../ttnn/cpp/ttnn/tensor/host_buffer/functions.hpp:148: tt::exception
E       info:
E       Tensor must have OwnedStorage

Close this issue when the failing test case passes with the correct output.

esmalTT commented 3 days ago

cc @cglagovich