[Bug]: torch compilation never seems to finish for LoKr

[orange-banal-pole]

What happened?

LoKr training speed is orders of magnitude slower than LoRA training.

torch keeps creating files under /tmp/torchinductor_${USER}. I currently have ~98k of them:

  21438 .json
  10721 .ttir
  10721 .ttgir
  10721 .source
  10721 .ptx
  10721 .llir
  10719 .cubin
   7272 .py
   4697 .best_config

Training data and logs: https://github.com/orange-banal-pole/onetrainer-issue

What did you expect would happen?

torch compiles all the needed kernels in a few steps, then the speed becomes closer to that of LoRA training.

Relevant log output

Generate and upload debug_report.log

=== System Information ===
OS: Linux 6.18.33-gentoo-r1-x86_64
Version: 'Gentoo Linux'

=== Hardware Information ===
CPU: 12th Gen Intel(R) Core(TM) i5-12400F (Cores: 2)
Total RAM: 78.41 GB

=== GPU Information ===
NVIDIA GPU (Index 0): NVIDIA RTX PRO 4000 Blackwell [NVIDIA]
Driver version: 610.43.02
Power Limit: 145.00 W

=== Python Environment ===
Global Python Version: 3.13.14
Python Executable Path: OneTrainer/venv/bin/python
PyTorch Info: torch==2.12.0+cu130
pip freeze output:
absl-py==2.5.0
accelerate==1.12.0
adv_optm==2.2.3
aiodns==4.0.4
aiohappyeyeballs==2.7.1
aiohttp==3.14.1
aiohttp-retry==2.9.1
aiosignal==1.4.0
annotated-doc==0.0.4
annotated-types==0.7.0
antlr4-python3-runtime==4.9.3
anyio==4.14.1
astunparse==1.6.3
attrs==26.1.0
av==16.1.0
backoff==2.2.1
backports.zstd==1.6.0
bcrypt==5.0.0
bitsandbytes==0.49.1
boto3==1.43.40
botocore==1.43.40
brotli==1.2.0
certifi==2026.6.17
cffi==2.0.0
charset-normalizer==3.4.7
click==8.4.2
cloudpickle==3.1.2
colorama==0.4.6
coloredlogs==15.0.1
contourpy==1.3.3
cryptography==45.0.7
cuda-bindings==13.3.1
cuda-pathfinder==1.5.6
cuda-toolkit==13.0.2
customtkinter==5.2.2
cycler==0.12.1
dadaptation==3.2
darkdetect==0.8.0
decorator==5.3.1
deepdiff==8.6.2
Deprecated==1.3.1
detect-installer==0.1.0
-e git+https://github.com/huggingface/diffusers.git@1ffa4236a48af8d1d2eaeed6c84027af19acfc39#egg=diffusers
dnspython==2.8.0
email-validator==2.3.0
fabric==3.2.2
fastapi==0.139.0
fastapi-cli==0.0.28
fastapi-cloud-cli==0.22.1
fastar==0.11.0
filelock==3.29.5
flatbuffers==25.12.19
fonttools==4.63.0
frozenlist==1.8.0
fsspec==2026.6.0
ftfy==6.3.1
gast==0.7.0
gguf==0.17.1
google-pasta==0.2.0
grpcio==1.81.1
h11==0.16.0
h5py==3.14.0
hf-xet==1.5.1
httpcore==1.0.9
httptools==0.8.0
httpx==0.28.1
huggingface_hub==1.22.0
humanfriendly==10.0
idna==3.18
imagesize==1.4.1
importlib_metadata==9.0.0
inquirerpy==0.3.4
invisible-watermark==0.2.0
invoke==3.0.3
itsdangerous==2.2.0
Jinja2==3.1.6
jmespath==1.1.0
keras==3.15.0
kiwisolver==1.5.0
libclang==18.1.1
lion-pytorch==0.2.3
Markdown==3.10.2
markdown-it-py==4.2.0
MarkupSafe==3.0.3
matplotlib==3.10.3
mdurl==0.1.2
-e git+https://github.com/Nerogar/mgds.git@3a6994ad18a490acc4f3876658167525b7d88da7#egg=mgds
ml_dtypes==0.5.4
mpmath==1.3.0
multidict==6.7.1
-e git+https://github.com/KellerJordan/Muon.git@f90a42b28e00b8d9d2d05865fe90d9f39abcbcbd#egg=muon_optimizer
namex==0.1.0
networkx==3.6.1
numpy==2.2.6
nvidia-cublas==13.1.1.3
nvidia-cuda-cupti==13.0.85
nvidia-cuda-nvrtc==13.0.88
nvidia-cuda-runtime==13.0.96
nvidia-cudnn-cu13==9.20.0.48
nvidia-cufft==12.0.0.61
nvidia-cufile==1.15.1.6
nvidia-curand==10.4.0.35
nvidia-cusolver==12.0.4.66
nvidia-cusparse==12.6.3.3
nvidia-cusparselt-cu13==0.8.1
nvidia-ml-py==13.610.43
nvidia-nccl-cu13==2.29.7
nvidia-nvjitlink==13.0.88
nvidia-nvshmem-cu13==3.4.5
nvidia-nvtx==13.0.85
omegaconf==2.3.0
onnxruntime-gpu==1.23.2
open_clip_torch==2.32.0
opencv-python==4.11.0.86
opt_einsum==3.4.0
optree==0.19.1
orderly-set==5.5.0
packaging==26.2
paramiko==5.0.0
parse==1.20.2
pfzy==0.3.4
pillow==12.2.0
platformdirs==4.10.0
pooch==1.8.2
prettytable==3.18.0
prodigy-plus-schedule-free==2.0.1
prodigyopt==1.1.2
prompt_toolkit==3.0.52
propcache==0.5.2
protobuf==7.35.1
psutil==7.0.0
py-cpuinfo==9.0.0
pycares==5.0.1
pycparser==3.0
pydantic==2.13.4
pydantic-extra-types==2.11.1
pydantic-settings==2.14.2
pydantic_core==2.46.4
Pygments==2.20.0
PyNaCl==1.6.2
pyparsing==3.3.2
PySide6==6.11.1
PySide6_Addons==6.11.1
PySide6_Essentials==6.11.1
python-dateutil==2.9.0.post0
python-dotenv==1.2.2
python-multipart==0.0.32
pytorch_optimizer==3.6.0
PyWavelets==1.9.0
PyYAML==6.0.2
regex==2026.6.28
requests==2.33.0
rich==15.0.0
rich-toolkit==0.20.1
rignore==0.7.6
runpod==1.7.10
s3transfer==0.19.0
safetensors==0.8.0
scalene==1.5.51
scenedetect==0.7
schedulefree==1.4.1
scipy==1.15.3
sentencepiece==0.2.1
sentry-sdk==2.64.0
setuptools==81.0.0
shellingham==1.5.4
shiboken6==6.11.1
six==1.17.0
starlette==1.3.1
sympy==1.14.0
tensorboard==2.20.0
tensorboard-data-server==0.7.2
termcolor==3.3.0
timm==1.0.27
tokenizers==0.22.2
tomli==2.4.1
tomlkit==0.15.0
torch==2.12.0+cu130
torchvision==0.27.0+cu130
tqdm==4.67.1
tqdm-loggable==0.4.1
transformers==5.5.4
triton==3.7.0
typer==0.26.8
typing-inspection==0.4.2
typing_extensions==4.16.0
urllib3==2.7.0
uvicorn==0.50.0
uvloop==0.22.1
watchdog==6.0.0
watchfiles==1.2.0
wcwidth==0.8.2
websockets==16.0
Werkzeug==3.1.8
wheel==0.47.0
wrapt==2.2.2
yarl==1.24.2
yt-dlp==2026.7.4
zipp==4.1.0

