nunchaku-style kernel fusion - 20-30% faster training

modules/modelSampler/AnimaSampler.py

@@ -12,13 +12,13 @@
 from modules.util.enum.ModelType import ModelType
 from modules.util.enum.NoiseScheduler import NoiseScheduler
 from modules.util.enum.VideoFormat import VideoFormat
+from modules.util.tqdm_util import tqdm
 
 import torch
 
 from diffusers import VaeImageProcessor
 
 import numpy as np
-from tqdm import tqdm
 
 
 @factory.register(BaseModelSampler, ModelType.ANIMA)

modules/modelSampler/ChromaSampler.py

@@ -12,11 +12,10 @@
 from modules.util.enum.ModelType import ModelType
 from modules.util.enum.NoiseScheduler import NoiseScheduler
 from modules.util.enum.VideoFormat import VideoFormat
+from modules.util.tqdm_util import tqdm
 
 import torch
 
-from tqdm import tqdm
-
 
 @factory.register(BaseModelSampler, ModelType.CHROMA_1)
 class ChromaSampler(BaseModelSampler):

modules/modelSampler/ErnieSampler.py

@@ -11,12 +11,12 @@
 from modules.util.enum.ModelType import ModelType
 from modules.util.enum.NoiseScheduler import NoiseScheduler
 from modules.util.enum.VideoFormat import VideoFormat
+from modules.util.tqdm_util import tqdm
 
 import torch
 
 import numpy as np
 from PIL import Image as PILImage
-from tqdm import tqdm
 
 
 @factory.register(BaseModelSampler, ModelType.ERNIE)

modules/modelSampler/Flux2Sampler.py

@@ -12,13 +12,13 @@
 from modules.util.enum.ModelType import ModelType
 from modules.util.enum.NoiseScheduler import NoiseScheduler
 from modules.util.enum.VideoFormat import VideoFormat
+from modules.util.tqdm_util import tqdm
 
 import torch
 
 from diffusers.pipelines.flux2.pipeline_flux2 import compute_empirical_mu
 
 import numpy as np
-from tqdm import tqdm
 
 
 @factory.register(BaseModelSampler, ModelType.FLUX_2)

modules/modelSampler/FluxSampler.py

@@ -14,13 +14,12 @@
 from modules.util.enum.NoiseScheduler import NoiseScheduler
 from modules.util.enum.VideoFormat import VideoFormat
 from modules.util.image_util import load_image
+from modules.util.tqdm_util import tqdm
 
 import torch
 from torch import nn
 from torchvision.transforms import transforms
 
-from tqdm import tqdm
-
 
 @factory.register(BaseModelSampler, ModelType.FLUX_DEV_1)
 @factory.register(BaseModelSampler, ModelType.FLUX_FILL_DEV_1)

modules/modelSampler/HiDreamSampler.py

@@ -12,11 +12,10 @@
 from modules.util.enum.ModelType import ModelType
 from modules.util.enum.NoiseScheduler import NoiseScheduler
 from modules.util.enum.VideoFormat import VideoFormat
+from modules.util.tqdm_util import tqdm
 
 import torch
 
-from tqdm import tqdm
-
 
 @factory.register(BaseModelSampler, ModelType.HI_DREAM_FULL)
 class HiDreamSampler(BaseModelSampler):

modules/modelSampler/HunyuanVideoSampler.py

@@ -12,11 +12,11 @@
 from modules.util.enum.ModelType import ModelType
 from modules.util.enum.NoiseScheduler import NoiseScheduler
 from modules.util.enum.VideoFormat import VideoFormat
+from modules.util.tqdm_util import tqdm
 
 import torch
 
 from PIL import Image
-from tqdm import tqdm
 
 
 @factory.register(BaseModelSampler, ModelType.HUNYUAN_VIDEO)

modules/modelSampler/IdeogramSampler.py

@@ -11,14 +11,14 @@
 from modules.util.enum.ModelType import ModelType
 from modules.util.enum.NoiseScheduler import NoiseScheduler
 from modules.util.enum.VideoFormat import VideoFormat
+from modules.util.tqdm_util import tqdm
 
 import torch
 
 from diffusers.pipelines.ideogram4.pipeline_ideogram4 import _logit_normal_sigmas, _resolution_aware_mu
 
 import numpy as np
 from PIL import Image as PILImage
-from tqdm import tqdm
 
 
 @factory.register(BaseModelSampler, ModelType.IDEOGRAM_4)

modules/modelSampler/Krea2Sampler.py

@@ -13,13 +13,12 @@
 from modules.util.enum.ModelType import ModelType
 from modules.util.enum.NoiseScheduler import NoiseScheduler
 from modules.util.enum.VideoFormat import VideoFormat
+from modules.util.tqdm_util import tqdm
 
 import torch
 
 from diffusers import Krea2Pipeline
 
