Train only one gpu?

[huster-wugj]

Thanks for your impressive work!
Can train silk on my single gpu--80gb?

[gleize]

Hi @huster-wugj,

Hum, it might be possible on large GPUs. You can try changing this line to use GPU 0.

However, usually Jax will try to pre-allocate 75% of the memory from the GPUs they are assigned to, which means PyTorch and Jax might run into conflict and trigger an OOM error. So you will likely have to disable the preallocation by setting XLA_PYTHON_CLIENT_PREALLOCATE to false. I haven't tried that, but please keep us updated if this worked. I'll put your solution in the FAQ in case others have the same GPU constraint.

[huster-wugj]

Thanks for your reply.
I will try this solution.
Or, is there anyway to disable computation by Jax?

[gleize]

Hi @huster-wugj ,

No, there is no way to disable Jax as it is necessary to be able to compute the large cosine similarity matrix.