trlx icon indicating copy to clipboard operation
trlx copied to clipboard

Published 20 hours ago • verified publisher icon CarperAI

Add Jax support

Open Dahoas opened this issue 2 years ago • 3 comments

🚀 The feature, motivation, and pitch

Add jax support for RLHF on TPUs.

Alternatives

No response

Additional context

No response

Dahoas avatar Oct 21 '22 13:10 Dahoas

We should, rather than just add JAX support, add T5X support so it can be used with models that were trained via T5X.

LouisCastricato avatar Oct 21 '22 15:10 LouisCastricato

can it be trained via T5X now ?

joytianya avatar Jan 30 '23 10:01 joytianya

T5X support is something we're planning for late 2023 at the earliest.

LouisCastricato avatar Jan 30 '23 10:01 LouisCastricato