Farm: DeepSpeed integration for training from scratch

[tholor]

Context
We already added in FARM 0.4.1:

• basic support for LM training from scratch
• full checkpointing & SageMaker integration to allow saving ~ 70% of costs via spot instances

We are currently working on:

• a fast streaming data silo to handle large datasets that don't fit in memory (https://github.com/deepset-ai/FARM/pull/239).

A logical next step, I see:

• optimizing the model speed during training

Current baseline
Right now we get:

• ~ 1.575 Mio samples per hour (15k steps / hour with batch size 105)

With:

• p3.8xlarge with 4x V100 GPUs
• model = bert-base
• batch size = 105
• gradient accumulation steps = 9
• max seq len = 128
• no automatic mixed precision (Amp)
• torch's DataParallel

Potential of DeepSpeed
I see a ton of options to optimize, but one interesting direction could be Microsoft's DeepSpeed / ZeRO seems to be an exciting opportunity: https://github.com/microsoft/DeepSpeed
Their PyTorch wrapper seems to be rather lightweight and the speed-up they achieved is significant.

Integration sketch

They initialize model and optimizer in a way that is already pretty similar to FARM style.

DeepSpeed:

model_engine, optimizer, _, _ = deepspeed.initialize(args=cmd_args,
                                                     model=model,
                                                     model_parameters=params)

FARM:
https://github.com/deepset-ai/FARM/blob/49019c1d074ce97002fc9cd990b8afa6524d1728/examples/train_from_scratch.py#L77-L85

The main work is probably to:

• figure out the cmd_args that are required
• investigate if their model_engine & optimizer and behave like standard PyTorch objects or if we have to deal with them in a special way (e.g. saving / loading and LR schedule)
• how to set up all dependencies (incl the required config file) in a docker

Any help from our community on this would be highly appreciated :heart:

@BramVanroy, I know you are keen on efficient optimization & parallelization. Might this be a topic that you are interested in :) ?

[BramVanroy]

I am _definitely_ interested and I wish this was something I could work full-time on. However, I am currently pre-occupied with other deadlines, probably until May at least. I can do code reviews or benchmark tests (I have a local (non-cloud) 4x V100 server available.)

DeepSpeed is a very promising (and necessary) development in Deep Learning. Not only its scalability and speed improvements, but also other features such as their LRRT.

I'd have to dig deeper into the FARM API and how different it is from DeepSpeed and/or how easily transferable it is. Do you want to build on top of DeepSpeed, or make it optional? That is important to decide how to go from there.

[tholor]

Oh nice, hadn't seen the LRRT - very useful!

Do you want to build on top of DeepSpeed, or make it optional?

I see it as an optional way of training in FARM

However, I am currently pre-occupied with other deadlines, probably until May at least. I can do code reviews or benchmark tests (I have a local (non-cloud) 4x V100 server available.)

What a pity! We have also a few other topics on our plate that we need to focus on first (e.g. scaling & extending QA with haystack), but we will let you know once we start working on this. Any support on code reviews / benchmarks would be great :)

[BramVanroy]

Making things optional is indeed more user-friendly, but it makes the implementation a lot harder I think. I am interested to see how this develops further, but I don't have the time to look into this now.

By the way, it might be useful to have a talk with the huggingface people, they're also working on their Q&A pipeline. Working together might yield cool ideas! (cf https://github.com/huggingface/transformers/issues/547#issuecomment-587957033)

[tholor]

Good point. With haystack we are more aiming for a large scale QA/semantic search system (incl. retriever, elasticsearch integration, re-ranking models ...). I feel this is quite different to the scope of Transformers (and FARM). However, we made sure that it's easy to load Transformers QA models into haystack. For this, we leverage the pipeline object and we will try to collaborate more on this - especially to support more model architectures.

[stale[bot]]

This issue has been automatically marked as stale because it has not had recent activity. It will be closed in 14 days if no further activity occurs.