[deepspeed] partial ZeRO-3 support

[stas00]

Fixes: #1865

Here is what's happening here:

When Deepspeed ZeRO-3 is used in examples/text_to_image/train_text_to_image.py the 3 models (text_encoder, vae, unet) get partitioned with zero.Init but only unet gets accelerate.prepare so Deepspeed doesn't know to automatically gather the already partitioned weights before each forward call for text_encoder and vae. That's the source of the problem.

Now, to provide full ZeRO-3 support you have to contact https://github.com/huggingface/accelerate and ask them to extend support to this situation. Specifically:

Currently Accelerate doesn't know how to handle multiple models.

For this to work properly all models must be run through accelerate.prepare. But Accelerate will try to assign the same optimizer with the same weights to all models during deepspeed.initialize, which of course doesn't work. It needs to assign the correct weights to the correct optimizer, so that the param_groups are populated with the weights that are inside the model itself. In the case of frozen models there should be no optimizer at all (or param_groups should be empty).

It will need to be able to support something like the following in accelerate.prepare:

# frozen models
deepspeed.initialize(model=text_encoder, model_parameters=None, optimizer=None, ...)
deepspeed.initialize(model=vae, model_parameters=None, optimizer=None, ...)
# normal model
deepspeed.initialize(model=unet, model_parameters=unet_model_parameters, optimizer=optimizer, ...)

For now this PR will partially support Deepspeed ZeRO-3, by excluding the 2 frozen models from being partitioned during zero.Init which gets called during from_pretrained. This is done for examples/text_to_image/train_text_to_image.py only, but can be replicated to other scripts if needed.

So CLIPTextModel and AutoencoderKL will not enjoy the parameter sharding across multiple gpus and only UNet2DConditionModel will get zero sharded, and so memory savings will only apply to the latter. The previous 2 will have full copies on each gpu.

---

to reviewers: As I'm not familiar with this project's structure I will leave to maintainers to choose where to put the helper code inside the example I modified to work. I just put it in one place for visibility.

---

If you want full ZeRO-3 support for optimal training efficiency, here is what you need to do

1. resolve the issue with Accelerate as discussed in this OP
2. then will have to ask Deepspeed to provide support for multiple models - but we can't do that until stage 1 is done.

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[stas00]

the ci failures look to be unrelated - some timeout.

[uygnef]

Thank a lot.

Fixes: #1865

Here is what's happening here:

When Deepspeed ZeRO-3 is used in examples/text_to_image/train_text_to_image.py the 3 models (text_encoder, vae, unet) get partitioned with zero.Init but only unet gets accelerate.prepare so Deepspeed doesn't know to automatically gather the already partitioned weights before each forward call for text_encoder and vae. That's the source of the problem.

Now, to provide full ZeRO-3 support you have to contact https://github.com/huggingface/accelerate and ask them to extend support to this situation. Specifically:

Currently Accelerate doesn't know how to handle multiple models.

For this to work properly all models must be run through accelerate.prepare. But Accelerate will try to assign the same optimizer with the same weights to all models during deepspeed.initialize, which of course doesn't work. It needs to assign the correct weights to the correct optimizer, so that the param_groups are populated with the weights that are inside the model itself. In the case of frozen models there should be no optimizer at all (or param_groups should be empty).

It will need to be able to support something like the following in accelerate.prepare:

# frozen models
deepspeed.initialize(model=text_encoder, model_parameters=None, optimizer=None, ...)
deepspeed.initialize(model=vae, model_parameters=None, optimizer=None, ...)
# normal model
deepspeed.initialize(model=unet, model_parameters=unet_model_parameters, optimizer=optimizer, ...)

For now this PR will partially support Deepspeed ZeRO-3, by excluding the 2 frozen models from being partitioned during zero.Init which gets called during from_pretrained. This is done for examples/text_to_image/train_text_to_image.py only, but can be replicated to other scripts if needed.

