fix(backend): revert non-blocking device transfer

[psychedelicious]

Summary

In #6490 we enabled non-blocking torch device transfers throughout the model manager's memory management code. When using this torch feature, torch attempts to wait until the tensor transfer has completed before allowing any access to the tensor. Theoretically, that should make this a safe feature to use.

This provides a small performance improvement but causes race conditions in some situations. Specific platforms/systems are affected, and complicated data dependencies can make this unsafe.

• Intermittent black images on MPS devices - reported on discord and [bug]: black image when using IP adapter after updating developer install #6545, fixed with special handling in fix(backend): mps should not use non_blocking #6549.
• Intermittent OOMs and black images on a P4000 GPU on Windows - reported in [bug]: VRAM not being released #6613, fixed in this commit.

On my system, I haven't experience any issues with generation, but targeted testing of non-blocking ops did expose a race condition when moving tensors from CUDA to CPU.

One workaround is to use torch streams with manual sync points. Our application logic is complicated enough that this would be a lot of work and feels ripe for edge cases and missed spots.

Much safer is to fully revert non-blocking - which is what this change does.

Test script demonstrating CUDA -> CPU race condition

This script induces the race condition. The tensor is different immediately after a device transfer and after waiting a couple seconds for torhc to sync. For me, I reliably get an inconsistency on the second GPU -> CPU transfer with non_blocking=True.

import time

import torch

if __name__ == "__main__":
    seed = 0
    torch.cuda.empty_cache()
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

    for (from_device, to_device) in (("cpu", "cuda"), ("cuda", "cpu")):
        x = torch.zeros(32, 256, 256, 256).to(from_device)
        y = torch.ones(32, 256, 256, 256).to(from_device)

        for non_blocking in (True, False):
            print(f"\033[94mfrom {from_device} to {to_device} non_blocking={non_blocking}\033[0m")
            for i in range(3):
                print(f"\033[96m{i}\033[0m")
                t = (x.max() + y.max()).to(torch.device(to_device), non_blocking=non_blocking)
                print("before waiting 2s:", t)
                old_t = t.clone()
                time.sleep(2.0)
                print("after waiting 2s:", t)
                if old_t != t:
                    print("\033[91minconsistency!\033[0m")

Related Issues / Discussions

Closes #6613

QA Instructions

I have tried these combinations of models and had no issues. I don't think this change introduces any changes to behaviours. It's a one-shot revert of #6490 and #6549.

SDXL

• ControlNet
• LoRA
• TI
• IP Adapter

SD1.5

• ControlNet
• LoRA
• TI
• IP Adapter

Merge Plan

We'll do a bugfix release with this once merged.

Checklist

• The PR has a short but descriptive title, suitable for a changelog
• Tests added / updated (if applicable)
• Documentation added / updated (if applicable)

[RyanJDick]

Thank you for the thorough investigation.

The changes look good to me. I double-checked that there are no remaining references to non_blocking, and ran a quick smoke test.