=== Git Information ===
Repo: Nerogar/OneTrainer
Branch: master
Commit: 192c6ac
No deleted, unmerged, or modified files relative to origin/master.

=== Network Connectivity ===
PyPI (https://pypi.org/): HTTPS to pypi.org failed: connection error: HTTPSConnectionPool(host='pypi.org', port=443): Max retries exceeded with url: / (Caused by NewConnectionError("HTTPSConnection(host='pypi.org', port=443): Failed to establish a new connection: [Errno 101] Network is unreachable")); Ping to pypi.org failed: ping: connect: Network is unreachable
HuggingFace (https://huggingface.co): HTTPS to huggingface.co succeeded: Status code 200; Ping to huggingface.co succeeded: Packet loss 0%
Google (https://www.google.com): HTTPS to www.google.com succeeded: Status code 200; Ping to www.google.com succeeded: Packet loss 0%

=== Intel Microcode Information ===
CPU is not detected as 13th or 14th Gen Intel - microcode info not applicable.

[dxqb]

@Koratahiu

[orange-banal-pole]

FWIW, with another dataset (100 images), AI Toolkit with the experimental compile feature enabled writes 6,562 files to /tmp/torchinductor_$USERNAME and then stops creating new ones. OneTrainer, on the other hand, keeps creating new files indefinitely.

[orange-banal-pole]

With the changes below (suggested by Claude), the recompilation madness stopped and the speed returned to LoRA levels. I can't vouch for their correctness, though.

--- a/modules/module/LoRAModule.py
+++ b/modules/module/LoRAModule.py
@@ -509,10 +509,10 @@ class LoKrModule(PeftBase):
                 x_reshaped = x.reshape(-1, self.in_m, self.in_n)
 
                 # Calculate delta = x @ (W1 x W2).T
-                delta_output = torch.einsum(
-                    'bmn, lm, kn -> blk',
-                    x_reshaped, w1.to(x.dtype), w2.to(x.dtype)
-                )
+                # sequential two-operand contractions — avoids opt_einsum's contract_path
+                # specialization that a single 3-operand einsum triggers on batch*seq size
+                intermediate = torch.einsum('bmn, lm -> bln', x_reshaped, w1.to(x.dtype))
+                delta_output = torch.einsum('bln, kn -> blk', intermediate, w2.to(x.dtype))
 
                 # Reshape back to [Batch, ..., Out_Features]
                 delta_output = delta_output.reshape(*x_shape[:-1], -1) * scale
diff --git a/modules/util/checkpointing_util.py b/modules/util/checkpointing_util.py
index 1f669e9a..0acef001 100644
--- a/modules/util/checkpointing_util.py
+++ b/modules/util/checkpointing_util.py
@@ -203,8 +203,7 @@ def create_checkpoint(
     else:
         if compile:
             layer = CheckpointLayer(orig_module=orig_module, orig_forward=None, train_device=train_device, checkpointing=checkpoin>
-            #do compile the checkpointing layer - slightly faster
-            layer.compile(fullgraph=True)
+            orig_module.compile(fullgraph=True)   # compile the block, not the wrapper
             return layer
         else:
             layer = CheckpointLayer(orig_module=None, orig_forward=orig_module.forward, train_device=train_device, checkpointing=c>

[Koratahiu]

This appears to be a torch.compile issue with the vec trick that only happens on specific models (like krea 2) when specializing for dynamic sizes.
The Claude diff seems correct as a workaround (only the LoRAModule change, with a slight performance gain).