So CLIPTextModel and AutoencoderKL will not enjoy the parameter sharding across multiple gpus and only UNet2DConditionModel will get zero sharded, and so memory savings will only apply to the latter. The previous 2 will have full copies on each gpu.

to reviewers: As I'm not familiar with this project's structure I will leave to maintainers to choose where to put the helper code inside the example I modified to work. I just put it in one place for visibility.

If you want full ZeRO-3 support for optimal training efficiency, here is what you need to do

1. resolve the issue with Accelerate as discussed in this OP
2. then will have to ask Deepspeed to provide support for multiple models - but we can't do that until stage 1 is done.

Thanks a lot. It works.
Just curious
"then will have to ask Deepspeed to provide support for multiple models - but we can't do that until stage 1 is done." If stage 1 for multiple models hasn't been completed yet, how does DeepSpeed-RLHF use zero-optimal?

[stas00]

Last I checked it had troubles with multiple models, but perhaps it was in a particular situation, and it might already be working. So if you'd like to take a lead on that, please file a feature request at accelerate to solve what I described in the OP. And once this is done we can then work on the deepspeed side.

If you do that please tag me on your feature request and link to this PR for context.

Thank you.

[stas00]

ok, so I was just talking to @tjruwase and he shared that the newly released DeepSpeed-Chat https://github.com/microsoft/DeepSpeedExamples/tree/master/applications/DeepSpeed-Chat which coincided with Deepspeed 0.9.0 release now supports multiple DS engines, e.g. see:

https://github.com/microsoft/DeepSpeedExamples/blob/d22b82b7d167f3e432234795da7fdf7965db8891/applications/DeepSpeed-Chat/training/step3_rlhf_finetuning/rlhf_engine.py#L116-L119

https://github.com/microsoft/DeepSpeedExamples/blob/d22b82b7d167f3e432234795da7fdf7965db8891/applications/DeepSpeed-Chat/training/step3_rlhf_finetuning/rlhf_engine.py#L145-L146

Therefore, let's ping @pacman100 - Sourab, please kindly have a look at this OP, we need to have accelerate support multiple DS engines, each with a different setup (in particular optimizer and optionally lack of), to make it work with diffusion models, which have at least 3 independent models. Not sure, perhaps let's ask @sgugger as well for an opinion.

To quote the relevant parts from the OP

It will need to be able to support something like the following in accelerate.prepare:

# frozen models
deepspeed.initialize(model=text_encoder, model_parameters=None,                  optimizer=None, ...)
deepspeed.initialize(model=vae,          model_parameters=None,                  optimizer=None, ...)
# normal model
deepspeed.initialize(model=unet,         model_parameters=unet_model_parameters, optimizer=optimizer, ...)

but of course it should also do the right thing for non-deepspeed situations. Probably be a no-op then? Not sure of the best way.

and what if unlike in the case of diffusers the user may want more than one model to be trainable - i.e. with a distinct optimizer?

[uygnef]

hi @stas00
I have published the feature request at #3098.

[pacman100]

Hello @stas00,

The thing is that there are multiple challenges with this which need discussion:

1. When there are multiple models, and if one of the models required zero-3 while other require zero-2 (for example actor model requires zero-3 but the small reward model needs zero-2). it will activate zero-init when remaining don't require it and it will involve a lot of juggling to use the disabling context manager as you did above ...
2. Accelerate currently takes as input via accelerate config questionnaire only single set of config, but different models will require different config setting. In the above DeepSpeedChat they are having helper functions to get different ds_config for each of the components such as actor and critic models. Accelerate can't have helper functions just for this usage as it won't be generic.
3. Even if the multiple config is somehow supported, big question about auto values and how to fill them, if user passes a few of the values via cmd args, then which of the multiple configs will it belong to? And then the remaining values being filled with prepare call, different models might need different values for gradient accumulation, train batch size per gpu ... So, how would prepare work in that case as it can't distinguish this dataloader is meant for this model and not that model ...

[tjruwase]

