Merge PR Nerogar#1481 (always enable mixed-precision / autocasting) into preview

Author: dxqb

modules/modelSetup/BaseChromaSetup.py

@@ -17,7 +17,6 @@
 )
 from modules.util.config.TrainConfig import TrainConfig
 from modules.util.dtype_util import create_autocast_context, disable_fp16_autocast_context
-from modules.util.enum.TrainingMethod import TrainingMethod
 from modules.util.quantization_util import quantize_layers
 from modules.util.torch_util import torch_gc
 from modules.util.TrainProgress import TrainProgress
@@ -56,24 +55,14 @@ def setup_optimizations(
                 model.text_encoder_offload_conductor = \
                     enable_checkpointing_for_t5_encoder_layers(model.text_encoder, config)
 
-        model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [
-            config.weight_dtypes().transformer,
-            config.weight_dtypes().text_encoder,
-            config.weight_dtypes().vae,
-            config.weight_dtypes().lora if config.training_method == TrainingMethod.LORA else None,
-            config.weight_dtypes().embedding if config.train_any_embedding() else None,
-        ], config.enable_autocast_cache)
+        model.autocast_context, model.train_dtype = create_autocast_context(
+            self.train_device, config.train_dtype, config.enable_autocast_cache)
 
         model.text_encoder_autocast_context, model.text_encoder_train_dtype = \
             disable_fp16_autocast_context(
                 self.train_device,
                 config.train_dtype,
                 config.fallback_train_dtype,
-                [
-                    config.weight_dtypes().text_encoder,
-                    config.weight_dtypes().lora if config.training_method == TrainingMethod.LORA else None,
-                    config.weight_dtypes().embedding if config.train_any_embedding() else None,
-                ],
                 config.enable_autocast_cache,
             )
 

modules/modelSetup/BaseErnieSetup.py

@@ -12,7 +12,6 @@
 from modules.util.checkpointing_util import enable_checkpointing_for_ernie_transformer
 from modules.util.config.TrainConfig import TrainConfig
 from modules.util.dtype_util import create_autocast_context, disable_fp16_autocast_context
-from modules.util.enum.TrainingMethod import TrainingMethod
 from modules.util.quantization_util import quantize_layers
 from modules.util.torch_util import torch_gc
 from modules.util.TrainProgress import TrainProgress
@@ -46,22 +45,14 @@ def setup_optimizations(
             model.transformer_offload_conductor = \
                 enable_checkpointing_for_ernie_transformer(model.transformer, config)
 
-        model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [
-            config.weight_dtypes().transformer,
-            config.weight_dtypes().text_encoder,
-            config.weight_dtypes().vae,
-            config.weight_dtypes().lora if config.training_method == TrainingMethod.LORA else None,
-        ], config.enable_autocast_cache)
+        model.autocast_context, model.train_dtype = create_autocast_context(
+            self.train_device, config.train_dtype, config.enable_autocast_cache)
 
         model.text_encoder_autocast_context, model.text_encoder_train_dtype = \
             disable_fp16_autocast_context(
                 self.train_device,
                 config.train_dtype,
                 config.fallback_train_dtype,
-                [
-                    config.weight_dtypes().text_encoder,
-                    config.weight_dtypes().lora if config.training_method == TrainingMethod.LORA else None,
-                ],
                 config.enable_autocast_cache,
             )
 

modules/modelSetup/BaseFlux2Setup.py

@@ -17,7 +17,6 @@
 )
 from modules.util.config.TrainConfig import TrainConfig
 from modules.util.dtype_util import create_autocast_context, disable_fp16_autocast_context
-from modules.util.enum.TrainingMethod import TrainingMethod
 from modules.util.quantization_util import quantize_layers
 from modules.util.torch_util import torch_gc
 from modules.util.TrainProgress import TrainProgress
@@ -56,22 +55,14 @@ def setup_optimizations(
                     model.text_encoder_offload_conductor = \
                         enable_checkpointing_for_qwen3_encoder_layers(model.text_encoder, config)
 
-        model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [
-            config.weight_dtypes().transformer,
-            config.weight_dtypes().text_encoder,
-            config.weight_dtypes().vae,
-            config.weight_dtypes().lora if config.training_method == TrainingMethod.LORA else None,
-        ], config.enable_autocast_cache)
+        model.autocast_context, model.train_dtype = create_autocast_context(
+            self.train_device, config.train_dtype, config.enable_autocast_cache)
 
         model.text_encoder_autocast_context, model.text_encoder_train_dtype = \
             disable_fp16_autocast_context(
                 self.train_device,
                 config.train_dtype,
                 config.fallback_train_dtype,
-                [
-                    config.weight_dtypes().text_encoder,
-                    config.weight_dtypes().lora if config.training_method == TrainingMethod.LORA else None,
-                ],
                 config.enable_autocast_cache,
             )
 

