search

BlinkDL / RWKV-LM

RWKV is an RNN with transformer-level LLM performance. It can be directly trained like a GPT (parallelizable). So it's combining the best of RNN and transformer - great performance, fast inference, saves VRAM, fast training, "infinite" ctx_len, and free sentence embedding.
Apache License 2.0
12.7k stars 867 forks source link

RWKV on TPU? #58

Open leejason opened 1 year ago

leejason commented 1 year ago

Is it possible to train RWKV on TPU? If not, what would be needed?

BlinkDL commented 1 year ago

On Nvidia GPUs you need a CUDA kernel for efficient training. Wonder if we can do similar things for TPUs