@pacman100, these are great questions that you raise and essentially the same ones we wrestled with in the DeepSpeed Chat project. I would like to think that our solutions are general, but they did require a re-think of how DeepSpeed worked.

For example, for (1), we do support actor and critic models using different zero stages as shown in our 1.3b actor on single GPU example. We also observed that some code refactoring could simplify the required context manager juggling.

If I sound a bit bullish, it is not because I want to diminish these challenges, but rather because I am excited to be tackling these difficult yet unavoidable problems to support the next generation of ML models.

[stas00]

Sourab, thank you for listing the challenges you see and I'm sure there are even more issues that we haven't yet thought about.

When we did the initial integration in HF Trainer and when you forked it to Accelerate there was always one model - as DS didn't support more than one model. But now it does so what I propose is that we design a new solution that accommodates multiple models each potentially having its own needs.

Wrt accelerate config - you already know my stand on it - it's a great approach for a very simplistic beginner get started quickly way, but it doesn't work for advanced uses. Advanced uses require custom config files for each unique situation and relying on a single global config file makes no sense.

Bottom line, do you want to meet and brainstorm on how to solve this new need, or we can of course do it in an Issue - basically we need an RFC with someone writing the initial design.

[stas00]

hi @stas00 I have published the feature request at #3098.

Thank you, @uygnef, but it needs to be moved to: https://github.com/huggingface/accelerate/issues

Diffusers is a user of Accelerate, it can't change Accelerate.

[tjruwase]

Bottom line, do you want to meet and brainstorm on how to solve this new need, or we can of course do it in an Issue - basically we need an RFC with someone writing the initial design.

If it makes sense, I would love to be a part of such a brainstorming effort.

[pacman100]

Hello @stas00 and @tjruwase,

If I sound a bit bullish, it is not because I want to diminish these challenges, but rather because I am excited to be tackling these difficult yet unavoidable problems to support the next generation of ML models.

Completely agree with this!

Bottom line, do you want to meet and brainstorm on how to solve this new need, or we can of course do it in an Issue - basically we need an RFC with someone writing the initial design.

I can take this up. Will try to put together a doc with the initial design and raise an RFC issue. ETA for this is Friday (21 April)

If it makes sense, I would love to be a part of such a brainstorming effort.

Yes, that would be great, we can even have a quick sync up post having an initial design. Thank you!

[stas00]

That sounds like a plan, Sourab. Thank you for working on it!

[patrickvonplaten]

Thanks for the fixes here @stas00 !

This looks good to me! @patil-suraj @sayakpaul @williamberman could you also take a look?

[patrickvonplaten]

works for me!

[patrickvonplaten]

Actually transformers is still not a hard dependency (some training scripts don't need transformers)

[patrickvonplaten]

Suggested change

	with transformers.utils.ContextManagers(context_manager):
	context_manager = deepspeed.zero.GatheredParameters(param, modifier_rank=None) if is_deepspeed_enabled else contextlib.nullcontext
	with context_manager:

[patrickvonplaten]

In general, this looks good to me? Would be nice to not force-add transformers in training_utils.py as it's not a hard dependency. Ideally we can just have some nice deepspeed tools also in diffusers? Also cc @patil-suraj

[sayakpaul]

The solution looks minimal, clean, and good to me. More so, because we're aware of the limitations too (thanks @stas00 for that). Also thanks to @pacman100 and @tjruwase for their inputs and being open to resolving the current challenges :)

I don't have additional comments on top of what @patrickvonplaten already shared.

[stas00]

@patrickvonplaten, thank you for the review - please feel free to adapt it to any way you like it as I'm not involved with this project. I made the key changes to make it work but it's totally up to you how you want it to be integrated.

Additionally, we are waiting for @pacman100 to write a spec to how Accelerate plans to handle multiple concurrent deepspeed model engines, which is already supported by Deepspeed, so Accelerate is the bottleneck. Once this is in place this PR can be replaced with a proper solution that is no longer a workaround.

[patrickvonplaten]

Tests are passing, let's merge :-)