-from tqdm import tqdm
-
 
 @factory.register(BaseModelSampler, ModelType.KREA_2)
 class Krea2Sampler(BaseModelSampler):

modules/modelSampler/PixArtAlphaSampler.py

@@ -11,11 +11,10 @@
 from modules.util.enum.ModelType import ModelType
 from modules.util.enum.NoiseScheduler import NoiseScheduler
 from modules.util.enum.VideoFormat import VideoFormat
+from modules.util.tqdm_util import tqdm
 
 import torch
 
-from tqdm import tqdm
-
 
 @factory.register(BaseModelSampler, ModelType.PIXART_ALPHA)
 @factory.register(BaseModelSampler, ModelType.PIXART_SIGMA)

modules/modelSampler/QwenSampler.py

@@ -13,11 +13,10 @@
 from modules.util.enum.ModelType import ModelType
 from modules.util.enum.NoiseScheduler import NoiseScheduler
 from modules.util.enum.VideoFormat import VideoFormat
+from modules.util.tqdm_util import tqdm
 
 import torch
 
-from tqdm import tqdm
-
 
 @factory.register(BaseModelSampler, ModelType.QWEN)
 class QwenSampler(BaseModelSampler):

modules/modelSampler/SanaSampler.py

@@ -12,11 +12,10 @@
 from modules.util.enum.ModelType import ModelType
 from modules.util.enum.NoiseScheduler import NoiseScheduler
 from modules.util.enum.VideoFormat import VideoFormat
+from modules.util.tqdm_util import tqdm
 
 import torch
 
-from tqdm import tqdm
-
 
 @factory.register(BaseModelSampler, ModelType.SANA)
 class SanaSampler(BaseModelSampler):

modules/modelSampler/StableDiffusion3Sampler.py

@@ -12,11 +12,10 @@
 from modules.util.enum.ModelType import ModelType
 from modules.util.enum.NoiseScheduler import NoiseScheduler
 from modules.util.enum.VideoFormat import VideoFormat
+from modules.util.tqdm_util import tqdm
 
 import torch
 
-from tqdm import tqdm
-
 
 @factory.register(BaseModelSampler, ModelType.STABLE_DIFFUSION_3)
 @factory.register(BaseModelSampler, ModelType.STABLE_DIFFUSION_35)

modules/modelSampler/StableDiffusionSampler.py

@@ -12,13 +12,12 @@
 from modules.util.enum.NoiseScheduler import NoiseScheduler
 from modules.util.enum.VideoFormat import VideoFormat
 from modules.util.image_util import load_image
+from modules.util.tqdm_util import tqdm
 
 import torch
 from torch import nn
 from torchvision.transforms import transforms
 
-from tqdm import tqdm
-
 
 @factory.register(BaseModelSampler, ModelType.STABLE_DIFFUSION_15)
 @factory.register(BaseModelSampler, ModelType.STABLE_DIFFUSION_15_INPAINTING)

modules/modelSampler/StableDiffusionXLSampler.py

@@ -12,13 +12,12 @@
 from modules.util.enum.NoiseScheduler import NoiseScheduler
 from modules.util.enum.VideoFormat import VideoFormat
 from modules.util.image_util import load_image
+from modules.util.tqdm_util import tqdm
 
 import torch
 from torch import nn
 from torchvision.transforms import transforms
 
-from tqdm import tqdm
-
 
 @factory.register(BaseModelSampler, ModelType.STABLE_DIFFUSION_XL_10_BASE)
 @factory.register(BaseModelSampler, ModelType.STABLE_DIFFUSION_XL_10_BASE_INPAINTING)

modules/modelSampler/WuerstchenSampler.py

@@ -11,11 +11,11 @@
 from modules.util.enum.ModelType import ModelType
 from modules.util.enum.NoiseScheduler import NoiseScheduler
 from modules.util.enum.VideoFormat import VideoFormat
+from modules.util.tqdm_util import tqdm
 
 import torch
 
 from PIL import Image
-from tqdm import tqdm
 
 
 @factory.register(BaseModelSampler, ModelType.WUERSTCHEN_2)

modules/modelSampler/ZImageSampler.py

@@ -12,11 +12,10 @@
 from modules.util.enum.ModelType import ModelType
 from modules.util.enum.NoiseScheduler import NoiseScheduler
 from modules.util.enum.VideoFormat import VideoFormat
+from modules.util.tqdm_util import tqdm
 
 import torch
 
-from tqdm import tqdm
-
 
 @factory.register(BaseModelSampler, ModelType.Z_IMAGE)
 class ZImageSampler(BaseModelSampler):

modules/module/BaseImageCaptionModel.py

@@ -6,9 +6,9 @@
 
 from modules.util import path_util
 from modules.util.image_util import load_image
+from modules.util.tqdm_util import tqdm
 
 from PIL import Image
