search

tenstorrent / tt-metal

:metal: TT-NN operator library, and TT-Metalium low level kernel programming model.
https://docs.tenstorrent.com/ttnn/latest/index.html
Apache License 2.0
483 stars 78 forks source link

Batch normalization support #12253

Open jdh8 opened 2 months ago

jdh8 commented 2 months ago

To compile aten._native_batch_norm_legit*, we need a batch norm op in tt-metal.

  1. Training mode requires native support (to update parameters on the fly)
  2. Probable performance gain in inference mode, compared with a handful of pointwise ops (add, mul, rsqrt)
mrakitaTT commented 1 week ago

Hi @eyonland, just to check is this already on your radar and is somebody working on it?

We can decompose batchnorm in compiler into smaller ops, but we would prefer to avoid that work if ttnn will soon have native support for batchnorm.