TPU Integration
Open
feature request
- Dominant language
- Python
- Stars
- 4.8k
- Forks
- 487
- PR merge metrics
- No merged PRs in 30d
Description
### 🚀 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.
Contributor guide
Assessment
This issue has not been assessed yet.