huggingface / huggingface/transformers

[trainer] new in pytorch: `torch.optim._multi_tensor` faster optimizers

Open
#9,965 8 comments 8 reactions 1 assignee Claimed by @stas00 View on GitHub
Performance WIP
Dominant language
Python
Stars
166k
Forks
34.6k
Avg merge
3d 8h
Merged PRs (30d)
276

Description

Back in September pytorch introduced `torch.optim._multi_tensor` https://github.com/pytorch/pytorch/pull/43507 which should be much more efficient for situations with lots of small feature tensors (`transformers`) and thus should show an appreciable speed up in training. If someone is interested in the progress of this project here is the stack to track: https://github.com/pytorch/pytorch/pull/48223

This feature is currently an alpha stage, so users can try to use it by simply replacing `torch.optim` with `torch.optim._multi_tensor` in HF Trainer or their own trainer.

Eventually it'll replace `torch.optim` so there is nothing that we need to do otherwise.

@blefaudeux who alerted me to this improvement suggested it should have good speed ups for the DDP/Sharded DDP training.

If resources allow it'd be good to run some benchmarks. Please feel free to beat me to it.

Thanks to @blefaudeux for the heads up, and @izdeby for working on this enhancement and clarifying where things are at.

heads up to: @sgugger, @patrickvonplaten - nothing else that needs to be done.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.