Repository: huggingface/diffusers
Pull request: [deepspeed] partial ZeRO-3 support
Fixes: https://github.com/huggingface/diffusers/issues/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.

Problem statement:
using zero stage3 finetune sd2, dimension error occurs
### Describe the bug

An error is reported when using deepspeed's zero stage3 finetune diffusers/examples/text_to_image/train_text_to_image.py script. My machine's GPU is 4*2080ti, and because a single GPU cannot accommodate all SD2 parameters, the deepspeed zero stage3 strategy must be used.

### Reproduction

accelerate.yaml
```
compute_environment: LOCAL_MACHINE
deepspeed_config:
 deepspeed_config_file: /home/kas/zero_stage3_offload_config.json
 zero3_init_flag: true
distributed_type: DEEPSPEED
fsdp_config: {}
machine_rank: 0
main_process_ip: null
main_process_port: null
main_training_function: main
mixed_precision: fp16
num_machines: 1
num_processes: 4
use_cpu: false
```

/home/kas/zero_stage3_offload_config.json
```
{
  "train_micro_batch_size_per_gpu": 16,
"gradient_accumulation_steps":2,
  "train_batch_size":128,
  "steps_per_print": 2,
  "gradient_clipping": 1,
  "zero_optimization": {
    "stage": 3,
    "allgather_partitions": false,
    "allgather_bucket_size": 2e8,
    "overlap_comm": true,
    "reduce_scatter": true,
    "reduce_bucket_size": 2e8,
    "contiguous_gradients": true,
    "offload_optimizer": {
      "device": "cpu",
      "pin_memory": true
    },
    "offload_param": {
      "device": "cpu",
      "pin_memory": true
    },
    "stage3_max_live_parameters" : 2e8,
    "stage3_max_reuse_distance" : 2e8,
    "stage3_prefetch_bucket_size": 2e8,
    "stage3_param_persistence_threshold": 2e8,
    "sub_group_size" : 2e8,
    "round_robin_gradients": true
  },
  "bf16": {
    "enabled": true
  }
}
```

launch script
```
pip install deepspeed
export MODEL_NAME="stabilityai/stable-diffusion-2"
export dataset_name="lambdalabs/pokemon-blip-captions"

accelerate launch --config_file ./accelerate.yaml --mixed_precision="fp16"  train_text_to_image.py  \
  --pretrained_model_name_or_path=$MODEL_NAME \
  --dataset_name=$dataset_name \
  --use_ema \
  --resolution=224 --center_crop --random_flip \
  --train_batch_size=16 \
  --gradient_accumulation_steps=2 \
  --gradient_checkpointing \
  --max_train_steps=500 \
  --learning_rate=6e-5 \
  --max_grad_norm=1 \
  --lr_scheduler="constant_with_warmup" --lr_warmup_steps=0 \
  --output_dir="sd-pokemon-model"
```

### Logs

```shell
`0%| | 0/500 [00:00<?, ?it/s] Steps: 0%| | 0/500 [00:00<?, ?it/s]Traceback (most recent call last):
File "train_text_to_image.py ", line 718, in <module>
main()
File "train_text_to_image.py ", line 648, in main
encoder_hidden_states = text_encoder(batch["input_ids"])[0]
File "/opt/miniconda3/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
return forward_call(*input, **kwargs)
File "/opt/miniconda3/lib/python3.7/site-packages/transformers/models/clip/modeling_clip.py", line 739, in forward
return_dict=return_dict,
File "/opt/miniconda3/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
return forward_call(*input, **kwargs)
File "/opt/miniconda3/lib/python3.7/site-packages/transformers/models/clip/modeling_clip.py", line 636, in forward
hidden_states = self.embeddings(input_ids=input_ids, position_ids=position_ids)
File "/opt/miniconda3/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
return forward_call(*input, **kwargs)
File "/opt/miniconda3/lib/python3.7/site-packages/transformers/models/clip/modeling_clip.py", line 165, in forward
inputs_embeds = self.token_embedding(input_ids)
File "/opt/miniconda3/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
return forward_call(*input, **kwargs)
File "/opt/miniconda3/lib/python3.7/site-packages/torch/nn/modules/sparse.py", line 160, in forward
self.norm_type, self.scale_grad_by_freq, self.sparse)
File "/opt/miniconda3/lib/python3.7/site-packages/torch/nn/functional.py", line 2183, in embedding
return torch.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse)
RuntimeError: 'weight' must be 2-D`
```


