Can unlimiformer be trained on mutiple gpus?

abertsch72 / unlimiformer

Public repo for the NeurIPS 2023 paper "Unlimiformer: Long-Range Transformers with Unlimited Length Input"

MIT License

1.05k stars 80 forks source link

Can unlimiformer be trained on mutiple gpus? #16

Open Muxv opened 1 year ago

Muxv commented 1 year ago

Is there ant option I can modify to support mutiple gpus training?

urialon commented 1 year ago

Hi @Muxv , Thank you for your interest in our work!

It is currently not fully supported. I tried to run with multiple GPUs once, but I got errors that the "memory" tensors are not on the same GPU as the query tensors. We would love to adopt contributions.

Best, Uri