A summary of primitives provided by DeepSpeed to reduce memory?

[g-karthik]

What are some primitives (other than "zero_optimization": true and "fp16": { "enabled": true }) that are provided in DeepSpeed to help reduce memory consumption? Also, what other relevant primitives will be available in future releases? Could you share tentative dates for those releases?

Thanks!

[samyam]

@g-karthik the ones you listed are currently the main ones. You can reduce the memory footprint for the activation using activation checkpointing in addition to those. For example, see the Megatron Tutorial, but it is not a primitive inside Deepspeed..

There are a few relevant primitives we are working on for future releases. ZeRO Stage 2 is a big one that will reduce the memory footprint of the model memory by another 2x. We don't have an exact date for the release, but we are targeting for a month or two from now.

Are there scenarios that you are running into where memory is a bottleneck. If so, can you help us understand those?

[samyam]

I forgot to mention, you could also use gradient accumulation in DeepSpeed to limit the size of the activation memory.

[g-karthik]

@samyam Thanks for the quick response! To be clear, do you mean doing gradient checkpointing in the underlying nn.Module sub-class when you say "activation checkpointing"? If so, yes, I am already doing it.

And yep, I am using gradient accumulation as well - I excluded that from the original list of primitives because it assumes that I am able to fit a single example in memory.

Are there some fancier sub-parameters/args one can use along with enabling FP16 that would further reduce memory consumption?

I am currently doing single-node multi-GPU fine-tuning. Even when DeepSpeed is enabled, I'm running into CUDA out-of-memory issues when trying to fine-tune GPT2DoubleHeadsModel from the Hugging Face transformers library with certain datasets where the max sequence length is ~ 500. Gradient checkpointing turned out to be the savior for me finally and the same job ran successfully when gradient checkpointing was enabled in addition to DeepSpeed. The wall clock time for fine-tuning took a significant hit, however, as you'd expect.

Having said that, for my use-case, the max sequence length is a parameter under my control, and I would like to be able to use all 1024 available positions in GPT-2 during fine-tuning. Do you have any suggestions/feedback on how I can accomplish this?

[samyam]

@g-karthik sorry for the late response.

I do mean gradient checkpointing or re-materialization. I prefer the term activation checkpointing because we are literally saving a few activation in memory for re-computation during backprop.

How much of a performance hit did you notice? We expect an about 30% perf drop from using gradient checkpointing.

I would imagine that for a 1.5B GPT-2 model with enough GPUs, to spread your optimizer states across with DeepSpeed, you should be able to run at least a single batch size without gradient checkpointing. But if you cannot run a batch size of 1, with zero enabled, gradient checkpointing is the only solution I can think of at the moment.

Are you not able to fit a sequence length of 1024 even with gradient checkpointing enabled?

[g-karthik]

@samyam I didn't measure the increase in train time, but it certainly took much longer to converge than without gradient checkpointing.

I would imagine that for a 1.5B GPT-2 model with enough GPUs, to spread your optimizer states across with DeepSpeed, you should be able to run at least a single batch size without gradient checkpointing.

Unfortunately I'm not able to do this, even with a max sequence length of about 500.

Are you not able to fit a sequence length of 1024 even with gradient checkpointing enabled?

I didn't try to fit a sequence length of 1024, because even 500 was not working without gradient checkpointing. But given that I was able to get 500 to work with gradient checkpointing, I suppose 1024 might be doable as well.

[g-karthik]

@samyam An unrelated question: in your experience, does using the Lamb optimizer with DeepSpeed typically reduce training time in comparison to Adam? If so, by how much?