はじめに

Tenstorrent の TT-Buda の Model の生成およびその Model のローディングの部分を見ていきます。

TensTorrent Device Image (TTI): Saving/Loading

User GuideのTensTorrent Device Image (TTI): Saving/Loadingに Tenstorrent Deveci Image (TTI) を呼び出す部分です。pybuda.TTDeviceにて、アーキテクチャとデバイスタイプを選択して、TTIを獲得します。このTTI を place_module にて配置した後に、compile_to_image でイメージを生成しています。

tt0 = pybuda.TTDevice("tt0",arch=BackendDevice.Grayskull, devtype=BackendType.Silicon)
tt0.place_module(...)
device_img: TTDeviceImage = tt0.compile_to_image(
    img_path="device_images/tt0.tti",
    training=training,
    sample_inputs=(...),
)

生成されたイメージは下記のようなZIPファイルになっているようです。

/unzipped_tti_directory
├── device.json # json file capturing device state
├── <module-name>.yaml # netlist yaml
├── backend_build_binaries # backend build files from tt_build/<test>
│   ├── blob_init
│   ├── brisc
│   ├── ...
├── *tensor*.pkl # pickled constant/parameter tensors
├── *module*.pkl # pickled PyBuda module object

生成したイメージは、下記のようなコードでロードします。

device_img: TTDeviceImage = TTDeviceImage.load_from_disk("device_images/tt0.tti")

読み込んだイメージの情報は下記のコマンドで出力されます。

device_img.info()

Image Info...
- Version Info:
      - pybuda_version: 0.1.220624+dev.f63c9d32
      - pybuda_commit: 7def2987
      - buda_backend_commit: f2fd0fa3
- Device Name: tt0

Device Info...
- arch: BackendDevice.Grayskull
- chip_ids: [0]
- backend device type: BackendType.Silicon
- grid size: [10, 12]
- harvested rows: 0

Compilation Graph State...
- training: False
- modules: ['bert_encoder']
- ordered input shapes: [[1, 128, 128], [1, 1, 128, 128]]
- ordered targets shapes: []

下記のコマンドにて、イメージファイルを読み込み、推論することができます。。。。

img = TTDeviceImage.load_from_disk(img_path="device_images/tt0.tti")
device = TTDevice.load_image(img=img)

inputs = [torch.rand(shape) for shape in img.get_input_shapes()] # create tensors using shape info from img
device.push_to_inputs(inputs) # push newly created input activation tensors to device
output_q = pybuda.run_inference()

下記のように、pybuda.set_configuration_optionsにて、一部のRowのTileを使わないようにイメージファイルを生成することもできます。harvedted_rows にて、1 と 11 を指定しています。

pybuda.set_configuration_options(harvested_rows=[1,11]) #manually harvest row 1 and 11

tt0 = pybuda.TTDevice("tt0",arch=BackendDevice.Grayskull, devtype=BackendType.Silicon)
tt0.place_module(...)
device_img: TTDeviceImage = tt0.compile_to_image(
    img_path="device_images/tt0.tti",
    training=training,
    sample_inputs=(...),
)

生成したイメージをロードした後に、

device_img.info()

を実行すると、harvested rowsのところに、最初と11番目に1が立っていますね。

- arch: BackendDevice.Grayskull
- chip_ids: [0]
- backend device type: BackendType.Silicon
- grid size: [8, 12]    # 2 rows are harvested from 10 rows
- harvested rows: 2050  # indicates row 1 and 11 are harvested (in binary, 100000000010)