Add support for torch.compile

[p1atdev]

Added:

• New options, --torch_compile and --dynamo_backend
  • --torch_compile: Enables torch.compile. Default is False.
    • This option is currently incompatible with --xformers. Please use --sdpa option instead.
  • --dynamo_backend: The backend used with torch.compile. Default is "inductor". "eager", "aot_eager", "inductor", "aot_ts_nvfuser", "nvprims_nvfuser", "cudagraphs", "ofi", "fx2trt", "onnxrt" are avaiable, but most are not tested.
    • inductor and eager were worked.

Changed:

• Bumped the einops version from 0.6.0 to 0.6.1 due to be compatible with torch.compile. (more information)

Related:

• Add accelerate torch.compile() support for faster training on Pytorch 2.0 #65

[FurkanGozukara]

Torch compile not available on Windows right?

Also what improvements / changes it brings?

[p1atdev]

Yes, torch.compile does not work on Windows, but it works on WSL.

In my small experiment, training with options --sdpa, --torch_compile and --dynamo_backend eager was faster than --xformers only. (RTX 3070Ti)

wandb: https://wandb.ai/p1atdev/sd-scripts-torch_compile/workspace?workspace=user-p1atdev

The training is very slow while a few steps after torch.compile but then gets faster. Also, torch.compile is expected to be high performance with a modern NVIDIA GPU (like H100, A100, or V100).

https://pytorch.org/tutorials/intermediate/torch_compile_tutorial.html

Therefore, using torch.compile may be faster than just using xformers for large trainings.

[FurkanGozukara]

Yes, torch.compile does not work on Windows, but it works on WSL.

In my small experiment, training with options --sdpa, --torch_compile and --dynamo_backend eager was faster than --xformers only. (RTX 3070Ti)

wandb: https://wandb.ai/p1atdev/sd-scripts-torch_compile/workspace?workspace=user-p1atdev

The training is very slow while a few steps after torch.compile but then gets faster. Also, torch.compile is expected to be high performance with a modern NVIDIA GPU (like H100, A100, or V100).

https://pytorch.org/tutorials/intermediate/torch_compile_tutorial.html

Therefore, using torch.compile may be faster than just using xformers for large trainings.

ty can you tell the it / s difference?

by the way best looking example is xformers

[p1atdev]

The following are screenshots taken during longer trainings:

• First time of --sdpa, --torch_compile and --dynamo_backend eager
  

• First time of --xformers
  

• Second time of --sdpa, --torch_compile and --dynamo_backend eager
  

wandb: https://wandb.ai/p1atdev/pvc-torch_compile

I'm not familiar with torch.compile so I don't know the exact details, but I think torch.compile requires a certain steps of warm-up, so the second training with torch.compile is faster than xformers.

[kohya-ss]

Sorry for the delay. Thank you so much for this great PR! I don't use Linux/WSL personally, but this is really nice!

[sdbds]

Yes, torch.compile does not work on Windows, but it works on WSL.

In my small experiment, training with options --sdpa, --torch_compile and --dynamo_backend eager was faster than --xformers only. (RTX 3070Ti)

wandb: https://wandb.ai/p1atdev/sd-scripts-torch_compile/workspace?workspace=user-p1atdev

The training is very slow while a few steps after torch.compile but then gets faster. Also, torch.compile is expected to be high performance with a modern NVIDIA GPU (like H100, A100, or V100).

https://pytorch.org/tutorials/intermediate/torch_compile_tutorial.html

Therefore, using torch.compile may be faster than just using xformers for large trainings.

I noticed that this is not sd-script's accelerator 0.0.23 but 0.0.25.
We should upgrade the dependency version otherwise a lot of options won't work.

[kohya-ss]

I updated accelerate to 0.0.25. I hope this makes this PR to work.

[FurkanGozukara]

p1atdev

what does second time of training means?

[dill-shower]

The following are screenshots taken during longer trainings

Can you please test speed with inductor backend?

[kohya-ss]

@p1atdev What version of PyTorch do you recommend, 2.1 or does 2.0 work fine? I would like to mention it in the documentation when updating.

[p1atdev]

I tested with PyTorch version 2.1.2+cu118 and it worked. Also according to the PyTorch release notes, torch.compile is more stable in version 2.1 or later.

https://github.com/pytorch/pytorch/releases/tag/v2.1.0

[kohya-ss]

Thank you for clarification!

[feffy380]

@p1atdev What training script did you test with? With train_network.py (SD1.x lora) it crashes in the sdpa forward function:

torch._dynamo.exc.TorchRuntimeError: Failed running call_function <function rearrange at 0x775e95df74c0>(*(FakeTensor(..., device='cuda:0', size=(4, s1, 320), dtype=torch.float16,
           grad_fn=<CloneBackward0>), 'b n (h d) -> b h n d'), **{'h': 8}):
unhashable type: non-singleton SymInt

from user code:
   File "/home/hope/src/sd/sd-scripts/library/original_unet.py", line 741, in resume_in_forward_sdpa_at_739
    q, k, v = map(lambda t: rearrange(t, "b n (h d) -> b h n d", h=h), (q_in, k_in, v_in))
  File "/home/hope/src/sd/sd-scripts/library/original_unet.py", line 741, in <lambda>
    q, k, v = map(lambda t: rearrange(t, "b n (h d) -> b h n d", h=h), (q_in, k_in, v_in))

(However, I am using pytorch-rocm 2.3 nightly and don't know if torch.compile is fully supported on AMD in the first place)

[jdack41]

I've got same error on mac, and also got same error on wsl with cuda. Mac's torch version is 2.2.0, and wsl has 2.1.2.

[jdack41]

@p1atdev @kohya-ss
It seems there is issue on training with sdxl_train.py. sdxl_train_network.py could run.

[ultranationalism]

I've got same error on mac, and also got same error on wsl with cuda. Mac's torch version is 2.2.0, and wsl has 2.1.2.

same error on torch 2.2.0 cu118

[ultranationalism]

@p1atdev @kohya-ss It seems there is issue on training with sdxl_train.py. sdxl_train_network.py could run.

try to upgrade your einops to the lastest version on torch2.1.2+cu118

[iamargentum]

i updated my torch to 2.1.1 and tried training with the torch_compile flag, but it keeps failing with this error - "LayerNormKernelImpl" not implemented for 'Half'
any idea about what this is and how it could be fixed?

[jdack41]

@p1atdev @kohya-ss It seems there is issue on training with sdxl_train.py. sdxl_train_network.py could run.

try to upgrade your einops to the lastest version on torch2.1.2+cu118

Thank you for reply. This solved non hashable error both Mac and wsl(updated einops to 0.7.0).
But saved weights causes NansException on Automatic1111(not lora, finetuned weights).

[jdack41]

https://discuss.pytorch.org/t/how-to-save-load-a-model-with-torch-compile/179739/2
According to this thread, torch.compile will add a prefix ‘_orig_mod.’ to state_dict() of the model.
So removing ‘_orig_mod.’ when saving weights solved NansException.
like this

    def update_sd(prefix, sd):
        for k, v in sd.items():
            key = prefix + k.replace('_orig_mod.', '')
            if save_dtype is not None:
                v = v.detach().clone().to("cpu").to(save_dtype)
            state_dict[key] = v

[jdack41]

Also saving Lora probably needs to do the same or it is structurally broken.