huggingface / optimum-neuron

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

Apache License 2.0

196 stars 59 forks source link

Optimize checkpointing #505

Closed dacorvo closed 7 months ago

dacorvo commented 7 months ago

What does this PR do?

This modifies the decoder models export code to reduce the disk usage when creating checkpoints:

use torch_dtype = auto when loading the model to avoid casting weigths to float32 (the default),
do not use snapshot_download before exporting to avoid downloading both pytorch and safetensors weights.

HuggingFaceDocBuilderDev commented 7 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.