[Community] Testing Stable Diffusion is hard 🥵

[patrickvonplaten]

It's really difficult to test stable diffusion due to the following:

• 1. Continous output: Diffusion models take float values as input and output float values. This is different from NLP models which tend to take int64 as inputs and int64 as outputs.
• 2. Output dimensions are huge. If an image has a output size of (1, 512, 512, 3) this means that there are 512 * 512 * 3 ~ 800,000 values that need to be within a given range. Say if you want to test for a max difference of (pred - ref).abs() < 1e-3 we have roughly a million values where this has to hold true. This is quite different in NLP where we rather test things like text generation or final logit layers which usually aren't bigger then a dozen or so tensors of size 768 or 1024.
• 3. Error propagation: We cannot simple test one forward pass for stable diffusion because in practice people use 50 forward passes. Error propagation becomes a real problem in this case. This again is different from say generation in NLP because in generation at every generation step errors can be somewhat "smoothed" out since a "argmax" of "softmax" operation is used after each step
• 4. Composite systems: Stable Diffusion has three main components for inference: A Unet, a scheduler and a VAE decoder. The UNet and Scheduler are very entangled during the forward pass. Just because we know the forward pass of both the scheduler and unet work independently, it doesn't mean that using them together works.

=> Therefore, we need to do full integration tests, meaning we need to make sure that the output of a full denoising process stays within a given error range. At the moment, we're having quite some problems though to get full reproducible of results on different GPUs, CUDA versions etc... (especially for FP16).

That being said, it is extremely important to test stable diffusion to avoid issues like this in the future: #902 whereas we should still be able to improve speed with PRs like this: #371

At the moment, we're running multiple integration tests for all 50 diffusion steps every time a PR is merged to master, see:

• diffusers/tests/pipelines/stable_diffusion/test_stable_diffusion.py

  Line 518 in 31af4d1

  class PipelineIntegrationTests(unittest.TestCase):

• diffusers/tests/pipelines/stable_diffusion/test_stable_diffusion_img2img.py

  Line 465 in 31af4d1

  class PipelineIntegrationTests(unittest.TestCase):

• diffusers/tests/pipelines/stable_diffusion/test_stable_diffusion_inpaint.py

  Line 262 in 31af4d1

  class PipelineIntegrationTests(unittest.TestCase):

Nevertheless, the tests weren't sufficient to detect: #902

Testing Puzzle 🧩: How can we find the best trade-off between fast & in-expensive test and best possible test coverage taking into account the above points?

We already looked quite a bit into: https://pytorch.org/docs/stable/notes/randomness.html

[patrickvonplaten]

Some ideas for improvement:

1. At the moment, we're not comparing against numpy float outputs but against the generated images:

   diffusers/tests/pipelines/stable_diffusion/test_stable_diffusion_inpaint.py

   Line 282 in 31af4d1

   expected_image = np.array(expected_image, dtype=np.float32) / 255.0

   => here we automatically loose precision of our reference image. We should instead compare directly to numpy matrices

2. We should probably set up a better "layered" testing structure:

   • a.) Have a test just for the UNet forward pass with a very low tolerance (maybe 1e-5 or 1e-4), same for VAE
   • b.) Run full integration tests on CPU as CPU has much less randomness than GPU. Problem here is obviously that the test will take very long - should we maybe set up nightly tests for this?
   • c.) (expensive - not sure if it's worth it) Test on both A100 and V100 and not just V100
   • d.) Just test more samples instead of just one. Let's maybe test Stable Diffusion on 5 different prompts (of different length) and each with 3 different seeds. Those tests should then run once per day only maybe
   • e.) Have at least one full integration test running at every push to master

[patrickvonplaten]

Pinging some people that might have good ideas (maybe big NLP models have similar problems) - cc @stas00 @sgugger @ydshieh @thomasw21 @anton-l

[anton-l]

Unit-testing individual parts of the pipeline + nightly CPU integration tests seems like a solid path to me at the moment.
For now we can check how spotty our line coverage is when disabling the slow integration tests: lots of room for improvement there.

[ydshieh]

Regarding 2), I see in transformers, we test just a slice of the outputs in some cases, something like

self.assertTrue(torch.allclose(logits[0, :3, :3, :3], expected_slice, atol=1e-4))