### System Info

- `diffusers` version: 0.11.1
- Platform: Linux-4.15.0-29-generic-x86_64-with-debian-buster-sid
- Python version: 3.7.7
- PyTorch version (GPU?): 1.11.0+cu113 (True)
- Huggingface_hub version: 0.10.1
- Transformers version: 4.23.1
- Using GPU in script?: yes
- Using distributed or parallel set-up in script?: yes

Hints from issue discussion:
cc @patil-suraj @williamberman @pcuenca 
I don't really know much about deepspeed maybe @williamberman  knows more here.
As far as I have tried, the script works well with stage 2, using CPU offloading which should help fit the model on 2080ti.  
I don't know off the top of my head. Someone opened an issue with the same error message using deepspeed stage 3 with transformers. https://github.com/microsoft/DeepSpeed/issues/2746 Maybe best to see if the deepspeed team knows before digging in :)

Commit message: cleanup

Patch to review:
diff --git a/examples/text_to_image/train_text_to_image.py b/examples/text_to_image/train_text_to_image.py
index 4bbf4706f01c..de9196157818 100644
--- a/examples/text_to_image/train_text_to_image.py
+++ b/examples/text_to_image/train_text_to_image.py
@@ -461,10 +461,43 @@ def main():
     tokenizer = CLIPTokenizer.from_pretrained(
         args.pretrained_model_name_or_path, subfolder="tokenizer", revision=args.revision
     )
-    text_encoder = CLIPTextModel.from_pretrained(
-        args.pretrained_model_name_or_path, subfolder="text_encoder", revision=args.revision
-    )
-    vae = AutoencoderKL.from_pretrained(args.pretrained_model_name_or_path, subfolder="vae", revision=args.revision)
+
+    from accelerate.state import AcceleratorState
+    from transformers.utils import ContextManagers
+
+    def get_deepspeed_plugin():
+        if accelerate.state.is_initialized():
+            return AcceleratorState().deepspeed_plugin
+        else:
+            return None
+
+    def deepspeed_zero_init_disabled_context_manager():
+        """
+        returns either a context list that includes one that will disable zero.Init or an empty context list
+        """
+        deepspeed_plugin = get_deepspeed_plugin()
+        if deepspeed_plugin is None:
+            return []
+
+        return [deepspeed_plugin.zero3_init_context_manager(enable=False)]
+
+    # Currently Accelerate doesn't know how to handle multiple models under Deepspeed ZeRO stage 3.
+    # 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.
+    #
+    # For now the following workaround will partially support Deepspeed ZeRO-3, by excluding the 2
+    # frozen models from being partitioned during `zero.Init` which gets called during
+    # `from_pretrained` So CLIPTextModel and AutoencoderKL will not enjoy the parameter sharding
+    # across multiple gpus and only UNet2DConditionModel will get ZeRO sharded.
+    with ContextManagers(deepspeed_zero_init_disabled_context_manager()):
+        text_encoder = CLIPTextModel.from_pretrained(
+            args.pretrained_model_name_or_path, subfolder="text_encoder", revision=args.revision
+        )
+        vae = AutoencoderKL.from_pretrained(
+            args.pretrained_model_name_or_path, subfolder="vae", revision=args.revision
+        )
+
     unet = UNet2DConditionModel.from_pretrained(
         args.pretrained_model_name_or_path, subfolder="unet", revision=args.non_ema_revision
     )
diff --git a/src/diffusers/training_utils.py b/src/diffusers/training_utils.py
index 340b96e29ac5..f0379d6ce5f0 100644
--- a/src/diffusers/training_utils.py
+++ b/src/diffusers/training_utils.py
@@ -5,6 +5,7 @@
 
 import numpy as np
 import torch
+import transformers
 
 from .utils import deprecate
 
