The device of model.w["emb.weight"] is in CPU

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.47k stars 847 forks source link

The device of model.w["emb.weight"] is in CPU #249

Closed MarshtompCS closed 2 months ago

MarshtompCS commented 3 months ago

Hi, I am using RWKV-6-world-7b and find the emb.weight is at CPU. I wonder why it is not put to the CUDA. Would it be faster in the GPU?

BlinkDL commented 3 months ago

Hi because there is almost no computation in emb layer, it's fine to put it on CPU, which saves VRAM.