(well, this approach is probably not what diffusers want to take)

• Question: in diffusers, is it possible to test against smaller input images (i.e. creating smaller models that accept smaller image inputs)?

Regarding different GPUs: I think we can focus on one type of GPU at this moment. In transformers, we only run on T4 so far. Once the testing become better/complete (i.e. when we are happy enough), we can enhance it.

[ydshieh]

If the pipeline used in that issue is tested in diffusers, it's probably a good idea to run that related test in the 2 mentioned diffusers versions and see if both pass. If so, we can try to figure out what the testing lacks to detect the regression in #902.

[thomasw21]

Continous output:

Maybe I'm naive, but in theory everything in discretized to bits. So if you can guarantee determinism (probably a bunch of flags in pytorch + blacklisting some operations), output can be compared bitwise, ie have perfect equality. That does mean that your test become somewhat setup specific (torch version, type of hardware etc ...) We have things like that on some unit tests in Megatron-DeepSpeed: https://github.com/bigscience-workshop/Megatron-DeepSpeed/blob/4255845d0685dd59cb70c0026db45f92b7ae4848/tests/test_activations.py#L30

Output dimensions are huge

If you have a very strict test (typically torch_assert_equal), unless you have some inductive bias where some values are not affected, you could randomly sample elements (like 0.1% for example) and run those assertions there? Though I'm not sure I understand what's the issue with comparing 1M values, those operations are pretty fast on GPUs for example.

Error propagation

I think the best comparison with NLP is rather depth of a model, and indeed in this case NLP is much shallower usually. However concerning the argument of 50 passes instead of 1, I'd say it's just a matter of making your test more strict. If you expect that after 50 forward passes the result is EXACTLY the same as explained #902 then you should probably make sure that a single pass is EXACTLY the same. Concerning the tradeoff with speed, I guess it's fine to update the test if you're willing to accept the breaking change that the increase speed implies. In fact it's usually the thing you would do manually:

• Make an optimization to speed up the process
• Check if the output changes
• If the output changes, figure out if the changes are acceptable

Composite systems

Is the issue that in the case where one of your test fails, it's hard to figure out where the bug is?

Concerning the issue where you're seeing issues with FP16. Maybe the issue is you're trying to ensure the model returns the same output regardless of setups (frameworks, hardware etc ...) I don't see how that's easily doable (how can you guarantee that LayerNorm implementations are the same across frameworks). My personal opinion is that you should rely on the APIs you're trying to compute the same thing?

[msaroufim]

My 2c on this, ideally you can have some fast unit tests that will cover some baseline correctness issues but for the vast majority of issues you're highlighting you need time-consuming runs so there's 2 things that help dramatically in my experience.

1. You can make benchmarks triggerable on some keyword in a PR if the PR's focus is performance improvements Add nightly benchmark and trigger in keyword meta-pytorch/data#740 which removes the need for community members of having to show you logs of perf improvements. The idea would be if something makes things mostly green on your benchmark suite then it should be safe to merge. The alternative is nuanced performance PRs stay open for a long time which will discourage repeat contributions.
2. The other problem though is some tests may be too time-consuming to run on every PR so you can run them at most once a day without spending way too much money. These kinds of tests are essential but they also change expectations because issues are detected late and may often be nuanced to debug, reverts will be a much bigger part of the workflow and that's not fun for anyone but unfortunately necessary which means you need to actively discourage all but the most dedicated contributors to find something to do besides improving performance

Finally, a lot of these problems become easier to handle if the models are faster either by changing the models or by using training compilers so focus on speed and test time will be crucial to make this process sane.

[patrickvonplaten]

Thanks for the great feedback here!

[keturn]

And testing is logistically difficult for another reason that I'm rediscovering with this PR for InvokeAI: we have to do secret token management in order for tests to be able to run with the only known useful diffusion model.

and secrets management + public PRs really don't mix well on GitHub.

I suspect you've worked around that problem here by self-hosting your test runners, but are there any options for other projects that would like to avoid that burden?

Are there any releases of text-to-image models under a creative commons or open source license, so we can pass them around without worrying about redistribution clauses?

They wouldn't need to have the breadth of subjects of the big Stable Diffusion models, or as refined. Just something that works well enough that we can tell if schedulers and guidance functions and whatnot are behaving.