Omnigen 2 Group Offloading

[mamei16]

Hi, I'd like to use group offloading on the new Omnigen 2 model, but it appears that applying it drastically changes the output, regardless of whether block_level or leaf_level offloading is used. It seems that some other 2D transformers don't support group offloading. Is this true for Omnigen 2 also? And if so why?

[a-r-r-o-w]

When you use group offloading, have you compared the results with use_stream=False and use_stream=True? Does one of these settings give you coherent results or both don't? It will help me try to pin down the major issue behind this because I've seen it happen at times too on specific models, but it is not obvious where the possible bug might be, or if it's even a bug.

Since you are able to run group offloading, it is unrelated to the other cases you linked to. Those models are skipped because of a call to F.multi_head_attention_forward in the forward pass, which directly takes in the weight/bias params, and so the hooks to trigger weight onloading/offloading are never called. Support for this case can be added in the near future once I find some time.

What might be happening in the case of OmniGen 2 is some kind of race condition or improper synchronization, but I don't know where to start with debugging this issue at the moment. Will first try to add native support for Omnigen 2 so that it's easier to get started for us.

[mamei16]

Thanks a lot for the answer! Both use_stream=False and use_stream=True result in incoherent results.

[mamei16]

Okay, I have narrowed the issue down: It only appears to happen if a reference image is provided. Without a reference image, the output with group offloading applied is the same as without.

[mamei16]

Thanks to your comment regarding the forward pass of the skipped models, I took a closer look at the forward method of the OmniGen2Transformer2DModel and discovered that a kwarg was used, which was not being included when group offloading was applied. By simply converting it to a regular argument, I was able to get it working!
Sorry for the comment spam and thanks again for your time.

[a-r-r-o-w]

Awesome, glad to know this specific issue could be resolved with a simple fix :)

Could you kindly point me to the kwarg part of the implementation and so we could have a linked reference in this discussion? It will help me better understand what's going on and maybe we could do something about those cases on our end. Thanks!

[mamei16]

Sure! The part in question can be found on the lines following pipeline_omnigen2.py#L715 to which I applied the following change:

 
         batch_size, num_channels_latents, height, width = latents.shape
         
-        optional_kwargs = {}
-        if 'ref_image_hidden_states' in set(inspect.signature(self.transformer.forward).parameters.keys()):
-            optional_kwargs['ref_image_hidden_states'] = ref_image_hidden_states
-        
         model_pred = self.transformer(
             latents,
             timestep,
             prompt_embeds,
             freqs_cis,
             prompt_attention_mask,
-            **optional_kwargs
+            ref_image_hidden_states

         )
         return model_pred

And as a result, I of course had to change transformer_omnigen2.py#L505:

         text_hidden_states: torch.Tensor,
         freqs_cis: torch.Tensor,
         text_attention_mask: torch.Tensor,
-        ref_image_hidden_states: Optional[List[List[torch.Tensor]]] = None,
+        ref_image_hidden_states: Optional[List[List[torch.Tensor]]],
         attention_kwargs: Optional[Dict[str, Any]] = None,
         return_dict: bool = False,
     ) -> Union[torch.Tensor, Transformer2DModelOutput]:

I also ran into the following error when applying group offloading:

 File "OmniGen2/omnigen2/pipelines/omnigen2/pipeline_omnigen2.py", line 597, in __call__
    latents = self.prepare_latents(
              ^^^^^^^^^^^^^^^^^^^^^
  File "OmniGen2/omnigen2/pipelines/omnigen2/pipeline_omnigen2.py", line 235, in prepare_latents
    latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "env/lib/python3.11/site-packages/diffusers/utils/torch_utils.py", line 67, in randn_tensor
    raise ValueError(f"Cannot generate a {device} tensor from a generator of type {gen_device_type}.")
ValueError: Cannot generate a cpu tensor from a generator of type cuda.

Which I "solved" by hard-coding the device set in the call method:

         else:
             batch_size = prompt_embeds.shape[0]
 
-        device = self._execution_device
+        device = torch.device("cuda")
 
         # 3. Encode input prompt
         (