modules/modelSetup/BaseFluxSetup.py

@@ -18,7 +18,6 @@
 )
 from modules.util.config.TrainConfig import TrainConfig
 from modules.util.dtype_util import create_autocast_context, disable_fp16_autocast_context
-from modules.util.enum.TrainingMethod import TrainingMethod
 from modules.util.quantization_util import quantize_layers
 from modules.util.torch_util import torch_gc
 from modules.util.TrainProgress import TrainProgress
@@ -58,25 +57,14 @@ def setup_optimizations(
                 model.text_encoder_2_offload_conductor = \
                     enable_checkpointing_for_t5_encoder_layers(model.text_encoder_2, config)
 
-        model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [
-            config.weight_dtypes().transformer,
-            config.weight_dtypes().text_encoder,
-            config.weight_dtypes().text_encoder_2,
-            config.weight_dtypes().vae,
-            config.weight_dtypes().lora if config.training_method == TrainingMethod.LORA else None,
-            config.weight_dtypes().embedding if config.train_any_embedding() else None,
-        ], config.enable_autocast_cache)
+        model.autocast_context, model.train_dtype = create_autocast_context(
+            self.train_device, config.train_dtype, config.enable_autocast_cache)
 
         model.text_encoder_2_autocast_context, model.text_encoder_2_train_dtype = \
             disable_fp16_autocast_context(
                 self.train_device,
                 config.train_dtype,
                 config.fallback_train_dtype,
-                [
-                    config.weight_dtypes().text_encoder_2,
-                    config.weight_dtypes().lora if config.training_method == TrainingMethod.LORA else None,
-                    config.weight_dtypes().embedding if config.train_any_embedding() else None,
-                ],
                 config.enable_autocast_cache,
             )
 

modules/modelSetup/BaseHiDreamSetup.py

@@ -18,7 +18,6 @@
 )
 from modules.util.config.TrainConfig import TrainConfig
 from modules.util.dtype_util import create_autocast_context, disable_fp16_autocast_context
-from modules.util.enum.TrainingMethod import TrainingMethod
 from modules.util.quantization_util import quantize_layers
 from modules.util.torch_util import torch_gc
 from modules.util.TrainProgress import TrainProgress
@@ -63,27 +62,14 @@ def setup_optimizations(
                 model.text_encoder_4_offload_conductor = \
                     enable_checkpointing_for_llama_encoder_layers(model.text_encoder_4, config)
 
-        model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [
-            config.weight_dtypes().transformer,
-            config.weight_dtypes().text_encoder,
-            config.weight_dtypes().text_encoder_2,
-            config.weight_dtypes().text_encoder_3,
-            config.weight_dtypes().text_encoder_4,
-            config.weight_dtypes().vae,
-            config.weight_dtypes().lora if config.training_method == TrainingMethod.LORA else None,
-            config.weight_dtypes().embedding if config.train_any_embedding() else None,
-        ], config.enable_autocast_cache)
+        model.autocast_context, model.train_dtype = create_autocast_context(
+            self.train_device, config.train_dtype, config.enable_autocast_cache)
 
         model.text_encoder_3_autocast_context, model.text_encoder_3_train_dtype = \
             disable_fp16_autocast_context(
                 self.train_device,
                 config.train_dtype,
                 config.fallback_train_dtype,
-                [
-                    config.weight_dtypes().text_encoder_3,
-                    config.weight_dtypes().lora if config.training_method == TrainingMethod.LORA else None,
-                    config.weight_dtypes().embedding if config.train_any_embedding() else None,
-                ],
                 config.enable_autocast_cache,
             )
 
@@ -92,11 +78,6 @@ def setup_optimizations(
                 self.train_device,
                 config.train_dtype,
                 config.fallback_train_dtype,
-                [
-                    config.weight_dtypes().transformer,
-                    config.weight_dtypes().lora if config.training_method == TrainingMethod.LORA else None,
-                    config.weight_dtypes().embedding if config.train_any_embedding() else None,
-                ],
                 config.enable_autocast_cache,
             )
 

