[tests] tighten compilation tests for quantization

[sayakpaul]

What does this PR do?

1. When not using any kind of offloading but just torch.compile with fullgraph=True on quantized models, we don't want any recompilations to get in the way of performance. This PR ensures that.
2. When using offloading, regional compilation with fullgraph=True is better in terms of cold-start and also overall execution time. When _repeated_blocks is available for a model class, we make use of compile_repeated_blocks() instead of compile().

[sayakpaul]

@matthewdouglas I get:

- 0/0: expected type of 'module._modules['norm_out']._modules['linear']._parameters['weight'].CB' to be a tensor type, ' but found <class 'NoneType'>

[matthewdouglas]

For the time being I'm not sure that we can do a whole lot to avoid this for bnb int8. At the very least it is not a high priority for us. Not 100% sure but it's possible you could get around this by making a forward pass through the model prior to compiling it.

[sayakpaul]

I will note this down then in the xfail reason.

[sayakpaul]

@anijain2305 I am planning to switch to compile_repeated_blocks(fullgraph=True) for:

diffusers/tests/quantization/test_torch_compile_utils.py

Line 95 in 2841504

def test_torch_compile_with_group_offload_leaf(self, use_stream=False):

This is to get rid of:

diffusers/tests/quantization/test_torch_compile_utils.py

Line 70 in 2841504

torch._dynamo.config.cache_size_limit = 1000

However, doing so with pytest tests/quantization/gguf/test_gguf.py::GGUFCompileTests::test_torch_compile_with_group_offload_leaf results into:

E               torch._dynamo.exc.Unsupported: Attempted to call function marked as skipped
E                 Explanation: Dynamo developers have intentionally marked that the function `current_accelerator` in file `/fsx/sayak/miniconda3/envs/diffusers/lib/python3.10/site-packages/torch/accelerator/__init__.py` should not be traced.
E                 Hint: Avoid calling the function `current_accelerator`.
E                 Hint: Apply `@torch._dynamo.dont_skip_tracing` to the function `current_accelerator` to force tracing into the function. More graph breaks may occur as a result of attempting to trace into the function.
E                 Hint: Please file an issue to PyTorch.
E               
E                 Developer debug context: module: torch.accelerator, qualname: current_accelerator, skip reason: <missing reason>
E               
E               
E               from user code:
E                  File "/fsx/sayak/diffusers/src/diffusers/models/transformers/transformer_flux.py", line 448, in forward
E                   norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb)
E                 File "/fsx/sayak/diffusers/src/diffusers/models/normalization.py", line 168, in forward
E                   emb = self.linear(self.silu(emb))
E                 File "/fsx/sayak/miniconda3/envs/diffusers/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1778, in _call_impl
E                   return forward_call(*args, **kwargs)
E                 File "/fsx/sayak/diffusers/src/diffusers/hooks/hooks.py", line 188, in new_forward
E                   args, kwargs = function_reference.pre_forward(module, *args, **kwargs)
E                 File "/fsx/sayak/diffusers/src/diffusers/hooks/group_offloading.py", line 339, in pre_forward
E                   self.group.onload_()
E                 File "/fsx/sayak/diffusers/src/diffusers/hooks/group_offloading.py", line 213, in onload_
E                   getattr(torch, torch.accelerator.current_accelerator().type)
E               
E               Set TORCHDYNAMO_VERBOSE=1 for the internal stack trace (please do this especially if you're reporting a bug to PyTorch). For even more developer context, set TORCH_LOGS="+dynamo"

I also had to comment out:

diffusers/src/diffusers/hooks/group_offloading.py

Line 210 in 2841504

@torch.compiler.disable()

Do you have any recommendations?