huggingface / optimum-neuron

Easy, fast and very cheap training and inference on AWS Trainium and Inferentia chips.

Apache License 2.0

176 stars 53 forks source link

Init on the `xla` device #521

Closed michaelbenayoun closed 3 months ago

michaelbenayoun commented 3 months ago

What does this PR do?

Initializes parallel layers on the xla device directly, instead of the cpu device. It saves a lot of host memory.
It does the same for the weights that are initialized during lazy load.

HuggingFaceDocBuilderDev commented 3 months ago

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.