modules/modelSetup/BaseHunyuanVideoSetup.py

@@ -18,7 +18,6 @@
 )
 from modules.util.config.TrainConfig import TrainConfig
 from modules.util.dtype_util import create_autocast_context, disable_fp16_autocast_context
-from modules.util.enum.TrainingMethod import TrainingMethod
 from modules.util.quantization_util import quantize_layers
 from modules.util.torch_util import torch_gc
 from modules.util.TrainProgress import TrainProgress
@@ -58,25 +57,14 @@ def setup_optimizations(
             if model.text_encoder_2 is not None:
                 enable_checkpointing_for_clip_encoder_layers(model.text_encoder_2, config)
 
-        model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [
-            config.weight_dtypes().transformer,
-            config.weight_dtypes().text_encoder,
-            config.weight_dtypes().text_encoder_2,
-            config.weight_dtypes().vae,
-            config.weight_dtypes().lora if config.training_method == TrainingMethod.LORA else None,
-            config.weight_dtypes().embedding if config.train_any_embedding() else None,
-        ], config.enable_autocast_cache)
+        model.autocast_context, model.train_dtype = create_autocast_context(
+            self.train_device, config.train_dtype, config.enable_autocast_cache)
 
         model.transformer_autocast_context, model.transformer_train_dtype = \
             disable_fp16_autocast_context(
                 self.train_device,
                 config.train_dtype,
                 config.fallback_train_dtype,
-                [
-                    config.weight_dtypes().transformer,
-                    config.weight_dtypes().lora if config.training_method == TrainingMethod.LORA else None,
-                    config.weight_dtypes().embedding if config.train_any_embedding() else None,
-                ],
                 config.enable_autocast_cache,
             )
 

modules/modelSetup/BasePixArtAlphaSetup.py

@@ -17,7 +17,6 @@
 )
 from modules.util.config.TrainConfig import TrainConfig
 from modules.util.dtype_util import create_autocast_context, disable_fp16_autocast_context
-from modules.util.enum.TrainingMethod import TrainingMethod
 from modules.util.quantization_util import quantize_layers
 from modules.util.torch_util import torch_gc
 from modules.util.TrainProgress import TrainProgress
@@ -58,23 +57,13 @@ def setup_optimizations(
             model.text_encoder_offload_conductor = \
                 enable_checkpointing_for_t5_encoder_layers(model.text_encoder, config)
 
-        model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [
-            config.weight_dtypes().transformer,
-            config.weight_dtypes().text_encoder,
-            config.weight_dtypes().vae,
-            config.weight_dtypes().lora if config.training_method == TrainingMethod.LORA else None,
-            config.weight_dtypes().embedding if config.train_any_embedding() else None,
-        ], config.enable_autocast_cache)
+        model.autocast_context, model.train_dtype = create_autocast_context(
+            self.train_device, config.train_dtype, config.enable_autocast_cache)
 
         model.text_encoder_autocast_context, model.text_encoder_train_dtype = disable_fp16_autocast_context(
             self.train_device,
             config.train_dtype,
             config.fallback_train_dtype,
-            [
-                config.weight_dtypes().text_encoder,
-                config.weight_dtypes().lora if config.training_method == TrainingMethod.LORA else None,
-                config.weight_dtypes().embedding if config.train_any_embedding() else None,
-            ],
             config.enable_autocast_cache,
         )
 

modules/modelSetup/BaseQwenSetup.py

@@ -15,7 +15,6 @@
 )
 from modules.util.config.TrainConfig import TrainConfig
 from modules.util.dtype_util import create_autocast_context, disable_fp16_autocast_context
-from modules.util.enum.TrainingMethod import TrainingMethod
 from modules.util.quantization_util import quantize_layers
 from modules.util.torch_util import torch_gc
 from modules.util.TrainProgress import TrainProgress
