Updated ramtorch

- Redid implementation, seems to work better, but be aware still not one to one with when it's not in use

Author: 67372a

flux_train_network.py

@@ -9,6 +9,7 @@
 
 from library.device_utils import clean_memory_on_device, init_ipex
 from library.strategy_flux import move_vision_encoder_to_device
+from ramtorch.helpers import replace_linear_with_ramtorch
 
 init_ipex()
 
@@ -141,16 +142,13 @@ def load_target_model(self, args, weight_dtype, accelerator):
 
         ae = flux_utils.load_ae(args.ae, weight_dtype, "cpu", disable_mmap=args.disable_mmap_load_safetensors)
 
+ 
         if args.use_ramtorch:
-            try:
-                from library.ramtorch_util import replace_linear_with_ramtorch_linear
-                logger.info("Applying RamTorch to model and Text Encoders for memory efficiency...")
-                replace_linear_with_ramtorch_linear(model, accelerator.device)
-                replace_linear_with_ramtorch_linear(clip_l, accelerator.device)
-                replace_linear_with_ramtorch_linear(t5xxl, accelerator.device)
-                logger.info("RamTorch applied successfully.")
-            except ImportError as e:
-                logger.error(f"Failed to apply RamTorch: {e}")
+            logger.info("Applying RamTorch to FLUX models (DiT, T5-XXL, CLIP-L, AE).")
+            model = replace_linear_with_ramtorch(model, accelerator.device)
+            clip_l = replace_linear_with_ramtorch(clip_l, accelerator.device)
+            t5xxl = replace_linear_with_ramtorch(t5xxl, accelerator.device)
+            ae = replace_linear_with_ramtorch(ae, accelerator.device)
 
         return flux_utils.MODEL_VERSION_FLUX_V1, [clip_l, t5xxl], ae, model
 

library/ramtorch_util.py

@@ -32,7 +32,7 @@ prodigy-plus-schedule-free~=2.0.1
 kornia~=0.8.1
 PyWavelets~=1.9.0
 hf-xet~=1.1.10
-RamTorch~=0.2.0
+RamTorch~=0.2.1
 adv_optm~=1.1.4
 # for kohya_ss library
 -e .

requirements.txt

@@ -11,6 +11,7 @@
 import train_network
 from library.utils import setup_logging
 from tools.stochastic_copy import to_stochastic
+from ramtorch.helpers import replace_linear_with_ramtorch
 
 setup_logging()
 import logging
@@ -54,22 +55,18 @@ def load_target_model(self, args, weight_dtype, accelerator):
         self.logit_scale = logit_scale
         self.ckpt_info = ckpt_info
 
+        if args.use_ramtorch:
+            logger.info("Applying RamTorch to SDXL UNet, VAE, and Text Encoders.")
+            unet = replace_linear_with_ramtorch(unet, accelerator.device)
+            vae = replace_linear_with_ramtorch(vae, accelerator.device)
+            text_encoder1 = replace_linear_with_ramtorch(text_encoder1, accelerator.device)
+            text_encoder2 = replace_linear_with_ramtorch(text_encoder2, accelerator.device)
+
         # モデルに xformers とか memory efficient attention を組み込む
         train_util.replace_unet_modules(unet, args.mem_eff_attn, args.xformers, args.sdpa)
         if torch.__version__ >= "2.0.0":  # PyTorch 2.0.0 以上対応のxformersなら以下が使える
             vae.set_use_memory_efficient_attention_xformers(args.xformers)
 
-        if args.use_ramtorch:
-            try:
-                from library.ramtorch_util import replace_linear_with_ramtorch_linear
-                logger.info("Applying RamTorch to U-Net and Text Encoders for memory efficiency...")
-                replace_linear_with_ramtorch_linear(unet, accelerator.device)
-                replace_linear_with_ramtorch_linear(text_encoder1, accelerator.device)
-                replace_linear_with_ramtorch_linear(text_encoder2, accelerator.device)
-                logger.info("RamTorch applied successfully.")
-            except ImportError as e:
-                logger.error(f"Failed to apply RamTorch: {e}")
-
         return sdxl_model_util.MODEL_VERSION_SDXL_BASE_V1_0, [text_encoder1, text_encoder2], vae, unet
 
     def get_tokenize_strategy(self, args):