Torch2 breaks training capabilities on low vram gpus?

[dpsalvatierra]

Recently I updated to torch==2.0.0+cu118, I own an RTX4070i that didn't have any problems training models until this change. The problem seems to be stemming from diffusers being broken by Torch Dynamo .

Have tried everything, nuking my venv, downloading a fresh copy of Automatic1111 and sd_dreambooth, and nothing.

Automatic1111's base project sets strict requirements on outdated dependencies.
If an extension is using a newer version, the dependency is uninstalled and reinstalled twice every startup.

[+] xformers version 0.0.17 installed.
[+] torch version 2.0.0+cu118 installed.
[+] torchvision version 0.15.1+cu118 installed.
[+] accelerate version 0.18.0 installed.
[+] diffusers version 0.14.0 installed.
[+] transformers version 4.26.1 installed.
[+] bitsandbytes version 0.35.4 installed.

Launching Web UI with arguments: --xformers --medvram
2023-05-11 17:17:14.872053: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: AVX2 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
Loading weights [88ecb78256] from /home/dsalvat1/stable-diffusion-webui/models/Stable-diffusion/v2-1_512-ema-pruned.ckpt
Creating model from config: /home/dsalvat1/stable-diffusion-webui/models/Stable-diffusion/v2-1_512-ema-pruned.yaml
LatentDiffusion: Running in eps-prediction mode
DiffusionWrapper has 865.91 M params.
Applying xformers cross attention optimization.
Textual inversion embeddings loaded(0):
Model loaded in 3.9s (load weights from disk: 2.4s, create model: 0.1s, apply weights to model: 0.8s, apply half(): 0.4s, move model to device: 0.1s).
-----------------------------------------------------------
[2023-05-11 17:05:20,147] torch._inductor.utils: [WARNING] not enough cuda cores to use max_autotune mode
--------------------------------------------------------------
  File "/home/dsalvat1/stable-diffusion-webui/venv/lib/python3.10/site-packages/bitsandbytes/optim/optimizer.py", line 393, in init_state
    state["state1"] = torch.zeros_like(
torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 20.00 MiB (GPU 0; 11.73 GiB total capacity; 11.37 GiB already allocated; 11.75 MiB free; 11.44 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF
Steps:   0%|                                                                                                                                                                                   | 0/3600 [00:02<?, ?it/s]
Traceback (most recent call last):
  File "/home/dsalvat1/stable-diffusion-webui/extensions/sd_dreambooth_extension/dreambooth/ui_functions.py", line 727, in start_training
    result = main(class_gen_method=class_gen_method)
  File "/home/dsalvat1/stable-diffusion-webui/extensions/sd_dreambooth_extension/dreambooth/train_dreambooth.py", line 1371, in main
    return inner_loop()
  File "/home/dsalvat1/stable-diffusion-webui/extensions/sd_dreambooth_extension/dreambooth/memory.py", line 117, in decorator
    raise RuntimeError("No executable batch size found, reached zero.")
RuntimeError: No executable batch size found, reached zero.
Restored system models.

++++++++++++++++++++++++++++++++++++++++++++

Webui.sh

#added by ds
# https://github.com/d8ahazard/sd_dreambooth_extension/pull/1186#issuecomment-1518694203
import os
if os.name == "posix":
    os.environ["TORCHDYNAMO_DISABLE"] = "1"
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128"

import torch._dynamo
torch._dynamo.config.suppress_errors = True
torch._dynamo.config.verbose = True

So the question is, is it broken? Any workaround or recommendation to get to a working state?