test-time-training / ttt-lm-pytorch

Official PyTorch implementation of Learning to (Learn at Test Time): RNNs with Expressive Hidden States
MIT License
1.01k stars 56 forks source link

Extremely slow performance ? #12

Closed junphine closed 3 months ago

junphine commented 3 months ago

V100 1B 0.06it/s

karan-dalal commented 3 months ago

We do not recommend training with this codebase, because it is written in pure PyTorch without any systems optimization, so training will be slow, especially when the per-device batch size is small. For faster training code, or to replicate results from our paper, please view our JAX codebase.