-from tqdm import tqdm
 
 
 class CaptionSample:

modules/module/BaseImageMaskModel.py

@@ -5,13 +5,13 @@
 
 from modules.util import path_util
 from modules.util.image_util import load_image
+from modules.util.tqdm_util import tqdm
 
 import torch
 from torch import Tensor
 from torchvision.transforms import transforms
 
 from PIL import Image
-from tqdm import tqdm
 
 
 class MaskSample:

modules/module/FusedModule.py

@@ -95,6 +95,11 @@ def __init__(self, prefix: str, leaves: list[nn.Module], klass, additional_args:
     # recompose the base weight itself (delta_forward returns None) keep going through the slower,
     # generic self.module.forward(x) path.
     def _leaf_output(self, leaf_index: int, start: int, end: int, x, *args, **kwargs):
+        # a leaf that can fold the adapter into its own base matmul returns the finished output, so
+        # neither the full fused delta nor the separate add happens for it
+        fused = self.module.fused_leaf_forward(self.leaves[leaf_index], x, start, end)
+        if fused is not None:
+            return fused
         delta = self.module.delta_forward(x, *args, **kwargs)
         if delta is None:
             return self.module.forward(x)[..., start:end]

modules/module/GenerateLossesModel.py

@@ -8,12 +8,11 @@
 from modules.util import create
 from modules.util.config.TrainConfig import QuantizationConfig, TrainConfig
 from modules.util.torch_util import torch_gc
+from modules.util.tqdm_util import tqdm
 from modules.util.TrainProgress import TrainProgress
 
 import torch
 
-from tqdm import tqdm
-
 
 class GenerateLossesModel:
     """Based on train args, writes a JSON instead of a model with filenames mapped to losses,

modules/module/LoRAModule.py

@@ -7,7 +7,7 @@
 
 from modules.module.FusedModule import FusedModuleGroup, check_fusion_match, discover_fused_groups
 from modules.module.oft_utils import OFTRotationModule
-from modules.module.quantized.LinearSVD import BaseLinearSVD
+from modules.module.quantized.mixin.LoRAFusableLinearMixin import LoRAFusableLinearMixin
 from modules.util.config.TrainConfig import TrainConfig
 from modules.util.enum.ModelType import PeftType
 from modules.util.lokr_utils import factorization, make_kron, rebuild_tucker
@@ -115,6 +115,13 @@ def delta_forward(self, x, *args, **kwargs) -> Tensor | None:
         # going through the fused forward.
         return None
 
+    def fused_leaf_forward(self, leaf: nn.Module, x, start: int, end: int) -> Tensor | None:
+        # Returns one leaf's complete output (base + this adapter's contribution) when the leaf can
+        # fold the adapter into its own base matmul, so no delta is computed or added separately.
+        # start/end are the leaf's row range in the fused output. Returns None (the default) when no
+        # such path exists, and the caller falls back to delta_forward.
+        return None
+
     @property
     def orig_module(self) -> nn.Module:
         assert self._orig_module is not None
@@ -570,11 +577,18 @@ def check_initialized(self):
 
     def forward(self, x, *args, **kwargs):
         self.check_initialized()
-        if isinstance(self.orig_module, BaseLinearSVD):
-            return self.orig_module.forward_with_lora(x, self.lora_down, self.lora_up, self.dropout, self.alpha)
+        if isinstance(self.orig_module, LoRAFusableLinearMixin):
+            return self.orig_module.forward_with_lora(x, self.lora_down.weight, self.lora_up.weight, self.dropout, self.alpha)
 
         return self.orig_forward(x) + self.delta_forward(x, *args, **kwargs)
 
+    def fused_leaf_forward(self, leaf: nn.Module, x, start: int, end: int) -> Tensor | None:
+        self.check_initialized()
+        if not isinstance(leaf, LoRAFusableLinearMixin):
+            return None
+        #the fused adapter's up spans all leaves' concatenated outputs; narrow it to this leaf's rows
+        return leaf.forward_with_lora(x, self.lora_down.weight, self.lora_up.weight[start:end], self.dropout, self.alpha)
+
     def delta_forward(self, x, *args, **kwargs) -> Tensor | None:
         self.check_initialized()
         ld = self.lora_up(self.dropout(self.lora_down(x)))
@@ -781,6 +795,11 @@ def delta_forward(self, x, *args, **kwargs) -> Tensor | None:
         # DoRA scales the recomposed weight, so there is no delta term; back to None from LoRAModule's.
         return None
 
+    def fused_leaf_forward(self, leaf: nn.Module, x, start: int, end: int) -> Tensor | None:
+        # same reason as delta_forward: the leaf folds an additive term into its own base matmul,
+        # and DoRA recomposes the weight instead
+        return None
+
     def forward(self, x, *args, **kwargs):
         self.check_initialized()
         A = self.lora_down.weight