Use consistent random number generation across hardware

[antoche]

Is your feature request related to a problem? Please describe.

pytorch.randn is not consistent across hardware devices (See pytorch/pytorch#84234).

diffusers calls torch.randn on the device computation is run on (typically 'cuda'). As a result, results produced with the exact same parameters will differ across machines.

Describe the solution you'd like

Until the issue is resolved in pytorch itself, diffusers should use a deterministic RNG so results can be consistent across hardware.

One possible workaround is to keep using torch.rng while enforcing generation to happen on the cpu, which currently seems consistent no matter the hardware.

Here is an example solution:

def randn(size, generator=None, device=None, **kwargs):
    """
    Wrapper around torch.randn providing proper reproducibility.

    Generation is done on the given generator's device, then moved to the
    given ``device``.

    Args:
        size: tensor size
        generator (torch.Generator): RNG generator
        device (torch.device): Target device for the resulting tensor
    """
    # FIXME: generator RNG device is ignored and needs to be passed to torch.randn (torch issue #62451)
    rng_device = generator.device if generator is not None else device
    image = torch.randn(size, generator=generator, device=rng_device, **kwargs)
    image = image.to(device=device)
    return image


def randn_like(tensor, generator=None, **kwargs):
    return randn(tensor.shape, layout=tensor.layout, generator=generator, device=tensor.device, **kwargs)

Calling these functions instead of the torch ones, with a generator whose device is cpu, gives deterministic results and still allows for the rest of the computations to run on cuda.

This would also simplify and speed up all the tests, which can simply use cpu-bound generators and leave device to be cuda even for those relying on RNG.

Describe alternatives you've considered

It's also possisble to switch to numpy's RNG, which is deterministic. The above solution is more torch-native.

[patrickvonplaten]

Linking pytorch/pytorch#87992 here

Generally this might actually be a good idea to always have the generator and all random operations on CPU . It might slow down the forward pass of DDPM though:

diffusers/src/diffusers/schedulers/scheduling_ddpm.py

Line 314 in 769f0be

variance_noise = torch.randn(

Since we would always move between CPU and GPU here. Essentially all random schedulers ("_ancestral") ones would be affected. But the reproducibility across devices might be worth it!

Maybe we could also add a set_deterministic() function to pipelines and schedulers that would decide whether to use CPU or GPU when creating random tensors.

[patrickvonplaten]

@patil-suraj @anton-l @pcuenca @williamberman curious to hear your opinion here.

[antoche]

Agree that this may be a performance concern for some.

With the proposal I suggested, no change in the interface or new function is needed. Users get deterministic results by passing a generator using device="cpu", and more performant, less deterministic results (i.e, current behaviour) by passing a generator using device="cuda" (which I believe is the default).

Note that this also addresses an issue in pytorch where the generator's device is sometimes ignored (pytorch/pytorch#62451)

[pcuenca]

Sorry for the delay in reacting to this.

I agree that random number generation is a mess and it's worthwhile taking a closer look. In addition, we have ugly workarounds for platform-specific bugs, such as mps.

I think many users won't care about reproducibility. The opt-in mechanism to use a cpu-backed Generator looks a bit obscure to me, but it may be enough for those who do. So, two main options here:

• Adopt this solution as is, and create a documentation page about how to achieve reproducible results.
• Create a helper as Patrick suggested. I'd call it something like enable_reproducibility(), as set_deterministic() sounds to me like "always produce the same results given the same inputs".

Since reproducibility is more of an advanced feature, I think it would suffice to use the first approach for now.

In any case (whether we add a helper or not), I'm in favor of adopting the proposed functions and use them throughout the code base. We could take advantage of them in our tests, and they could absorb the workarounds for mps too.

[patil-suraj]

Agree with Patrick here, IMO it's better to add a method which can enable reproducibility. I'm usually hesitant to create such wrappers around framework functions as it might make it a bit harder to go through the code. For example, if we replace torch.randn by custom randn then users who are going through the code might wonder that something different is happening here.

[williamberman]

Agreed reproducibility is really nice. Just for my own context, outside of our test suite, what's the general use case for reproducibility? In my limited experience so far most of the diffusion model codebases I've worked on have just used the standard torch rng functionality.

[antoche]

Reproducibility/determinism is not just a "nice thing, it's vitally important for multiple reasons:

• For research, if results are not reproducible, then other researchers cannot compare their results against previous work. Reproducibility is a cornerstone of the scientific process.
• For contributors, if tests aren't reproducible on their machines, then it is a major hindrance in contributing to a project. If the tests don't even pass on a contributor's machine without any changes, they don't know what their changes might be breaking.
• For production software, reproducibility is a requirement. If a developer writes software relying on a library that is not deterministic, they can't make any guarantees that their software will be reliable, and they can't write reliable tests for their own code.
• For users, reproducibility is an important feature. The same inputs ought to result in the same results, so users can explore and refine their work, build on other users' inputs, as well as recover it if the outputs have been lost.