@@ -197,11 +198,17 @@ def step(self, parameters: Iterable[torch.nn.Parameter]):
         self.cur_decay_value = decay
         one_minus_decay = 1 - decay
 
+        if transformers.deepspeed.is_deepspeed_zero3_enabled():
+            import deepspeed
+        context_manager = []
         for s_param, param in zip(self.shadow_params, parameters):
-            if param.requires_grad:
-                s_param.sub_(one_minus_decay * (s_param - param))
-            else:
-                s_param.copy_(param)
+            if transformers.deepspeed.is_deepspeed_zero3_enabled():
+                context_manager = [deepspeed.zero.GatheredParameters(param, modifier_rank=None)]
+            with transformers.utils.ContextManagers(context_manager):
+                if param.requires_grad:
+                    s_param.sub_(one_minus_decay * (s_param - param))
+                else:
+                    s_param.copy_(param)
 
     def copy_to(self, parameters: Iterable[torch.nn.Parameter]) -> None:
         """


Code review comments:
diff_hunk: @@ -461,10 +461,43 @@ def main():
     tokenizer = CLIPTokenizer.from_pretrained(
         args.pretrained_model_name_or_path, subfolder="tokenizer", revision=args.revision
     )
-    text_encoder = CLIPTextModel.from_pretrained(
-        args.pretrained_model_name_or_path, subfolder="text_encoder", revision=args.revision
-    )
-    vae = AutoencoderKL.from_pretrained(args.pretrained_model_name_or_path, subfolder="vae", revision=args.revision)
+
+    from accelerate.state import AcceleratorState
+    from transformers.utils import ContextManagers
+
+    def get_deepspeed_plugin():
+        if accelerate.state.is_initialized():
+            return AcceleratorState().deepspeed_plugin
+        else:
+            return None
+
+    def deepspeed_zero_init_disabled_context_manager():
+        """
+        returns either a context list that includes one that will disable zero.Init or an empty context list
+        """
+        deepspeed_plugin = get_deepspeed_plugin()
path: examples/text_to_image/train_text_to_image.py
text: @user1:
```suggestion
        deepspeed_plugin = AcceleratorState().deepspeed_plugin if accelerate.state.is_initialized() else None
```

and remove function above? Think function is only used once and this way the example is more concise, readable.

diff_hunk: @@ -5,6 +5,7 @@
 
 import numpy as np
 import torch
+import transformers
path: src/diffusers/training_utils.py
text: @user1:
Actually `transformers` is still not a hard dependency (some training scripts don't need `transformers`)

diff_hunk: @@ -197,11 +198,17 @@ def step(self, parameters: Iterable[torch.nn.Parameter]):
         self.cur_decay_value = decay
         one_minus_decay = 1 - decay
 
+        if transformers.deepspeed.is_deepspeed_zero3_enabled():
+            import deepspeed
+        context_manager = []
         for s_param, param in zip(self.shadow_params, parameters):
-            if param.requires_grad:
-                s_param.sub_(one_minus_decay * (s_param - param))
-            else:
-                s_param.copy_(param)
+            if transformers.deepspeed.is_deepspeed_zero3_enabled():
+                context_manager = [deepspeed.zero.GatheredParameters(param, modifier_rank=None)]
+            with transformers.utils.ContextManagers(context_manager):
path: src/diffusers/training_utils.py
text: @user1:
```suggestion
            context_manager = deepspeed.zero.GatheredParameters(param, modifier_rank=None) if is_deepspeed_enabled else contextlib.nullcontext
            with context_manager:
```

diff_hunk: @@ -461,10 +461,43 @@ def main():
     tokenizer = CLIPTokenizer.from_pretrained(
         args.pretrained_model_name_or_path, subfolder="tokenizer", revision=args.revision
     )
-    text_encoder = CLIPTextModel.from_pretrained(
-        args.pretrained_model_name_or_path, subfolder="text_encoder", revision=args.revision
-    )
-    vae = AutoencoderKL.from_pretrained(args.pretrained_model_name_or_path, subfolder="vae", revision=args.revision)
+
+    from accelerate.state import AcceleratorState
path: examples/text_to_image/train_text_to_image.py
text: @user1:
can we move those imports to the top?