@@ -53,22 +52,14 @@ def setup_optimizations(
                 model.text_encoder_offload_conductor = \
                     enable_checkpointing_for_qwen25vl_encoder_layers(model.text_encoder, config)
 
-        model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [
-            config.weight_dtypes().transformer,
-            config.weight_dtypes().text_encoder,
-            config.weight_dtypes().vae,
-            config.weight_dtypes().lora if config.training_method == TrainingMethod.LORA else None,
-        ], config.enable_autocast_cache)
+        model.autocast_context, model.train_dtype = create_autocast_context(
+            self.train_device, config.train_dtype, config.enable_autocast_cache)
 
         model.text_encoder_autocast_context, model.text_encoder_train_dtype = \
             disable_fp16_autocast_context(
                 self.train_device,
                 config.train_dtype,
                 config.fallback_train_dtype,
-                [
-                    config.weight_dtypes().text_encoder,
-                    config.weight_dtypes().lora if config.training_method == TrainingMethod.LORA else None,
-                ],
                 config.enable_autocast_cache,
             )
 

modules/modelSetup/BaseSanaSetup.py

@@ -17,7 +17,6 @@
 )
 from modules.util.config.TrainConfig import TrainConfig
 from modules.util.dtype_util import create_autocast_context, disable_fp16_autocast_context
-from modules.util.enum.TrainingMethod import TrainingMethod
 from modules.util.quantization_util import quantize_layers
 from modules.util.torch_util import torch_gc
 from modules.util.TrainProgress import TrainProgress
@@ -59,33 +58,20 @@ def setup_optimizations(
             model.text_encoder_offload_conductor = \
                 enable_checkpointing_for_gemma_layers(model.text_encoder, config)
 
-        model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [
-            config.weight_dtypes().transformer,
-            config.weight_dtypes().text_encoder,
-            config.weight_dtypes().vae,
-            config.weight_dtypes().lora if config.training_method == TrainingMethod.LORA else None,
-            config.weight_dtypes().embedding if config.train_any_embedding() else None,
-        ], config.enable_autocast_cache)
+        model.autocast_context, model.train_dtype = create_autocast_context(
+            self.train_device, config.train_dtype, config.enable_autocast_cache)
 
         model.text_encoder_autocast_context, model.text_encoder_train_dtype = disable_fp16_autocast_context(
             self.train_device,
             config.train_dtype,
             config.fallback_train_dtype,
-            [
-                config.weight_dtypes().text_encoder,
-                config.weight_dtypes().lora if config.training_method == TrainingMethod.LORA else None,
-                config.weight_dtypes().embedding if config.train_any_embedding() else None,
-            ],
             config.enable_autocast_cache,
         )
 
         model.vae_autocast_context, model.vae_train_dtype = disable_fp16_autocast_context(
             self.train_device,
             config.train_dtype,
             config.fallback_train_dtype,
-            [
-                config.weight_dtypes().vae,
-            ],
             config.enable_autocast_cache,
         )
 

modules/modelSetup/BaseStableDiffusion3Setup.py

@@ -18,7 +18,6 @@
 )
 from modules.util.config.TrainConfig import TrainConfig
 from modules.util.dtype_util import create_autocast_context, disable_fp16_autocast_context
-from modules.util.enum.TrainingMethod import TrainingMethod
 from modules.util.quantization_util import quantize_layers
 from modules.util.torch_util import torch_gc
 from modules.util.TrainProgress import TrainProgress
@@ -59,26 +58,14 @@ def setup_optimizations(
                 model.text_encoder_3_offload_conductor = \
                     enable_checkpointing_for_t5_encoder_layers(model.text_encoder_3, config)
 
-        model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [
-            config.weight_dtypes().transformer,
-            config.weight_dtypes().text_encoder,
-            config.weight_dtypes().text_encoder_2,
-            config.weight_dtypes().text_encoder_3,
-            config.weight_dtypes().vae,
-            config.weight_dtypes().lora if config.training_method == TrainingMethod.LORA else None,
-            config.weight_dtypes().embedding if config.train_any_embedding() else None,
-        ], config.enable_autocast_cache)
+        model.autocast_context, model.train_dtype = create_autocast_context(
+            self.train_device, config.train_dtype, config.enable_autocast_cache)
 
         model.text_encoder_3_autocast_context, model.text_encoder_3_train_dtype = \
             disable_fp16_autocast_context(
                 self.train_device,
                 config.train_dtype,
                 config.fallback_train_dtype,
-                [
-                    config.weight_dtypes().text_encoder_3,
-                    config.weight_dtypes().lora if config.training_method == TrainingMethod.LORA else None,
-                    config.weight_dtypes().embedding if config.train_any_embedding() else None,
-                ],
                 config.enable_autocast_cache,
             )