Fix nondeterministic tests for GPU runs

[anton-l]

This ensures that:

• Stable Diffusion fast tests only run on CPU due to a device-dependent Generator
• Training tests only run on CPU, to avoid setting extra determinism environment flags for CUBLAS and CUDA>10.2 which might affect other tests

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[patrickvonplaten]

@anton-l what do you think about instead of skipping those tests when a GPU is available as is now simply forcing the device to be on CPU? This should be a bit safer no? Also I think we can be sure all GPUs have a CPU

[anton-l]

@patrickvonplaten if removing autocast is ok with you, then yes, that was my first thought too :D

[patrickvonplaten]

@patrickvonplaten if removing autocast is ok with you, then yes, that was my first thought too :D

Sure (sorry I copy-pasted that yesterday)

[patrickvonplaten]

Thanks for making the change!

[patrickvonplaten]

Thanks for fixing!

[antoche]

Just wading into this after the fact as I've hit similar issues in the past and this is affecting me too...

Unfortunately the device arg in the Generator is ignored by pytorch (see bug 62451), so just using a pre-seeded, cpu-bound generator is not enough to ensure determinism.

The following code allows you to get deterministic tests by enforcing random number generation to actually happen on the CPU, while still allowing model evaluation to happen on the GPU:

import torch

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)

Replacing all uses of torch.randn and torch.randn_like with this code lets you get deterministic results when using a generator such as generator = torch.Generator(device='cpu').manual_seed(0), and you can keep device to "cuda" everywhere else.

[patrickvonplaten]

Just wading into this after the fact as I've hit similar issues in the past and this is affecting me too...

Unfortunately the device arg in the Generator is ignored by pytorch (see bug 62451), so just using a pre-seeded, cpu-bound generator is not enough to ensure determinism.

The following code allows you to get deterministic tests by enforcing random number generation to actually happen on the CPU, while still allowing model evaluation to happen on the GPU:

import torch

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)

Replacing all uses of torch.randn and torch.randn_like with this code lets you get deterministic results when using a generator such as generator = torch.Generator(device='cpu').manual_seed(0), and you can keep device to "cuda" everywhere else.

Thanks - that looks very interesting! @anton-l - let's try it out no?