Memory Efficient FP 16 Training

[rajgar114]

🚀 Feature request

Fairseq uses memory efficient FP 16 training as explained in https://arxiv.org/pdf/1904.01038.pdf.

Motivation

Generally the model requires high end GPU's to fine-tune on larger length datasets. Using memory efficient FP 16 we can reduce the need of high GPU's and thus models can be fine-tune without OOM problems.

[LysandreJik]

Is "memory efficient" fp16 different to the fp16 available in the training scripts?

[rajgar114]

Yes, both are different.
facebookresearch/fairseq#2907

[LysandreJik]

Thanks for the link! cc @sgugger @stas00

[stas00]

I think you must be referring to this comment facebookresearch/fairseq#2907 (comment)

--memory-efficient-fp16 gets rid of the FP32 model copy and only maintains FP32 momentum in the optimizer. Thus you'll see 0.5x memory usage from the model weights, 0.5x memory usage in the forward/backward, and 1.0x memory usage in the optimizer (relative to FP32).

correct?

I think this is the implementation: https://github.com/pytorch/fairseq/blob/f8a7c93440cd925f70979a6082c18f830b39e44b/fairseq/optim/fp16_optimizer.py#L456

Appears to be added 2 years ago. And you quoted an old paper from 2019. Do you think it's actually something that's worth investigating? Somehow I'd expect for it to adopted by other projects if it were to work great, so it'd be good to ask someone experienced with it whether it's actually good.

[stas00]

Not sure how to evaluate the effectiveness this proposal, it would be helpful to have some case studies that show the actual improvements.

I asked around and someone reported that someone mentioned this was useful for certain models, but it'd help to know which and how they were trained so that there is an actual proven setup to work with.

I found pytext included a variation of it here: facebookresearch/pytext@c6d13ac, but I can't find any other references via google, which is not very encouraging. But we don't know whether other implementations use the same name.

But also let's ask this, @rajgar114, would you like to try to work on this and submit a PR when you have something working?

[rajgar114]

I think you must be referring to this comment pytorch/fairseq#2907 (comment) ?

Yes, @stas00 I was referring to this comment pytorch/fairseq#2907 only.

Appears to be added 2 years ago.

No doubt the paper is quite old but I have found some instances that memory efficient fp16 training worked for people having low end GPU's. Here is an example:
https://bleepcoder.com/fairseq/551167214/oom-while-trying-to-train-bart

Do you think it's actually something that's worth investigating?

I am not completely sure that how much impact it can cause on the reduction in GPU memory consumption. We should definitely ask some experienced guys and also try to compare and analyze the results with & without using memory efficient fp16 version.

Would you like to try to work on this and submit a PR when you have something working?

@stas00 Thanks for giving this wonderful opportunity. I would love to work on open source projects. But I can't because of my busy schedule, I would not be able to spend time on this project. Sorry for that.

[stas00]

Thank you for the feedback, @rajgar114 and letting us know that you'd love to work on that but reality won't allow that at the moment.

I have added it to #12126 so it won't get lost and hopefully it'd find a champion.

From your part what would help is to find specific models/setups where it has been found to be converging well despite the limitations it imposed on itself. So that whoever works on this will have a good chance of succeeding. Thank you.

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

[ksopyla]

I do not know exactly if my case fits this issue, but I have trained seq2seq model with fairseq (v0.10.1) and current transformer 4.9. In both cases I have used encoder-decoder and BART architecture and saw that Fariseq allows me to use a greater batch size 120- 160 (sentences) around 5000 tokens. Transformers library with deepspeed integration handles max 48(sentences).
I have used 3x Geforce 3090 with 24 GB ram. In both cases model size ~70M parameters (6 layers for encoder and decoder, hidden_size=512, ff=2048)
Concludes, fairseq training is almost 3 times faster (measured by the number of tokens seen during training in a fixed time budget).

[stas00]

Have you tried using activation checkpointing? That should save a lot of memory and enable much larger batch sizes.

Also it might be good to try both Deepspeed zero-2 and zero-3 stages - I don't know which one you were optimizing with.