Skip to content

always enable mixed-precision / autocasting #1481

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
dxqb merged 5 commits into from
Jul 5, 2026
Merged

always enable mixed-precision / autocasting #1481

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
d624556
fix: gate autocast on train dtype instead of dtype-set uniformity
dxqb May 27, 2026
bcc7d5d
raise instead of silently degrading float32 fallback on non-CUDA
dxqb May 28, 2026
2609822
Merge remote-tracking branch 'Nerogar/master' into fix/mixed-precisio…
dxqb Jul 4, 2026
0873225
Merge remote-tracking branch 'private/merge' into fix/mixed-precision…
dxqb Jul 5, 2026
abfe6a1
Update Anima/Krea2/Ideogram to new create_autocast_context/disable_fp…
dxqb Jul 5, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 2 additions & 11 deletions modules/modelSetup/BaseAnimaSetup.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -50,22 +49,14 @@ def setup_optimizations(
model.transformer_offload_conductor = enable_checkpointing_for_qwen_transformer(model.transformer, config, config.transformer)
model.text_encoder_offload_conductor = enable_checkpointing_for_qwen3_encoder_layers(model.text_encoder, config, config.text_encoder)

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,
)

Expand Down
15 changes: 2 additions & 13 deletions modules/modelSetup/BaseChromaSetup.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -52,24 +51,14 @@ def setup_optimizations(
model.transformer_offload_conductor = enable_checkpointing_for_chroma_transformer(model.transformer, config, config.transformer)
model.text_encoder_offload_conductor = enable_checkpointing_for_t5_encoder_layers(model.text_encoder, config, config.text_encoder)

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,
)

Expand Down
13 changes: 2 additions & 11 deletions modules/modelSetup/BaseErnieSetup.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -48,22 +47,14 @@ def setup_optimizations(
model.transformer_offload_conductor = enable_checkpointing_for_ernie_transformer(model.transformer, config, config.transformer)
model.text_encoder_offload_conductor = enable_checkpointing_for_mistral_encoder_layers(model.text_encoder, config, config.text_encoder)

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,
)

Expand Down
13 changes: 2 additions & 11 deletions modules/modelSetup/BaseFlux2Setup.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -51,22 +50,14 @@ def setup_optimizations(
else:
model.text_encoder_offload_conductor = enable_checkpointing_for_qwen3_encoder_layers(model.text_encoder, config, config.text_encoder)

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,
)

Expand Down
16 changes: 2 additions & 14 deletions modules/modelSetup/BaseFluxSetup.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -55,25 +54,14 @@ def setup_optimizations(
if model.text_encoder_2 is not None:
model.text_encoder_2_offload_conductor = enable_checkpointing_for_t5_encoder_layers(model.text_encoder_2, config, config.text_encoder_2)

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,
)

Expand Down
23 changes: 2 additions & 21 deletions modules/modelSetup/BaseHiDreamSetup.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -59,27 +58,14 @@ def setup_optimizations(
if model.text_encoder_4 is not None:
model.text_encoder_4_offload_conductor = enable_checkpointing_for_llama_encoder_layers(model.text_encoder_4, config, config.text_encoder_4)

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,
)

Expand All @@ -88,11 +74,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,
)

Expand Down
16 changes: 2 additions & 14 deletions modules/modelSetup/BaseHunyuanVideoSetup.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -55,25 +54,14 @@ def setup_optimizations(
if model.text_encoder_2 is not None:
enable_checkpointing_for_clip_encoder_layers(model.text_encoder_2, config, config.text_encoder_2)

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,
)

Expand Down
13 changes: 2 additions & 11 deletions modules/modelSetup/BaseIdeogramSetup.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -57,22 +56,14 @@ def setup_optimizations(

model.text_encoder_offload_conductor = enable_checkpointing_for_qwen3vl_encoder_layers(model.text_encoder, config, config.text_encoder)

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,
)

Expand Down
13 changes: 2 additions & 11 deletions modules/modelSetup/BaseKrea2Setup.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -50,22 +49,14 @@ def setup_optimizations(
model.transformer_offload_conductor = enable_checkpointing_for_krea2_transformer(model.transformer, config, config.transformer)
model.text_encoder_offload_conductor = enable_checkpointing_for_qwen3vl_encoder_layers(model.text_encoder, config, config.text_encoder)

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,
)

Expand Down
15 changes: 2 additions & 13 deletions modules/modelSetup/BasePixArtAlphaSetup.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -54,23 +53,13 @@ def setup_optimizations(
model.transformer_offload_conductor = enable_checkpointing_for_basic_transformer_blocks(model.transformer, config, config.transformer, offload_enabled=True)
model.text_encoder_offload_conductor = enable_checkpointing_for_t5_encoder_layers(model.text_encoder, config, config.text_encoder)

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,
)

Expand Down
Loading