chenyu 18e854cdbf shrink MLB on sharded axis (#3255) ...

* shrink MLB on sharded axis

use onehot structure to store the real partition. goal is unsynced batchnorm2d that can be run on multigpu for training.

draft version in https://github.com/chenyuxyz/tinygrad/pull/109

* SYNCBN flag

* test unclean shrinks

* UnsyncedBatchNorm reuses BatchNorm

* more robust pad arg check

* better types

* more tests!

* 6 gpus in benchmark

* disable slow GPUS=6 benchmark

2024-01-31 21:48:25 -05:00

..

workflows

shrink MLB on sharded axis (#3255)

2024-01-31 21:48:25 -05:00