TPU Integration

[steventk-g]

🚀 The feature, motivation, and pitch

trlX uses HuggingFace accelerate under the hood. Accelerate has the capability to leverage Google's TPUs for faster training. I'm interested in supporting trlX on TPUs. This can be user configurable using the TPU fields in the accelerate config.

Alternatives

No response

Additional context

I'm a SWE on the PyTorch/XLA team at Google, and I work with HF to test and support their libraries on TPU. I would be happy to help test and develop this feature as well.

[steventk-g]

@reciprocated feel free to assign this me, just wanted to put it on your radar.

[joytianya]

does it support tpu?

[leejason]

Any progress on TPU? Thanks for information.