Performance regression in torch 2.0 with deterministic algorithms

username: antoche
cooked: <p>Hi,</p>
<p>I’ve noticed a significant performance slowdown in torch 2.0 when enabling determinism.</p>
<p>Here is a simple example using the diffusers library:</p>
<pre><code class="lang-python">import os
import sys
from datetime import timedelta
import time
import torch
from diffusers import UNet2DModel

import torch

torch.backends.cuda.matmul.allow_tf32 = True

def set_deterministic(mode=True):
    torch.backends.cudnn.benchmark = not mode
    torch.backends.cudnn.deterministic = mode
    torch.use_deterministic_algorithms(mode, warn_only=True)
    if mode:
        os.putenv("CUBLAS_WORKSPACE_CONFIG", ":4096:8")
    else:
        os.unsetenv("CUBLAS_WORKSPACE_CONFIG")
    print(f"Deterministic: {mode}")

def go():
    scaler = torch.cuda.amp.GradScaler()
    batch_size = 8
    channels = 3
    sample_size = 64
    n = 20
    device = torch.device("cuda")
    model = UNet2DModel(
        sample_size=sample_size,
        in_channels=channels, out_channels=channels,
        layers_per_block=2,
        block_out_channels=(128, 128, 256, 256, 512, 512),
        norm_num_groups=32,
        down_block_types=("DownBlock2D", "DownBlock2D", "DownBlock2D", "DownBlock2D", "AttnDownBlock2D", "DownBlock2D"),
        up_block_types=("UpBlock2D", "AttnUpBlock2D", "UpBlock2D", "UpBlock2D", "UpBlock2D", "UpBlock2D"))

    model = model.to(device=device)
    model.train()
    start = time.time()
    rng = torch.Generator(device="cuda").manual_seed(0)
    for step in range(n):
        input = torch.randn((batch_size, channels, sample_size, sample_size), device=device)
        target = torch.randn((batch_size, channels, sample_size, sample_size), device=device)
        bs = input.shape[0]
        timestep = torch.randint(0, 1000, (bs,), generator=rng, device=device)
        with torch.autocast(device_type='cuda', dtype=torch.float16):
            output = model(input, timestep=timestep)
            loss = torch.nn.functional.mse_loss(output.sample, target, reduction="none").mean()
        scaler.scale(loss).backward()

    duration = timedelta(seconds=time.time() - start)

    print(f"Train duration {duration} ({n/duration.total_seconds():.02f} it/s)")

    model = model.to(dtype=torch.float16)
    model.eval()
    start = time.time()
    with torch.no_grad():
        for i in range(n):
            input = torch.randn((batch_size, channels, sample_size, sample_size), device=model.device, dtype=model.dtype)
            timestep = torch.randint(0, 1000, (batch_size,), device=model.device, dtype=model.dtype)
            output = model(input, timestep=timestep)

    duration = timedelta(seconds=time.time() - start)

    print(f"Eval duration {duration} ({n/duration.total_seconds():.02f} it/s)")

def main(mode):
    print(f"Torch version: {torch.__version__}")
    set_deterministic(mode)
    go()

if __name__ == "__main__":
    main(bool(int(sys.argv[1])))
</code></pre>
<p>With pytorch-1.13, performance is roughly equal whether determinism is enabled or not:</p>
<pre data-code-wrap="txt"><code class="lang-plaintext">Torch version: 1.13.0a0+git49444c3
Deterministic: False
Train duration 0:00:02.445595 (8.18 it/s)
Eval duration 0:00:00.488221 (40.97 it/s)
Torch version: 1.13.0a0+git49444c3
Deterministic: True
Train duration 0:00:02.433920 (8.22 it/s)
Eval duration 0:00:00.484679 (41.26 it/s)
</code></pre>
<p>But with pytorch-2.0, performance degrades by 2-4x (or even worse on more complex cases):</p>
<pre data-code-wrap="txt"><code class="lang-plaintext">Torch version: 2.0.0a0+gite9ebda2
Deterministic: False
Train duration 0:00:02.245691 (8.91 it/s)
Eval duration 0:00:00.477144 (41.92 it/s)
Torch version: 2.0.0a0+gite9ebda2
Deterministic: True
Train duration 0:00:05.969440 (3.35 it/s)
Eval duration 0:00:01.809939 (11.05 it/s)
</code></pre>
<p>The difference also happens without using mixed precision, but it is especially visible when using it. GPU usage goes from 100% in non-deterministic mode to &lt;50% in deterministic mode, making me think some operations might be running on the CPU.</p>
<p>Given that determinism did not degrade performance in 1.13, I would expect similar results in 2.0. Did something change in 2.0 to explain this result? Does determinism need to be enabled differently?</p>
<p>This is using cuda-11.2.2 and libcudnn-8.9.4 in both cases.</p>
<p>Thanks,<br>
A.</p>

username: ptrblck
cooked: <p>CUDA operations are executed asynchronously so you would need to synchronize the code before starting and stopping the timers. Your current profile is thus invalid.</p>
<aside class="quote no-group" data-username="antoche" data-post="1" data-topic="188690">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/antoche/48/26241_2.png" class="avatar"> antoche:</div>
<blockquote>
<p>Given that determinism did not degrade performance in 1.13, I would expect similar results in 2.0.</p>
</blockquote>
</aside>
<p>That’s a wrong expectation since forcing deterministic algorithms is expected to potentially downgrade the performance.</p>

username: antoche
cooked: <aside class="quote no-group" data-username="ptrblck" data-post="2" data-topic="188690">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>CUDA operations are executed asynchronously so you would need to synchronize the code before starting and stopping the timers. Your current profile is thus invalid.</p>
</blockquote>
</aside>
<p>I’ve added calls to <code>torch.cuda.synchronize(device)</code> before my <code>time.time()</code> calls and that made no difference to the timings.</p>
<aside class="quote no-group" data-username="ptrblck" data-post="2" data-topic="188690">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>That’s a wrong expectation since forcing deterministic algorithms is expected to potentially downgrade the performance.</p>
</blockquote>
</aside>
<p>I understand that there are no guarantees, but given that 1.13 showed no degradation, and torch 2.0 is branded as “faster” and being the same as 1.13 but with (optional) compilation on top, it would be great to have at least some explanation as to where the degradation is coming from.</p>