Groupoffloading introduce bad results

[seed93]

Describe the bug

I am using groupoffloading for saving gpu memory. I got worse results with a cosine similarity aboud 0.934 on A800, which is unexpected. And I got results with a cosine similarity about 0.78 on 4090, which is worse.
Could anyone give me any suggestions to fix the precision?

Reproduction

apply_group_offloading(
    transformer,
    onload_device=torch.device(f"cuda:{self.local_rank}"),
    offload_device=torch.device("cpu"),
    offload_type="block_level",
    num_blocks_per_group=1,
    non_blocking=True,
    use_stream=True,
)

### Logs

```shell

System Info

I tried diffusers 0.33.1 and 0.34.

Who can help?

No response

[a-r-r-o-w]

Hi @seed93 , we've come across different/worse results in atleast one other case (reported by Sayak from the team) with block level offloading. It is most likely related to faulty synchronisation of streams, but we haven't been able to identify the issue yet. Ideally, group offloading will be lossless once the potential bug is fixed

Do you notice any problems with leaf level offloading?

[seed93]

Hi @seed93 , we've come across different/worse results in atleast one other case (reported by Sayak from the team) with block level offloading. It is most likely related to faulty synchronisation of streams, but we haven't been able to identify the issue yet. Ideally, group offloading will be lossless once the potential bug is fixed

Do you notice any problems with leaf level offloading?

Thank you for your quick reply. I tried leaf_level on a800 just minites ago and it worked good.

[seed93]

One more thing, I failed to run leaf_level on 4090 with an error like this:

Failed running call_function <method 'set_' of 'torch._C.TensorBase' objects>(*(Parameter(FakeTensor(..., device='cuda:0', size=(4096, 4096), dtype=torch.bfloat16,
           requires_grad=True)), FakeTensor(..., size=(4096, 4096), dtype=torch.bfloat16)), **{}):
Unhandled FakeTensor Device Propagation for aten.set_.source_Tensor, found two different devices cuda:0, cpu

from user code:
   File "/usr/local/lib/python3.12/dist-packages/diffusers/hooks/hooks.py", line 149, in torch_dynamo_resume_in_new_forward_at_147
    return function_reference.post_forward(module, output)
  File "/usr/local/lib/python3.12/dist-packages/diffusers/hooks/group_offloading.py", line 230, in post_forward
    self.group.offload_()
  File "/usr/local/lib/python3.12/dist-packages/diffusers/hooks/group_offloading.py", line 171, in offload_
    param.data = self.cpu_param_dict[param]

[a-r-r-o-w]

Can you share the minimal reproducible snippet fro the FakeTensor/torch.compile error?

[seed93]

Can you share the minimal reproducible snippet fro the FakeTensor/torch.compile error?

#11960 (comment)
You could try this code and change the block_level to leaf_level. It worked for A800 but not for 4090.

[seed93]

I managed to implement my own code for offloading finally. I have some notes for you.

1. Different gpu may have different behavior of random even using the same random seed.
2. Groupoffloading of diffuser has two bugs:
   • double register when using compile as I made a PR [fix] compile for groupoffloading #11960
   • precision loss
3. Stream for every block may be too time consuming, I only use 1/n blocks for offloading, which may be a balance between memory and speed.
4. Prefetch for the first block when computing the last block.

[a-r-r-o-w]

Hi, I've found some corner cases that are not handled by group offloading and I'm trying to fix it.

Regarding (1), I think we still have some race condition and improper synchronization in specific code examples. I've found one such minimal example and shared below.

Regarding (2), I've opened a PR (#11987) for compile compatibility with fullgraph=False and added some changes to your code to make it work. Precision loss is a separate issue that I share about below.

Regarding (3), I will add support for opting-out of some blocks using a regex pattern. I think that's the cleanest way of supporting having some blocks on the device, while some blocks dynamically onloaded/offloaded as required.

Regarding (4), looping the prefetching seems do-able but I haven't been able to look into it yet. I'm working on debugging (1) for now to get to the bottom of the potential synchronization issues.

---

When a layer that is not of type X (let's call it L) is followed by a layer of type X in the model, we see that the final outputs are incorrect when L contains any parameter/buffer.

Here, X corresponds to either one of torch.nn.Conv1d, torch.nn.Conv2d, torch.nn.Conv3d, torch.nn.ConvTranspose1d, torch.nn.ConvTranspose2d, torch.nn.ConvTranspose3d, torch.nn.Linear.

The following example shows that when we set elementwise_affine=False, everything works well and as expected. When elementwise_affine=True, the final outputs end up being incorrect.

You mentioned that block_level produces incorrect results for you, while leaf_level worked okay. It was probably just very lucky because I can reproduce incorrect results with both when elementwise_affine=True

Code

import torch
import torch.nn as nn
import torch.nn.functional as F
from diffusers.hooks import apply_group_offloading, ModelHook, HookRegistry
from diffusers.models import ModelMixin
from diffusers.utils.logging import set_verbosity_debug

set_verbosity_debug()


class LayerOutputTrackerHook(ModelHook):
    def __init__(self):
        super().__init__()
        self.outputs = []
    
    def post_forward(self, module, output):
        self.outputs.append(output)
        return output


class Model(ModelMixin):
    def __init__(self, d_model=1024, num_layers=1):
        super().__init__()
        self.d_model = d_model
        
        self.input_proj = nn.Linear(1024, d_model)
        self.blocks = nn.ModuleList([nn.Linear(d_model, d_model) for _ in range(num_layers)])
        
        # This is problematic
        self.norm = nn.LayerNorm(d_model, elementwise_affine=True)
        # This works
        # self.norm = nn.LayerNorm(d_model, elementwise_affine=False)
        
        self.output_proj = nn.Linear(d_model, 1024)
    
    def forward(self, x):
        x = self.input_proj(x)
        for block in self.blocks:
            x = block(x)
            x = F.relu(x)
        x = self.norm(x)
        x = self.output_proj(x)
        return x


def apply_layer_output_tracker_hook(model: Model):
    for name, module in model.named_modules():
        if not isinstance(module, torch.nn.Linear):
            continue
        registry = HookRegistry.check_if_exists_or_initialize(module)
        hook = LayerOutputTrackerHook()
        registry.register_hook(hook, "layer_output_tracker")


def print_output_diffs(ref_model: Model, model: Model):
    for (ref_name, ref_module), (name, module) in zip(ref_model.named_modules(), model.named_modules()):
        assert ref_name == name
        if not isinstance(ref_module, torch.nn.Linear):
            continue
        ref_outputs = HookRegistry.check_if_exists_or_initialize(ref_module).get_hook("layer_output_tracker").outputs
        outputs = HookRegistry.check_if_exists_or_initialize(module).get_hook("layer_output_tracker").outputs
        cumulated_absmax = 0.0
        for i in range(len(outputs)):
            diff = ref_outputs[0] - outputs[i]
            absdiff = diff.abs()
            absmax = absdiff.max().item()
            cumulated_absmax += absmax
            if ref_name == "output_proj":
                print(f"{ref_name} absmax {i}: {absmax}")
        print(f"{name}: cumulated_absmax={cumulated_absmax:.5f}, num_outputs={len(outputs)}")


torch.manual_seed(42)
model_ref = Model()
model1 = Model()
model2 = Model()

model1.load_state_dict(model_ref.state_dict())
model2.load_state_dict(model_ref.state_dict())

model_ref.eval()
model1.eval()
model2.eval()

onload_device = torch.device("cuda:0")
offload_device = torch.device("cpu")

model_ref = model_ref.to(onload_device)
apply_group_offloading(
    model1,
    onload_device=onload_device,
    offload_device=offload_device,
    offload_type="block_level",
    num_blocks_per_group=1,
    use_stream=True,
)
apply_group_offloading(
    model2,
    onload_device=onload_device,
    offload_device=offload_device,
    offload_type="leaf_level",
    use_stream=True,
)

apply_layer_output_tracker_hook(model_ref)
apply_layer_output_tracker_hook(model1)
apply_layer_output_tracker_hook(model2)

x = torch.randn(1, 512, 1024).to("cuda")
out_ref = model_ref(x)

def compare_outputs(out1, out2):
    diff = out1 - out2
    absdiff = diff.abs()
    absmax = absdiff.max()
    mae = absdiff.mean()
    mse = (absdiff ** 2).mean()
    cossim = F.cosine_similarity(out1.flatten(), out2.flatten(), dim=0)
    print(f"{absmax=:.5f}, {mae=:.5f}, {mse=:.5f}, {cossim=:.5f}")

for _ in range(2):
    model1(x)
    print("=" * 80)
    model2(x)

with torch.inference_mode():
    for i in range(10):
        output1 = model1(x)
        output2 = model2(x)

        print(i)
        compare_outputs(out_ref, output1)
        compare_outputs(out_ref, output2)
        print()


print_output_diffs(model_ref, model1)
print()
print_output_diffs(model_ref, model2)

Output on A100

(nightly-venv) aryan@hf-dgx-01:~/work/diffusers$ CUDA_VISIBLE_DEVICES=3 python3 dump12.py
Adding blocks.0 to the execution order
================================================================================
Adding input_proj to the execution order
Adding blocks.0 to the execution order
Adding norm to the execution order
Adding output_proj to the execution order
Applying lazy prefetch group offloading from input_proj to blocks.0
Applying lazy prefetch group offloading from blocks.0 to norm
Applying lazy prefetch group offloading from norm to output_proj
================================================================================
0
absmax=0.71945, mae=0.08477, mse=0.01412, cossim=0.97881
absmax=3.15842, mae=0.09697, mse=0.06828, cossim=0.89760

1
absmax=0.00000, mae=0.00000, mse=0.00000, cossim=1.00000
absmax=3.15842, mae=0.09018, mse=0.06186, cossim=0.90723

2
absmax=0.00000, mae=0.00000, mse=0.00000, cossim=1.00000
absmax=3.15842, mae=0.10204, mse=0.07219, cossim=0.89171

3
absmax=0.00000, mae=0.00000, mse=0.00000, cossim=1.00000
absmax=3.15842, mae=0.09956, mse=0.07020, cossim=0.89471

4
absmax=0.00000, mae=0.00000, mse=0.00000, cossim=1.00000
absmax=3.15842, mae=0.09767, mse=0.06892, cossim=0.89663

5
absmax=0.00000, mae=0.00000, mse=0.00000, cossim=1.00000
absmax=3.15842, mae=0.09435, mse=0.06566, cossim=0.90152

6
absmax=0.00000, mae=0.00000, mse=0.00000, cossim=1.00000
absmax=3.40971, mae=0.11625, mse=0.08446, cossim=0.87340

7
absmax=0.00000, mae=0.00000, mse=0.00000, cossim=1.00000
absmax=3.15842, mae=0.09477, mse=0.06620, cossim=0.90073

8
absmax=0.00000, mae=0.00000, mse=0.00000, cossim=1.00000
absmax=3.15842, mae=0.09691, mse=0.06826, cossim=0.89763

9
absmax=0.00000, mae=0.00000, mse=0.00000, cossim=1.00000
absmax=3.40971, mae=0.14465, mse=0.11311, cossim=0.83038

input_proj: cumulated_absmax=0.00000, num_outputs=12
blocks.0: cumulated_absmax=1.57468, num_outputs=12
output_proj absmax 0: 0.0
output_proj absmax 1: 0.24564117193222046
output_proj absmax 2: 0.7194460034370422
output_proj absmax 3: 0.0
output_proj absmax 4: 0.0
output_proj absmax 5: 0.0
output_proj absmax 6: 0.0
output_proj absmax 7: 0.0
output_proj absmax 8: 0.0
output_proj absmax 9: 0.0
output_proj absmax 10: 0.0
output_proj absmax 11: 0.0
output_proj: cumulated_absmax=0.96509, num_outputs=12

input_proj: cumulated_absmax=0.00000, num_outputs=12
blocks.0: cumulated_absmax=0.00000, num_outputs=12
output_proj absmax 0: 0.0
output_proj absmax 1: 0.0
output_proj absmax 2: 3.1584205627441406
output_proj absmax 3: 3.1584205627441406
output_proj absmax 4: 3.1584205627441406
output_proj absmax 5: 3.1584205627441406
output_proj absmax 6: 3.1584205627441406
output_proj absmax 7: 3.1584205627441406
output_proj absmax 8: 3.4097068309783936
output_proj absmax 9: 3.1584205627441406
output_proj absmax 10: 3.1584205627441406
output_proj absmax 11: 3.4097068309783936
output_proj: cumulated_absmax=32.08678, num_outputs=12

The code is run with PR #11990 but the problem can be reproduced on main as well.

[seed93]

Thank you so much for this great work.

[sayakpaul]

@seed93 thanks so much!

@a-r-r-o-w prioritizing the sync issue is definitely a good idea!

I can try looking into 4 in case #11990 doesn't cover it.