From e432c14d7f2fe56713433b93a997e4ec9dfcdb5f Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Mon, 19 Jan 2026 00:11:13 +0300 Subject: [PATCH 01/84] initial --- modules/ui/OptimizerParamsWindow.py | 1 + modules/util/config/TrainConfig.py | 2 + modules/util/create.py | 34 +++++++++++-- modules/util/optimizer/muon_util.py | 79 ++++++++++++++++++++++++++++- modules/util/optimizer_util.py | 10 ++-- 5 files changed, 113 insertions(+), 13 deletions(-) diff --git a/modules/ui/OptimizerParamsWindow.py b/modules/ui/OptimizerParamsWindow.py index ed50e9f57..95fef5cd1 100644 --- a/modules/ui/OptimizerParamsWindow.py +++ b/modules/ui/OptimizerParamsWindow.py @@ -199,6 +199,7 @@ def create_dynamic_ui( 'kappa_p': {'title': 'Lion-K P-value', 'tooltip': 'Controls the Lp-norm geometry for the Lion update. 1.0 = Standard Lion (Sign update, coordinate-wise), best for Transformers. 2.0 = Spherical Lion (Normalized update, rotational invariant), best for Conv2d layers (in unet models). Values between 1.0 and 2.0 interpolate behavior between the two.', 'type': 'float'}, 'auto_kappa_p': {'title': 'Auto Lion-K', 'tooltip': 'Automatically determines the optimal P-value based on layer dimensions. Uses p=2.0 (Spherical) for 4D (Conv) tensors for stability and rotational invariance, and p=1.0 (Sign) for 2D (Linear) tensors for sparsity. Overrides the manual P-value. Recommend for unet models.', 'type': 'bool'}, 'compile': {'title': 'Compiled Optimizer', 'tooltip': 'Enables PyTorch compilation for the optimizer internal step logic. This is intended to improve performance by allowing PyTorch to fuse operations and optimize the computational graph.', 'type': 'bool'}, + 'spectral_normalization': {'title': 'Spectral Scaling', 'tooltip': 'Enables explicit Spectral Normalization to automatically rescale the update magnitude and Weight Decay based on layer dimensions. This allows hyperparameters to transfer seamlessly from small to large models without retuning, while making the optimizer highly robust to a wide range of learning rates. This ensures consistent performance across different model sizes, adapter methods, and ranks. ', 'type': 'bool'}, } # @formatter:on diff --git a/modules/util/config/TrainConfig.py b/modules/util/config/TrainConfig.py index b8184e1c2..8b06719ba 100644 --- a/modules/util/config/TrainConfig.py +++ b/modules/util/config/TrainConfig.py @@ -144,6 +144,7 @@ class TrainOptimizerConfig(BaseConfig): kappa_p: float auto_kappa_p: False compile: False + spectral_normalization: False def __init__(self, data: list[(str, Any, type, bool)]): super().__init__(data) @@ -263,6 +264,7 @@ def default_values(): data.append(("kappa_p", None, float, True)) data.append(("auto_kappa_p", False, bool, False)) data.append(("compile", False, bool, False)) + data.append(("spectral_normalization", False, bool, False)) return TrainOptimizerConfig(data) diff --git a/modules/util/create.py b/modules/util/create.py index a03ce3067..e0a71860e 100644 --- a/modules/util/create.py +++ b/modules/util/create.py @@ -159,7 +159,7 @@ from modules.util.optimizer.adafactor_extensions import patch_adafactor from modules.util.optimizer.adam_extensions import patch_adam from modules.util.optimizer.adamw_extensions import patch_adamw -from modules.util.optimizer.muon_util import split_parameters_for_muon +from modules.util.optimizer.muon_util import calculate_muon_n_layers, split_parameters_for_muon from modules.util.TrainProgress import TrainProgress from modules.zluda import ZLUDA @@ -520,7 +520,7 @@ def create_optimizer( parameter_group_collection: NamedParameterGroupCollection, state_dict: dict | None, config: TrainConfig, - layer_key_fn: dict[int, str] | None = None, + model: BaseModel | None = None, ) -> torch.optim.Optimizer | None: optimizer = None optimizer_config = config.optimizer @@ -1237,7 +1237,18 @@ def create_optimizer( from adv_optm import Muon_adv - params_for_optimizer, MuonWithAuxAdam = split_parameters_for_muon(parameters, layer_key_fn, config) + params_for_optimizer, MuonWithAuxAdam = split_parameters_for_muon(model, parameters, config) + + if optimizer_config.spectral_normalization: + # Calculate n_layers for spectral normalization + n_layers_map = calculate_muon_n_layers(model) + + for group in params_for_optimizer: + group_name = group.get('name') + if group_name in n_layers_map: + group['n_layers'] = n_layers_map[group_name] + else: + group['n_layers'] = n_layers_map.get('default', 1) # Prepare Adam-specific keyword arguments from the config adam_kwargs = {} @@ -1278,6 +1289,7 @@ def create_optimizer( orthogonal_gradient=optimizer_config.orthogonal_gradient if optimizer_config.orthogonal_gradient is not None else False, approx_mars=optimizer_config.approx_mars if optimizer_config.approx_mars is not None else False, compiled_optimizer=optimizer_config.compile if optimizer_config.compile is not None else False, + spectral_normalization=optimizer_config.spectral_normalization if optimizer_config.spectral_normalization is not None else False, **adam_kwargs ) @@ -1287,7 +1299,18 @@ def create_optimizer( from adv_optm import AdaMuon_adv - params_for_optimizer, MuonWithAuxAdam = split_parameters_for_muon(parameters, layer_key_fn, config) + params_for_optimizer, MuonWithAuxAdam = split_parameters_for_muon(model, parameters, config) + + if optimizer_config.spectral_normalization: + # Calculate n_layers for spectral normalization + n_layers_map = calculate_muon_n_layers(model) + + for group in params_for_optimizer: + group_name = group.get('name') + if group_name in n_layers_map: + group['n_layers'] = n_layers_map[group_name] + else: + group['n_layers'] = n_layers_map.get('default', 1) # Prepare Adam-specific keyword arguments from the config adam_kwargs = {} @@ -1332,6 +1355,7 @@ def create_optimizer( orthogonal_gradient=optimizer_config.orthogonal_gradient if optimizer_config.orthogonal_gradient is not None else False, approx_mars=optimizer_config.approx_mars if optimizer_config.approx_mars is not None else False, compiled_optimizer=optimizer_config.compile if optimizer_config.compile is not None else False, + spectral_normalization=optimizer_config.spectral_normalization if optimizer_config.spectral_normalization is not None else False, **adam_kwargs ) @@ -1340,7 +1364,7 @@ def create_optimizer( from muon import MuonWithAuxAdam, SingleDeviceMuonWithAuxAdam - params_for_optimizer, ___ = split_parameters_for_muon(parameters, layer_key_fn, config) + params_for_optimizer, ___ = split_parameters_for_muon(model, parameters, config) final_param_groups = [] for group in params_for_optimizer: diff --git a/modules/util/optimizer/muon_util.py b/modules/util/optimizer/muon_util.py index 74c7e0430..8cf9e4814 100644 --- a/modules/util/optimizer/muon_util.py +++ b/modules/util/optimizer/muon_util.py @@ -1,3 +1,4 @@ +import re from collections.abc import Callable from modules.model.BaseModel import BaseModel, TrainConfig @@ -8,6 +9,80 @@ import torch +def calculate_muon_n_layers(model: BaseModel) -> dict[str, int]: + """ + Calculates the number of residual layers (the depth) in each component of the model. + Used for Muon optimizer spectral normalization scaling. + """ + match model.model_type: + case (ModelType.STABLE_DIFFUSION_15 | ModelType.STABLE_DIFFUSION_15_INPAINTING | + ModelType.STABLE_DIFFUSION_20_BASE | ModelType.STABLE_DIFFUSION_20_INPAINTING | + ModelType.STABLE_DIFFUSION_20 | ModelType.STABLE_DIFFUSION_21 | + ModelType.STABLE_DIFFUSION_21_BASE | ModelType.STABLE_DIFFUSION_XL_10_BASE | + ModelType.STABLE_DIFFUSION_XL_10_BASE_INPAINTING | ModelType.STABLE_CASCADE_1 | + ModelType.WUERSTCHEN_2): + default_patterns = ['transformer_blocks', 'resnets', 'layers'] + case (ModelType.STABLE_DIFFUSION_3 | ModelType.STABLE_DIFFUSION_35 | ModelType.SANA | + ModelType.FLUX_DEV_1 | ModelType.CHROMA_1 | ModelType.QWEN | + ModelType.PIXART_ALPHA | ModelType.PIXART_SIGMA): + default_patterns = ['transformer_blocks', 'encoder.block'] + case ModelType.HI_DREAM_FULL: + default_patterns = ['double_stream_blocks', 'single_stream_blocks'] + case ModelType.Z_IMAGE: + default_patterns = [ + 'layers', + 'refiner', + ] + case _: + raise NotImplementedError(f"Muon optimizer spectral normalization is not implemented for model type: {model.model_type}") + + # Build the regex pattern dynamically + joined_patterns = "|".join([re.escape(p) for p in default_patterns]) + pattern = re.compile(rf'(?:^|\.)(?:{joined_patterns})\.\d+$') + + layer_counts = {} + + # Iterate over model components (e.g., 'unet', 'text_encoder', 'transformer') + for attr_name, module in vars(model).items(): + + # Identify the 'Ground Truth' blocks in this component. + target_module = module + if isinstance(module, LoRAModuleWrapper): + target_module = module.orig_module + + valid_component_blocks = set() + if isinstance(target_module, torch.nn.Module): + for name, _ in target_module.named_modules(): + if pattern.search(name): + valid_component_blocks.add(name) + + if not valid_component_blocks: + continue + + active_component_blocks = set() + + # Filter: Only count blocks that are actually being trained. + if isinstance(module, LoRAModuleWrapper): + # For LoRA, we check if the active leaves reside inside a valid block. + for layer_name in module.lora_modules: + parts = layer_name.split('.') + for i in range(len(parts), 0, -1): + candidate = ".".join(parts[:i]) + if candidate in valid_component_blocks: + active_component_blocks.add(candidate) + + elif isinstance(module, torch.nn.Module): + # For standard full-finetuning, all valid blocks are active. + active_component_blocks = valid_component_blocks + + count = len(active_component_blocks) + if count > 0: + print(f"[MuonUtil] Component '{attr_name}': detected {count} residual layers.") + layer_counts[attr_name] = count + + return layer_counts + + def build_muon_adam_key_fn( model: BaseModel, config: TrainConfig, @@ -105,8 +180,8 @@ def get_optim_type(param_name: str, p: torch.nn.Parameter) -> str: return param_map def split_parameters_for_muon( + model: BaseModel, parameters: list[dict], - layer_key_fn: dict[int, str], config: TrainConfig, ) -> tuple[list[dict], bool]: """ @@ -116,6 +191,8 @@ def split_parameters_for_muon( optimizer_config = config.optimizer has_adam_params = False + layer_key_fn = build_muon_adam_key_fn(model, config) + if layer_key_fn: for group in parameters: for p in group['params']: diff --git a/modules/util/optimizer_util.py b/modules/util/optimizer_util.py index f879bb5d9..52689dd65 100644 --- a/modules/util/optimizer_util.py +++ b/modules/util/optimizer_util.py @@ -4,7 +4,6 @@ from modules.util.config.TrainConfig import TrainConfig, TrainOptimizerConfig from modules.util.enum.Optimizer import Optimizer from modules.util.NamedParameterGroup import NamedParameterGroupCollection -from modules.util.optimizer.muon_util import build_muon_adam_key_fn from modules.util.torch_util import optimizer_to_device_ import torch @@ -59,13 +58,8 @@ def init_model_parameters( #to be safe, do that before the optimizer is created because the optimizer could take copies multi.broadcast_parameters(parameters.parameters(), train_device) - layer_key_fn = None - if model.train_config.optimizer.MuonWithAuxAdam: - print("INFO: Creating layer keys for MuonWithAuxAdam.") - layer_key_fn = build_muon_adam_key_fn(model, model.train_config) - model.optimizer = create.create_optimizer( - parameters, model.optimizer_state_dict, model.train_config, layer_key_fn + parameters, model.optimizer_state_dict, model.train_config, model=model ) if model.optimizer is not None: @@ -596,6 +590,7 @@ def init_model_parameters( "low_rank_ortho": False, "ortho_rank": 128, "rms_rescaling": True, + "spectral_normalization": False, "nnmf_factor": False, "stochastic_rounding": True, "compile": False, @@ -627,6 +622,7 @@ def init_model_parameters( "low_rank_ortho": False, "ortho_rank": 128, "rms_rescaling": True, + "spectral_normalization": False, "nnmf_factor": False, "stochastic_rounding": True, "compile": False, From b359a156b9c913829f74604d784174880db97c22 Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Mon, 19 Jan 2026 00:19:26 +0300 Subject: [PATCH 02/84] dev1 --- requirements-global.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements-global.txt b/requirements-global.txt index 3438ce9ef..3f3406348 100644 --- a/requirements-global.txt +++ b/requirements-global.txt @@ -41,7 +41,7 @@ prodigyopt==1.1.2 # prodigy optimizer schedulefree==1.4.1 # schedule-free optimizers pytorch_optimizer==3.6.0 # pytorch optimizers prodigy-plus-schedule-free==2.0.1 # Prodigy plus optimizer -adv_optm==2.1.0 # advanced optimizers +adv_optm==2.2.dev1 # advanced optimizers -e git+https://github.com/KellerJordan/Muon.git@f90a42b#egg=muon-optimizer # Profiling From 972ee771af08a449640b12aee6baa028b86c62a7 Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Mon, 19 Jan 2026 22:04:56 +0300 Subject: [PATCH 03/84] dev2 --- requirements-global.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements-global.txt b/requirements-global.txt index 3f3406348..fc8ead00d 100644 --- a/requirements-global.txt +++ b/requirements-global.txt @@ -41,7 +41,7 @@ prodigyopt==1.1.2 # prodigy optimizer schedulefree==1.4.1 # schedule-free optimizers pytorch_optimizer==3.6.0 # pytorch optimizers prodigy-plus-schedule-free==2.0.1 # Prodigy plus optimizer -adv_optm==2.2.dev1 # advanced optimizers +adv_optm==2.2.dev2 # advanced optimizers -e git+https://github.com/KellerJordan/Muon.git@f90a42b#egg=muon-optimizer # Profiling From 1d491758411d010422bcfca69c6009896ebd72d8 Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Tue, 20 Jan 2026 18:50:46 +0300 Subject: [PATCH 04/84] dev3 --- requirements-global.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements-global.txt b/requirements-global.txt index fc8ead00d..2b118adc0 100644 --- a/requirements-global.txt +++ b/requirements-global.txt @@ -41,7 +41,7 @@ prodigyopt==1.1.2 # prodigy optimizer schedulefree==1.4.1 # schedule-free optimizers pytorch_optimizer==3.6.0 # pytorch optimizers prodigy-plus-schedule-free==2.0.1 # Prodigy plus optimizer -adv_optm==2.2.dev2 # advanced optimizers +adv_optm==2.2.dev3 # advanced optimizers -e git+https://github.com/KellerJordan/Muon.git@f90a42b#egg=muon-optimizer # Profiling From 6ca6f52532e3dca041264be1853e886363ec9934 Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Tue, 20 Jan 2026 19:12:58 +0300 Subject: [PATCH 05/84] dev4 --- requirements-global.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements-global.txt b/requirements-global.txt index 2b118adc0..f3309ad9b 100644 --- a/requirements-global.txt +++ b/requirements-global.txt @@ -41,7 +41,7 @@ prodigyopt==1.1.2 # prodigy optimizer schedulefree==1.4.1 # schedule-free optimizers pytorch_optimizer==3.6.0 # pytorch optimizers prodigy-plus-schedule-free==2.0.1 # Prodigy plus optimizer -adv_optm==2.2.dev3 # advanced optimizers +adv_optm==2.2.dev4 # advanced optimizers -e git+https://github.com/KellerJordan/Muon.git@f90a42b#egg=muon-optimizer # Profiling From e40579acaa6ee9a9f6555139a32aaebdbf338d14 Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Mon, 26 Jan 2026 04:52:10 +0300 Subject: [PATCH 06/84] add Chroma residual filter --- modules/util/optimizer/muon_util.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modules/util/optimizer/muon_util.py b/modules/util/optimizer/muon_util.py index 8cf9e4814..ece9e2f31 100644 --- a/modules/util/optimizer/muon_util.py +++ b/modules/util/optimizer/muon_util.py @@ -25,7 +25,7 @@ def calculate_muon_n_layers(model: BaseModel) -> dict[str, int]: case (ModelType.STABLE_DIFFUSION_3 | ModelType.STABLE_DIFFUSION_35 | ModelType.SANA | ModelType.FLUX_DEV_1 | ModelType.CHROMA_1 | ModelType.QWEN | ModelType.PIXART_ALPHA | ModelType.PIXART_SIGMA): - default_patterns = ['transformer_blocks', 'encoder.block'] + default_patterns = ['transformer_blocks', 'single_transformer_blocks', 'encoder.block'] case ModelType.HI_DREAM_FULL: default_patterns = ['double_stream_blocks', 'single_stream_blocks'] case ModelType.Z_IMAGE: From 2bc6ae1b98f899c9ad474d6b590c8e6a86bf316d Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Sat, 31 Jan 2026 21:35:47 +0300 Subject: [PATCH 07/84] stable 2.2 and edit rms tooltip --- modules/ui/OptimizerParamsWindow.py | 2 +- requirements-global.txt | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/modules/ui/OptimizerParamsWindow.py b/modules/ui/OptimizerParamsWindow.py index 95fef5cd1..124d99b9d 100644 --- a/modules/ui/OptimizerParamsWindow.py +++ b/modules/ui/OptimizerParamsWindow.py @@ -187,7 +187,7 @@ def create_dynamic_ui( 'muon_adam_lr': {'title': 'Auxiliary Adam LR', 'tooltip': 'Learning rate for the auxiliary AdamW optimizer. If empty, it will use the main learning rate.', 'type': 'float'}, 'muon_te1_adam_lr': {'title': 'AuxAdam TE1 LR', 'tooltip': 'Learning rate for the auxiliary AdamW optimizer for the first text encoder. If empty, it will use the Auxiliary Adam LR.', 'type': 'float'}, 'muon_te2_adam_lr': {'title': 'AuxAdam TE2 LR', 'tooltip': 'Learning rate for the auxiliary AdamW optimizer for the second text encoder. If empty, it will use the Auxiliary Adam LR.', 'type': 'float'}, - 'rms_rescaling': {'title': 'RMS Rescaling', 'tooltip': 'Muon already scales its updates to approximate and use the same learning rate (LR) as Adam. This option integrates a more accurate method to match the Adam LR, but it is slower.', 'type': 'bool'}, + 'rms_rescaling': {'title': 'RMS Rescaling', 'tooltip': 'Normalizes Muon update magnitudes to align with Adam. This allows to reuse standard "Adam-style" learning rates instead of specialized Muon scales.', 'type': 'bool'}, 'normuon_variant': {'title': 'NorMuon Variant', 'tooltip': 'Enables the NorMuon optimizer variant, which combines Muon orthogonalization with per-neuron adaptive learning rates for better convergence and balanced parameter updates. Costs only one scalar state buffer per parameter group, size few KBs, maintaining high memory efficiency.', 'type': 'bool'}, 'beta2_normuon': {'title': 'NorMuon Beta2', 'tooltip': 'Exponential decay rate for the neuron-wise second-moment estimator in NorMuon (analogous to Adams beta2). Controls how past squared updates influence current normalization.', 'type': 'float'}, 'normuon_eps': {'title': 'NorMuon EPS', 'tooltip': 'Epsilon for NorMuon normalization stability.', 'type': 'float'}, diff --git a/requirements-global.txt b/requirements-global.txt index f3309ad9b..42b9adc5b 100644 --- a/requirements-global.txt +++ b/requirements-global.txt @@ -41,7 +41,7 @@ prodigyopt==1.1.2 # prodigy optimizer schedulefree==1.4.1 # schedule-free optimizers pytorch_optimizer==3.6.0 # pytorch optimizers prodigy-plus-schedule-free==2.0.1 # Prodigy plus optimizer -adv_optm==2.2.dev4 # advanced optimizers +adv_optm==2.2.0 # advanced optimizers -e git+https://github.com/KellerJordan/Muon.git@f90a42b#egg=muon-optimizer # Profiling From 31287b2c708aee28931563aa394bd22ddc84f933 Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Sat, 31 Jan 2026 21:43:49 +0300 Subject: [PATCH 08/84] remove the print --- modules/util/optimizer/muon_util.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/modules/util/optimizer/muon_util.py b/modules/util/optimizer/muon_util.py index ece9e2f31..58a84b1c4 100644 --- a/modules/util/optimizer/muon_util.py +++ b/modules/util/optimizer/muon_util.py @@ -43,7 +43,7 @@ def calculate_muon_n_layers(model: BaseModel) -> dict[str, int]: layer_counts = {} # Iterate over model components (e.g., 'unet', 'text_encoder', 'transformer') - for attr_name, module in vars(model).items(): + for _attr_name, module in vars(model).values(): # Identify the 'Ground Truth' blocks in this component. target_module = module @@ -75,10 +75,6 @@ def calculate_muon_n_layers(model: BaseModel) -> dict[str, int]: # For standard full-finetuning, all valid blocks are active. active_component_blocks = valid_component_blocks - count = len(active_component_blocks) - if count > 0: - print(f"[MuonUtil] Component '{attr_name}': detected {count} residual layers.") - layer_counts[attr_name] = count return layer_counts From 44cca26009cc1b8e340ebdb8cfc6bedf7fc528ec Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Sat, 31 Jan 2026 21:45:21 +0300 Subject: [PATCH 09/84] use .values() --- modules/util/optimizer/muon_util.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modules/util/optimizer/muon_util.py b/modules/util/optimizer/muon_util.py index 58a84b1c4..1d15a85ca 100644 --- a/modules/util/optimizer/muon_util.py +++ b/modules/util/optimizer/muon_util.py @@ -43,7 +43,7 @@ def calculate_muon_n_layers(model: BaseModel) -> dict[str, int]: layer_counts = {} # Iterate over model components (e.g., 'unet', 'text_encoder', 'transformer') - for _attr_name, module in vars(model).values(): + for module in vars(model).values(): # Identify the 'Ground Truth' blocks in this component. target_module = module From bee4b86a8e2cd5ab0a4a7520b86424a62ee2e805 Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Sun, 22 Feb 2026 23:37:33 +0300 Subject: [PATCH 10/84] initial --- modules/module/LoRAModule.py | 4 ++++ modules/ui/OptimizerParamsWindow.py | 1 + modules/util/config/TrainConfig.py | 2 ++ modules/util/create.py | 5 +++++ modules/util/optimizer_util.py | 5 +++++ 5 files changed, 17 insertions(+) diff --git a/modules/module/LoRAModule.py b/modules/module/LoRAModule.py index c2a3a34b9..5539acc33 100644 --- a/modules/module/LoRAModule.py +++ b/modules/module/LoRAModule.py @@ -314,6 +314,8 @@ def initialize_weights(self): self.lora_down, self.lora_up = self.create_layer() nn.init.kaiming_uniform_(self.lora_down.weight, a=math.sqrt(5)) nn.init.zeros_(self.lora_up.weight) + self.lora_down.weight._is_lora_A = True + self.lora_up.weight._is_lora_B = True def check_initialized(self): super().check_initialized() @@ -427,6 +429,7 @@ def initialize_weights(self): ) nn.init.zeros_(self.oft_R.weight) + self.oft_R.weight._is_oft = True def forward(self, x, *args, **kwargs): self.check_initialized() @@ -519,6 +522,7 @@ def initialize_weights(self): .to(device=self.orig_module.weight.device) ) + self.dora_scale._is_dora_scale = True del orig_weight def check_initialized(self): diff --git a/modules/ui/OptimizerParamsWindow.py b/modules/ui/OptimizerParamsWindow.py index 124d99b9d..e40bb44c1 100644 --- a/modules/ui/OptimizerParamsWindow.py +++ b/modules/ui/OptimizerParamsWindow.py @@ -200,6 +200,7 @@ def create_dynamic_ui( 'auto_kappa_p': {'title': 'Auto Lion-K', 'tooltip': 'Automatically determines the optimal P-value based on layer dimensions. Uses p=2.0 (Spherical) for 4D (Conv) tensors for stability and rotational invariance, and p=1.0 (Sign) for 2D (Linear) tensors for sparsity. Overrides the manual P-value. Recommend for unet models.', 'type': 'bool'}, 'compile': {'title': 'Compiled Optimizer', 'tooltip': 'Enables PyTorch compilation for the optimizer internal step logic. This is intended to improve performance by allowing PyTorch to fuse operations and optimize the computational graph.', 'type': 'bool'}, 'spectral_normalization': {'title': 'Spectral Scaling', 'tooltip': 'Enables explicit Spectral Normalization to automatically rescale the update magnitude and Weight Decay based on layer dimensions. This allows hyperparameters to transfer seamlessly from small to large models without retuning, while making the optimizer highly robust to a wide range of learning rates. This ensures consistent performance across different model sizes, adapter methods, and ranks. ', 'type': 'bool'}, + 'scaled_optm': {'title': 'Scaled Optimizer', 'tooltip': 'Automatically rescale the update magnitude and Weight Decay based on layer dimensions and type. This allows hyperparameters to transfer seamlessly from small to large models without retuning. This ensures consistent performance across different model sizes, adapter methods, and ranks. For LoRAs set alpha=rank. Using this, LoRA shares the same LR as full finetuning for all ranks.', 'type': 'bool'}, } # @formatter:on diff --git a/modules/util/config/TrainConfig.py b/modules/util/config/TrainConfig.py index db57322b3..1a6d6c57a 100644 --- a/modules/util/config/TrainConfig.py +++ b/modules/util/config/TrainConfig.py @@ -144,6 +144,7 @@ class TrainOptimizerConfig(BaseConfig): auto_kappa_p: False compile: False spectral_normalization: False + scaled_optm: False def __init__(self, data: list[(str, Any, type, bool)]): super().__init__(data) @@ -263,6 +264,7 @@ def default_values(): data.append(("auto_kappa_p", False, bool, False)) data.append(("compile", False, bool, False)) data.append(("spectral_normalization", False, bool, False)) + data.append(("scaled_optm", False, bool, False)) return TrainOptimizerConfig(data) diff --git a/modules/util/create.py b/modules/util/create.py index 6fe2175c9..fa65c2658 100644 --- a/modules/util/create.py +++ b/modules/util/create.py @@ -686,6 +686,7 @@ def create_optimizer( kourkoutas_beta=optimizer_config.kourkoutas_beta if optimizer_config.kourkoutas_beta is not None else False, k_warmup_steps=optimizer_config.k_warmup_steps if optimizer_config.k_warmup_steps is not None else 0, compiled_optimizer=optimizer_config.compile if optimizer_config.compile is not None else False, + scaled_optm=optimizer_config.scaled_optm if optimizer_config.scaled_optm is not None else False, ) # ADOPT_ADV Optimizer @@ -713,6 +714,7 @@ def create_optimizer( kourkoutas_beta=optimizer_config.kourkoutas_beta if optimizer_config.kourkoutas_beta is not None else False, k_warmup_steps=optimizer_config.k_warmup_steps if optimizer_config.k_warmup_steps is not None else 0, compiled_optimizer=optimizer_config.compile if optimizer_config.compile is not None else False, + scaled_optm=optimizer_config.scaled_optm if optimizer_config.scaled_optm is not None else False, ) # PRODIGY_ADV Optimizer @@ -770,6 +772,7 @@ def create_optimizer( kourkoutas_beta=optimizer_config.kourkoutas_beta if optimizer_config.kourkoutas_beta is not None else False, k_warmup_steps=optimizer_config.k_warmup_steps if optimizer_config.k_warmup_steps is not None else 0, compiled_optimizer=optimizer_config.compile if optimizer_config.compile is not None else False, + scaled_optm=optimizer_config.scaled_optm if optimizer_config.scaled_optm is not None else False, ) # SignSGD_ADV Optimizer @@ -787,6 +790,7 @@ def create_optimizer( compiled_optimizer=optimizer_config.compile if optimizer_config.compile is not None else False, Simplified_AdEMAMix=optimizer_config.Simplified_AdEMAMix if optimizer_config.Simplified_AdEMAMix is not None else False, alpha_grad=optimizer_config.alpha_grad if optimizer_config.alpha_grad is not None else 100, + scaled_optm=optimizer_config.scaled_optm if optimizer_config.scaled_optm is not None else False, ) # LION_ADV Optimizer @@ -807,6 +811,7 @@ def create_optimizer( kappa_p=optimizer_config.kappa_p if optimizer_config.kappa_p is not None else 1.0, auto_kappa_p=optimizer_config.auto_kappa_p if optimizer_config.auto_kappa_p is not None else False, compiled_optimizer=optimizer_config.compile if optimizer_config.compile is not None else False, + scaled_optm=optimizer_config.scaled_optm if optimizer_config.scaled_optm is not None else False, ) # LION_PRODIGY_ADV Optimizer diff --git a/modules/util/optimizer_util.py b/modules/util/optimizer_util.py index fa026bd7b..0c4fa31f7 100644 --- a/modules/util/optimizer_util.py +++ b/modules/util/optimizer_util.py @@ -462,6 +462,7 @@ def init_model_parameters( "alpha": 5, "kourkoutas_beta": False, "k_warmup_steps": None, + "scaled_optm": False, }, Optimizer.ADOPT_ADV: { "beta1": 0.9, @@ -484,6 +485,7 @@ def init_model_parameters( "alpha_grad": 100.0, "kourkoutas_beta": False, "k_warmup_steps": None, + "scaled_optm": False, }, Optimizer.PRODIGY_ADV: { "beta1": 0.9, @@ -531,6 +533,7 @@ def init_model_parameters( "orthogonal_gradient": False, "kourkoutas_beta": False, "k_warmup_steps": None, + "scaled_optm": False, }, Optimizer.SIGNSGD_ADV: { "momentum": 0.95, @@ -543,6 +546,7 @@ def init_model_parameters( "orthogonal_gradient": False, "Simplified_AdEMAMix": False, "alpha_grad": 100.0, + "scaled_optm": False, }, Optimizer.LION_ADV: { "beta1": 0.9, @@ -558,6 +562,7 @@ def init_model_parameters( "orthogonal_gradient": False, "kappa_p": 1.0, "auto_kappa_p": True, + "scaled_optm": False, }, Optimizer.LION_PRODIGY_ADV: { "beta1": 0.9, From 1ebe54d62baa479abcc6ddb80c43b114b2ddaef6 Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Thu, 26 Feb 2026 01:32:38 +0300 Subject: [PATCH 11/84] initial cwd, signed --- modules/ui/OptimizerParamsWindow.py | 4 ++++ modules/util/config/TrainConfig.py | 10 +++++++++- modules/util/create.py | 16 ++++++++++++++++ modules/util/enum/CenteredWDMode.py | 7 +++++++ modules/util/optimizer_util.py | 17 +++++++++++++++++ 5 files changed, 53 insertions(+), 1 deletion(-) create mode 100644 modules/util/enum/CenteredWDMode.py diff --git a/modules/ui/OptimizerParamsWindow.py b/modules/ui/OptimizerParamsWindow.py index e40bb44c1..d9173d9ef 100644 --- a/modules/ui/OptimizerParamsWindow.py +++ b/modules/ui/OptimizerParamsWindow.py @@ -201,6 +201,10 @@ def create_dynamic_ui( 'compile': {'title': 'Compiled Optimizer', 'tooltip': 'Enables PyTorch compilation for the optimizer internal step logic. This is intended to improve performance by allowing PyTorch to fuse operations and optimize the computational graph.', 'type': 'bool'}, 'spectral_normalization': {'title': 'Spectral Scaling', 'tooltip': 'Enables explicit Spectral Normalization to automatically rescale the update magnitude and Weight Decay based on layer dimensions. This allows hyperparameters to transfer seamlessly from small to large models without retuning, while making the optimizer highly robust to a wide range of learning rates. This ensures consistent performance across different model sizes, adapter methods, and ranks. ', 'type': 'bool'}, 'scaled_optm': {'title': 'Scaled Optimizer', 'tooltip': 'Automatically rescale the update magnitude and Weight Decay based on layer dimensions and type. This allows hyperparameters to transfer seamlessly from small to large models without retuning. This ensures consistent performance across different model sizes, adapter methods, and ranks. For LoRAs set alpha=rank. Using this, LoRA shares the same LR as full finetuning for all ranks.', 'type': 'bool'}, + 'freeze_on_flip': {'title': 'Freeze-on-Flip', 'tooltip': 'Projected One-hit freeze. Masks updates for coordinates where the gradient sign flips compared to the previous step. Over the steps, this makes the signed update semi-continuous.', 'type': 'bool'}, + 'l1_adaptive': {'title': 'L1 Adaptive', 'tooltip': 'Scales the learning rate dynamically by the L1 norm of the gradient to handle gradient heterogeneity. This makes the signed-optimizers semi-adaptive.', 'type': 'bool'}, + 'centered_wd': {'title': 'Centered Weight Decay', 'tooltip': 'Centered Weight Decay coefficient. Instead of decaying weights toward zero, they are decayed toward their initial values (anchors). This can be used together with standard weight decay.', 'type': 'float'}, + 'centered_wd_mode': {'title': 'Centered WD Mode', 'tooltip': """The quantization format used to store the anchor weights to save VRAM. Options include: 'full': Stores anchors in the original parameter's precision. 'float8': Uses torch.float8_e4m3fn for a balance of precision and memory. 'int8': Uses 8-bit block-wise quantization. 'int4': Uses 4-bit block-wise quantization.""", 'type': 'str'}, } # @formatter:on diff --git a/modules/util/config/TrainConfig.py b/modules/util/config/TrainConfig.py index 1a6d6c57a..263e66411 100644 --- a/modules/util/config/TrainConfig.py +++ b/modules/util/config/TrainConfig.py @@ -30,7 +30,7 @@ from modules.util.ModelNames import EmbeddingName, ModelNames from modules.util.ModelWeightDtypes import ModelWeightDtypes from modules.util.torch_util import default_device - +from modules.util.enum.CenteredWDMode import CenteredWDMode class TrainOptimizerConfig(BaseConfig): optimizer: Optimizer @@ -145,6 +145,10 @@ class TrainOptimizerConfig(BaseConfig): compile: False spectral_normalization: False scaled_optm: False + freeze_on_flip: False + l1_adaptive: False + centered_wd: float + centered_wd_mode: CenteredWDMode def __init__(self, data: list[(str, Any, type, bool)]): super().__init__(data) @@ -265,6 +269,10 @@ def default_values(): data.append(("compile", False, bool, False)) data.append(("spectral_normalization", False, bool, False)) data.append(("scaled_optm", False, bool, False)) + data.append(("freeze_on_flip", False, bool, False)) + data.append(("l1_adaptive", False, bool, False)) + data.append(("centered_wd", 0.0, float, False)) + data.append(("centered_wd_mode", CenteredWDMode.FLOAT8, CenteredWDMode, False)) return TrainOptimizerConfig(data) diff --git a/modules/util/create.py b/modules/util/create.py index fa65c2658..77e4895c7 100644 --- a/modules/util/create.py +++ b/modules/util/create.py @@ -687,6 +687,10 @@ def create_optimizer( k_warmup_steps=optimizer_config.k_warmup_steps if optimizer_config.k_warmup_steps is not None else 0, compiled_optimizer=optimizer_config.compile if optimizer_config.compile is not None else False, scaled_optm=optimizer_config.scaled_optm if optimizer_config.scaled_optm is not None else False, + centered_wd=optimizer_config.centered_wd if optimizer_config.centered_wd is not None else 0.0, + centered_wd_mode=optimizer_config.centered_wd_mode if optimizer_config.centered_wd_mode is not None else "full", + centered_wd=optimizer_config.centered_wd if optimizer_config.centered_wd is not None else 0.0, + centered_wd_mode=optimizer_config.centered_wd_mode if optimizer_config.centered_wd_mode is not None else "full", ) # ADOPT_ADV Optimizer @@ -715,6 +719,8 @@ def create_optimizer( k_warmup_steps=optimizer_config.k_warmup_steps if optimizer_config.k_warmup_steps is not None else 0, compiled_optimizer=optimizer_config.compile if optimizer_config.compile is not None else False, scaled_optm=optimizer_config.scaled_optm if optimizer_config.scaled_optm is not None else False, + centered_wd=optimizer_config.centered_wd if optimizer_config.centered_wd is not None else 0.0, + centered_wd_mode=optimizer_config.centered_wd_mode if optimizer_config.centered_wd_mode is not None else "full", ) # PRODIGY_ADV Optimizer @@ -773,6 +779,8 @@ def create_optimizer( k_warmup_steps=optimizer_config.k_warmup_steps if optimizer_config.k_warmup_steps is not None else 0, compiled_optimizer=optimizer_config.compile if optimizer_config.compile is not None else False, scaled_optm=optimizer_config.scaled_optm if optimizer_config.scaled_optm is not None else False, + centered_wd=optimizer_config.centered_wd if optimizer_config.centered_wd is not None else 0.0, + centered_wd_mode=optimizer_config.centered_wd_mode if optimizer_config.centered_wd_mode is not None else "full", ) # SignSGD_ADV Optimizer @@ -791,6 +799,10 @@ def create_optimizer( Simplified_AdEMAMix=optimizer_config.Simplified_AdEMAMix if optimizer_config.Simplified_AdEMAMix is not None else False, alpha_grad=optimizer_config.alpha_grad if optimizer_config.alpha_grad is not None else 100, scaled_optm=optimizer_config.scaled_optm if optimizer_config.scaled_optm is not None else False, + freeze_on_flip=optimizer_config.freeze_on_flip if optimizer_config.freeze_on_flip is not None else False, + l1_adaptive=optimizer_config.l1_adaptive if optimizer_config.l1_adaptive is not None else False, + centered_wd=optimizer_config.centered_wd if optimizer_config.centered_wd is not None else 0.0, + centered_wd_mode=optimizer_config.centered_wd_mode if optimizer_config.centered_wd_mode is not None else "full", ) # LION_ADV Optimizer @@ -811,7 +823,11 @@ def create_optimizer( kappa_p=optimizer_config.kappa_p if optimizer_config.kappa_p is not None else 1.0, auto_kappa_p=optimizer_config.auto_kappa_p if optimizer_config.auto_kappa_p is not None else False, compiled_optimizer=optimizer_config.compile if optimizer_config.compile is not None else False, + freeze_on_flip=optimizer_config.freeze_on_flip if optimizer_config.freeze_on_flip is not None else False, + l1_adaptive=optimizer_config.l1_adaptive if optimizer_config.l1_adaptive is not None else False, scaled_optm=optimizer_config.scaled_optm if optimizer_config.scaled_optm is not None else False, + centered_wd=optimizer_config.centered_wd if optimizer_config.centered_wd is not None else 0.0, + centered_wd_mode=optimizer_config.centered_wd_mode if optimizer_config.centered_wd_mode is not None else "full", ) # LION_PRODIGY_ADV Optimizer diff --git a/modules/util/enum/CenteredWDMode.py b/modules/util/enum/CenteredWDMode.py new file mode 100644 index 000000000..9eb62470f --- /dev/null +++ b/modules/util/enum/CenteredWDMode.py @@ -0,0 +1,7 @@ +from enum import Enum + +class CenteredWDMode(str, Enum): + FULL = "full" + FLOAT8 = "float8" + INT8 = "int8" + INT4 = "int4" \ No newline at end of file diff --git a/modules/util/optimizer_util.py b/modules/util/optimizer_util.py index 0c4fa31f7..2d4bd699a 100644 --- a/modules/util/optimizer_util.py +++ b/modules/util/optimizer_util.py @@ -5,6 +5,7 @@ from modules.util.enum.Optimizer import Optimizer from modules.util.NamedParameterGroup import NamedParameterGroupCollection from modules.util.torch_util import optimizer_to_device_ +from modules.util.enum.CenteredWDMode import CenteredWDMode import torch @@ -463,6 +464,8 @@ def init_model_parameters( "kourkoutas_beta": False, "k_warmup_steps": None, "scaled_optm": False, + "centered_wd": 0.0, + "centered_wd_mode": CenteredWDMode.FLOAT8, }, Optimizer.ADOPT_ADV: { "beta1": 0.9, @@ -486,6 +489,8 @@ def init_model_parameters( "kourkoutas_beta": False, "k_warmup_steps": None, "scaled_optm": False, + "centered_wd": 0.0, + "centered_wd_mode": CenteredWDMode.FLOAT8, }, Optimizer.PRODIGY_ADV: { "beta1": 0.9, @@ -515,6 +520,8 @@ def init_model_parameters( "alpha_grad": 100.0, "kourkoutas_beta": False, "k_warmup_steps": None, + "centered_wd": 0.0, + "centered_wd_mode": CenteredWDMode.FLOAT8, }, Optimizer.SIMPLIFIED_AdEMAMix: { "beta1": 0.99, @@ -534,6 +541,8 @@ def init_model_parameters( "kourkoutas_beta": False, "k_warmup_steps": None, "scaled_optm": False, + "centered_wd": 0.0, + "centered_wd_mode": CenteredWDMode.FLOAT8, }, Optimizer.SIGNSGD_ADV: { "momentum": 0.95, @@ -546,7 +555,11 @@ def init_model_parameters( "orthogonal_gradient": False, "Simplified_AdEMAMix": False, "alpha_grad": 100.0, + "freeze_on_flip": False, + "l1_adaptive": False, "scaled_optm": False, + "centered_wd": 0.0, + "centered_wd_mode": CenteredWDMode.FLOAT8, }, Optimizer.LION_ADV: { "beta1": 0.9, @@ -562,7 +575,11 @@ def init_model_parameters( "orthogonal_gradient": False, "kappa_p": 1.0, "auto_kappa_p": True, + "freeze_on_flip": False, + "l1_adaptive": False, "scaled_optm": False, + "centered_wd": 0.0, + "centered_wd_mode": CenteredWDMode.FLOAT8, }, Optimizer.LION_PRODIGY_ADV: { "beta1": 0.9, From 30f7b2843d208819bf07670ba4517b248eb67266 Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Thu, 26 Feb 2026 01:32:59 +0300 Subject: [PATCH 12/84] dev1 --- requirements-global.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements-global.txt b/requirements-global.txt index 7c3254adc..b37f91069 100644 --- a/requirements-global.txt +++ b/requirements-global.txt @@ -42,7 +42,7 @@ prodigyopt==1.1.2 # prodigy optimizer schedulefree==1.4.1 # schedule-free optimizers pytorch_optimizer==3.6.0 # pytorch optimizers prodigy-plus-schedule-free==2.0.1 # Prodigy plus optimizer -adv_optm==2.2.3 # advanced optimizers +adv_optm==2.4.dev1 # advanced optimizers -e git+https://github.com/KellerJordan/Muon.git@f90a42b#egg=muon-optimizer # Profiling From d72c03dfdceb69ff0b564e3e8da0b6fcb8f9b76b Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Thu, 26 Feb 2026 01:41:57 +0300 Subject: [PATCH 13/84] add factored_2nd --- modules/ui/OptimizerParamsWindow.py | 1 + modules/util/config/TrainConfig.py | 2 ++ modules/util/create.py | 4 ++++ modules/util/optimizer_util.py | 4 ++++ 4 files changed, 11 insertions(+) diff --git a/modules/ui/OptimizerParamsWindow.py b/modules/ui/OptimizerParamsWindow.py index d9173d9ef..7e80c706b 100644 --- a/modules/ui/OptimizerParamsWindow.py +++ b/modules/ui/OptimizerParamsWindow.py @@ -205,6 +205,7 @@ def create_dynamic_ui( 'l1_adaptive': {'title': 'L1 Adaptive', 'tooltip': 'Scales the learning rate dynamically by the L1 norm of the gradient to handle gradient heterogeneity. This makes the signed-optimizers semi-adaptive.', 'type': 'bool'}, 'centered_wd': {'title': 'Centered Weight Decay', 'tooltip': 'Centered Weight Decay coefficient. Instead of decaying weights toward zero, they are decayed toward their initial values (anchors). This can be used together with standard weight decay.', 'type': 'float'}, 'centered_wd_mode': {'title': 'Centered WD Mode', 'tooltip': """The quantization format used to store the anchor weights to save VRAM. Options include: 'full': Stores anchors in the original parameter's precision. 'float8': Uses torch.float8_e4m3fn for a balance of precision and memory. 'int8': Uses 8-bit block-wise quantization. 'int4': Uses 4-bit block-wise quantization.""", 'type': 'str'}, + 'factored_2nd': {'title': 'Factored 2nd', 'tooltip': 'Whether to keep the first moment uncompressed (dense), while only factorizing the second moment. This makes the optimizer highly robust to a wide range of LRs, mimicking high-order optimization.', 'type': 'bool'}, } # @formatter:on diff --git a/modules/util/config/TrainConfig.py b/modules/util/config/TrainConfig.py index 263e66411..69ea71de9 100644 --- a/modules/util/config/TrainConfig.py +++ b/modules/util/config/TrainConfig.py @@ -149,6 +149,7 @@ class TrainOptimizerConfig(BaseConfig): l1_adaptive: False centered_wd: float centered_wd_mode: CenteredWDMode + factored_2nd: False def __init__(self, data: list[(str, Any, type, bool)]): super().__init__(data) @@ -273,6 +274,7 @@ def default_values(): data.append(("l1_adaptive", False, bool, False)) data.append(("centered_wd", 0.0, float, False)) data.append(("centered_wd_mode", CenteredWDMode.FLOAT8, CenteredWDMode, False)) + data.append(("factored_2nd", False, bool, False)) return TrainOptimizerConfig(data) diff --git a/modules/util/create.py b/modules/util/create.py index 77e4895c7..910b4160c 100644 --- a/modules/util/create.py +++ b/modules/util/create.py @@ -673,6 +673,7 @@ def create_optimizer( eps=optimizer_config.eps if optimizer_config.eps is not None else 1e-8, weight_decay=optimizer_config.weight_decay if optimizer_config.weight_decay is not None else 0.0, use_bias_correction=optimizer_config.use_bias_correction if optimizer_config.use_bias_correction is not None else True, + factored_2nd=optimizer_config.factored_2nd if optimizer_config.factored_2nd is not None else False, nnmf_factor=optimizer_config.nnmf_factor if optimizer_config.nnmf_factor is not None else False, cautious_wd=optimizer_config.cautious_wd if optimizer_config.cautious_wd is not None else False, stochastic_rounding=optimizer_config.stochastic_rounding, @@ -704,6 +705,7 @@ def create_optimizer( eps=optimizer_config.eps if optimizer_config.eps is not None else 1e-6, weight_decay=optimizer_config.weight_decay if optimizer_config.weight_decay is not None else 0.0, nnmf_factor=optimizer_config.nnmf_factor if optimizer_config.nnmf_factor is not None else False, + factored_2nd=optimizer_config.factored_2nd if optimizer_config.factored_2nd is not None else False, cautious_wd=optimizer_config.cautious_wd if optimizer_config.cautious_wd is not None else False, stochastic_rounding=optimizer_config.stochastic_rounding, use_atan2=optimizer_config.use_atan2 if optimizer_config.use_atan2 is not None else False, @@ -735,6 +737,7 @@ def create_optimizer( eps=optimizer_config.eps if optimizer_config.eps is not None else 1e-8, weight_decay=optimizer_config.weight_decay if optimizer_config.weight_decay is not None else 0.0, nnmf_factor=optimizer_config.nnmf_factor if optimizer_config.nnmf_factor is not None else False, + factored_2nd=optimizer_config.factored_2nd if optimizer_config.factored_2nd is not None else False, cautious_wd=optimizer_config.cautious_wd if optimizer_config.cautious_wd is not None else False, stochastic_rounding=optimizer_config.stochastic_rounding, d0=optimizer_config.d0 if optimizer_config.d0 is not None else 1e-6, @@ -772,6 +775,7 @@ def create_optimizer( min_beta1=optimizer_config.min_beta1 if optimizer_config.min_beta1 is not None else 0.9, use_bias_correction=optimizer_config.use_bias_correction if optimizer_config.use_bias_correction is not None else True, nnmf_factor=optimizer_config.nnmf_factor if optimizer_config.nnmf_factor is not None else False, + factored_2nd=optimizer_config.factored_2nd if optimizer_config.factored_2nd is not None else False, cautious_wd=optimizer_config.cautious_wd if optimizer_config.cautious_wd is not None else False, stochastic_rounding=optimizer_config.stochastic_rounding, orthogonal_gradient=optimizer_config.orthogonal_gradient if optimizer_config.orthogonal_gradient is not None else False, diff --git a/modules/util/optimizer_util.py b/modules/util/optimizer_util.py index 2d4bd699a..91502311b 100644 --- a/modules/util/optimizer_util.py +++ b/modules/util/optimizer_util.py @@ -451,6 +451,7 @@ def init_model_parameters( "weight_decay": 0.0, "use_bias_correction": True, "nnmf_factor": False, + "factored_2nd": False, "stochastic_rounding": True, "compile": False, "fused_back_pass": False, @@ -474,6 +475,7 @@ def init_model_parameters( "cautious_wd": False, "weight_decay": 0.0, "nnmf_factor": False, + "factored_2nd": False, "stochastic_rounding": True, "compile": False, "fused_back_pass": False, @@ -500,6 +502,7 @@ def init_model_parameters( "cautious_wd": False, "weight_decay": 0.0, "nnmf_factor": False, + "factored_2nd": False, "stochastic_rounding": True, "compile": False, "fused_back_pass": False, @@ -534,6 +537,7 @@ def init_model_parameters( "min_beta1": 0.9, "use_bias_correction": True, "nnmf_factor": False, + "factored_2nd": False, "stochastic_rounding": True, "compile": False, "fused_back_pass": False, From 69f24178255faa106c94eac8ca710ab798d96230 Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Thu, 26 Feb 2026 01:56:55 +0300 Subject: [PATCH 14/84] pre-commit --- modules/util/config/TrainConfig.py | 3 ++- modules/util/create.py | 2 -- modules/util/enum/CenteredWDMode.py | 3 ++- modules/util/optimizer_util.py | 2 +- 4 files changed, 5 insertions(+), 5 deletions(-) diff --git a/modules/util/config/TrainConfig.py b/modules/util/config/TrainConfig.py index 69ea71de9..6cd9bc23b 100644 --- a/modules/util/config/TrainConfig.py +++ b/modules/util/config/TrainConfig.py @@ -10,6 +10,7 @@ from modules.util.config.SampleConfig import SampleConfig from modules.util.config.SecretsConfig import SecretsConfig from modules.util.enum.AudioFormat import AudioFormat +from modules.util.enum.CenteredWDMode import CenteredWDMode from modules.util.enum.ConfigPart import ConfigPart from modules.util.enum.DataType import DataType from modules.util.enum.EMAMode import EMAMode @@ -30,7 +31,7 @@ from modules.util.ModelNames import EmbeddingName, ModelNames from modules.util.ModelWeightDtypes import ModelWeightDtypes from modules.util.torch_util import default_device -from modules.util.enum.CenteredWDMode import CenteredWDMode + class TrainOptimizerConfig(BaseConfig): optimizer: Optimizer diff --git a/modules/util/create.py b/modules/util/create.py index 910b4160c..0fcc247d9 100644 --- a/modules/util/create.py +++ b/modules/util/create.py @@ -690,8 +690,6 @@ def create_optimizer( scaled_optm=optimizer_config.scaled_optm if optimizer_config.scaled_optm is not None else False, centered_wd=optimizer_config.centered_wd if optimizer_config.centered_wd is not None else 0.0, centered_wd_mode=optimizer_config.centered_wd_mode if optimizer_config.centered_wd_mode is not None else "full", - centered_wd=optimizer_config.centered_wd if optimizer_config.centered_wd is not None else 0.0, - centered_wd_mode=optimizer_config.centered_wd_mode if optimizer_config.centered_wd_mode is not None else "full", ) # ADOPT_ADV Optimizer diff --git a/modules/util/enum/CenteredWDMode.py b/modules/util/enum/CenteredWDMode.py index 9eb62470f..a3ebb19d1 100644 --- a/modules/util/enum/CenteredWDMode.py +++ b/modules/util/enum/CenteredWDMode.py @@ -1,7 +1,8 @@ from enum import Enum + class CenteredWDMode(str, Enum): FULL = "full" FLOAT8 = "float8" INT8 = "int8" - INT4 = "int4" \ No newline at end of file + INT4 = "int4" diff --git a/modules/util/optimizer_util.py b/modules/util/optimizer_util.py index 91502311b..16baaa342 100644 --- a/modules/util/optimizer_util.py +++ b/modules/util/optimizer_util.py @@ -2,10 +2,10 @@ from modules.model.BaseModel import BaseModel from modules.util import create from modules.util.config.TrainConfig import TrainConfig, TrainOptimizerConfig +from modules.util.enum.CenteredWDMode import CenteredWDMode from modules.util.enum.Optimizer import Optimizer from modules.util.NamedParameterGroup import NamedParameterGroupCollection from modules.util.torch_util import optimizer_to_device_ -from modules.util.enum.CenteredWDMode import CenteredWDMode import torch From 06c9e6e1fd5a536e59528fdaed534b3225ac2aa4 Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Thu, 26 Feb 2026 02:09:00 +0300 Subject: [PATCH 15/84] fix CenteredWDMode --- modules/ui/OptimizerParamsWindow.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/modules/ui/OptimizerParamsWindow.py b/modules/ui/OptimizerParamsWindow.py index 7e80c706b..e6f274ee8 100644 --- a/modules/ui/OptimizerParamsWindow.py +++ b/modules/ui/OptimizerParamsWindow.py @@ -3,6 +3,7 @@ from modules.ui.MuonAdamWindow import MUON_AUX_ADAM_DEFAULTS, MuonAdamWindow from modules.util.config.TrainConfig import TrainConfig, TrainOptimizerConfig +from modules.util.enum.CenteredWDMode import CenteredWDMode from modules.util.enum.Optimizer import Optimizer from modules.util.optimizer_util import ( OPTIMIZER_DEFAULT_PARAMETERS, @@ -204,7 +205,7 @@ def create_dynamic_ui( 'freeze_on_flip': {'title': 'Freeze-on-Flip', 'tooltip': 'Projected One-hit freeze. Masks updates for coordinates where the gradient sign flips compared to the previous step. Over the steps, this makes the signed update semi-continuous.', 'type': 'bool'}, 'l1_adaptive': {'title': 'L1 Adaptive', 'tooltip': 'Scales the learning rate dynamically by the L1 norm of the gradient to handle gradient heterogeneity. This makes the signed-optimizers semi-adaptive.', 'type': 'bool'}, 'centered_wd': {'title': 'Centered Weight Decay', 'tooltip': 'Centered Weight Decay coefficient. Instead of decaying weights toward zero, they are decayed toward their initial values (anchors). This can be used together with standard weight decay.', 'type': 'float'}, - 'centered_wd_mode': {'title': 'Centered WD Mode', 'tooltip': """The quantization format used to store the anchor weights to save VRAM. Options include: 'full': Stores anchors in the original parameter's precision. 'float8': Uses torch.float8_e4m3fn for a balance of precision and memory. 'int8': Uses 8-bit block-wise quantization. 'int4': Uses 4-bit block-wise quantization.""", 'type': 'str'}, + 'centered_wd_mode': {'title': 'Centered WD Mode', 'tooltip': """The quantization format used to store the anchor weights to save VRAM. Options include: 'full': Stores anchors in the original parameter's precision. 'float8': Uses torch.float8_e4m3fn for a balance of precision and memory. 'int8': Uses 8-bit block-wise quantization. 'int4': Uses 4-bit block-wise quantization.""", 'type': 'CenteredWDMode'}, 'factored_2nd': {'title': 'Factored 2nd', 'tooltip': 'Whether to keep the first moment uncompressed (dense), while only factorizing the second moment. This makes the optimizer highly robust to a wide range of LRs, mimicking high-order optimization.', 'type': 'bool'}, } # @formatter:on @@ -242,6 +243,9 @@ def create_dynamic_ui( tooltip="Configure the auxiliary AdamW_adv optimizer", width=20, padx=5 ) self.toggle_muon_adam_button() + elif type == 'CenteredWDMode': + components.options(master, row, col + 1, [x.value for x in CenteredWDMode], self.optimizer_ui_state, key, + command=self.update_user_pref) elif type != 'bool': components.entry(master, row, col + 1, self.optimizer_ui_state, key, command=self.update_user_pref) From 6d623733a8248df8be3f9f9965e0e443a3db201a Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Thu, 26 Feb 2026 23:16:32 +0300 Subject: [PATCH 16/84] maybe fix --- modules/ui/OptimizerParamsWindow.py | 2 +- modules/util/enum/CenteredWDMode.py | 3 +++ 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/modules/ui/OptimizerParamsWindow.py b/modules/ui/OptimizerParamsWindow.py index e6f274ee8..755a40876 100644 --- a/modules/ui/OptimizerParamsWindow.py +++ b/modules/ui/OptimizerParamsWindow.py @@ -244,7 +244,7 @@ def create_dynamic_ui( width=20, padx=5 ) self.toggle_muon_adam_button() elif type == 'CenteredWDMode': - components.options(master, row, col + 1, [x.value for x in CenteredWDMode], self.optimizer_ui_state, key, + components.options(master, row, col + 1, [str(x) for x in CenteredWDMode], self.optimizer_ui_state, key, command=self.update_user_pref) elif type != 'bool': components.entry(master, row, col + 1, self.optimizer_ui_state, key, diff --git a/modules/util/enum/CenteredWDMode.py b/modules/util/enum/CenteredWDMode.py index a3ebb19d1..69e926720 100644 --- a/modules/util/enum/CenteredWDMode.py +++ b/modules/util/enum/CenteredWDMode.py @@ -6,3 +6,6 @@ class CenteredWDMode(str, Enum): FLOAT8 = "float8" INT8 = "int8" INT4 = "int4" + + def __str__(self): + return self.name \ No newline at end of file From 32ef49a5d290aa26ded45193537db6b61cc8f53a Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Fri, 27 Feb 2026 00:43:21 +0300 Subject: [PATCH 17/84] dev2 --- requirements-global.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements-global.txt b/requirements-global.txt index b37f91069..b23d950a5 100644 --- a/requirements-global.txt +++ b/requirements-global.txt @@ -42,7 +42,7 @@ prodigyopt==1.1.2 # prodigy optimizer schedulefree==1.4.1 # schedule-free optimizers pytorch_optimizer==3.6.0 # pytorch optimizers prodigy-plus-schedule-free==2.0.1 # Prodigy plus optimizer -adv_optm==2.4.dev1 # advanced optimizers +adv_optm==2.4.dev2 # advanced optimizers -e git+https://github.com/KellerJordan/Muon.git@f90a42b#egg=muon-optimizer # Profiling From 852f389ec0bc32b61022aa37a76b031c8d1e907e Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Fri, 27 Feb 2026 00:44:12 +0300 Subject: [PATCH 18/84] pre-commit --- modules/util/enum/CenteredWDMode.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modules/util/enum/CenteredWDMode.py b/modules/util/enum/CenteredWDMode.py index 69e926720..9f1c280b7 100644 --- a/modules/util/enum/CenteredWDMode.py +++ b/modules/util/enum/CenteredWDMode.py @@ -8,4 +8,4 @@ class CenteredWDMode(str, Enum): INT4 = "int4" def __str__(self): - return self.name \ No newline at end of file + return self.name From add186dee7e5f385586cf213b6f89242ee7d5345 Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Sat, 28 Feb 2026 03:58:48 +0300 Subject: [PATCH 19/84] fix and remove CenteredWDMode enum --- modules/ui/OptimizerParamsWindow.py | 3 +-- modules/util/config/TrainConfig.py | 5 ++--- modules/util/enum/CenteredWDMode.py | 11 ----------- modules/util/optimizer_util.py | 13 ++++++------- 4 files changed, 9 insertions(+), 23 deletions(-) delete mode 100644 modules/util/enum/CenteredWDMode.py diff --git a/modules/ui/OptimizerParamsWindow.py b/modules/ui/OptimizerParamsWindow.py index 755a40876..ded081821 100644 --- a/modules/ui/OptimizerParamsWindow.py +++ b/modules/ui/OptimizerParamsWindow.py @@ -3,7 +3,6 @@ from modules.ui.MuonAdamWindow import MUON_AUX_ADAM_DEFAULTS, MuonAdamWindow from modules.util.config.TrainConfig import TrainConfig, TrainOptimizerConfig -from modules.util.enum.CenteredWDMode import CenteredWDMode from modules.util.enum.Optimizer import Optimizer from modules.util.optimizer_util import ( OPTIMIZER_DEFAULT_PARAMETERS, @@ -244,7 +243,7 @@ def create_dynamic_ui( width=20, padx=5 ) self.toggle_muon_adam_button() elif type == 'CenteredWDMode': - components.options(master, row, col + 1, [str(x) for x in CenteredWDMode], self.optimizer_ui_state, key, + components.options(master, row, col + 1, ["full", "float8", "int8", "int4"], self.optimizer_ui_state, key, command=self.update_user_pref) elif type != 'bool': components.entry(master, row, col + 1, self.optimizer_ui_state, key, diff --git a/modules/util/config/TrainConfig.py b/modules/util/config/TrainConfig.py index 6cd9bc23b..300487870 100644 --- a/modules/util/config/TrainConfig.py +++ b/modules/util/config/TrainConfig.py @@ -10,7 +10,6 @@ from modules.util.config.SampleConfig import SampleConfig from modules.util.config.SecretsConfig import SecretsConfig from modules.util.enum.AudioFormat import AudioFormat -from modules.util.enum.CenteredWDMode import CenteredWDMode from modules.util.enum.ConfigPart import ConfigPart from modules.util.enum.DataType import DataType from modules.util.enum.EMAMode import EMAMode @@ -149,7 +148,7 @@ class TrainOptimizerConfig(BaseConfig): freeze_on_flip: False l1_adaptive: False centered_wd: float - centered_wd_mode: CenteredWDMode + centered_wd_mode: str factored_2nd: False def __init__(self, data: list[(str, Any, type, bool)]): @@ -274,7 +273,7 @@ def default_values(): data.append(("freeze_on_flip", False, bool, False)) data.append(("l1_adaptive", False, bool, False)) data.append(("centered_wd", 0.0, float, False)) - data.append(("centered_wd_mode", CenteredWDMode.FLOAT8, CenteredWDMode, False)) + data.append(("centered_wd_mode", "float8", str, False)) data.append(("factored_2nd", False, bool, False)) return TrainOptimizerConfig(data) diff --git a/modules/util/enum/CenteredWDMode.py b/modules/util/enum/CenteredWDMode.py deleted file mode 100644 index 9f1c280b7..000000000 --- a/modules/util/enum/CenteredWDMode.py +++ /dev/null @@ -1,11 +0,0 @@ -from enum import Enum - - -class CenteredWDMode(str, Enum): - FULL = "full" - FLOAT8 = "float8" - INT8 = "int8" - INT4 = "int4" - - def __str__(self): - return self.name diff --git a/modules/util/optimizer_util.py b/modules/util/optimizer_util.py index 16baaa342..9ff538844 100644 --- a/modules/util/optimizer_util.py +++ b/modules/util/optimizer_util.py @@ -2,7 +2,6 @@ from modules.model.BaseModel import BaseModel from modules.util import create from modules.util.config.TrainConfig import TrainConfig, TrainOptimizerConfig -from modules.util.enum.CenteredWDMode import CenteredWDMode from modules.util.enum.Optimizer import Optimizer from modules.util.NamedParameterGroup import NamedParameterGroupCollection from modules.util.torch_util import optimizer_to_device_ @@ -466,7 +465,7 @@ def init_model_parameters( "k_warmup_steps": None, "scaled_optm": False, "centered_wd": 0.0, - "centered_wd_mode": CenteredWDMode.FLOAT8, + "centered_wd_mode": "float8", }, Optimizer.ADOPT_ADV: { "beta1": 0.9, @@ -492,7 +491,7 @@ def init_model_parameters( "k_warmup_steps": None, "scaled_optm": False, "centered_wd": 0.0, - "centered_wd_mode": CenteredWDMode.FLOAT8, + "centered_wd_mode": "float8", }, Optimizer.PRODIGY_ADV: { "beta1": 0.9, @@ -524,7 +523,7 @@ def init_model_parameters( "kourkoutas_beta": False, "k_warmup_steps": None, "centered_wd": 0.0, - "centered_wd_mode": CenteredWDMode.FLOAT8, + "centered_wd_mode": "float8", }, Optimizer.SIMPLIFIED_AdEMAMix: { "beta1": 0.99, @@ -546,7 +545,7 @@ def init_model_parameters( "k_warmup_steps": None, "scaled_optm": False, "centered_wd": 0.0, - "centered_wd_mode": CenteredWDMode.FLOAT8, + "centered_wd_mode": "float8", }, Optimizer.SIGNSGD_ADV: { "momentum": 0.95, @@ -563,7 +562,7 @@ def init_model_parameters( "l1_adaptive": False, "scaled_optm": False, "centered_wd": 0.0, - "centered_wd_mode": CenteredWDMode.FLOAT8, + "centered_wd_mode": "float8", }, Optimizer.LION_ADV: { "beta1": 0.9, @@ -583,7 +582,7 @@ def init_model_parameters( "l1_adaptive": False, "scaled_optm": False, "centered_wd": 0.0, - "centered_wd_mode": CenteredWDMode.FLOAT8, + "centered_wd_mode": "float8", }, Optimizer.LION_PRODIGY_ADV: { "beta1": 0.9, From 8e55171ffeb30f0075d46e4d9bdefb3a1efc7691 Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Sun, 1 Mar 2026 18:04:41 +0300 Subject: [PATCH 20/84] dev4 --- requirements-global.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements-global.txt b/requirements-global.txt index b23d950a5..36bead3f7 100644 --- a/requirements-global.txt +++ b/requirements-global.txt @@ -42,7 +42,7 @@ prodigyopt==1.1.2 # prodigy optimizer schedulefree==1.4.1 # schedule-free optimizers pytorch_optimizer==3.6.0 # pytorch optimizers prodigy-plus-schedule-free==2.0.1 # Prodigy plus optimizer -adv_optm==2.4.dev2 # advanced optimizers +adv_optm==2.4.dev4 # advanced optimizers -e git+https://github.com/KellerJordan/Muon.git@f90a42b#egg=muon-optimizer # Profiling From d15d91577694dfa8c6b72ece3ecabc7f5f7e4c3e Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Mon, 16 Mar 2026 19:21:54 +0300 Subject: [PATCH 21/84] Dev5: Add Fisher WD to Adam-variants --- modules/ui/OptimizerParamsWindow.py | 1 + modules/util/config/TrainConfig.py | 2 ++ modules/util/create.py | 3 +++ modules/util/optimizer_util.py | 3 +++ requirements-global.txt | 2 +- 5 files changed, 10 insertions(+), 1 deletion(-) diff --git a/modules/ui/OptimizerParamsWindow.py b/modules/ui/OptimizerParamsWindow.py index ded081821..e7910bb39 100644 --- a/modules/ui/OptimizerParamsWindow.py +++ b/modules/ui/OptimizerParamsWindow.py @@ -206,6 +206,7 @@ def create_dynamic_ui( 'centered_wd': {'title': 'Centered Weight Decay', 'tooltip': 'Centered Weight Decay coefficient. Instead of decaying weights toward zero, they are decayed toward their initial values (anchors). This can be used together with standard weight decay.', 'type': 'float'}, 'centered_wd_mode': {'title': 'Centered WD Mode', 'tooltip': """The quantization format used to store the anchor weights to save VRAM. Options include: 'full': Stores anchors in the original parameter's precision. 'float8': Uses torch.float8_e4m3fn for a balance of precision and memory. 'int8': Uses 8-bit block-wise quantization. 'int4': Uses 4-bit block-wise quantization.""", 'type': 'CenteredWDMode'}, 'factored_2nd': {'title': 'Factored 2nd', 'tooltip': 'Whether to keep the first moment uncompressed (dense), while only factorizing the second moment. This makes the optimizer highly robust to a wide range of LRs, mimicking high-order optimization.', 'type': 'bool'}, + 'fisher_wd': {'title': 'Fisher Weight Decay', 'tooltip': 'Applies adaptive, scale-invariant weight-decay regularization based on the Fisher Information Matrix (approximated by Adam\'s second moment). It reduces penalty for "important" high-curvature weights while accelerating decay for "useless" weights in flat regions. Leading to improved convergence and better final performance.', 'type': 'bool'}, } # @formatter:on diff --git a/modules/util/config/TrainConfig.py b/modules/util/config/TrainConfig.py index 300487870..3d678dc78 100644 --- a/modules/util/config/TrainConfig.py +++ b/modules/util/config/TrainConfig.py @@ -150,6 +150,7 @@ class TrainOptimizerConfig(BaseConfig): centered_wd: float centered_wd_mode: str factored_2nd: False + fisher_wd: False def __init__(self, data: list[(str, Any, type, bool)]): super().__init__(data) @@ -275,6 +276,7 @@ def default_values(): data.append(("centered_wd", 0.0, float, False)) data.append(("centered_wd_mode", "float8", str, False)) data.append(("factored_2nd", False, bool, False)) + data.append(("fisher_wd", False, bool, False)) return TrainOptimizerConfig(data) diff --git a/modules/util/create.py b/modules/util/create.py index 0fcc247d9..26ce8eb48 100644 --- a/modules/util/create.py +++ b/modules/util/create.py @@ -675,6 +675,7 @@ def create_optimizer( use_bias_correction=optimizer_config.use_bias_correction if optimizer_config.use_bias_correction is not None else True, factored_2nd=optimizer_config.factored_2nd if optimizer_config.factored_2nd is not None else False, nnmf_factor=optimizer_config.nnmf_factor if optimizer_config.nnmf_factor is not None else False, + fisher_wd=optimizer_config.fisher_wd if optimizer_config.fisher_wd is not None else False, cautious_wd=optimizer_config.cautious_wd if optimizer_config.cautious_wd is not None else False, stochastic_rounding=optimizer_config.stochastic_rounding, use_atan2=optimizer_config.use_atan2 if optimizer_config.use_atan2 is not None else False, @@ -704,6 +705,7 @@ def create_optimizer( weight_decay=optimizer_config.weight_decay if optimizer_config.weight_decay is not None else 0.0, nnmf_factor=optimizer_config.nnmf_factor if optimizer_config.nnmf_factor is not None else False, factored_2nd=optimizer_config.factored_2nd if optimizer_config.factored_2nd is not None else False, + fisher_wd=optimizer_config.fisher_wd if optimizer_config.fisher_wd is not None else False, cautious_wd=optimizer_config.cautious_wd if optimizer_config.cautious_wd is not None else False, stochastic_rounding=optimizer_config.stochastic_rounding, use_atan2=optimizer_config.use_atan2 if optimizer_config.use_atan2 is not None else False, @@ -736,6 +738,7 @@ def create_optimizer( weight_decay=optimizer_config.weight_decay if optimizer_config.weight_decay is not None else 0.0, nnmf_factor=optimizer_config.nnmf_factor if optimizer_config.nnmf_factor is not None else False, factored_2nd=optimizer_config.factored_2nd if optimizer_config.factored_2nd is not None else False, + fisher_wd=optimizer_config.fisher_wd if optimizer_config.fisher_wd is not None else False, cautious_wd=optimizer_config.cautious_wd if optimizer_config.cautious_wd is not None else False, stochastic_rounding=optimizer_config.stochastic_rounding, d0=optimizer_config.d0 if optimizer_config.d0 is not None else 1e-6, diff --git a/modules/util/optimizer_util.py b/modules/util/optimizer_util.py index 9ff538844..1c6d9a6d1 100644 --- a/modules/util/optimizer_util.py +++ b/modules/util/optimizer_util.py @@ -446,6 +446,7 @@ def init_model_parameters( "beta1": 0.9, "beta2": 0.99, "eps": 1e-8, + "fisher_wd": False, "cautious_wd": False, "weight_decay": 0.0, "use_bias_correction": True, @@ -471,6 +472,7 @@ def init_model_parameters( "beta1": 0.9, "beta2": 0.9999, "eps": 1e-6, + "fisher_wd": False, "cautious_wd": False, "weight_decay": 0.0, "nnmf_factor": False, @@ -498,6 +500,7 @@ def init_model_parameters( "beta2": 0.99, "beta3": None, "eps": 1e-8, + "fisher_wd": False, "cautious_wd": False, "weight_decay": 0.0, "nnmf_factor": False, diff --git a/requirements-global.txt b/requirements-global.txt index 36bead3f7..0b54f9560 100644 --- a/requirements-global.txt +++ b/requirements-global.txt @@ -42,7 +42,7 @@ prodigyopt==1.1.2 # prodigy optimizer schedulefree==1.4.1 # schedule-free optimizers pytorch_optimizer==3.6.0 # pytorch optimizers prodigy-plus-schedule-free==2.0.1 # Prodigy plus optimizer -adv_optm==2.4.dev4 # advanced optimizers +adv_optm==2.4.dev5 # advanced optimizers -e git+https://github.com/KellerJordan/Muon.git@f90a42b#egg=muon-optimizer # Profiling From b9ece3443ce26f63494d71861ca03378bb721344 Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Mon, 16 Mar 2026 20:46:21 +0300 Subject: [PATCH 22/84] depth_calculator --- modules/util/create.py | 5 ++ modules/util/optimizer/depth_calculator.py | 74 ++++++++++++++++++++++ 2 files changed, 79 insertions(+) create mode 100644 modules/util/optimizer/depth_calculator.py diff --git a/modules/util/create.py b/modules/util/create.py index 26ce8eb48..36351a338 100644 --- a/modules/util/create.py +++ b/modules/util/create.py @@ -33,6 +33,7 @@ from modules.util.optimizer.adafactor_extensions import patch_adafactor from modules.util.optimizer.adam_extensions import patch_adam from modules.util.optimizer.adamw_extensions import patch_adamw +from modules.util.optimizer.depth_calculator import inject_depth_into_param_groups from modules.util.optimizer.muon_util import calculate_muon_n_layers, split_parameters_for_muon from modules.util.TrainProgress import TrainProgress from modules.zluda import ZLUDA @@ -140,6 +141,10 @@ def create_optimizer( parameters = parameter_group_collection.parameters_for_optimizer(config) + if optimizer_config.scaled_optm or optimizer_config.spectral_normalization: + # _adv optimizers O(1) depth scaling. + inject_depth_into_param_groups(model, parameters) + match config.optimizer.optimizer: # SGD Optimizer diff --git a/modules/util/optimizer/depth_calculator.py b/modules/util/optimizer/depth_calculator.py new file mode 100644 index 000000000..ed22fc040 --- /dev/null +++ b/modules/util/optimizer/depth_calculator.py @@ -0,0 +1,74 @@ +import re + +from modules.model.BaseModel import BaseModel +from modules.module.LoRAModule import LoRAModuleWrapper +from modules.util.enum.ModelType import ModelType + +import torch + + +def calculate_n_layers(model: BaseModel) -> dict[str, int]: + """ + Calculates the number of residual layers (the depth) in each component of the model. + """ + match model.model_type: + case (ModelType.STABLE_DIFFUSION_15 | ModelType.STABLE_DIFFUSION_15_INPAINTING | + ModelType.STABLE_DIFFUSION_20_BASE | ModelType.STABLE_DIFFUSION_20_INPAINTING | + ModelType.STABLE_DIFFUSION_20 | ModelType.STABLE_DIFFUSION_21 | + ModelType.STABLE_DIFFUSION_21_BASE | ModelType.STABLE_DIFFUSION_XL_10_BASE | + ModelType.STABLE_DIFFUSION_XL_10_BASE_INPAINTING | ModelType.STABLE_CASCADE_1 | + ModelType.WUERSTCHEN_2): + default_patterns = ['transformer_blocks', 'resnets', 'layers'] + case (ModelType.STABLE_DIFFUSION_3 | ModelType.STABLE_DIFFUSION_35 | ModelType.SANA | + ModelType.FLUX_DEV_1 | ModelType.FLUX_2 | ModelType.CHROMA_1 | ModelType.QWEN | + ModelType.PIXART_ALPHA | ModelType.PIXART_SIGMA): + default_patterns = ['transformer_blocks', 'single_transformer_blocks', 'encoder.block'] + case ModelType.HI_DREAM_FULL: + default_patterns = ['double_stream_blocks', 'single_stream_blocks'] + case ModelType.Z_IMAGE: + default_patterns = [ + 'layers', + 'refiner', + ] + case _: + raise NotImplementedError(f"Scaled Optimizer is not implemented for model type: {model.model_type}") + + # Build the regex pattern + joined_patterns = "|".join([re.escape(p) for p in default_patterns]) + pattern = re.compile(rf'(?:^|\.)(?:{joined_patterns})\.\d+$') + + layer_counts = {} + + # Iterate over model components (e.g., 'unet', 'text_encoder', 'transformer') + for attr_name, module in vars(model).items(): + # Identify the 'Ground Truth' blocks in this component. + target_module = module + if isinstance(module, LoRAModuleWrapper): + target_module = module.orig_module + valid_component_blocks = set() + if isinstance(target_module, torch.nn.Module): + for name, _ in target_module.named_modules(): + if pattern.search(name): + valid_component_blocks.add(name) + if not valid_component_blocks: + continue + count = len(valid_component_blocks) + if count > 0: + layer_counts[attr_name] = count + + return layer_counts + + +def inject_depth_into_param_groups(model: BaseModel, parameters): + """ + Calculates the model depth and injects the 'n_layers' key directly + into the optimizer parameter groups. + """ + n_layers_map = calculate_n_layers(model) + + for group in parameters: + group_name = group.get('name') + if group_name in n_layers_map: + group['n_layers'] = n_layers_map[group_name] + else: + group['n_layers'] = n_layers_map.get('default', 1) From 89ce64fc0463dd5ad23b60a5f7961443cf143c02 Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Mon, 16 Mar 2026 20:48:21 +0300 Subject: [PATCH 23/84] remove calculate_muon_n_layers --- modules/util/create.py | 22 +-------- modules/util/optimizer/muon_util.py | 71 ----------------------------- 2 files changed, 1 insertion(+), 92 deletions(-) diff --git a/modules/util/create.py b/modules/util/create.py index 36351a338..7aea42818 100644 --- a/modules/util/create.py +++ b/modules/util/create.py @@ -34,7 +34,7 @@ from modules.util.optimizer.adam_extensions import patch_adam from modules.util.optimizer.adamw_extensions import patch_adamw from modules.util.optimizer.depth_calculator import inject_depth_into_param_groups -from modules.util.optimizer.muon_util import calculate_muon_n_layers, split_parameters_for_muon +from modules.util.optimizer.muon_util import split_parameters_for_muon from modules.util.TrainProgress import TrainProgress from modules.zluda import ZLUDA @@ -875,16 +875,6 @@ def create_optimizer( params_for_optimizer, MuonWithAuxAdam = split_parameters_for_muon(model, parameters, config) - if optimizer_config.spectral_normalization: - # Calculate n_layers for spectral normalization - n_layers_map = calculate_muon_n_layers(model) - - for group in params_for_optimizer: - group_name = group.get('name') - if group_name in n_layers_map: - group['n_layers'] = n_layers_map[group_name] - else: - group['n_layers'] = n_layers_map.get('default', 1) # Prepare Adam-specific keyword arguments from the config adam_kwargs = {} @@ -939,16 +929,6 @@ def create_optimizer( params_for_optimizer, MuonWithAuxAdam = split_parameters_for_muon(model, parameters, config) - if optimizer_config.spectral_normalization: - # Calculate n_layers for spectral normalization - n_layers_map = calculate_muon_n_layers(model) - - for group in params_for_optimizer: - group_name = group.get('name') - if group_name in n_layers_map: - group['n_layers'] = n_layers_map[group_name] - else: - group['n_layers'] = n_layers_map.get('default', 1) # Prepare Adam-specific keyword arguments from the config adam_kwargs = {} diff --git a/modules/util/optimizer/muon_util.py b/modules/util/optimizer/muon_util.py index 1c0230a8a..e230b6f7d 100644 --- a/modules/util/optimizer/muon_util.py +++ b/modules/util/optimizer/muon_util.py @@ -1,4 +1,3 @@ -import re from collections.abc import Callable from modules.model.BaseModel import BaseModel, TrainConfig @@ -9,76 +8,6 @@ import torch -def calculate_muon_n_layers(model: BaseModel) -> dict[str, int]: - """ - Calculates the number of residual layers (the depth) in each component of the model. - Used for Muon optimizer spectral normalization scaling. - """ - match model.model_type: - case (ModelType.STABLE_DIFFUSION_15 | ModelType.STABLE_DIFFUSION_15_INPAINTING | - ModelType.STABLE_DIFFUSION_20_BASE | ModelType.STABLE_DIFFUSION_20_INPAINTING | - ModelType.STABLE_DIFFUSION_20 | ModelType.STABLE_DIFFUSION_21 | - ModelType.STABLE_DIFFUSION_21_BASE | ModelType.STABLE_DIFFUSION_XL_10_BASE | - ModelType.STABLE_DIFFUSION_XL_10_BASE_INPAINTING | ModelType.STABLE_CASCADE_1 | - ModelType.WUERSTCHEN_2): - default_patterns = ['transformer_blocks', 'resnets', 'layers'] - case (ModelType.STABLE_DIFFUSION_3 | ModelType.STABLE_DIFFUSION_35 | ModelType.SANA | - ModelType.FLUX_DEV_1 | ModelType.CHROMA_1 | ModelType.QWEN | - ModelType.PIXART_ALPHA | ModelType.PIXART_SIGMA): - default_patterns = ['transformer_blocks', 'single_transformer_blocks', 'encoder.block'] - case ModelType.HI_DREAM_FULL: - default_patterns = ['double_stream_blocks', 'single_stream_blocks'] - case ModelType.Z_IMAGE: - default_patterns = [ - 'layers', - 'refiner', - ] - case _: - raise NotImplementedError(f"Muon optimizer spectral normalization is not implemented for model type: {model.model_type}") - - # Build the regex pattern dynamically - joined_patterns = "|".join([re.escape(p) for p in default_patterns]) - pattern = re.compile(rf'(?:^|\.)(?:{joined_patterns})\.\d+$') - - layer_counts = {} - - # Iterate over model components (e.g., 'unet', 'text_encoder', 'transformer') - for module in vars(model).values(): - - # Identify the 'Ground Truth' blocks in this component. - target_module = module - if isinstance(module, LoRAModuleWrapper): - target_module = module.orig_module - - valid_component_blocks = set() - if isinstance(target_module, torch.nn.Module): - for name, _ in target_module.named_modules(): - if pattern.search(name): - valid_component_blocks.add(name) - - if not valid_component_blocks: - continue - - active_component_blocks = set() - - # Filter: Only count blocks that are actually being trained. - if isinstance(module, LoRAModuleWrapper): - # For LoRA, we check if the active leaves reside inside a valid block. - for layer_name in module.lora_modules: - parts = layer_name.split('.') - for i in range(len(parts), 0, -1): - candidate = ".".join(parts[:i]) - if candidate in valid_component_blocks: - active_component_blocks.add(candidate) - - elif isinstance(module, torch.nn.Module): - # For standard full-finetuning, all valid blocks are active. - active_component_blocks = valid_component_blocks - - - return layer_counts - - def build_muon_adam_key_fn( model: BaseModel, config: TrainConfig, From fd669d3c72bccd33f81f85a67a82085e9c28baf9 Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Tue, 17 Mar 2026 21:23:36 +0300 Subject: [PATCH 24/84] dev6: add scaled eps --- requirements-global.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements-global.txt b/requirements-global.txt index 0b54f9560..c17c3a9e3 100644 --- a/requirements-global.txt +++ b/requirements-global.txt @@ -42,7 +42,7 @@ prodigyopt==1.1.2 # prodigy optimizer schedulefree==1.4.1 # schedule-free optimizers pytorch_optimizer==3.6.0 # pytorch optimizers prodigy-plus-schedule-free==2.0.1 # Prodigy plus optimizer -adv_optm==2.4.dev5 # advanced optimizers +adv_optm==2.4.dev6 # advanced optimizers -e git+https://github.com/KellerJordan/Muon.git@f90a42b#egg=muon-optimizer # Profiling From 7db66b51f4620f0e79dd98ec8348556bf1d3b17c Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Sun, 12 Apr 2026 00:22:52 +0300 Subject: [PATCH 25/84] add StatePrecision --- modules/ui/OptimizerParamsWindow.py | 4 ++++ modules/util/config/TrainConfig.py | 2 ++ modules/util/create.py | 3 +++ 3 files changed, 9 insertions(+) diff --git a/modules/ui/OptimizerParamsWindow.py b/modules/ui/OptimizerParamsWindow.py index e7910bb39..0112ad2f1 100644 --- a/modules/ui/OptimizerParamsWindow.py +++ b/modules/ui/OptimizerParamsWindow.py @@ -207,6 +207,7 @@ def create_dynamic_ui( 'centered_wd_mode': {'title': 'Centered WD Mode', 'tooltip': """The quantization format used to store the anchor weights to save VRAM. Options include: 'full': Stores anchors in the original parameter's precision. 'float8': Uses torch.float8_e4m3fn for a balance of precision and memory. 'int8': Uses 8-bit block-wise quantization. 'int4': Uses 4-bit block-wise quantization.""", 'type': 'CenteredWDMode'}, 'factored_2nd': {'title': 'Factored 2nd', 'tooltip': 'Whether to keep the first moment uncompressed (dense), while only factorizing the second moment. This makes the optimizer highly robust to a wide range of LRs, mimicking high-order optimization.', 'type': 'bool'}, 'fisher_wd': {'title': 'Fisher Weight Decay', 'tooltip': 'Applies adaptive, scale-invariant weight-decay regularization based on the Fisher Information Matrix (approximated by Adam\'s second moment). It reduces penalty for "important" high-curvature weights while accelerating decay for "useless" weights in flat regions. Leading to improved convergence and better final performance.', 'type': 'bool'}, + 'state_precision': {'title': 'state_precision', 'tooltip': """The quantization format used to store the optimizer states to save VRAM. Options include: 'auto': Stores the states in the original parameter's precision. 'factored': Enables a memory-efficient mode by applying fast low-rank factorization to the optimizers states. It combines factorization for magnitudes with 1-bit compression for signs, drastically reducing VRAM usage and allowing for larger models or batch sizes. 'bf16_sr': Uses BF16 with stochastic rounding for a balance of precision and memory. 'int8_sr': Uses 8-bit block-wise quantization with stochastic rounding. 'fp8_sr': Uses torch.float8_e4m3fn with stochastic rounding.""", 'type': 'StatePrecision'}, } # @formatter:on @@ -246,6 +247,9 @@ def create_dynamic_ui( elif type == 'CenteredWDMode': components.options(master, row, col + 1, ["full", "float8", "int8", "int4"], self.optimizer_ui_state, key, command=self.update_user_pref) + elif type == 'StatePrecision': + components.options(master, row, col + 1, ["auto", "factored", "bf16_sr", "int8_sr", "fp8_sr"], self.optimizer_ui_state, key, + command=self.update_user_pref) elif type != 'bool': components.entry(master, row, col + 1, self.optimizer_ui_state, key, command=self.update_user_pref) diff --git a/modules/util/config/TrainConfig.py b/modules/util/config/TrainConfig.py index 3d678dc78..7d9f1957d 100644 --- a/modules/util/config/TrainConfig.py +++ b/modules/util/config/TrainConfig.py @@ -151,6 +151,7 @@ class TrainOptimizerConfig(BaseConfig): centered_wd_mode: str factored_2nd: False fisher_wd: False + state_precision: str def __init__(self, data: list[(str, Any, type, bool)]): super().__init__(data) @@ -277,6 +278,7 @@ def default_values(): data.append(("centered_wd_mode", "float8", str, False)) data.append(("factored_2nd", False, bool, False)) data.append(("fisher_wd", False, bool, False)) + data.append(("state_precision", "float8", str, False)) return TrainOptimizerConfig(data) diff --git a/modules/util/create.py b/modules/util/create.py index 7aea42818..0f0a1dfe6 100644 --- a/modules/util/create.py +++ b/modules/util/create.py @@ -696,6 +696,7 @@ def create_optimizer( scaled_optm=optimizer_config.scaled_optm if optimizer_config.scaled_optm is not None else False, centered_wd=optimizer_config.centered_wd if optimizer_config.centered_wd is not None else 0.0, centered_wd_mode=optimizer_config.centered_wd_mode if optimizer_config.centered_wd_mode is not None else "full", + state_precision=optimizer_config.state_precision if optimizer_config.state_precision is not None else "auto", ) # ADOPT_ADV Optimizer @@ -728,6 +729,7 @@ def create_optimizer( scaled_optm=optimizer_config.scaled_optm if optimizer_config.scaled_optm is not None else False, centered_wd=optimizer_config.centered_wd if optimizer_config.centered_wd is not None else 0.0, centered_wd_mode=optimizer_config.centered_wd_mode if optimizer_config.centered_wd_mode is not None else "full", + state_precision=optimizer_config.state_precision if optimizer_config.state_precision is not None else "auto", ) # PRODIGY_ADV Optimizer @@ -813,6 +815,7 @@ def create_optimizer( l1_adaptive=optimizer_config.l1_adaptive if optimizer_config.l1_adaptive is not None else False, centered_wd=optimizer_config.centered_wd if optimizer_config.centered_wd is not None else 0.0, centered_wd_mode=optimizer_config.centered_wd_mode if optimizer_config.centered_wd_mode is not None else "full", + state_precision=optimizer_config.state_precision if optimizer_config.state_precision is not None else "auto", ) # LION_ADV Optimizer From 2499da1a885eeedd5e792e62f7deabc5dbdba646 Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Sun, 12 Apr 2026 01:04:59 +0300 Subject: [PATCH 26/84] add SGD_ADV, various changes --- modules/ui/MuonAdamWindow.py | 1 - modules/ui/OptimizerParamsWindow.py | 3 +- modules/util/config/TrainConfig.py | 6 ++-- modules/util/create.py | 46 +++++++++++++++++++++-------- modules/util/enum/Optimizer.py | 2 ++ modules/util/optimizer_util.py | 40 ++++++++++++++++++------- 6 files changed, 71 insertions(+), 27 deletions(-) diff --git a/modules/ui/MuonAdamWindow.py b/modules/ui/MuonAdamWindow.py index 2f1682f36..a665a10b7 100644 --- a/modules/ui/MuonAdamWindow.py +++ b/modules/ui/MuonAdamWindow.py @@ -73,7 +73,6 @@ def create_adam_params_ui(self, master): 'use_bias_correction': {'title': 'Bias Correction', 'tooltip': 'Turn on Adam\'s bias correction.', 'type': 'bool'}, 'weight_decay': {'title': 'Weight Decay', 'tooltip': 'Regularization to prevent overfitting.', 'type': 'float'}, 'use_orthograd': {'title': 'use_orthograd', 'tooltip': 'Use orthograd method', 'type': 'bool'}, - 'nnmf_factor': {'title': 'Factored Optimizer', 'tooltip': 'Enables a memory-efficient mode by applying fast low-rank factorization to the optimizers states. It combines factorization for magnitudes with 1-bit compression for signs, drastically reducing VRAM usage and allowing for larger models or batch sizes. This is an approximation which may slightly alter training dynamics.', 'type': 'bool'}, 'orthogonal_gradient': {'title': 'OrthoGrad', 'tooltip': 'Reduces overfitting by removing the gradient component parallel to the weight, thus improving generalization.', 'type': 'bool'}, 'use_atan2': {'title': 'Atan2 Scaling', 'tooltip': 'A robust replacement for eps, which also incorporates gradient clipping, bounding and stabilizing the optimizer updates.', 'type': 'bool'}, 'cautious_mask': {'title': 'Cautious Variant', 'tooltip': 'Applies a mask to dampen or zero-out momentum components that disagree with the current gradients direction.', 'type': 'bool'}, diff --git a/modules/ui/OptimizerParamsWindow.py b/modules/ui/OptimizerParamsWindow.py index 0112ad2f1..b6e2bb8e7 100644 --- a/modules/ui/OptimizerParamsWindow.py +++ b/modules/ui/OptimizerParamsWindow.py @@ -200,7 +200,6 @@ def create_dynamic_ui( 'auto_kappa_p': {'title': 'Auto Lion-K', 'tooltip': 'Automatically determines the optimal P-value based on layer dimensions. Uses p=2.0 (Spherical) for 4D (Conv) tensors for stability and rotational invariance, and p=1.0 (Sign) for 2D (Linear) tensors for sparsity. Overrides the manual P-value. Recommend for unet models.', 'type': 'bool'}, 'compile': {'title': 'Compiled Optimizer', 'tooltip': 'Enables PyTorch compilation for the optimizer internal step logic. This is intended to improve performance by allowing PyTorch to fuse operations and optimize the computational graph.', 'type': 'bool'}, 'spectral_normalization': {'title': 'Spectral Scaling', 'tooltip': 'Enables explicit Spectral Normalization to automatically rescale the update magnitude and Weight Decay based on layer dimensions. This allows hyperparameters to transfer seamlessly from small to large models without retuning, while making the optimizer highly robust to a wide range of learning rates. This ensures consistent performance across different model sizes, adapter methods, and ranks. ', 'type': 'bool'}, - 'scaled_optm': {'title': 'Scaled Optimizer', 'tooltip': 'Automatically rescale the update magnitude and Weight Decay based on layer dimensions and type. This allows hyperparameters to transfer seamlessly from small to large models without retuning. This ensures consistent performance across different model sizes, adapter methods, and ranks. For LoRAs set alpha=rank. Using this, LoRA shares the same LR as full finetuning for all ranks.', 'type': 'bool'}, 'freeze_on_flip': {'title': 'Freeze-on-Flip', 'tooltip': 'Projected One-hit freeze. Masks updates for coordinates where the gradient sign flips compared to the previous step. Over the steps, this makes the signed update semi-continuous.', 'type': 'bool'}, 'l1_adaptive': {'title': 'L1 Adaptive', 'tooltip': 'Scales the learning rate dynamically by the L1 norm of the gradient to handle gradient heterogeneity. This makes the signed-optimizers semi-adaptive.', 'type': 'bool'}, 'centered_wd': {'title': 'Centered Weight Decay', 'tooltip': 'Centered Weight Decay coefficient. Instead of decaying weights toward zero, they are decayed toward their initial values (anchors). This can be used together with standard weight decay.', 'type': 'float'}, @@ -208,6 +207,8 @@ def create_dynamic_ui( 'factored_2nd': {'title': 'Factored 2nd', 'tooltip': 'Whether to keep the first moment uncompressed (dense), while only factorizing the second moment. This makes the optimizer highly robust to a wide range of LRs, mimicking high-order optimization.', 'type': 'bool'}, 'fisher_wd': {'title': 'Fisher Weight Decay', 'tooltip': 'Applies adaptive, scale-invariant weight-decay regularization based on the Fisher Information Matrix (approximated by Adam\'s second moment). It reduces penalty for "important" high-curvature weights while accelerating decay for "useless" weights in flat regions. Leading to improved convergence and better final performance.', 'type': 'bool'}, 'state_precision': {'title': 'state_precision', 'tooltip': """The quantization format used to store the optimizer states to save VRAM. Options include: 'auto': Stores the states in the original parameter's precision. 'factored': Enables a memory-efficient mode by applying fast low-rank factorization to the optimizers states. It combines factorization for magnitudes with 1-bit compression for signs, drastically reducing VRAM usage and allowing for larger models or batch sizes. 'bf16_sr': Uses BF16 with stochastic rounding for a balance of precision and memory. 'int8_sr': Uses 8-bit block-wise quantization with stochastic rounding. 'fp8_sr': Uses torch.float8_e4m3fn with stochastic rounding.""", 'type': 'StatePrecision'}, + 'sinkhorn': {'title': 'Multi-Normed Sinkhorn', 'tooltip': 'Applies fast iterative row and column normalization to distribute weight updates evenly while giving some room to the original gradient\'s direction. This serves as a "soft," smoothed alternative to orthogonalization, which is theoretically better.', 'type': 'bool'}, + 'sinkhorn_iterations': {'title': 'Sinkhorn Iterations', 'tooltip': 'Controls the number of iterations for Multi-Normed Sinkhorn. While 1 iteration is often sufficient for convergence and 3 offers a slight refinement, 5 is the default.', 'type': 'int'}, } # @formatter:on diff --git a/modules/util/config/TrainConfig.py b/modules/util/config/TrainConfig.py index 7d9f1957d..10a8595cf 100644 --- a/modules/util/config/TrainConfig.py +++ b/modules/util/config/TrainConfig.py @@ -144,7 +144,6 @@ class TrainOptimizerConfig(BaseConfig): auto_kappa_p: False compile: False spectral_normalization: False - scaled_optm: False freeze_on_flip: False l1_adaptive: False centered_wd: float @@ -152,6 +151,8 @@ class TrainOptimizerConfig(BaseConfig): factored_2nd: False fisher_wd: False state_precision: str + sinkhorn: False + sinkhorn_iterations: int def __init__(self, data: list[(str, Any, type, bool)]): super().__init__(data) @@ -271,7 +272,6 @@ def default_values(): data.append(("auto_kappa_p", False, bool, False)) data.append(("compile", False, bool, False)) data.append(("spectral_normalization", False, bool, False)) - data.append(("scaled_optm", False, bool, False)) data.append(("freeze_on_flip", False, bool, False)) data.append(("l1_adaptive", False, bool, False)) data.append(("centered_wd", 0.0, float, False)) @@ -279,6 +279,8 @@ def default_values(): data.append(("factored_2nd", False, bool, False)) data.append(("fisher_wd", False, bool, False)) data.append(("state_precision", "float8", str, False)) + data.append(("sinkhorn", False, bool, False)) + data.append(("sinkhorn_iterations", None, int, True)) return TrainOptimizerConfig(data) diff --git a/modules/util/create.py b/modules/util/create.py index 0f0a1dfe6..980bdfd68 100644 --- a/modules/util/create.py +++ b/modules/util/create.py @@ -141,7 +141,7 @@ def create_optimizer( parameters = parameter_group_collection.parameters_for_optimizer(config) - if optimizer_config.scaled_optm or optimizer_config.spectral_normalization: + if optimizer_config.spectral_normalization: # _adv optimizers O(1) depth scaling. inject_depth_into_param_groups(model, parameters) @@ -667,6 +667,27 @@ def create_optimizer( quant_block_size=optimizer_config.quant_block_size if optimizer_config.quant_block_size is not None else 2048 ) + # SGD_ADV Optimizer + case Optimizer.SGD_ADV: + from adv_optm import SGD_adv + optimizer = SGD_adv( + params=parameters, + lr=config.learning_rate, + momentum=optimizer_config.momentum if optimizer_config.beta1 is not None else 0, + weight_decay=optimizer_config.weight_decay if optimizer_config.weight_decay is not None else 0.0, + nesterov=optimizer_config.nesterov if optimizer_config.nesterov is not None else True, + cautious_wd=optimizer_config.cautious_wd if optimizer_config.cautious_wd is not None else False, + stochastic_rounding=optimizer_config.stochastic_rounding, + orthogonal_gradient=optimizer_config.orthogonal_gradient if optimizer_config.orthogonal_gradient is not None else False, + compiled_optimizer=optimizer_config.compile if optimizer_config.compile is not None else False, + sinkhorn=optimizer_config.sinkhorn if optimizer_config.sinkhorn is not None else False, + sinkhorn_iterations=optimizer_config.sinkhorn_iterations if optimizer_config.sinkhorn_iterations is not None else 5, + spectral_normalization=optimizer_config.spectral_normalization if optimizer_config.spectral_normalization is not None else False, + centered_wd=optimizer_config.centered_wd if optimizer_config.centered_wd is not None else 0.0, + centered_wd_mode=optimizer_config.centered_wd_mode if optimizer_config.centered_wd_mode is not None else "full", + state_precision=optimizer_config.state_precision if optimizer_config.state_precision is not None else "auto", + ) + # ADAMW_ADV Optimizer case Optimizer.ADAMW_ADV: from adv_optm import AdamW_adv @@ -679,7 +700,6 @@ def create_optimizer( weight_decay=optimizer_config.weight_decay if optimizer_config.weight_decay is not None else 0.0, use_bias_correction=optimizer_config.use_bias_correction if optimizer_config.use_bias_correction is not None else True, factored_2nd=optimizer_config.factored_2nd if optimizer_config.factored_2nd is not None else False, - nnmf_factor=optimizer_config.nnmf_factor if optimizer_config.nnmf_factor is not None else False, fisher_wd=optimizer_config.fisher_wd if optimizer_config.fisher_wd is not None else False, cautious_wd=optimizer_config.cautious_wd if optimizer_config.cautious_wd is not None else False, stochastic_rounding=optimizer_config.stochastic_rounding, @@ -693,7 +713,7 @@ def create_optimizer( kourkoutas_beta=optimizer_config.kourkoutas_beta if optimizer_config.kourkoutas_beta is not None else False, k_warmup_steps=optimizer_config.k_warmup_steps if optimizer_config.k_warmup_steps is not None else 0, compiled_optimizer=optimizer_config.compile if optimizer_config.compile is not None else False, - scaled_optm=optimizer_config.scaled_optm if optimizer_config.scaled_optm is not None else False, + spectral_normalization=optimizer_config.spectral_normalization if optimizer_config.spectral_normalization is not None else False, centered_wd=optimizer_config.centered_wd if optimizer_config.centered_wd is not None else 0.0, centered_wd_mode=optimizer_config.centered_wd_mode if optimizer_config.centered_wd_mode is not None else "full", state_precision=optimizer_config.state_precision if optimizer_config.state_precision is not None else "auto", @@ -709,7 +729,6 @@ def create_optimizer( optimizer_config.beta2 if optimizer_config.beta2 is not None else 0.9999), eps=optimizer_config.eps if optimizer_config.eps is not None else 1e-6, weight_decay=optimizer_config.weight_decay if optimizer_config.weight_decay is not None else 0.0, - nnmf_factor=optimizer_config.nnmf_factor if optimizer_config.nnmf_factor is not None else False, factored_2nd=optimizer_config.factored_2nd if optimizer_config.factored_2nd is not None else False, fisher_wd=optimizer_config.fisher_wd if optimizer_config.fisher_wd is not None else False, cautious_wd=optimizer_config.cautious_wd if optimizer_config.cautious_wd is not None else False, @@ -726,7 +745,7 @@ def create_optimizer( kourkoutas_beta=optimizer_config.kourkoutas_beta if optimizer_config.kourkoutas_beta is not None else False, k_warmup_steps=optimizer_config.k_warmup_steps if optimizer_config.k_warmup_steps is not None else 0, compiled_optimizer=optimizer_config.compile if optimizer_config.compile is not None else False, - scaled_optm=optimizer_config.scaled_optm if optimizer_config.scaled_optm is not None else False, + spectral_normalization=optimizer_config.spectral_normalization if optimizer_config.spectral_normalization is not None else False, centered_wd=optimizer_config.centered_wd if optimizer_config.centered_wd is not None else 0.0, centered_wd_mode=optimizer_config.centered_wd_mode if optimizer_config.centered_wd_mode is not None else "full", state_precision=optimizer_config.state_precision if optimizer_config.state_precision is not None else "auto", @@ -743,7 +762,6 @@ def create_optimizer( beta3=optimizer_config.beta3 if optimizer_config.beta3 is not None else None, eps=optimizer_config.eps if optimizer_config.eps is not None else 1e-8, weight_decay=optimizer_config.weight_decay if optimizer_config.weight_decay is not None else 0.0, - nnmf_factor=optimizer_config.nnmf_factor if optimizer_config.nnmf_factor is not None else False, factored_2nd=optimizer_config.factored_2nd if optimizer_config.factored_2nd is not None else False, fisher_wd=optimizer_config.fisher_wd if optimizer_config.fisher_wd is not None else False, cautious_wd=optimizer_config.cautious_wd if optimizer_config.cautious_wd is not None else False, @@ -782,7 +800,6 @@ def create_optimizer( beta1_warmup=optimizer_config.beta1_warmup if optimizer_config.beta1_warmup is not None else None, min_beta1=optimizer_config.min_beta1 if optimizer_config.min_beta1 is not None else 0.9, use_bias_correction=optimizer_config.use_bias_correction if optimizer_config.use_bias_correction is not None else True, - nnmf_factor=optimizer_config.nnmf_factor if optimizer_config.nnmf_factor is not None else False, factored_2nd=optimizer_config.factored_2nd if optimizer_config.factored_2nd is not None else False, cautious_wd=optimizer_config.cautious_wd if optimizer_config.cautious_wd is not None else False, stochastic_rounding=optimizer_config.stochastic_rounding, @@ -790,7 +807,7 @@ def create_optimizer( kourkoutas_beta=optimizer_config.kourkoutas_beta if optimizer_config.kourkoutas_beta is not None else False, k_warmup_steps=optimizer_config.k_warmup_steps if optimizer_config.k_warmup_steps is not None else 0, compiled_optimizer=optimizer_config.compile if optimizer_config.compile is not None else False, - scaled_optm=optimizer_config.scaled_optm if optimizer_config.scaled_optm is not None else False, + spectral_normalization=optimizer_config.spectral_normalization if optimizer_config.spectral_normalization is not None else False, centered_wd=optimizer_config.centered_wd if optimizer_config.centered_wd is not None else 0.0, centered_wd_mode=optimizer_config.centered_wd_mode if optimizer_config.centered_wd_mode is not None else "full", ) @@ -803,14 +820,13 @@ def create_optimizer( lr=config.learning_rate, momentum=optimizer_config.momentum if optimizer_config.momentum is not None else 0, weight_decay=optimizer_config.weight_decay if optimizer_config.weight_decay is not None else 0.0, - nnmf_factor=optimizer_config.nnmf_factor if optimizer_config.nnmf_factor is not None else False, cautious_wd=optimizer_config.cautious_wd if optimizer_config.cautious_wd is not None else False, stochastic_rounding=optimizer_config.stochastic_rounding, orthogonal_gradient=optimizer_config.orthogonal_gradient if optimizer_config.orthogonal_gradient is not None else False, compiled_optimizer=optimizer_config.compile if optimizer_config.compile is not None else False, Simplified_AdEMAMix=optimizer_config.Simplified_AdEMAMix if optimizer_config.Simplified_AdEMAMix is not None else False, alpha_grad=optimizer_config.alpha_grad if optimizer_config.alpha_grad is not None else 100, - scaled_optm=optimizer_config.scaled_optm if optimizer_config.scaled_optm is not None else False, + spectral_normalization=optimizer_config.spectral_normalization if optimizer_config.spectral_normalization is not None else False, freeze_on_flip=optimizer_config.freeze_on_flip if optimizer_config.freeze_on_flip is not None else False, l1_adaptive=optimizer_config.l1_adaptive if optimizer_config.l1_adaptive is not None else False, centered_wd=optimizer_config.centered_wd if optimizer_config.centered_wd is not None else 0.0, @@ -838,7 +854,7 @@ def create_optimizer( compiled_optimizer=optimizer_config.compile if optimizer_config.compile is not None else False, freeze_on_flip=optimizer_config.freeze_on_flip if optimizer_config.freeze_on_flip is not None else False, l1_adaptive=optimizer_config.l1_adaptive if optimizer_config.l1_adaptive is not None else False, - scaled_optm=optimizer_config.scaled_optm if optimizer_config.scaled_optm is not None else False, + spectral_normalization=optimizer_config.spectral_normalization if optimizer_config.spectral_normalization is not None else False, centered_wd=optimizer_config.centered_wd if optimizer_config.centered_wd is not None else 0.0, centered_wd_mode=optimizer_config.centered_wd_mode if optimizer_config.centered_wd_mode is not None else "full", ) @@ -905,7 +921,6 @@ def create_optimizer( ns_steps=optimizer_config.ns_steps if optimizer_config.ns_steps is not None else 5, weight_decay=optimizer_config.weight_decay if optimizer_config.weight_decay is not None else 0.0, rms_rescaling=optimizer_config.rms_rescaling if optimizer_config.rms_rescaling is not None else True, - nnmf_factor=optimizer_config.nnmf_factor if optimizer_config.nnmf_factor is not None else False, cautious_wd=optimizer_config.cautious_wd if optimizer_config.cautious_wd is not None else False, stochastic_rounding=optimizer_config.stochastic_rounding, nesterov=optimizer_config.nesterov if optimizer_config.nesterov is not None else True, @@ -921,6 +936,9 @@ def create_optimizer( Simplified_AdEMAMix=optimizer_config.Simplified_AdEMAMix if optimizer_config.Simplified_AdEMAMix is not None else False, alpha_grad=optimizer_config.alpha_grad if optimizer_config.alpha_grad is not None else 100, spectral_normalization=optimizer_config.spectral_normalization if optimizer_config.spectral_normalization is not None else False, + centered_wd=optimizer_config.centered_wd if optimizer_config.centered_wd is not None else 0.0, + centered_wd_mode=optimizer_config.centered_wd_mode if optimizer_config.centered_wd_mode is not None else "full", + state_precision=optimizer_config.state_precision if optimizer_config.state_precision is not None else "auto", **adam_kwargs ) @@ -962,7 +980,6 @@ def create_optimizer( ns_steps=optimizer_config.ns_steps if optimizer_config.ns_steps is not None else 5, rms_rescaling=optimizer_config.rms_rescaling if optimizer_config.rms_rescaling is not None else True, weight_decay=optimizer_config.weight_decay if optimizer_config.weight_decay is not None else 0.0, - nnmf_factor=optimizer_config.nnmf_factor if optimizer_config.nnmf_factor is not None else False, cautious_wd=optimizer_config.cautious_wd if optimizer_config.cautious_wd is not None else False, stochastic_rounding=optimizer_config.stochastic_rounding, nesterov=optimizer_config.nesterov if optimizer_config.nesterov is not None else True, @@ -977,6 +994,9 @@ def create_optimizer( approx_mars=optimizer_config.approx_mars if optimizer_config.approx_mars is not None else False, compiled_optimizer=optimizer_config.compile if optimizer_config.compile is not None else False, spectral_normalization=optimizer_config.spectral_normalization if optimizer_config.spectral_normalization is not None else False, + centered_wd=optimizer_config.centered_wd if optimizer_config.centered_wd is not None else 0.0, + centered_wd_mode=optimizer_config.centered_wd_mode if optimizer_config.centered_wd_mode is not None else "full", + state_precision=optimizer_config.state_precision if optimizer_config.state_precision is not None else "auto", **adam_kwargs ) diff --git a/modules/util/enum/Optimizer.py b/modules/util/enum/Optimizer.py index 58edf419c..2227772c6 100644 --- a/modules/util/enum/Optimizer.py +++ b/modules/util/enum/Optimizer.py @@ -42,6 +42,7 @@ class Optimizer(Enum): # 32 bit is torch and not bnb SGD = 'SGD' SGD_8BIT = 'SGD_8BIT' + SGD_ADV = 'SGD_ADV' SIGNSGD_ADV = 'SIGNSGD_ADV' # Schedule-free optimizers @@ -118,6 +119,7 @@ def supports_fused_back_pass(self): Optimizer.MUON_ADV, Optimizer.ADAMUON_ADV, Optimizer.SIGNSGD_ADV, + Optimizer.SGD_ADV, ] # Small helper for adjusting learning rates to adaptive optimizers. diff --git a/modules/util/optimizer_util.py b/modules/util/optimizer_util.py index 1c6d9a6d1..c7125a2b0 100644 --- a/modules/util/optimizer_util.py +++ b/modules/util/optimizer_util.py @@ -442,6 +442,22 @@ def init_model_parameters( "min_8bit_size": 16384, "quant_block_size": 2048 }, + Optimizer.SGD_ADV: { + "momentum": 0, + "cautious_wd": False, + "weight_decay": 0.0, + "sinkhorn": False, + "stochastic_rounding": True, + "compile": False, + "fused_back_pass": False, + "orthogonal_gradient": False, + "sinkhorn": False, + "sinkhorn_iterations": 5, + "spectral_normalization": False, + "centered_wd": 0.0, + "centered_wd_mode": "float8", + "state_precision": "auto", + }, Optimizer.ADAMW_ADV: { "beta1": 0.9, "beta2": 0.99, @@ -450,7 +466,6 @@ def init_model_parameters( "cautious_wd": False, "weight_decay": 0.0, "use_bias_correction": True, - "nnmf_factor": False, "factored_2nd": False, "stochastic_rounding": True, "compile": False, @@ -464,9 +479,10 @@ def init_model_parameters( "alpha": 5, "kourkoutas_beta": False, "k_warmup_steps": None, - "scaled_optm": False, + "spectral_normalization": False, "centered_wd": 0.0, "centered_wd_mode": "float8", + "state_precision": "auto", }, Optimizer.ADOPT_ADV: { "beta1": 0.9, @@ -475,7 +491,6 @@ def init_model_parameters( "fisher_wd": False, "cautious_wd": False, "weight_decay": 0.0, - "nnmf_factor": False, "factored_2nd": False, "stochastic_rounding": True, "compile": False, @@ -491,9 +506,10 @@ def init_model_parameters( "alpha_grad": 100.0, "kourkoutas_beta": False, "k_warmup_steps": None, - "scaled_optm": False, + "spectral_normalization": False, "centered_wd": 0.0, "centered_wd_mode": "float8", + "state_precision": "auto", }, Optimizer.PRODIGY_ADV: { "beta1": 0.9, @@ -546,7 +562,7 @@ def init_model_parameters( "orthogonal_gradient": False, "kourkoutas_beta": False, "k_warmup_steps": None, - "scaled_optm": False, + "spectral_normalization": False, "centered_wd": 0.0, "centered_wd_mode": "float8", }, @@ -554,7 +570,6 @@ def init_model_parameters( "momentum": 0.95, "cautious_wd": False, "weight_decay": 0.0, - "nnmf_factor": False, "stochastic_rounding": True, "compile": False, "fused_back_pass": False, @@ -563,9 +578,10 @@ def init_model_parameters( "alpha_grad": 100.0, "freeze_on_flip": False, "l1_adaptive": False, - "scaled_optm": False, + "spectral_normalization": False, "centered_wd": 0.0, "centered_wd_mode": "float8", + "state_precision": "auto", }, Optimizer.LION_ADV: { "beta1": 0.9, @@ -583,7 +599,7 @@ def init_model_parameters( "auto_kappa_p": True, "freeze_on_flip": False, "l1_adaptive": False, - "scaled_optm": False, + "spectral_normalization": False, "centered_wd": 0.0, "centered_wd_mode": "float8", }, @@ -619,7 +635,6 @@ def init_model_parameters( "ortho_rank": 128, "rms_rescaling": True, "spectral_normalization": False, - "nnmf_factor": False, "stochastic_rounding": True, "compile": False, "fused_back_pass": False, @@ -637,6 +652,9 @@ def init_model_parameters( "normuon_eps": 1e-8, "orthogonal_gradient": False, "approx_mars": False, + "centered_wd": 0.0, + "centered_wd_mode": "float8", + "state_precision": "auto", "muon_adam_config": {}, }, Optimizer.ADAMUON_ADV: { @@ -651,7 +669,6 @@ def init_model_parameters( "ortho_rank": 128, "rms_rescaling": True, "spectral_normalization": False, - "nnmf_factor": False, "stochastic_rounding": True, "compile": False, "fused_back_pass": False, @@ -668,6 +685,9 @@ def init_model_parameters( "normuon_variant": True, "orthogonal_gradient": False, "approx_mars": False, + "centered_wd": 0.0, + "centered_wd_mode": "float8", + "state_precision": "auto", "muon_adam_config": {}, }, Optimizer.ADABELIEF: { From 535d206e7a22bc5545a4d364d8fcbccea7e1c5a8 Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Sun, 12 Apr 2026 01:05:57 +0300 Subject: [PATCH 27/84] dev8 --- requirements-global.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements-global.txt b/requirements-global.txt index c17c3a9e3..3cdc2abcd 100644 --- a/requirements-global.txt +++ b/requirements-global.txt @@ -42,7 +42,7 @@ prodigyopt==1.1.2 # prodigy optimizer schedulefree==1.4.1 # schedule-free optimizers pytorch_optimizer==3.6.0 # pytorch optimizers prodigy-plus-schedule-free==2.0.1 # Prodigy plus optimizer -adv_optm==2.4.dev6 # advanced optimizers +adv_optm==2.4.dev8 # advanced optimizers -e git+https://github.com/KellerJordan/Muon.git@f90a42b#egg=muon-optimizer # Profiling From 16eef8f85748c20f426998e20afcf94e1bd29c2e Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Sun, 12 Apr 2026 01:13:23 +0300 Subject: [PATCH 28/84] fix --- modules/util/optimizer_util.py | 1 - 1 file changed, 1 deletion(-) diff --git a/modules/util/optimizer_util.py b/modules/util/optimizer_util.py index c7125a2b0..928990e9f 100644 --- a/modules/util/optimizer_util.py +++ b/modules/util/optimizer_util.py @@ -446,7 +446,6 @@ def init_model_parameters( "momentum": 0, "cautious_wd": False, "weight_decay": 0.0, - "sinkhorn": False, "stochastic_rounding": True, "compile": False, "fused_back_pass": False, From 753d265ae1390247e057829d3f54367bfd783592 Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Sun, 12 Apr 2026 01:14:22 +0300 Subject: [PATCH 29/84] add nesterov --- modules/util/optimizer_util.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/modules/util/optimizer_util.py b/modules/util/optimizer_util.py index 928990e9f..e3ed5196a 100644 --- a/modules/util/optimizer_util.py +++ b/modules/util/optimizer_util.py @@ -443,7 +443,8 @@ def init_model_parameters( "quant_block_size": 2048 }, Optimizer.SGD_ADV: { - "momentum": 0, + "momentum": 0.95, + "nesterov": True, "cautious_wd": False, "weight_decay": 0.0, "stochastic_rounding": True, From 4ab3d4aa3bb279720f623a82c7101640d057aaca Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Sun, 12 Apr 2026 01:25:54 +0300 Subject: [PATCH 30/84] dev9 --- requirements-global.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements-global.txt b/requirements-global.txt index 3cdc2abcd..9bf0eb33f 100644 --- a/requirements-global.txt +++ b/requirements-global.txt @@ -42,7 +42,7 @@ prodigyopt==1.1.2 # prodigy optimizer schedulefree==1.4.1 # schedule-free optimizers pytorch_optimizer==3.6.0 # pytorch optimizers prodigy-plus-schedule-free==2.0.1 # Prodigy plus optimizer -adv_optm==2.4.dev8 # advanced optimizers +adv_optm==2.4.dev9 # advanced optimizers -e git+https://github.com/KellerJordan/Muon.git@f90a42b#egg=muon-optimizer # Profiling From 49b24a83dbb8120273f01325e30829383cd0b526 Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Sun, 12 Apr 2026 01:35:32 +0300 Subject: [PATCH 31/84] dev10 --- requirements-global.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements-global.txt b/requirements-global.txt index 9bf0eb33f..179a82ba5 100644 --- a/requirements-global.txt +++ b/requirements-global.txt @@ -42,7 +42,7 @@ prodigyopt==1.1.2 # prodigy optimizer schedulefree==1.4.1 # schedule-free optimizers pytorch_optimizer==3.6.0 # pytorch optimizers prodigy-plus-schedule-free==2.0.1 # Prodigy plus optimizer -adv_optm==2.4.dev9 # advanced optimizers +adv_optm==2.4.dev10 # advanced optimizers -e git+https://github.com/KellerJordan/Muon.git@f90a42b#egg=muon-optimizer # Profiling From dba147d1008a044c9ac22bbccc6aeee608698312 Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Sun, 12 Apr 2026 01:38:35 +0300 Subject: [PATCH 32/84] dev11 --- requirements-global.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements-global.txt b/requirements-global.txt index 179a82ba5..a941d41d1 100644 --- a/requirements-global.txt +++ b/requirements-global.txt @@ -42,7 +42,7 @@ prodigyopt==1.1.2 # prodigy optimizer schedulefree==1.4.1 # schedule-free optimizers pytorch_optimizer==3.6.0 # pytorch optimizers prodigy-plus-schedule-free==2.0.1 # Prodigy plus optimizer -adv_optm==2.4.dev10 # advanced optimizers +adv_optm==2.4.dev11 # advanced optimizers -e git+https://github.com/KellerJordan/Muon.git@f90a42b#egg=muon-optimizer # Profiling From b092fbbb5fa8869c100499ad23e6ab800a348f7f Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Tue, 28 Apr 2026 11:32:48 +0300 Subject: [PATCH 33/84] Change sgd to sinksgd and add orthogonal sinkhorn --- modules/ui/OptimizerParamsWindow.py | 2 +- modules/util/config/TrainConfig.py | 4 ++-- modules/util/create.py | 10 +++++----- modules/util/enum/Optimizer.py | 4 ++-- modules/util/optimizer_util.py | 6 +++--- requirements-global.txt | 2 +- 6 files changed, 14 insertions(+), 14 deletions(-) diff --git a/modules/ui/OptimizerParamsWindow.py b/modules/ui/OptimizerParamsWindow.py index b6e2bb8e7..c927505d1 100644 --- a/modules/ui/OptimizerParamsWindow.py +++ b/modules/ui/OptimizerParamsWindow.py @@ -207,7 +207,7 @@ def create_dynamic_ui( 'factored_2nd': {'title': 'Factored 2nd', 'tooltip': 'Whether to keep the first moment uncompressed (dense), while only factorizing the second moment. This makes the optimizer highly robust to a wide range of LRs, mimicking high-order optimization.', 'type': 'bool'}, 'fisher_wd': {'title': 'Fisher Weight Decay', 'tooltip': 'Applies adaptive, scale-invariant weight-decay regularization based on the Fisher Information Matrix (approximated by Adam\'s second moment). It reduces penalty for "important" high-curvature weights while accelerating decay for "useless" weights in flat regions. Leading to improved convergence and better final performance.', 'type': 'bool'}, 'state_precision': {'title': 'state_precision', 'tooltip': """The quantization format used to store the optimizer states to save VRAM. Options include: 'auto': Stores the states in the original parameter's precision. 'factored': Enables a memory-efficient mode by applying fast low-rank factorization to the optimizers states. It combines factorization for magnitudes with 1-bit compression for signs, drastically reducing VRAM usage and allowing for larger models or batch sizes. 'bf16_sr': Uses BF16 with stochastic rounding for a balance of precision and memory. 'int8_sr': Uses 8-bit block-wise quantization with stochastic rounding. 'fp8_sr': Uses torch.float8_e4m3fn with stochastic rounding.""", 'type': 'StatePrecision'}, - 'sinkhorn': {'title': 'Multi-Normed Sinkhorn', 'tooltip': 'Applies fast iterative row and column normalization to distribute weight updates evenly while giving some room to the original gradient\'s direction. This serves as a "soft," smoothed alternative to orthogonalization, which is theoretically better.', 'type': 'bool'}, + 'orthogonal_sinkhorn': {'title': 'Orthogonal Sinkhorn', 'tooltip': 'Applies iterative row and column orthogonal projection to make the updates orthogonal to the current weight, leading to robust regularization and better generalization.', 'type': 'bool'}, 'sinkhorn_iterations': {'title': 'Sinkhorn Iterations', 'tooltip': 'Controls the number of iterations for Multi-Normed Sinkhorn. While 1 iteration is often sufficient for convergence and 3 offers a slight refinement, 5 is the default.', 'type': 'int'}, } # @formatter:on diff --git a/modules/util/config/TrainConfig.py b/modules/util/config/TrainConfig.py index 10a8595cf..7774da2af 100644 --- a/modules/util/config/TrainConfig.py +++ b/modules/util/config/TrainConfig.py @@ -151,7 +151,7 @@ class TrainOptimizerConfig(BaseConfig): factored_2nd: False fisher_wd: False state_precision: str - sinkhorn: False + orthogonal_sinkhorn: False sinkhorn_iterations: int def __init__(self, data: list[(str, Any, type, bool)]): @@ -279,7 +279,7 @@ def default_values(): data.append(("factored_2nd", False, bool, False)) data.append(("fisher_wd", False, bool, False)) data.append(("state_precision", "float8", str, False)) - data.append(("sinkhorn", False, bool, False)) + data.append(("orthogonal_sinkhorn", False, bool, False)) data.append(("sinkhorn_iterations", None, int, True)) return TrainOptimizerConfig(data) diff --git a/modules/util/create.py b/modules/util/create.py index 980bdfd68..ea25acb46 100644 --- a/modules/util/create.py +++ b/modules/util/create.py @@ -667,10 +667,10 @@ def create_optimizer( quant_block_size=optimizer_config.quant_block_size if optimizer_config.quant_block_size is not None else 2048 ) - # SGD_ADV Optimizer - case Optimizer.SGD_ADV: - from adv_optm import SGD_adv - optimizer = SGD_adv( + # SINKSGD_ADV Optimizer + case Optimizer.SINKSGD_ADV: + from adv_optm import SinkSGD_adv + optimizer = SinkSGD_adv( params=parameters, lr=config.learning_rate, momentum=optimizer_config.momentum if optimizer_config.beta1 is not None else 0, @@ -680,7 +680,7 @@ def create_optimizer( stochastic_rounding=optimizer_config.stochastic_rounding, orthogonal_gradient=optimizer_config.orthogonal_gradient if optimizer_config.orthogonal_gradient is not None else False, compiled_optimizer=optimizer_config.compile if optimizer_config.compile is not None else False, - sinkhorn=optimizer_config.sinkhorn if optimizer_config.sinkhorn is not None else False, + orthogonal_sinkhorn=optimizer_config.orthogonal_sinkhorn if optimizer_config.orthogonal_sinkhorn is not None else False, sinkhorn_iterations=optimizer_config.sinkhorn_iterations if optimizer_config.sinkhorn_iterations is not None else 5, spectral_normalization=optimizer_config.spectral_normalization if optimizer_config.spectral_normalization is not None else False, centered_wd=optimizer_config.centered_wd if optimizer_config.centered_wd is not None else 0.0, diff --git a/modules/util/enum/Optimizer.py b/modules/util/enum/Optimizer.py index 2227772c6..cb17c6c48 100644 --- a/modules/util/enum/Optimizer.py +++ b/modules/util/enum/Optimizer.py @@ -42,7 +42,7 @@ class Optimizer(Enum): # 32 bit is torch and not bnb SGD = 'SGD' SGD_8BIT = 'SGD_8BIT' - SGD_ADV = 'SGD_ADV' + SINKSGD_ADV = 'SINKSGD_ADV' SIGNSGD_ADV = 'SIGNSGD_ADV' # Schedule-free optimizers @@ -119,7 +119,7 @@ def supports_fused_back_pass(self): Optimizer.MUON_ADV, Optimizer.ADAMUON_ADV, Optimizer.SIGNSGD_ADV, - Optimizer.SGD_ADV, + Optimizer.SINKSGD_ADV, ] # Small helper for adjusting learning rates to adaptive optimizers. diff --git a/modules/util/optimizer_util.py b/modules/util/optimizer_util.py index e3ed5196a..9811353d4 100644 --- a/modules/util/optimizer_util.py +++ b/modules/util/optimizer_util.py @@ -442,17 +442,17 @@ def init_model_parameters( "min_8bit_size": 16384, "quant_block_size": 2048 }, - Optimizer.SGD_ADV: { + Optimizer.SINKSGD_ADV: { "momentum": 0.95, "nesterov": True, + "sinkhorn_iterations": 5, "cautious_wd": False, "weight_decay": 0.0, "stochastic_rounding": True, "compile": False, "fused_back_pass": False, + "orthogonal_sinkhorn": False, "orthogonal_gradient": False, - "sinkhorn": False, - "sinkhorn_iterations": 5, "spectral_normalization": False, "centered_wd": 0.0, "centered_wd_mode": "float8", diff --git a/requirements-global.txt b/requirements-global.txt index a941d41d1..dc6c21af2 100644 --- a/requirements-global.txt +++ b/requirements-global.txt @@ -42,7 +42,7 @@ prodigyopt==1.1.2 # prodigy optimizer schedulefree==1.4.1 # schedule-free optimizers pytorch_optimizer==3.6.0 # pytorch optimizers prodigy-plus-schedule-free==2.0.1 # Prodigy plus optimizer -adv_optm==2.4.dev11 # advanced optimizers +adv_optm==2.4.dev12 # advanced optimizers -e git+https://github.com/KellerJordan/Muon.git@f90a42b#egg=muon-optimizer # Profiling From f0ed83d8a073a29d287fa29a4058a569ec96fc1f Mon Sep 17 00:00:00 2001 From: Koratahiu Date: Wed, 29 Apr 2026 19:48:52 +0300 Subject: [PATCH 34/84] dev13: sinksgd bugfixes --- requirements-global.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements-global.txt b/requirements-global.txt index dc6c21af2..145eae03d 100644 --- a/requirements-global.txt +++ b/requirements-global.txt @@ -42,7 +42,7 @@ prodigyopt==1.1.2 # prodigy optimizer schedulefree==1.4.1 # schedule-free optimizers pytorch_optimizer==3.6.0 # pytorch optimizers prodigy-plus-schedule-free==2.0.1 # Prodigy plus optimizer -adv_optm==2.4.dev12 # advanced optimizers +adv_optm==2.4.dev13 # advanced optimizers -e git+https://github.com/KellerJordan/Muon.git@f90a42b#egg=muon-optimizer # Profiling From 081e26035a80b43c083a2a6e702679a55949b292 Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Sat, 9 May 2026 22:02:08 +0300 Subject: [PATCH 35/84] add stochastic_sign --- modules/ui/OptimizerParamsWindow.py | 3 +-- modules/util/config/TrainConfig.py | 6 ++---- modules/util/create.py | 6 ++---- modules/util/optimizer_util.py | 6 ++---- 4 files changed, 7 insertions(+), 14 deletions(-) diff --git a/modules/ui/OptimizerParamsWindow.py b/modules/ui/OptimizerParamsWindow.py index 50fd6a10a..1c700b681 100644 --- a/modules/ui/OptimizerParamsWindow.py +++ b/modules/ui/OptimizerParamsWindow.py @@ -195,8 +195,7 @@ def create_dynamic_ui( 'auto_kappa_p': {'title': 'Auto Lion-K', 'tooltip': 'Automatically determines the optimal P-value based on layer dimensions. Uses p=2.0 (Spherical) for 4D (Conv) tensors for stability and rotational invariance, and p=1.0 (Sign) for 2D (Linear) tensors for sparsity. Overrides the manual P-value. Recommend for unet models.', 'type': 'bool'}, 'compile': {'title': 'Compiled Optimizer', 'tooltip': 'Enables PyTorch compilation for the optimizer internal step logic. This is intended to improve performance by allowing PyTorch to fuse operations and optimize the computational graph.', 'type': 'bool'}, 'spectral_normalization': {'title': 'Spectral Scaling', 'tooltip': 'Enables explicit Spectral Normalization to automatically rescale the update magnitude and Weight Decay based on layer dimensions. This allows hyperparameters to transfer seamlessly from small to large models without retuning, while making the optimizer highly robust to a wide range of learning rates. This ensures consistent performance across different model sizes, adapter methods, and ranks. ', 'type': 'bool'}, - 'freeze_on_flip': {'title': 'Freeze-on-Flip', 'tooltip': 'Projected One-hit freeze. Masks updates for coordinates where the gradient sign flips compared to the previous step. Over the steps, this makes the signed update semi-continuous.', 'type': 'bool'}, - 'l1_adaptive': {'title': 'L1 Adaptive', 'tooltip': 'Scales the learning rate dynamically by the L1 norm of the gradient to handle gradient heterogeneity. This makes the signed-optimizers semi-adaptive.', 'type': 'bool'}, + 'stochastic_sign': {'title': 'Adaptive Sign', 'tooltip': 'Applies Adaptive Sign operation, that respects the geometry and direction of the original gradient, and scales the learning rate dynamically by the L1 norm of the gradient. This makes the signed-optimizers adaptive and more robust.', 'type': 'bool'}, 'centered_wd': {'title': 'Centered Weight Decay', 'tooltip': 'Centered Weight Decay coefficient. Instead of decaying weights toward zero, they are decayed toward their initial values (anchors). This can be used together with standard weight decay.', 'type': 'float'}, 'centered_wd_mode': {'title': 'Centered WD Mode', 'tooltip': """The quantization format used to store the anchor weights to save VRAM. Options include: 'full': Stores anchors in the original parameter's precision. 'float8': Uses torch.float8_e4m3fn for a balance of precision and memory. 'int8': Uses 8-bit block-wise quantization. 'int4': Uses 4-bit block-wise quantization.""", 'type': 'CenteredWDMode'}, 'factored_2nd': {'title': 'Factored 2nd', 'tooltip': 'Whether to keep the first moment uncompressed (dense), while only factorizing the second moment. This makes the optimizer highly robust to a wide range of LRs, mimicking high-order optimization.', 'type': 'bool'}, diff --git a/modules/util/config/TrainConfig.py b/modules/util/config/TrainConfig.py index e78161eda..f4f97a2c0 100644 --- a/modules/util/config/TrainConfig.py +++ b/modules/util/config/TrainConfig.py @@ -139,8 +139,7 @@ class TrainOptimizerConfig(BaseConfig): auto_kappa_p: False compile: False spectral_normalization: False - freeze_on_flip: False - l1_adaptive: False + stochastic_sign: False centered_wd: float centered_wd_mode: str factored_2nd: False @@ -262,8 +261,7 @@ def default_values(): data.append(("auto_kappa_p", False, bool, False)) data.append(("compile", False, bool, False)) data.append(("spectral_normalization", False, bool, False)) - data.append(("freeze_on_flip", False, bool, False)) - data.append(("l1_adaptive", False, bool, False)) + data.append(("stochastic_sign", False, bool, False)) data.append(("centered_wd", 0.0, float, False)) data.append(("centered_wd_mode", "float8", str, False)) data.append(("factored_2nd", False, bool, False)) diff --git a/modules/util/create.py b/modules/util/create.py index 1fbe86807..0a13df541 100644 --- a/modules/util/create.py +++ b/modules/util/create.py @@ -797,8 +797,7 @@ def create_optimizer( Simplified_AdEMAMix=optimizer_config.Simplified_AdEMAMix if optimizer_config.Simplified_AdEMAMix is not None else False, alpha_grad=optimizer_config.alpha_grad if optimizer_config.alpha_grad is not None else 100, spectral_normalization=optimizer_config.spectral_normalization if optimizer_config.spectral_normalization is not None else False, - freeze_on_flip=optimizer_config.freeze_on_flip if optimizer_config.freeze_on_flip is not None else False, - l1_adaptive=optimizer_config.l1_adaptive if optimizer_config.l1_adaptive is not None else False, + stochastic_sign=optimizer_config.stochastic_sign if optimizer_config.stochastic_sign is not None else False, centered_wd=optimizer_config.centered_wd if optimizer_config.centered_wd is not None else 0.0, centered_wd_mode=optimizer_config.centered_wd_mode if optimizer_config.centered_wd_mode is not None else "full", state_precision=optimizer_config.state_precision if optimizer_config.state_precision is not None else "auto", @@ -820,8 +819,7 @@ def create_optimizer( orthogonal_gradient=optimizer_config.orthogonal_gradient if optimizer_config.orthogonal_gradient is not None else False, auto_kappa_p=optimizer_config.auto_kappa_p if optimizer_config.auto_kappa_p is not None else False, compiled_optimizer=optimizer_config.compile if optimizer_config.compile is not None else False, - freeze_on_flip=optimizer_config.freeze_on_flip if optimizer_config.freeze_on_flip is not None else False, - l1_adaptive=optimizer_config.l1_adaptive if optimizer_config.l1_adaptive is not None else False, + stochastic_sign=optimizer_config.stochastic_sign if optimizer_config.stochastic_sign is not None else False, spectral_normalization=optimizer_config.spectral_normalization if optimizer_config.spectral_normalization is not None else False, centered_wd=optimizer_config.centered_wd if optimizer_config.centered_wd is not None else 0.0, centered_wd_mode=optimizer_config.centered_wd_mode if optimizer_config.centered_wd_mode is not None else "full", diff --git a/modules/util/optimizer_util.py b/modules/util/optimizer_util.py index 69dd24334..c5abc86e8 100644 --- a/modules/util/optimizer_util.py +++ b/modules/util/optimizer_util.py @@ -544,8 +544,7 @@ def init_model_parameters( "orthogonal_gradient": False, "Simplified_AdEMAMix": False, "alpha_grad": 100.0, - "freeze_on_flip": False, - "l1_adaptive": False, + "stochastic_sign": False, "spectral_normalization": False, "centered_wd": 0.0, "centered_wd_mode": "float8", @@ -563,8 +562,7 @@ def init_model_parameters( "fused_back_pass": False, "cautious_mask": False, "orthogonal_gradient": False, - "freeze_on_flip": False, - "l1_adaptive": False, + "stochastic_sign": False, "spectral_normalization": False, "centered_wd": 0.0, "centered_wd_mode": "float8", From 8d2b1de80451c6755730ed968c7a175cecc9b69c Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Sat, 9 May 2026 22:06:22 +0300 Subject: [PATCH 36/84] add normed_momentum (Normalization then Momentum) --- modules/ui/OptimizerParamsWindow.py | 1 + modules/util/config/TrainConfig.py | 2 ++ modules/util/create.py | 4 +++- modules/util/optimizer_util.py | 6 ++++-- 4 files changed, 10 insertions(+), 3 deletions(-) diff --git a/modules/ui/OptimizerParamsWindow.py b/modules/ui/OptimizerParamsWindow.py index 1c700b681..354aacd27 100644 --- a/modules/ui/OptimizerParamsWindow.py +++ b/modules/ui/OptimizerParamsWindow.py @@ -203,6 +203,7 @@ def create_dynamic_ui( 'state_precision': {'title': 'state_precision', 'tooltip': """The quantization format used to store the optimizer states to save VRAM. Options include: 'auto': Stores the states in the original parameter's precision. 'factored': Enables a memory-efficient mode by applying fast low-rank factorization to the optimizers states. It combines factorization for magnitudes with 1-bit compression for signs, drastically reducing VRAM usage and allowing for larger models or batch sizes. 'bf16_sr': Uses BF16 with stochastic rounding for a balance of precision and memory. 'int8_sr': Uses 8-bit block-wise quantization with stochastic rounding. 'fp8_sr': Uses torch.float8_e4m3fn with stochastic rounding.""", 'type': 'StatePrecision'}, 'orthogonal_sinkhorn': {'title': 'Orthogonal Sinkhorn', 'tooltip': 'Applies iterative row and column orthogonal projection to make the updates orthogonal to the current weight, leading to robust regularization and better generalization.', 'type': 'bool'}, 'sinkhorn_iterations': {'title': 'Sinkhorn Iterations', 'tooltip': 'Controls the number of iterations for Multi-Normed Sinkhorn. While 1 iteration is often sufficient for convergence and 3 offers a slight refinement, 5 is the default.', 'type': 'int'}, + 'normed_momentum': {'title': 'Normalization-then-Momentum', 'tooltip': 'Applies the momentum after the optimizer normalization. This makes the momentum scale invariant and tracks the true variance of the normalized gradients.', 'type': 'bool'}, } # @formatter:on diff --git a/modules/util/config/TrainConfig.py b/modules/util/config/TrainConfig.py index f4f97a2c0..43195cdd3 100644 --- a/modules/util/config/TrainConfig.py +++ b/modules/util/config/TrainConfig.py @@ -147,6 +147,7 @@ class TrainOptimizerConfig(BaseConfig): state_precision: str orthogonal_sinkhorn: False sinkhorn_iterations: int + normed_momentum: False def __init__(self, data: list[(str, Any, type, bool)]): super().__init__(data) @@ -269,6 +270,7 @@ def default_values(): data.append(("state_precision", "float8", str, False)) data.append(("orthogonal_sinkhorn", False, bool, False)) data.append(("sinkhorn_iterations", None, int, True)) + data.append(("normed_momentum", False, bool, False)) return TrainOptimizerConfig(data) diff --git a/modules/util/create.py b/modules/util/create.py index 0a13df541..172a67483 100644 --- a/modules/util/create.py +++ b/modules/util/create.py @@ -673,6 +673,7 @@ def create_optimizer( optimizer = SinkSGD_adv( params=parameters, lr=config.learning_rate, + normed_momentum=optimizer_config.normed_momentum if optimizer_config.normed_momentum is not None else False, momentum=optimizer_config.momentum if optimizer_config.beta1 is not None else 0, weight_decay=optimizer_config.weight_decay if optimizer_config.weight_decay is not None else 0.0, nesterov=optimizer_config.nesterov if optimizer_config.nesterov is not None else True, @@ -788,6 +789,8 @@ def create_optimizer( optimizer = SignSGD_adv( params=parameters, lr=config.learning_rate, + stochastic_sign=optimizer_config.stochastic_sign if optimizer_config.stochastic_sign is not None else False, + normed_momentum=optimizer_config.normed_momentum if optimizer_config.normed_momentum is not None else False, momentum=optimizer_config.momentum if optimizer_config.momentum is not None else 0, weight_decay=optimizer_config.weight_decay if optimizer_config.weight_decay is not None else 0.0, cautious_wd=optimizer_config.cautious_wd if optimizer_config.cautious_wd is not None else False, @@ -797,7 +800,6 @@ def create_optimizer( Simplified_AdEMAMix=optimizer_config.Simplified_AdEMAMix if optimizer_config.Simplified_AdEMAMix is not None else False, alpha_grad=optimizer_config.alpha_grad if optimizer_config.alpha_grad is not None else 100, spectral_normalization=optimizer_config.spectral_normalization if optimizer_config.spectral_normalization is not None else False, - stochastic_sign=optimizer_config.stochastic_sign if optimizer_config.stochastic_sign is not None else False, centered_wd=optimizer_config.centered_wd if optimizer_config.centered_wd is not None else 0.0, centered_wd_mode=optimizer_config.centered_wd_mode if optimizer_config.centered_wd_mode is not None else "full", state_precision=optimizer_config.state_precision if optimizer_config.state_precision is not None else "auto", diff --git a/modules/util/optimizer_util.py b/modules/util/optimizer_util.py index c5abc86e8..e0caeba92 100644 --- a/modules/util/optimizer_util.py +++ b/modules/util/optimizer_util.py @@ -443,6 +443,7 @@ def init_model_parameters( "quant_block_size": 2048 }, Optimizer.SINKSGD_ADV: { + "normed_momentum": False, "momentum": 0.95, "nesterov": True, "sinkhorn_iterations": 5, @@ -535,6 +536,8 @@ def init_model_parameters( "centered_wd_mode": "float8", }, Optimizer.SIGNSGD_ADV: { + "stochastic_sign": False, + "normed_momentum": False, "momentum": 0.95, "cautious_wd": False, "weight_decay": 0.0, @@ -544,13 +547,13 @@ def init_model_parameters( "orthogonal_gradient": False, "Simplified_AdEMAMix": False, "alpha_grad": 100.0, - "stochastic_sign": False, "spectral_normalization": False, "centered_wd": 0.0, "centered_wd_mode": "float8", "state_precision": "auto", }, Optimizer.LION_ADV: { + "stochastic_sign": False, "beta1": 0.9, "beta2": 0.99, "cautious_wd": False, @@ -562,7 +565,6 @@ def init_model_parameters( "fused_back_pass": False, "cautious_mask": False, "orthogonal_gradient": False, - "stochastic_sign": False, "spectral_normalization": False, "centered_wd": 0.0, "centered_wd_mode": "float8", From 5005b208c36b1b7447e7df2e3ea24bde50714211 Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Sat, 9 May 2026 22:07:45 +0300 Subject: [PATCH 37/84] remove kappa_p --- modules/ui/OptimizerParamsWindow.py | 1 - modules/util/config/TrainConfig.py | 2 -- modules/util/create.py | 1 - 3 files changed, 4 deletions(-) diff --git a/modules/ui/OptimizerParamsWindow.py b/modules/ui/OptimizerParamsWindow.py index 354aacd27..a08475d07 100644 --- a/modules/ui/OptimizerParamsWindow.py +++ b/modules/ui/OptimizerParamsWindow.py @@ -192,7 +192,6 @@ def create_dynamic_ui( 'accelerated_ns': {'title': 'Accelerated Newton-Schulz', 'tooltip': 'Applies an enhanced Newton-Schulz variant that replaces heuristic coefficients with optimal coefficients derived at each step. This improves performance and convergence by reducing the number of required operations.', 'type': 'bool'}, 'cautious_wd': {'title': 'Cautious Weight Decay', 'tooltip': 'Applies weight decay only to parameter coordinates whose signs align with the optimizer update direction. This preserves the original optimization objective while still benefiting from regularization effects, leading to improved convergence and better final performance.', 'type': 'bool'}, 'approx_mars': {'title': 'Approx MARS-M', 'tooltip': 'Enables Approximated MARS-M, a variance reduction technique. It uses the previous step\'s gradient to correct the current update, leading to lower losses and improved convergence stability. This requires additional state to store the previous gradient.', 'type': 'bool'}, - 'auto_kappa_p': {'title': 'Auto Lion-K', 'tooltip': 'Automatically determines the optimal P-value based on layer dimensions. Uses p=2.0 (Spherical) for 4D (Conv) tensors for stability and rotational invariance, and p=1.0 (Sign) for 2D (Linear) tensors for sparsity. Overrides the manual P-value. Recommend for unet models.', 'type': 'bool'}, 'compile': {'title': 'Compiled Optimizer', 'tooltip': 'Enables PyTorch compilation for the optimizer internal step logic. This is intended to improve performance by allowing PyTorch to fuse operations and optimize the computational graph.', 'type': 'bool'}, 'spectral_normalization': {'title': 'Spectral Scaling', 'tooltip': 'Enables explicit Spectral Normalization to automatically rescale the update magnitude and Weight Decay based on layer dimensions. This allows hyperparameters to transfer seamlessly from small to large models without retuning, while making the optimizer highly robust to a wide range of learning rates. This ensures consistent performance across different model sizes, adapter methods, and ranks. ', 'type': 'bool'}, 'stochastic_sign': {'title': 'Adaptive Sign', 'tooltip': 'Applies Adaptive Sign operation, that respects the geometry and direction of the original gradient, and scales the learning rate dynamically by the L1 norm of the gradient. This makes the signed-optimizers adaptive and more robust.', 'type': 'bool'}, diff --git a/modules/util/config/TrainConfig.py b/modules/util/config/TrainConfig.py index 43195cdd3..83aa34213 100644 --- a/modules/util/config/TrainConfig.py +++ b/modules/util/config/TrainConfig.py @@ -136,7 +136,6 @@ class TrainOptimizerConfig(BaseConfig): accelerated_ns: False cautious_wd: False approx_mars: False - auto_kappa_p: False compile: False spectral_normalization: False stochastic_sign: False @@ -259,7 +258,6 @@ def default_values(): data.append(("accelerated_ns", False, bool, False)) data.append(("cautious_wd", False, bool, False)) data.append(("approx_mars", False, bool, False)) - data.append(("auto_kappa_p", False, bool, False)) data.append(("compile", False, bool, False)) data.append(("spectral_normalization", False, bool, False)) data.append(("stochastic_sign", False, bool, False)) diff --git a/modules/util/create.py b/modules/util/create.py index 172a67483..3ee70d9a7 100644 --- a/modules/util/create.py +++ b/modules/util/create.py @@ -819,7 +819,6 @@ def create_optimizer( cautious_wd=optimizer_config.cautious_wd if optimizer_config.cautious_wd is not None else False, stochastic_rounding=optimizer_config.stochastic_rounding, orthogonal_gradient=optimizer_config.orthogonal_gradient if optimizer_config.orthogonal_gradient is not None else False, - auto_kappa_p=optimizer_config.auto_kappa_p if optimizer_config.auto_kappa_p is not None else False, compiled_optimizer=optimizer_config.compile if optimizer_config.compile is not None else False, stochastic_sign=optimizer_config.stochastic_sign if optimizer_config.stochastic_sign is not None else False, spectral_normalization=optimizer_config.spectral_normalization if optimizer_config.spectral_normalization is not None else False, From 147e6b9b59b64bc2af2e463060e7577ca6627b88 Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Sat, 9 May 2026 22:11:06 +0300 Subject: [PATCH 38/84] remove Simplified_AdEMAMix --- modules/ui/MuonAdamWindow.py | 2 -- modules/ui/OptimizerParamsWindow.py | 2 -- modules/util/config/TrainConfig.py | 4 ---- modules/util/create.py | 10 ---------- modules/util/optimizer_util.py | 10 ---------- 5 files changed, 28 deletions(-) diff --git a/modules/ui/MuonAdamWindow.py b/modules/ui/MuonAdamWindow.py index f043d290e..216813754 100644 --- a/modules/ui/MuonAdamWindow.py +++ b/modules/ui/MuonAdamWindow.py @@ -77,8 +77,6 @@ def create_adam_params_ui(self, master): 'use_atan2': {'title': 'Atan2 Scaling', 'tooltip': 'A robust replacement for eps, which also incorporates gradient clipping, bounding and stabilizing the optimizer updates.', 'type': 'bool'}, 'use_AdEMAMix': {'title': 'AdEMAMix EMA', 'tooltip': 'Adds a second, slow-moving EMA, which is combined with the primary momentum to stabilize updates, and accelerate the training.', 'type': 'bool'}, 'beta3_ema': {'title': 'Beta3 EMA', 'tooltip': 'Coefficient for slow-moving EMA of AdEMAMix.', 'type': 'float'}, - 'Simplified_AdEMAMix': {'title': 'Simplified AdEMAMix', 'tooltip': "Enables a simplified, single-EMA variant of AdEMAMix. Instead of blending two moving averages (fast and slow momentum), this version combines the raw current gradient (controlled by 'Grad α') directly with a single theory-based momentum. This makes the optimizer highly responsive to recent gradient information, which can accelerate training in all batch size scenarios when tuned correctly.", 'type': 'bool'}, - 'alpha_grad': {'title': 'Grad α', 'tooltip': 'Controls the mixing coefficient between raw gradients and momentum gradients in Simplified AdEMAMix. Higher values (e.g., 10-100) emphasize recent gradients, suitable for small batch sizes to reduce noise. Lower values (e.g., 0-1) emphasize historical gradients, suitable for large batch sizes for stability. Setting to 0 uses only momentum gradients without raw gradient contribution.', 'type': 'float'}, 'kourkoutas_beta': {'title': 'Kourkoutas Beta', 'tooltip': 'Enables a layer-wise dynamic β₂ adaptation. This feature makes the optimizer more responsive to "spiky" gradients by lowering β₂ during periods of high variance, and more stable during calm periods by raising β₂ towards its maximum. It can significantly improve training stability and final loss.', 'type': 'bool'}, } # @formatter:on diff --git a/modules/ui/OptimizerParamsWindow.py b/modules/ui/OptimizerParamsWindow.py index a08475d07..ae9984528 100644 --- a/modules/ui/OptimizerParamsWindow.py +++ b/modules/ui/OptimizerParamsWindow.py @@ -173,8 +173,6 @@ def create_dynamic_ui( 'beta3_ema': {'title': 'Beta3 EMA', 'tooltip': 'Coefficient for slow-moving EMA of AdEMAMix.', 'type': 'float'}, 'beta1_warmup': {'title': 'Beta1 Warmup Steps', 'tooltip': 'Number of warmup steps to gradually increase beta1 from Minimum Beta1 Value to its final value. During warmup, beta1 increases linearly. leave it empty to disable warmup and use constant beta1.', 'type': 'int'}, 'min_beta1': {'title': 'Minimum Beta1', 'tooltip': 'Starting beta1 value for warmup scheduling. Used only when beta1 warmup is enabled. Lower values allow faster initial adaptation, while higher values provide more smoothing. The final beta1 value is specified in the beta1 parameter.', 'type': 'float'}, - 'Simplified_AdEMAMix': {'title': 'Simplified AdEMAMix', 'tooltip': "Enables a simplified, single-EMA variant of AdEMAMix. Instead of blending two moving averages (fast and slow momentum), this version combines the raw current gradient (controlled by 'Grad α') directly with a single theory-based momentum. This makes the optimizer highly responsive to recent gradient information, which can accelerate training in all batch size scenarios when tuned correctly.", 'type': 'bool'}, - 'alpha_grad': {'title': 'Grad α', 'tooltip': 'Controls the mixing coefficient between raw gradients and momentum gradients in Simplified AdEMAMix. Higher values (e.g., 10-100) emphasize recent gradients, suitable for small batch sizes to reduce noise. Lower values (e.g., 0-1) emphasize historical gradients, suitable for large batch sizes for stability. Setting to 0 uses only momentum gradients without raw gradient contribution.', 'type': 'float'}, 'kourkoutas_beta': {'title': 'Kourkoutas Beta', 'tooltip': 'Enables a layer-wise dynamic β₂ adaptation. This feature makes the optimizer more responsive to "spiky" gradients by lowering β₂ during periods of high variance, and more stable during calm periods by raising β₂ towards its maximum. It can significantly improve training stability and final loss.', 'type': 'bool'}, 'schedulefree_c': {'title': 'Schedule free averaging strength', 'tooltip': 'Larger values = more responsive (shorter averaging window); smaller values = smoother (longer window). Set to 0 to disable and use the original Schedule-Free rule. Short small batches (≈6-12); long/large-batch (≈50-200).', 'type': 'float'}, 'ns_steps': {'title': 'Newton-Schulz Iterations', 'tooltip': 'Controls the number of iterations for update orthogonalization. Higher values improve the updates quality but make each step slower. Lower values are faster per step but may be less effective.', 'type': 'int'}, diff --git a/modules/util/config/TrainConfig.py b/modules/util/config/TrainConfig.py index 83aa34213..f8654ab8f 100644 --- a/modules/util/config/TrainConfig.py +++ b/modules/util/config/TrainConfig.py @@ -114,10 +114,8 @@ class TrainOptimizerConfig(BaseConfig): use_atan2: False use_AdEMAMix: False beta3_ema: float - alpha_grad: float beta1_warmup: int min_beta1: float - Simplified_AdEMAMix: False kourkoutas_beta: False schedulefree_c: float ns_steps: int @@ -236,10 +234,8 @@ def default_values(): data.append(("use_atan2", False, bool, False)) data.append(("use_AdEMAMix", False, bool, False)) data.append(("beta3_ema", None, float, True)) - data.append(("alpha_grad", None, float, True)) data.append(("beta1_warmup", None, int, True)) data.append(("min_beta1", None, float, True)) - data.append(("Simplified_AdEMAMix", False, bool, False)) data.append(("kourkoutas_beta", False, bool, False)) data.append(("schedulefree_c", None, float, True)) data.append(("ns_steps", None, int, True)) diff --git a/modules/util/create.py b/modules/util/create.py index 3ee70d9a7..f62bea542 100644 --- a/modules/util/create.py +++ b/modules/util/create.py @@ -736,8 +736,6 @@ def create_optimizer( use_AdEMAMix=optimizer_config.use_AdEMAMix if optimizer_config.use_AdEMAMix is not None else False, beta3_ema=optimizer_config.beta3 if optimizer_config.beta3 is not None else 0.9999, alpha=optimizer_config.alpha if optimizer_config.alpha is not None else 5, - Simplified_AdEMAMix=optimizer_config.Simplified_AdEMAMix if optimizer_config.Simplified_AdEMAMix is not None else False, - alpha_grad=optimizer_config.alpha_grad if optimizer_config.alpha_grad is not None else 100, kourkoutas_beta=optimizer_config.kourkoutas_beta if optimizer_config.kourkoutas_beta is not None else False, k_warmup_steps=(config.learning_rate_warmup_steps / config.gradient_accumulation_steps), compiled_optimizer=optimizer_config.compile if optimizer_config.compile is not None else False, @@ -773,8 +771,6 @@ def create_optimizer( use_AdEMAMix=optimizer_config.use_AdEMAMix if optimizer_config.use_AdEMAMix is not None else False, beta3_ema=optimizer_config.beta3_ema if optimizer_config.beta3_ema is not None else 0.9999, alpha=optimizer_config.alpha if optimizer_config.alpha is not None else 5, - Simplified_AdEMAMix=optimizer_config.Simplified_AdEMAMix if optimizer_config.Simplified_AdEMAMix is not None else False, - alpha_grad=optimizer_config.alpha_grad if optimizer_config.alpha_grad is not None else 100, kourkoutas_beta=optimizer_config.kourkoutas_beta if optimizer_config.kourkoutas_beta is not None else False, k_warmup_steps=(config.learning_rate_warmup_steps / config.gradient_accumulation_steps), compiled_optimizer=optimizer_config.compile if optimizer_config.compile is not None else False, @@ -797,8 +793,6 @@ def create_optimizer( stochastic_rounding=optimizer_config.stochastic_rounding, orthogonal_gradient=optimizer_config.orthogonal_gradient if optimizer_config.orthogonal_gradient is not None else False, compiled_optimizer=optimizer_config.compile if optimizer_config.compile is not None else False, - Simplified_AdEMAMix=optimizer_config.Simplified_AdEMAMix if optimizer_config.Simplified_AdEMAMix is not None else False, - alpha_grad=optimizer_config.alpha_grad if optimizer_config.alpha_grad is not None else 100, spectral_normalization=optimizer_config.spectral_normalization if optimizer_config.spectral_normalization is not None else False, centered_wd=optimizer_config.centered_wd if optimizer_config.centered_wd is not None else 0.0, centered_wd_mode=optimizer_config.centered_wd_mode if optimizer_config.centered_wd_mode is not None else "full", @@ -872,8 +866,6 @@ def create_optimizer( orthogonal_gradient=optimizer_config.orthogonal_gradient if optimizer_config.orthogonal_gradient is not None else False, approx_mars=optimizer_config.approx_mars if optimizer_config.approx_mars is not None else False, compiled_optimizer=optimizer_config.compile if optimizer_config.compile is not None else False, - Simplified_AdEMAMix=optimizer_config.Simplified_AdEMAMix if optimizer_config.Simplified_AdEMAMix is not None else False, - alpha_grad=optimizer_config.alpha_grad if optimizer_config.alpha_grad is not None else 100, spectral_normalization=optimizer_config.spectral_normalization if optimizer_config.spectral_normalization is not None else False, centered_wd=optimizer_config.centered_wd if optimizer_config.centered_wd is not None else 0.0, centered_wd_mode=optimizer_config.centered_wd_mode if optimizer_config.centered_wd_mode is not None else "full", @@ -923,8 +915,6 @@ def create_optimizer( stochastic_rounding=optimizer_config.stochastic_rounding, nesterov=optimizer_config.nesterov if optimizer_config.nesterov is not None else True, use_atan2=optimizer_config.use_atan2 if optimizer_config.use_atan2 is not None else False, - Simplified_AdEMAMix=optimizer_config.Simplified_AdEMAMix if optimizer_config.Simplified_AdEMAMix is not None else False, - alpha_grad=optimizer_config.alpha_grad if optimizer_config.alpha_grad is not None else 100, low_rank_ortho=optimizer_config.low_rank_ortho if optimizer_config.low_rank_ortho is not None else False, ortho_rank=optimizer_config.ortho_rank if optimizer_config.ortho_rank is not None else 128, normuon_variant=optimizer_config.normuon_variant if optimizer_config.normuon_variant is not None else False, diff --git a/modules/util/optimizer_util.py b/modules/util/optimizer_util.py index e0caeba92..234280cd4 100644 --- a/modules/util/optimizer_util.py +++ b/modules/util/optimizer_util.py @@ -497,8 +497,6 @@ def init_model_parameters( "use_AdEMAMix": False, "beta3_ema": 0.9999, "alpha": 5, - "Simplified_AdEMAMix": False, - "alpha_grad": 100.0, "kourkoutas_beta": False, "spectral_normalization": False, "centered_wd": 0.0, @@ -529,8 +527,6 @@ def init_model_parameters( "use_AdEMAMix": False, "beta3_ema": 0.9999, "alpha": 5, - "Simplified_AdEMAMix": False, - "alpha_grad": 100.0, "kourkoutas_beta": False, "centered_wd": 0.0, "centered_wd_mode": "float8", @@ -545,8 +541,6 @@ def init_model_parameters( "compile": False, "fused_back_pass": False, "orthogonal_gradient": False, - "Simplified_AdEMAMix": False, - "alpha_grad": 100.0, "spectral_normalization": False, "centered_wd": 0.0, "centered_wd_mode": "float8", @@ -589,8 +583,6 @@ def init_model_parameters( "muon_te1_adam_lr": None, "muon_te2_adam_lr": None, "nesterov": True, - "Simplified_AdEMAMix": False, - "alpha_grad": 100.0, "normuon_variant": True, "beta2_normuon": 0.95, "orthogonal_gradient": False, @@ -623,8 +615,6 @@ def init_model_parameters( "muon_te2_adam_lr": None, "nesterov": False, "use_atan2": False, - "Simplified_AdEMAMix": False, - "alpha_grad": 100.0, "normuon_variant": True, "orthogonal_gradient": False, "approx_mars": False, From bd515daac91aa8a7969896e43fcfd1bdb3fb7428 Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Sat, 9 May 2026 23:37:34 +0300 Subject: [PATCH 39/84] add nesterov and nesterov_coef --- modules/ui/OptimizerParamsWindow.py | 1 + modules/util/config/TrainConfig.py | 2 ++ modules/util/create.py | 11 +++++++++++ modules/util/optimizer_util.py | 15 +++++++++++++-- 4 files changed, 27 insertions(+), 2 deletions(-) diff --git a/modules/ui/OptimizerParamsWindow.py b/modules/ui/OptimizerParamsWindow.py index ae9984528..44c615d09 100644 --- a/modules/ui/OptimizerParamsWindow.py +++ b/modules/ui/OptimizerParamsWindow.py @@ -201,6 +201,7 @@ def create_dynamic_ui( 'orthogonal_sinkhorn': {'title': 'Orthogonal Sinkhorn', 'tooltip': 'Applies iterative row and column orthogonal projection to make the updates orthogonal to the current weight, leading to robust regularization and better generalization.', 'type': 'bool'}, 'sinkhorn_iterations': {'title': 'Sinkhorn Iterations', 'tooltip': 'Controls the number of iterations for Multi-Normed Sinkhorn. While 1 iteration is often sufficient for convergence and 3 offers a slight refinement, 5 is the default.', 'type': 'int'}, 'normed_momentum': {'title': 'Normalization-then-Momentum', 'tooltip': 'Applies the momentum after the optimizer normalization. This makes the momentum scale invariant and tracks the true variance of the normalized gradients.', 'type': 'bool'}, + 'nesterov_coef': {'title': 'Nesterov Coef', 'tooltip': 'Controls the mixing coefficient between momentum gradients and raw gradients in Nesterov momentum. For a factor of 0.8, the final update will be 80% of the momentum gradients and 20% raw gradient. Leaving it unset toggles the standard Nestrov behavior (where nesterov_coef = beta1 or momentum). Setting it to 0 cancels momentum contribution.', 'type': 'float'}, } # @formatter:on diff --git a/modules/util/config/TrainConfig.py b/modules/util/config/TrainConfig.py index f8654ab8f..a8374c87e 100644 --- a/modules/util/config/TrainConfig.py +++ b/modules/util/config/TrainConfig.py @@ -145,6 +145,7 @@ class TrainOptimizerConfig(BaseConfig): orthogonal_sinkhorn: False sinkhorn_iterations: int normed_momentum: False + nesterov_coef: float def __init__(self, data: list[(str, Any, type, bool)]): super().__init__(data) @@ -265,6 +266,7 @@ def default_values(): data.append(("orthogonal_sinkhorn", False, bool, False)) data.append(("sinkhorn_iterations", None, int, True)) data.append(("normed_momentum", False, bool, False)) + data.append(("nesterov_coef", None, float, True)) return TrainOptimizerConfig(data) diff --git a/modules/util/create.py b/modules/util/create.py index f62bea542..a1cb08aa2 100644 --- a/modules/util/create.py +++ b/modules/util/create.py @@ -677,6 +677,7 @@ def create_optimizer( momentum=optimizer_config.momentum if optimizer_config.beta1 is not None else 0, weight_decay=optimizer_config.weight_decay if optimizer_config.weight_decay is not None else 0.0, nesterov=optimizer_config.nesterov if optimizer_config.nesterov is not None else True, + nesterov_coef=optimizer_config.nesterov_coef if optimizer_config.nesterov_coef is not None else None, cautious_wd=optimizer_config.cautious_wd if optimizer_config.cautious_wd is not None else False, stochastic_rounding=optimizer_config.stochastic_rounding, orthogonal_gradient=optimizer_config.orthogonal_gradient if optimizer_config.orthogonal_gradient is not None else False, @@ -715,6 +716,8 @@ def create_optimizer( centered_wd=optimizer_config.centered_wd if optimizer_config.centered_wd is not None else 0.0, centered_wd_mode=optimizer_config.centered_wd_mode if optimizer_config.centered_wd_mode is not None else "full", state_precision=optimizer_config.state_precision if optimizer_config.state_precision is not None else "auto", + nesterov=optimizer_config.nesterov if optimizer_config.nesterov is not None else False, + nesterov_coef=optimizer_config.nesterov_coef if optimizer_config.nesterov_coef is not None else None, ) # ADOPT_ADV Optimizer @@ -743,6 +746,8 @@ def create_optimizer( centered_wd=optimizer_config.centered_wd if optimizer_config.centered_wd is not None else 0.0, centered_wd_mode=optimizer_config.centered_wd_mode if optimizer_config.centered_wd_mode is not None else "full", state_precision=optimizer_config.state_precision if optimizer_config.state_precision is not None else "auto", + nesterov=optimizer_config.nesterov if optimizer_config.nesterov is not None else False, + nesterov_coef=optimizer_config.nesterov_coef if optimizer_config.nesterov_coef is not None else None, ) # PRODIGY_ADV Optimizer @@ -777,6 +782,8 @@ def create_optimizer( spectral_normalization=optimizer_config.spectral_normalization if optimizer_config.spectral_normalization is not None else False, centered_wd=optimizer_config.centered_wd if optimizer_config.centered_wd is not None else 0.0, centered_wd_mode=optimizer_config.centered_wd_mode if optimizer_config.centered_wd_mode is not None else "full", + nesterov=optimizer_config.nesterov if optimizer_config.nesterov is not None else False, + nesterov_coef=optimizer_config.nesterov_coef if optimizer_config.nesterov_coef is not None else None, ) # SignSGD_ADV Optimizer @@ -797,6 +804,8 @@ def create_optimizer( centered_wd=optimizer_config.centered_wd if optimizer_config.centered_wd is not None else 0.0, centered_wd_mode=optimizer_config.centered_wd_mode if optimizer_config.centered_wd_mode is not None else "full", state_precision=optimizer_config.state_precision if optimizer_config.state_precision is not None else "auto", + nesterov=optimizer_config.nesterov if optimizer_config.nesterov is not None else False, + nesterov_coef=optimizer_config.nesterov_coef if optimizer_config.nesterov_coef is not None else None, ) # LION_ADV Optimizer @@ -858,6 +867,7 @@ def create_optimizer( cautious_wd=optimizer_config.cautious_wd if optimizer_config.cautious_wd is not None else False, stochastic_rounding=optimizer_config.stochastic_rounding, nesterov=optimizer_config.nesterov if optimizer_config.nesterov is not None else True, + nesterov_coef=optimizer_config.nesterov_coef if optimizer_config.nesterov_coef is not None else None, normuon_variant=optimizer_config.normuon_variant if optimizer_config.normuon_variant is not None else False, beta2_normuon=optimizer_config.beta2_normuon if optimizer_config.beta2_normuon is not None else 0.95, low_rank_ortho=optimizer_config.low_rank_ortho if optimizer_config.low_rank_ortho is not None else False, @@ -914,6 +924,7 @@ def create_optimizer( cautious_wd=optimizer_config.cautious_wd if optimizer_config.cautious_wd is not None else False, stochastic_rounding=optimizer_config.stochastic_rounding, nesterov=optimizer_config.nesterov if optimizer_config.nesterov is not None else True, + nesterov_coef=optimizer_config.nesterov_coef if optimizer_config.nesterov_coef is not None else None, use_atan2=optimizer_config.use_atan2 if optimizer_config.use_atan2 is not None else False, low_rank_ortho=optimizer_config.low_rank_ortho if optimizer_config.low_rank_ortho is not None else False, ortho_rank=optimizer_config.ortho_rank if optimizer_config.ortho_rank is not None else 128, diff --git a/modules/util/optimizer_util.py b/modules/util/optimizer_util.py index 234280cd4..aa36ff944 100644 --- a/modules/util/optimizer_util.py +++ b/modules/util/optimizer_util.py @@ -446,6 +446,7 @@ def init_model_parameters( "normed_momentum": False, "momentum": 0.95, "nesterov": True, + "nesterov_coef": None, "sinkhorn_iterations": 5, "cautious_wd": False, "weight_decay": 0.0, @@ -461,6 +462,8 @@ def init_model_parameters( }, Optimizer.ADAMW_ADV: { "beta1": 0.9, + "nesterov": False, + "nesterov_coef": None, "beta2": 0.99, "eps": 1e-8, "fisher_wd": False, @@ -483,6 +486,8 @@ def init_model_parameters( }, Optimizer.ADOPT_ADV: { "beta1": 0.9, + "nesterov": False, + "nesterov_coef": None, "beta2": 0.9999, "eps": 1e-6, "fisher_wd": False, @@ -505,6 +510,8 @@ def init_model_parameters( }, Optimizer.PRODIGY_ADV: { "beta1": 0.9, + "nesterov": False, + "nesterov_coef": None, "beta2": 0.99, "beta3": None, "eps": 1e-8, @@ -535,6 +542,8 @@ def init_model_parameters( "stochastic_sign": False, "normed_momentum": False, "momentum": 0.95, + "nesterov": False, + "nesterov_coef": None, "cautious_wd": False, "weight_decay": 0.0, "stochastic_rounding": True, @@ -565,6 +574,8 @@ def init_model_parameters( }, Optimizer.MUON_ADV: { "beta1": 0.9, + "nesterov": True, + "nesterov_coef": None, "cautious_wd": False, "weight_decay": 0.0, "accelerated_ns": False, @@ -582,7 +593,6 @@ def init_model_parameters( "muon_adam_lr": 1e-6, "muon_te1_adam_lr": None, "muon_te2_adam_lr": None, - "nesterov": True, "normuon_variant": True, "beta2_normuon": 0.95, "orthogonal_gradient": False, @@ -594,6 +604,8 @@ def init_model_parameters( }, Optimizer.ADAMUON_ADV: { "beta1": 0.95, + "nesterov": True, + "nesterov_coef": None, "beta2": 0.95, "eps": 1e-8, "cautious_wd": False, @@ -613,7 +625,6 @@ def init_model_parameters( "muon_adam_lr": 1e-6, "muon_te1_adam_lr": None, "muon_te2_adam_lr": None, - "nesterov": False, "use_atan2": False, "normuon_variant": True, "orthogonal_gradient": False, From 91e06da00fe232db1554069aedb4e2ed21130293 Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Sat, 9 May 2026 23:42:15 +0300 Subject: [PATCH 40/84] Scale invariant eps when eps=None --- modules/util/create.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/modules/util/create.py b/modules/util/create.py index a1cb08aa2..79abd770a 100644 --- a/modules/util/create.py +++ b/modules/util/create.py @@ -698,7 +698,7 @@ def create_optimizer( lr=config.learning_rate, betas=(optimizer_config.beta1 if optimizer_config.beta1 is not None else 0, optimizer_config.beta2 if optimizer_config.beta2 is not None else 0.99), - eps=optimizer_config.eps if optimizer_config.eps is not None else 1e-8, + eps=optimizer_config.eps if optimizer_config.eps is not None else None, weight_decay=optimizer_config.weight_decay if optimizer_config.weight_decay is not None else 0.0, factored_2nd=optimizer_config.factored_2nd if optimizer_config.factored_2nd is not None else False, fisher_wd=optimizer_config.fisher_wd if optimizer_config.fisher_wd is not None else False, @@ -728,7 +728,7 @@ def create_optimizer( lr=config.learning_rate, betas=(optimizer_config.beta1 if optimizer_config.beta1 is not None else 0, optimizer_config.beta2 if optimizer_config.beta2 is not None else 0.9999), - eps=optimizer_config.eps if optimizer_config.eps is not None else 1e-6, + eps=optimizer_config.eps if optimizer_config.eps is not None else None, weight_decay=optimizer_config.weight_decay if optimizer_config.weight_decay is not None else 0.0, factored_2nd=optimizer_config.factored_2nd if optimizer_config.factored_2nd is not None else False, fisher_wd=optimizer_config.fisher_wd if optimizer_config.fisher_wd is not None else False, @@ -759,7 +759,7 @@ def create_optimizer( betas=(optimizer_config.beta1 if optimizer_config.beta1 is not None else 0, optimizer_config.beta2 if optimizer_config.beta2 is not None else 0.99), beta3=optimizer_config.beta3 if optimizer_config.beta3 is not None else None, - eps=optimizer_config.eps if optimizer_config.eps is not None else 1e-8, + eps=optimizer_config.eps if optimizer_config.eps is not None else None, weight_decay=optimizer_config.weight_decay if optimizer_config.weight_decay is not None else 0.0, factored_2nd=optimizer_config.factored_2nd if optimizer_config.factored_2nd is not None else False, fisher_wd=optimizer_config.fisher_wd if optimizer_config.fisher_wd is not None else False, From 8aa14e143d8e6525cdee58ea84ba0c97d8f02d26 Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Sat, 9 May 2026 23:53:19 +0300 Subject: [PATCH 41/84] update create_adam_params_ui --- modules/ui/MuonAdamWindow.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/modules/ui/MuonAdamWindow.py b/modules/ui/MuonAdamWindow.py index 216813754..c757e5cd3 100644 --- a/modules/ui/MuonAdamWindow.py +++ b/modules/ui/MuonAdamWindow.py @@ -78,6 +78,10 @@ def create_adam_params_ui(self, master): 'use_AdEMAMix': {'title': 'AdEMAMix EMA', 'tooltip': 'Adds a second, slow-moving EMA, which is combined with the primary momentum to stabilize updates, and accelerate the training.', 'type': 'bool'}, 'beta3_ema': {'title': 'Beta3 EMA', 'tooltip': 'Coefficient for slow-moving EMA of AdEMAMix.', 'type': 'float'}, 'kourkoutas_beta': {'title': 'Kourkoutas Beta', 'tooltip': 'Enables a layer-wise dynamic β₂ adaptation. This feature makes the optimizer more responsive to "spiky" gradients by lowering β₂ during periods of high variance, and more stable during calm periods by raising β₂ towards its maximum. It can significantly improve training stability and final loss.', 'type': 'bool'}, + 'nesterov_coef': {'title': 'Nesterov Coef', 'tooltip': 'Controls the mixing coefficient between momentum gradients and raw gradients in Nesterov momentum. For a factor of 0.8, the final update will be 80% of the momentum gradients and 20% raw gradient. Leaving it unset toggles the standard Nestrov behavior (where nesterov_coef = beta1 or momentum). Setting it to 0 cancels momentum contribution.', 'type': 'float'}, + 'factored_2nd': {'title': 'Factored 2nd', 'tooltip': 'Whether to keep the first moment uncompressed (dense), while only factorizing the second moment. This makes the optimizer highly robust to a wide range of LRs, mimicking high-order optimization.', 'type': 'bool'}, + 'fisher_wd': {'title': 'Fisher Weight Decay', 'tooltip': 'Applies adaptive, scale-invariant weight-decay regularization based on the Fisher Information Matrix (approximated by Adam\'s second moment). It reduces penalty for "important" high-curvature weights while accelerating decay for "useless" weights in flat regions. Leading to improved convergence and better final performance.', 'type': 'bool'}, + 'state_precision': {'title': 'state_precision', 'tooltip': """The quantization format used to store the optimizer states to save VRAM. Options include: 'auto': Stores the states in the original parameter's precision. 'factored': Enables a memory-efficient mode by applying fast low-rank factorization to the optimizers states. It combines factorization for magnitudes with 1-bit compression for signs, drastically reducing VRAM usage and allowing for larger models or batch sizes. 'bf16_sr': Uses BF16 with stochastic rounding for a balance of precision and memory. 'int8_sr': Uses 8-bit block-wise quantization with stochastic rounding. 'fp8_sr': Uses torch.float8_e4m3fn with stochastic rounding.""", 'type': 'StatePrecision'}, } # @formatter:on From 0f8d59ee2fff74942c641b9fc0a1b8c1eb9a51fe Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Sun, 10 May 2026 01:28:19 +0300 Subject: [PATCH 42/84] dev14 --- requirements-global.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements-global.txt b/requirements-global.txt index f6e292bc0..b69996f7b 100644 --- a/requirements-global.txt +++ b/requirements-global.txt @@ -41,7 +41,7 @@ prodigyopt==1.1.2 # prodigy optimizer schedulefree==1.4.1 # schedule-free optimizers pytorch_optimizer==3.6.0 # pytorch optimizers prodigy-plus-schedule-free==2.0.1 # Prodigy plus optimizer -adv_optm==2.4.dev13 # advanced optimizers +adv_optm==2.4.dev14 # advanced optimizers -e git+https://github.com/KellerJordan/Muon.git@f90a42b#egg=muon-optimizer # Profiling From 7889d14862e204eb45d272619805a5d56961cc5f Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Sun, 10 May 2026 19:35:55 +0300 Subject: [PATCH 43/84] dev15 --- requirements-global.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements-global.txt b/requirements-global.txt index b69996f7b..f775e254f 100644 --- a/requirements-global.txt +++ b/requirements-global.txt @@ -41,7 +41,7 @@ prodigyopt==1.1.2 # prodigy optimizer schedulefree==1.4.1 # schedule-free optimizers pytorch_optimizer==3.6.0 # pytorch optimizers prodigy-plus-schedule-free==2.0.1 # Prodigy plus optimizer -adv_optm==2.4.dev14 # advanced optimizers +adv_optm==2.4.dev15 # advanced optimizers -e git+https://github.com/KellerJordan/Muon.git@f90a42b#egg=muon-optimizer # Profiling From cfa876e5b1c5dc8bf1d9fcdb2d13685a3ee302af Mon Sep 17 00:00:00 2001 From: Koratahiu Date: Mon, 11 May 2026 18:47:13 +0300 Subject: [PATCH 44/84] Update requirements-global.txt --- requirements-global.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements-global.txt b/requirements-global.txt index f775e254f..aa9484aab 100644 --- a/requirements-global.txt +++ b/requirements-global.txt @@ -41,7 +41,7 @@ prodigyopt==1.1.2 # prodigy optimizer schedulefree==1.4.1 # schedule-free optimizers pytorch_optimizer==3.6.0 # pytorch optimizers prodigy-plus-schedule-free==2.0.1 # Prodigy plus optimizer -adv_optm==2.4.dev15 # advanced optimizers +adv_optm==2.4.dev16 # advanced optimizers -e git+https://github.com/KellerJordan/Muon.git@f90a42b#egg=muon-optimizer # Profiling From 2d13b17a69175a6e157c27dd947c07cc95c3fe86 Mon Sep 17 00:00:00 2001 From: Koratahiu Date: Tue, 12 May 2026 14:26:20 +0300 Subject: [PATCH 45/84] Update requirements-global.txt --- requirements-global.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements-global.txt b/requirements-global.txt index aa9484aab..8fa1e9cad 100644 --- a/requirements-global.txt +++ b/requirements-global.txt @@ -41,7 +41,7 @@ prodigyopt==1.1.2 # prodigy optimizer schedulefree==1.4.1 # schedule-free optimizers pytorch_optimizer==3.6.0 # pytorch optimizers prodigy-plus-schedule-free==2.0.1 # Prodigy plus optimizer -adv_optm==2.4.dev16 # advanced optimizers +adv_optm==2.4.dev17 # advanced optimizers -e git+https://github.com/KellerJordan/Muon.git@f90a42b#egg=muon-optimizer # Profiling From 46fbf73a406b778d71735b13d596499160db4b08 Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Sun, 24 May 2026 05:54:10 +0300 Subject: [PATCH 46/84] fix SinkSGD typo --- modules/util/create.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modules/util/create.py b/modules/util/create.py index 79abd770a..814197479 100644 --- a/modules/util/create.py +++ b/modules/util/create.py @@ -674,7 +674,7 @@ def create_optimizer( params=parameters, lr=config.learning_rate, normed_momentum=optimizer_config.normed_momentum if optimizer_config.normed_momentum is not None else False, - momentum=optimizer_config.momentum if optimizer_config.beta1 is not None else 0, + momentum=optimizer_config.momentum if optimizer_config.momentum is not None else 0, weight_decay=optimizer_config.weight_decay if optimizer_config.weight_decay is not None else 0.0, nesterov=optimizer_config.nesterov if optimizer_config.nesterov is not None else True, nesterov_coef=optimizer_config.nesterov_coef if optimizer_config.nesterov_coef is not None else None, From 137f4cd1682bc8ee20f1c440b9ddf4aeb7efd1da Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Mon, 25 May 2026 14:04:09 +0300 Subject: [PATCH 47/84] add geometric_wd, centered_vt --- modules/ui/OptimizerParamsWindow.py | 2 ++ modules/util/config/TrainConfig.py | 4 ++++ modules/util/create.py | 2 ++ modules/util/optimizer_util.py | 2 ++ requirements-global.txt | 2 +- 5 files changed, 11 insertions(+), 1 deletion(-) diff --git a/modules/ui/OptimizerParamsWindow.py b/modules/ui/OptimizerParamsWindow.py index 44c615d09..63977d4ba 100644 --- a/modules/ui/OptimizerParamsWindow.py +++ b/modules/ui/OptimizerParamsWindow.py @@ -202,6 +202,8 @@ def create_dynamic_ui( 'sinkhorn_iterations': {'title': 'Sinkhorn Iterations', 'tooltip': 'Controls the number of iterations for Multi-Normed Sinkhorn. While 1 iteration is often sufficient for convergence and 3 offers a slight refinement, 5 is the default.', 'type': 'int'}, 'normed_momentum': {'title': 'Normalization-then-Momentum', 'tooltip': 'Applies the momentum after the optimizer normalization. This makes the momentum scale invariant and tracks the true variance of the normalized gradients.', 'type': 'bool'}, 'nesterov_coef': {'title': 'Nesterov Coef', 'tooltip': 'Controls the mixing coefficient between momentum gradients and raw gradients in Nesterov momentum. For a factor of 0.8, the final update will be 80% of the momentum gradients and 20% raw gradient. Leaving it unset toggles the standard Nestrov behavior (where nesterov_coef = beta1 or momentum). Setting it to 0 cancels momentum contribution.', 'type': 'float'}, + 'centered_vt': {'title': 'Variance Preconditioner', 'tooltip': 'Applies centered, scale-invariant row and column confidence pre-conditioner. This is intended to be used with Normalization-then-Momentum.', 'type': 'bool'}, + 'geometric_wd': {'title': 'Geometric Weight Decay', 'tooltip': 'Regularizes weights based on the geometric structure of the optimizer. Compatible with cautious weight decay.', 'type': 'bool'}, } # @formatter:on diff --git a/modules/util/config/TrainConfig.py b/modules/util/config/TrainConfig.py index a8374c87e..2ff4c5983 100644 --- a/modules/util/config/TrainConfig.py +++ b/modules/util/config/TrainConfig.py @@ -146,6 +146,8 @@ class TrainOptimizerConfig(BaseConfig): sinkhorn_iterations: int normed_momentum: False nesterov_coef: float + centered_vt: False + geometric_wd: False def __init__(self, data: list[(str, Any, type, bool)]): super().__init__(data) @@ -267,6 +269,8 @@ def default_values(): data.append(("sinkhorn_iterations", None, int, True)) data.append(("normed_momentum", False, bool, False)) data.append(("nesterov_coef", None, float, True)) + data.append(("centered_vt", False, bool, False)) + data.append(("geometric_wd", False, bool, False)) return TrainOptimizerConfig(data) diff --git a/modules/util/create.py b/modules/util/create.py index 814197479..0d53512b9 100644 --- a/modules/util/create.py +++ b/modules/util/create.py @@ -688,6 +688,8 @@ def create_optimizer( centered_wd=optimizer_config.centered_wd if optimizer_config.centered_wd is not None else 0.0, centered_wd_mode=optimizer_config.centered_wd_mode if optimizer_config.centered_wd_mode is not None else "full", state_precision=optimizer_config.state_precision if optimizer_config.state_precision is not None else "auto", + centered_vt=optimizer_config.centered_vt if optimizer_config.centered_vt is not None else False, + geometric_wd=optimizer_config.geometric_wd if optimizer_config.geometric_wd is not None else False, ) # ADAMW_ADV Optimizer diff --git a/modules/util/optimizer_util.py b/modules/util/optimizer_util.py index aa36ff944..727aed7bd 100644 --- a/modules/util/optimizer_util.py +++ b/modules/util/optimizer_util.py @@ -449,7 +449,9 @@ def init_model_parameters( "nesterov_coef": None, "sinkhorn_iterations": 5, "cautious_wd": False, + "geometric_wd": False, "weight_decay": 0.0, + "centered_vt": False, "stochastic_rounding": True, "compile": False, "fused_back_pass": False, diff --git a/requirements-global.txt b/requirements-global.txt index 8fa1e9cad..81813de4b 100644 --- a/requirements-global.txt +++ b/requirements-global.txt @@ -41,7 +41,7 @@ prodigyopt==1.1.2 # prodigy optimizer schedulefree==1.4.1 # schedule-free optimizers pytorch_optimizer==3.6.0 # pytorch optimizers prodigy-plus-schedule-free==2.0.1 # Prodigy plus optimizer -adv_optm==2.4.dev17 # advanced optimizers +adv_optm==2.4.dev18 # advanced optimizers -e git+https://github.com/KellerJordan/Muon.git@f90a42b#egg=muon-optimizer # Profiling From 5f0a9e8e43a0ace8a6ee724ae5b31eff88c662f3 Mon Sep 17 00:00:00 2001 From: Koratahiu Date: Wed, 27 May 2026 05:08:04 +0300 Subject: [PATCH 48/84] dev19: sinksgd bugifxes --- requirements-global.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements-global.txt b/requirements-global.txt index 81813de4b..0713619bb 100644 --- a/requirements-global.txt +++ b/requirements-global.txt @@ -41,7 +41,7 @@ prodigyopt==1.1.2 # prodigy optimizer schedulefree==1.4.1 # schedule-free optimizers pytorch_optimizer==3.6.0 # pytorch optimizers prodigy-plus-schedule-free==2.0.1 # Prodigy plus optimizer -adv_optm==2.4.dev18 # advanced optimizers +adv_optm==2.4.dev19 # advanced optimizers -e git+https://github.com/KellerJordan/Muon.git@f90a42b#egg=muon-optimizer # Profiling From 59d90228f5cfa380933cc0f7914b058f9961579a Mon Sep 17 00:00:00 2001 From: Koratahiu Date: Wed, 27 May 2026 05:57:21 +0300 Subject: [PATCH 49/84] Update requirements-global.txt --- requirements-global.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements-global.txt b/requirements-global.txt index 0713619bb..adecbfbf4 100644 --- a/requirements-global.txt +++ b/requirements-global.txt @@ -41,7 +41,7 @@ prodigyopt==1.1.2 # prodigy optimizer schedulefree==1.4.1 # schedule-free optimizers pytorch_optimizer==3.6.0 # pytorch optimizers prodigy-plus-schedule-free==2.0.1 # Prodigy plus optimizer -adv_optm==2.4.dev19 # advanced optimizers +adv_optm==2.4.dev20 # advanced optimizers -e git+https://github.com/KellerJordan/Muon.git@f90a42b#egg=muon-optimizer # Profiling From 47c2539503a1a5298ac6ac2b0553ccf87737a5da Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Sat, 30 May 2026 23:40:32 +0300 Subject: [PATCH 50/84] centered_vt -> snr_cond, add it and geometric_wd to signsgd_adv --- modules/ui/OptimizerParamsWindow.py | 4 ++-- modules/util/config/TrainConfig.py | 4 ++-- modules/util/create.py | 4 +++- modules/util/optimizer_util.py | 4 +++- 4 files changed, 10 insertions(+), 6 deletions(-) diff --git a/modules/ui/OptimizerParamsWindow.py b/modules/ui/OptimizerParamsWindow.py index 0d12baca7..25ec44d37 100644 --- a/modules/ui/OptimizerParamsWindow.py +++ b/modules/ui/OptimizerParamsWindow.py @@ -200,9 +200,9 @@ def create_dynamic_ui( 'state_precision': {'title': 'state_precision', 'tooltip': """The quantization format used to store the optimizer states to save VRAM. Options include: 'auto': Stores the states in the original parameter's precision. 'factored': Enables a memory-efficient mode by applying fast low-rank factorization to the optimizers states. It combines factorization for magnitudes with 1-bit compression for signs, drastically reducing VRAM usage and allowing for larger models or batch sizes. 'bf16_sr': Uses BF16 with stochastic rounding for a balance of precision and memory. 'int8_sr': Uses 8-bit block-wise quantization with stochastic rounding. 'fp8_sr': Uses torch.float8_e4m3fn with stochastic rounding.""", 'type': 'StatePrecision'}, 'orthogonal_sinkhorn': {'title': 'Orthogonal Sinkhorn', 'tooltip': 'Applies iterative row and column orthogonal projection to make the updates orthogonal to the current weight, leading to robust regularization and better generalization.', 'type': 'bool'}, 'sinkhorn_iterations': {'title': 'Sinkhorn Iterations', 'tooltip': 'Controls the number of iterations for Multi-Normed Sinkhorn. While 1 iteration is often sufficient for convergence and 3 offers a slight refinement, 5 is the default.', 'type': 'int'}, - 'normed_momentum': {'title': 'Normalization-then-Momentum', 'tooltip': 'Applies the momentum after the optimizer normalization. This makes the momentum scale invariant and tracks the true variance of the normalized gradients.', 'type': 'bool'}, + 'normed_momentum': {'title': 'Normalization-then-Momentum (NtM)', 'tooltip': 'Applies the momentum after the optimizer normalization. This makes the momentum scale invariant and tracks the true variance of the normalized gradients.', 'type': 'bool'}, 'nesterov_coef': {'title': 'Nesterov Coef', 'tooltip': 'Controls the mixing coefficient between momentum gradients and raw gradients in Nesterov momentum. For a factor of 0.8, the final update will be 80% of the momentum gradients and 20% raw gradient. Leaving it unset toggles the standard Nestrov behavior (where nesterov_coef = beta1 or momentum). Setting it to 0 cancels momentum contribution.', 'type': 'float'}, - 'centered_vt': {'title': 'Variance Preconditioner', 'tooltip': 'Applies centered, scale-invariant row and column confidence pre-conditioner. This is intended to be used with Normalization-then-Momentum.', 'type': 'bool'}, + 'snr_cond': {'title': 'SNR Precondition', 'tooltip': 'Applies a Signal-to-Noise Ratio (SNR) precondition to reshape the optimization curve. It prioritizes high-confidence signals and dampens noise. On-the-fly math with zero memory overhead. Requires Normalization-then-Momentum (NtM). ', 'type': 'bool'}, 'geometric_wd': {'title': 'Geometric Weight Decay', 'tooltip': 'Regularizes weights based on the geometric structure of the optimizer. Compatible with cautious weight decay.', 'type': 'bool'}, } # @formatter:on diff --git a/modules/util/config/TrainConfig.py b/modules/util/config/TrainConfig.py index ac9932ca6..b759fd550 100644 --- a/modules/util/config/TrainConfig.py +++ b/modules/util/config/TrainConfig.py @@ -146,7 +146,7 @@ class TrainOptimizerConfig(BaseConfig): sinkhorn_iterations: int normed_momentum: False nesterov_coef: float - centered_vt: False + snr_cond: False geometric_wd: False def __init__(self, data: list[(str, Any, type, bool)]): @@ -269,7 +269,7 @@ def default_values(): data.append(("sinkhorn_iterations", None, int, True)) data.append(("normed_momentum", False, bool, False)) data.append(("nesterov_coef", None, float, True)) - data.append(("centered_vt", False, bool, False)) + data.append(("snr_cond", False, bool, False)) data.append(("geometric_wd", False, bool, False)) return TrainOptimizerConfig(data) diff --git a/modules/util/create.py b/modules/util/create.py index 0d53512b9..afad5107d 100644 --- a/modules/util/create.py +++ b/modules/util/create.py @@ -688,7 +688,7 @@ def create_optimizer( centered_wd=optimizer_config.centered_wd if optimizer_config.centered_wd is not None else 0.0, centered_wd_mode=optimizer_config.centered_wd_mode if optimizer_config.centered_wd_mode is not None else "full", state_precision=optimizer_config.state_precision if optimizer_config.state_precision is not None else "auto", - centered_vt=optimizer_config.centered_vt if optimizer_config.centered_vt is not None else False, + snr_cond=optimizer_config.snr_cond if optimizer_config.snr_cond is not None else False, geometric_wd=optimizer_config.geometric_wd if optimizer_config.geometric_wd is not None else False, ) @@ -808,6 +808,8 @@ def create_optimizer( state_precision=optimizer_config.state_precision if optimizer_config.state_precision is not None else "auto", nesterov=optimizer_config.nesterov if optimizer_config.nesterov is not None else False, nesterov_coef=optimizer_config.nesterov_coef if optimizer_config.nesterov_coef is not None else None, + snr_cond=optimizer_config.snr_cond if optimizer_config.snr_cond is not None else False, + geometric_wd=optimizer_config.geometric_wd if optimizer_config.geometric_wd is not None else False, ) # LION_ADV Optimizer diff --git a/modules/util/optimizer_util.py b/modules/util/optimizer_util.py index dc32d8512..c06ae6014 100644 --- a/modules/util/optimizer_util.py +++ b/modules/util/optimizer_util.py @@ -451,7 +451,7 @@ def init_model_parameters( "cautious_wd": False, "geometric_wd": False, "weight_decay": 0.0, - "centered_vt": False, + "snr_cond": False, "stochastic_rounding": True, "compile": False, "fused_back_pass": False, @@ -546,8 +546,10 @@ def init_model_parameters( "momentum": 0.95, "nesterov": False, "nesterov_coef": None, + "geometric_wd": False, "cautious_wd": False, "weight_decay": 0.0, + "snr_cond": False, "stochastic_rounding": True, "compile": False, "fused_back_pass": False, From 1c7ee1f746b2cb610975c3c3e9afbba8aec6d97c Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Sat, 30 May 2026 23:47:00 +0300 Subject: [PATCH 51/84] remove use_AdEMAMix and its parameters --- modules/ui/MuonAdamWindow.py | 2 -- modules/ui/OptimizerParamsWindow.py | 2 -- modules/util/config/TrainConfig.py | 4 ---- modules/util/create.py | 9 --------- modules/util/optimizer_util.py | 9 --------- 5 files changed, 26 deletions(-) diff --git a/modules/ui/MuonAdamWindow.py b/modules/ui/MuonAdamWindow.py index c757e5cd3..6cf58c9a1 100644 --- a/modules/ui/MuonAdamWindow.py +++ b/modules/ui/MuonAdamWindow.py @@ -75,8 +75,6 @@ def create_adam_params_ui(self, master): 'use_orthograd': {'title': 'use_orthograd', 'tooltip': 'Use orthograd method', 'type': 'bool'}, 'orthogonal_gradient': {'title': 'OrthoGrad', 'tooltip': 'Reduces overfitting by removing the gradient component parallel to the weight, thus improving generalization.', 'type': 'bool'}, 'use_atan2': {'title': 'Atan2 Scaling', 'tooltip': 'A robust replacement for eps, which also incorporates gradient clipping, bounding and stabilizing the optimizer updates.', 'type': 'bool'}, - 'use_AdEMAMix': {'title': 'AdEMAMix EMA', 'tooltip': 'Adds a second, slow-moving EMA, which is combined with the primary momentum to stabilize updates, and accelerate the training.', 'type': 'bool'}, - 'beta3_ema': {'title': 'Beta3 EMA', 'tooltip': 'Coefficient for slow-moving EMA of AdEMAMix.', 'type': 'float'}, 'kourkoutas_beta': {'title': 'Kourkoutas Beta', 'tooltip': 'Enables a layer-wise dynamic β₂ adaptation. This feature makes the optimizer more responsive to "spiky" gradients by lowering β₂ during periods of high variance, and more stable during calm periods by raising β₂ towards its maximum. It can significantly improve training stability and final loss.', 'type': 'bool'}, 'nesterov_coef': {'title': 'Nesterov Coef', 'tooltip': 'Controls the mixing coefficient between momentum gradients and raw gradients in Nesterov momentum. For a factor of 0.8, the final update will be 80% of the momentum gradients and 20% raw gradient. Leaving it unset toggles the standard Nestrov behavior (where nesterov_coef = beta1 or momentum). Setting it to 0 cancels momentum contribution.', 'type': 'float'}, 'factored_2nd': {'title': 'Factored 2nd', 'tooltip': 'Whether to keep the first moment uncompressed (dense), while only factorizing the second moment. This makes the optimizer highly robust to a wide range of LRs, mimicking high-order optimization.', 'type': 'bool'}, diff --git a/modules/ui/OptimizerParamsWindow.py b/modules/ui/OptimizerParamsWindow.py index 25ec44d37..b1d95bf7e 100644 --- a/modules/ui/OptimizerParamsWindow.py +++ b/modules/ui/OptimizerParamsWindow.py @@ -169,8 +169,6 @@ def create_dynamic_ui( 'nnmf_factor': {'title': 'Factored Optimizer', 'tooltip': 'Enables a memory-efficient mode by applying fast low-rank factorization to the optimizers states. It combines factorization for magnitudes with 1-bit compression for signs, drastically reducing VRAM usage and allowing for larger models or batch sizes. This is an approximation which may slightly alter training dynamics.', 'type': 'bool'}, 'orthogonal_gradient': {'title': 'OrthoGrad', 'tooltip': 'Reduces overfitting by removing the gradient component parallel to the weight, thus improving generalization.', 'type': 'bool'}, 'use_atan2': {'title': 'Atan2 Scaling', 'tooltip': 'A robust replacement for eps, which also incorporates gradient clipping, bounding and stabilizing the optimizer updates.', 'type': 'bool'}, - 'use_AdEMAMix': {'title': 'AdEMAMix EMA', 'tooltip': 'Adds a second, slow-moving EMA, which is combined with the primary momentum to stabilize updates, and accelerate the training.', 'type': 'bool'}, - 'beta3_ema': {'title': 'Beta3 EMA', 'tooltip': 'Coefficient for slow-moving EMA of AdEMAMix.', 'type': 'float'}, 'beta1_warmup': {'title': 'Beta1 Warmup Steps', 'tooltip': 'Number of warmup steps to gradually increase beta1 from Minimum Beta1 Value to its final value. During warmup, beta1 increases linearly. leave it empty to disable warmup and use constant beta1.', 'type': 'int'}, 'min_beta1': {'title': 'Minimum Beta1', 'tooltip': 'Starting beta1 value for warmup scheduling. Used only when beta1 warmup is enabled. Lower values allow faster initial adaptation, while higher values provide more smoothing. The final beta1 value is specified in the beta1 parameter.', 'type': 'float'}, 'kourkoutas_beta': {'title': 'Kourkoutas Beta', 'tooltip': 'Enables a layer-wise dynamic β₂ adaptation. This feature makes the optimizer more responsive to "spiky" gradients by lowering β₂ during periods of high variance, and more stable during calm periods by raising β₂ towards its maximum. It can significantly improve training stability and final loss.', 'type': 'bool'}, diff --git a/modules/util/config/TrainConfig.py b/modules/util/config/TrainConfig.py index b759fd550..6acb0faa1 100644 --- a/modules/util/config/TrainConfig.py +++ b/modules/util/config/TrainConfig.py @@ -112,8 +112,6 @@ class TrainOptimizerConfig(BaseConfig): nnmf_factor: False orthogonal_gradient: False use_atan2: False - use_AdEMAMix: False - beta3_ema: float beta1_warmup: int min_beta1: float kourkoutas_beta: False @@ -235,8 +233,6 @@ def default_values(): data.append(("nnmf_factor", False, bool, False)) data.append(("orthogonal_gradient", False, bool, False)) data.append(("use_atan2", False, bool, False)) - data.append(("use_AdEMAMix", False, bool, False)) - data.append(("beta3_ema", None, float, True)) data.append(("beta1_warmup", None, int, True)) data.append(("min_beta1", None, float, True)) data.append(("kourkoutas_beta", False, bool, False)) diff --git a/modules/util/create.py b/modules/util/create.py index afad5107d..2d4df595b 100644 --- a/modules/util/create.py +++ b/modules/util/create.py @@ -708,9 +708,6 @@ def create_optimizer( stochastic_rounding=optimizer_config.stochastic_rounding, use_atan2=optimizer_config.use_atan2 if optimizer_config.use_atan2 is not None else False, orthogonal_gradient=optimizer_config.orthogonal_gradient if optimizer_config.orthogonal_gradient is not None else False, - use_AdEMAMix=optimizer_config.use_AdEMAMix if optimizer_config.use_AdEMAMix is not None else False, - beta3_ema=optimizer_config.beta3 if optimizer_config.beta3 is not None else 0.9999, - alpha=optimizer_config.alpha if optimizer_config.alpha is not None else 5, kourkoutas_beta=optimizer_config.kourkoutas_beta if optimizer_config.kourkoutas_beta is not None else False, k_warmup_steps=(config.learning_rate_warmup_steps / config.gradient_accumulation_steps), compiled_optimizer=optimizer_config.compile if optimizer_config.compile is not None else False, @@ -738,9 +735,6 @@ def create_optimizer( stochastic_rounding=optimizer_config.stochastic_rounding, use_atan2=optimizer_config.use_atan2 if optimizer_config.use_atan2 is not None else False, orthogonal_gradient=optimizer_config.orthogonal_gradient if optimizer_config.orthogonal_gradient is not None else False, - use_AdEMAMix=optimizer_config.use_AdEMAMix if optimizer_config.use_AdEMAMix is not None else False, - beta3_ema=optimizer_config.beta3 if optimizer_config.beta3 is not None else 0.9999, - alpha=optimizer_config.alpha if optimizer_config.alpha is not None else 5, kourkoutas_beta=optimizer_config.kourkoutas_beta if optimizer_config.kourkoutas_beta is not None else False, k_warmup_steps=(config.learning_rate_warmup_steps / config.gradient_accumulation_steps), compiled_optimizer=optimizer_config.compile if optimizer_config.compile is not None else False, @@ -775,9 +769,6 @@ def create_optimizer( d_limiter=optimizer_config.d_limiter if optimizer_config.d_limiter is not None else False, use_atan2=optimizer_config.use_atan2 if optimizer_config.use_atan2 is not None else False, orthogonal_gradient=optimizer_config.orthogonal_gradient if optimizer_config.orthogonal_gradient is not None else False, - use_AdEMAMix=optimizer_config.use_AdEMAMix if optimizer_config.use_AdEMAMix is not None else False, - beta3_ema=optimizer_config.beta3_ema if optimizer_config.beta3_ema is not None else 0.9999, - alpha=optimizer_config.alpha if optimizer_config.alpha is not None else 5, kourkoutas_beta=optimizer_config.kourkoutas_beta if optimizer_config.kourkoutas_beta is not None else False, k_warmup_steps=(config.learning_rate_warmup_steps / config.gradient_accumulation_steps), compiled_optimizer=optimizer_config.compile if optimizer_config.compile is not None else False, diff --git a/modules/util/optimizer_util.py b/modules/util/optimizer_util.py index c06ae6014..ee9acb108 100644 --- a/modules/util/optimizer_util.py +++ b/modules/util/optimizer_util.py @@ -477,9 +477,6 @@ def init_model_parameters( "fused_back_pass": False, "use_atan2": False, "orthogonal_gradient": False, - "use_AdEMAMix": False, - "beta3_ema": 0.9999, - "alpha": 5, "kourkoutas_beta": False, "spectral_normalization": False, "centered_wd": 0.0, @@ -501,9 +498,6 @@ def init_model_parameters( "fused_back_pass": False, "use_atan2": True, "orthogonal_gradient": False, - "use_AdEMAMix": False, - "beta3_ema": 0.9999, - "alpha": 5, "kourkoutas_beta": False, "spectral_normalization": False, "centered_wd": 0.0, @@ -533,9 +527,6 @@ def init_model_parameters( "d_limiter": False, "use_atan2": False, "orthogonal_gradient": False, - "use_AdEMAMix": False, - "beta3_ema": 0.9999, - "alpha": 5, "kourkoutas_beta": False, "centered_wd": 0.0, "centered_wd_mode": "float8", From b4b9484d61ef4e3f7cb754286949c762af7b01c1 Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Sat, 30 May 2026 23:51:56 +0300 Subject: [PATCH 52/84] remove fp8_sr and add fp32 to StatePrecision --- modules/ui/MuonAdamWindow.py | 4 +++- modules/ui/OptimizerParamsWindow.py | 4 ++-- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/modules/ui/MuonAdamWindow.py b/modules/ui/MuonAdamWindow.py index 6cf58c9a1..36569ded2 100644 --- a/modules/ui/MuonAdamWindow.py +++ b/modules/ui/MuonAdamWindow.py @@ -79,7 +79,7 @@ def create_adam_params_ui(self, master): 'nesterov_coef': {'title': 'Nesterov Coef', 'tooltip': 'Controls the mixing coefficient between momentum gradients and raw gradients in Nesterov momentum. For a factor of 0.8, the final update will be 80% of the momentum gradients and 20% raw gradient. Leaving it unset toggles the standard Nestrov behavior (where nesterov_coef = beta1 or momentum). Setting it to 0 cancels momentum contribution.', 'type': 'float'}, 'factored_2nd': {'title': 'Factored 2nd', 'tooltip': 'Whether to keep the first moment uncompressed (dense), while only factorizing the second moment. This makes the optimizer highly robust to a wide range of LRs, mimicking high-order optimization.', 'type': 'bool'}, 'fisher_wd': {'title': 'Fisher Weight Decay', 'tooltip': 'Applies adaptive, scale-invariant weight-decay regularization based on the Fisher Information Matrix (approximated by Adam\'s second moment). It reduces penalty for "important" high-curvature weights while accelerating decay for "useless" weights in flat regions. Leading to improved convergence and better final performance.', 'type': 'bool'}, - 'state_precision': {'title': 'state_precision', 'tooltip': """The quantization format used to store the optimizer states to save VRAM. Options include: 'auto': Stores the states in the original parameter's precision. 'factored': Enables a memory-efficient mode by applying fast low-rank factorization to the optimizers states. It combines factorization for magnitudes with 1-bit compression for signs, drastically reducing VRAM usage and allowing for larger models or batch sizes. 'bf16_sr': Uses BF16 with stochastic rounding for a balance of precision and memory. 'int8_sr': Uses 8-bit block-wise quantization with stochastic rounding. 'fp8_sr': Uses torch.float8_e4m3fn with stochastic rounding.""", 'type': 'StatePrecision'}, + 'state_precision': {'title': 'state_precision', 'tooltip': """The quantization format used to store the optimizer states to save VRAM. Options include: 'auto': Stores the states in the original parameter's precision. 'factored': Enables a memory-efficient mode by applying fast low-rank factorization to the optimizers states. It combines factorization for magnitudes with 1-bit compression for signs, drastically reducing VRAM usage and allowing for larger models or batch sizes. 'fp32': Uses full FP32. 'bf16_sr': Uses BF16 with stochastic rounding for a balance of precision and memory. 'int8_sr': Uses 8-bit block-wise quantization with stochastic rounding.""", 'type': 'StatePrecision'}, } # @formatter:on @@ -102,5 +102,7 @@ def create_adam_params_ui(self, master): if param_type != 'bool': components.entry(master, row, col + 1, self.adam_ui_state, key) + elif type == 'StatePrecision': + components.options(master, row, col + 1, ["auto", "factored", "fp32", "bf16_sr", "int8_sr"], self.optimizer_ui_state, key) else: components.switch(master, row, col + 1, self.adam_ui_state, key) diff --git a/modules/ui/OptimizerParamsWindow.py b/modules/ui/OptimizerParamsWindow.py index b1d95bf7e..7b6b3ef9f 100644 --- a/modules/ui/OptimizerParamsWindow.py +++ b/modules/ui/OptimizerParamsWindow.py @@ -195,7 +195,7 @@ def create_dynamic_ui( 'centered_wd_mode': {'title': 'Centered WD Mode', 'tooltip': """The quantization format used to store the anchor weights to save VRAM. Options include: 'full': Stores anchors in the original parameter's precision. 'float8': Uses torch.float8_e4m3fn for a balance of precision and memory. 'int8': Uses 8-bit block-wise quantization. 'int4': Uses 4-bit block-wise quantization.""", 'type': 'CenteredWDMode'}, 'factored_2nd': {'title': 'Factored 2nd', 'tooltip': 'Whether to keep the first moment uncompressed (dense), while only factorizing the second moment. This makes the optimizer highly robust to a wide range of LRs, mimicking high-order optimization.', 'type': 'bool'}, 'fisher_wd': {'title': 'Fisher Weight Decay', 'tooltip': 'Applies adaptive, scale-invariant weight-decay regularization based on the Fisher Information Matrix (approximated by Adam\'s second moment). It reduces penalty for "important" high-curvature weights while accelerating decay for "useless" weights in flat regions. Leading to improved convergence and better final performance.', 'type': 'bool'}, - 'state_precision': {'title': 'state_precision', 'tooltip': """The quantization format used to store the optimizer states to save VRAM. Options include: 'auto': Stores the states in the original parameter's precision. 'factored': Enables a memory-efficient mode by applying fast low-rank factorization to the optimizers states. It combines factorization for magnitudes with 1-bit compression for signs, drastically reducing VRAM usage and allowing for larger models or batch sizes. 'bf16_sr': Uses BF16 with stochastic rounding for a balance of precision and memory. 'int8_sr': Uses 8-bit block-wise quantization with stochastic rounding. 'fp8_sr': Uses torch.float8_e4m3fn with stochastic rounding.""", 'type': 'StatePrecision'}, + 'state_precision': {'title': 'state_precision', 'tooltip': """The quantization format used to store the optimizer states to save VRAM. Options include: 'auto': Stores the states in the original parameter's precision. 'factored': Enables a memory-efficient mode by applying fast low-rank factorization to the optimizers states. It combines factorization for magnitudes with 1-bit compression for signs, drastically reducing VRAM usage and allowing for larger models or batch sizes. 'fp32': Uses full FP32. 'bf16_sr': Uses BF16 with stochastic rounding for a balance of precision and memory. 'int8_sr': Uses 8-bit block-wise quantization with stochastic rounding.""", 'type': 'StatePrecision'}, 'orthogonal_sinkhorn': {'title': 'Orthogonal Sinkhorn', 'tooltip': 'Applies iterative row and column orthogonal projection to make the updates orthogonal to the current weight, leading to robust regularization and better generalization.', 'type': 'bool'}, 'sinkhorn_iterations': {'title': 'Sinkhorn Iterations', 'tooltip': 'Controls the number of iterations for Multi-Normed Sinkhorn. While 1 iteration is often sufficient for convergence and 3 offers a slight refinement, 5 is the default.', 'type': 'int'}, 'normed_momentum': {'title': 'Normalization-then-Momentum (NtM)', 'tooltip': 'Applies the momentum after the optimizer normalization. This makes the momentum scale invariant and tracks the true variance of the normalized gradients.', 'type': 'bool'}, @@ -242,7 +242,7 @@ def create_dynamic_ui( components.options(master, row, col + 1, ["full", "float8", "int8", "int4"], self.optimizer_ui_state, key, command=self.update_user_pref) elif type == 'StatePrecision': - components.options(master, row, col + 1, ["auto", "factored", "bf16_sr", "int8_sr", "fp8_sr"], self.optimizer_ui_state, key, + components.options(master, row, col + 1, ["auto", "factored", "fp32", "bf16_sr", "int8_sr"], self.optimizer_ui_state, key, command=self.update_user_pref) elif type != 'bool': components.entry(master, row, col + 1, self.optimizer_ui_state, key, From 5181601661782df0e112ff2c88d6697540e9c418 Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Sat, 30 May 2026 23:52:59 +0300 Subject: [PATCH 53/84] adam_ui_state --- modules/ui/MuonAdamWindow.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modules/ui/MuonAdamWindow.py b/modules/ui/MuonAdamWindow.py index 36569ded2..6c9a1b545 100644 --- a/modules/ui/MuonAdamWindow.py +++ b/modules/ui/MuonAdamWindow.py @@ -103,6 +103,6 @@ def create_adam_params_ui(self, master): if param_type != 'bool': components.entry(master, row, col + 1, self.adam_ui_state, key) elif type == 'StatePrecision': - components.options(master, row, col + 1, ["auto", "factored", "fp32", "bf16_sr", "int8_sr"], self.optimizer_ui_state, key) + components.options(master, row, col + 1, ["auto", "factored", "fp32", "bf16_sr", "int8_sr"], self.adam_ui_state, key) else: components.switch(master, row, col + 1, self.adam_ui_state, key) From 7b6d349afd71c5c90b6d51fadd2e801bebb23932 Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Sat, 30 May 2026 23:54:41 +0300 Subject: [PATCH 54/84] dev23 --- requirements-global.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements-global.txt b/requirements-global.txt index 81813de4b..8388d2fd6 100644 --- a/requirements-global.txt +++ b/requirements-global.txt @@ -41,7 +41,7 @@ prodigyopt==1.1.2 # prodigy optimizer schedulefree==1.4.1 # schedule-free optimizers pytorch_optimizer==3.6.0 # pytorch optimizers prodigy-plus-schedule-free==2.0.1 # Prodigy plus optimizer -adv_optm==2.4.dev18 # advanced optimizers +adv_optm==2.4.dev23 # advanced optimizers -e git+https://github.com/KellerJordan/Muon.git@f90a42b#egg=muon-optimizer # Profiling From f331379e3b65d3fdcda2c7458a4a0fdcba4399ae Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Sun, 31 May 2026 11:49:06 +0300 Subject: [PATCH 55/84] remove depth --- modules/util/create.py | 4 -- modules/util/optimizer/depth_calculator.py | 74 ---------------------- 2 files changed, 78 deletions(-) delete mode 100644 modules/util/optimizer/depth_calculator.py diff --git a/modules/util/create.py b/modules/util/create.py index 2d4df595b..b6c9f64c7 100644 --- a/modules/util/create.py +++ b/modules/util/create.py @@ -33,7 +33,6 @@ from modules.util.optimizer.adafactor_extensions import patch_adafactor from modules.util.optimizer.adam_extensions import patch_adam from modules.util.optimizer.adamw_extensions import patch_adamw -from modules.util.optimizer.depth_calculator import inject_depth_into_param_groups from modules.util.optimizer.muon_util import split_parameters_for_muon from modules.util.TrainProgress import TrainProgress from modules.zluda import ZLUDA @@ -141,9 +140,6 @@ def create_optimizer( parameters = parameter_group_collection.parameters_for_optimizer(config) - if optimizer_config.spectral_normalization: - # _adv optimizers O(1) depth scaling. - inject_depth_into_param_groups(model, parameters) match config.optimizer.optimizer: diff --git a/modules/util/optimizer/depth_calculator.py b/modules/util/optimizer/depth_calculator.py deleted file mode 100644 index ed22fc040..000000000 --- a/modules/util/optimizer/depth_calculator.py +++ /dev/null @@ -1,74 +0,0 @@ -import re - -from modules.model.BaseModel import BaseModel -from modules.module.LoRAModule import LoRAModuleWrapper -from modules.util.enum.ModelType import ModelType - -import torch - - -def calculate_n_layers(model: BaseModel) -> dict[str, int]: - """ - Calculates the number of residual layers (the depth) in each component of the model. - """ - match model.model_type: - case (ModelType.STABLE_DIFFUSION_15 | ModelType.STABLE_DIFFUSION_15_INPAINTING | - ModelType.STABLE_DIFFUSION_20_BASE | ModelType.STABLE_DIFFUSION_20_INPAINTING | - ModelType.STABLE_DIFFUSION_20 | ModelType.STABLE_DIFFUSION_21 | - ModelType.STABLE_DIFFUSION_21_BASE | ModelType.STABLE_DIFFUSION_XL_10_BASE | - ModelType.STABLE_DIFFUSION_XL_10_BASE_INPAINTING | ModelType.STABLE_CASCADE_1 | - ModelType.WUERSTCHEN_2): - default_patterns = ['transformer_blocks', 'resnets', 'layers'] - case (ModelType.STABLE_DIFFUSION_3 | ModelType.STABLE_DIFFUSION_35 | ModelType.SANA | - ModelType.FLUX_DEV_1 | ModelType.FLUX_2 | ModelType.CHROMA_1 | ModelType.QWEN | - ModelType.PIXART_ALPHA | ModelType.PIXART_SIGMA): - default_patterns = ['transformer_blocks', 'single_transformer_blocks', 'encoder.block'] - case ModelType.HI_DREAM_FULL: - default_patterns = ['double_stream_blocks', 'single_stream_blocks'] - case ModelType.Z_IMAGE: - default_patterns = [ - 'layers', - 'refiner', - ] - case _: - raise NotImplementedError(f"Scaled Optimizer is not implemented for model type: {model.model_type}") - - # Build the regex pattern - joined_patterns = "|".join([re.escape(p) for p in default_patterns]) - pattern = re.compile(rf'(?:^|\.)(?:{joined_patterns})\.\d+$') - - layer_counts = {} - - # Iterate over model components (e.g., 'unet', 'text_encoder', 'transformer') - for attr_name, module in vars(model).items(): - # Identify the 'Ground Truth' blocks in this component. - target_module = module - if isinstance(module, LoRAModuleWrapper): - target_module = module.orig_module - valid_component_blocks = set() - if isinstance(target_module, torch.nn.Module): - for name, _ in target_module.named_modules(): - if pattern.search(name): - valid_component_blocks.add(name) - if not valid_component_blocks: - continue - count = len(valid_component_blocks) - if count > 0: - layer_counts[attr_name] = count - - return layer_counts - - -def inject_depth_into_param_groups(model: BaseModel, parameters): - """ - Calculates the model depth and injects the 'n_layers' key directly - into the optimizer parameter groups. - """ - n_layers_map = calculate_n_layers(model) - - for group in parameters: - group_name = group.get('name') - if group_name in n_layers_map: - group['n_layers'] = n_layers_map[group_name] - else: - group['n_layers'] = n_layers_map.get('default', 1) From 38ab7e71fb145861cd9c25cd882d9b79c09126fa Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Sun, 31 May 2026 11:50:02 +0300 Subject: [PATCH 56/84] tag util --- modules/util/create.py | 4 ++++ modules/util/enum/Optimizer.py | 13 +++++++++++++ modules/util/optimizer/tag_util.py | 31 ++++++++++++++++++++++++++++++ 3 files changed, 48 insertions(+) create mode 100644 modules/util/optimizer/tag_util.py diff --git a/modules/util/create.py b/modules/util/create.py index b6c9f64c7..42ed28cda 100644 --- a/modules/util/create.py +++ b/modules/util/create.py @@ -33,6 +33,7 @@ from modules.util.optimizer.adafactor_extensions import patch_adafactor from modules.util.optimizer.adam_extensions import patch_adam from modules.util.optimizer.adamw_extensions import patch_adamw +from modules.util.optimizer.tag_util import tag_peft_parameters from modules.util.optimizer.muon_util import split_parameters_for_muon from modules.util.TrainProgress import TrainProgress from modules.zluda import ZLUDA @@ -140,6 +141,9 @@ def create_optimizer( parameters = parameter_group_collection.parameters_for_optimizer(config) + if config.optimizer.optimizer.is_adv: + # Tag PEFT parameters based on parameter names. + tag_peft_parameters(model) match config.optimizer.optimizer: diff --git a/modules/util/enum/Optimizer.py b/modules/util/enum/Optimizer.py index bbab87b4b..6326a1861 100644 --- a/modules/util/enum/Optimizer.py +++ b/modules/util/enum/Optimizer.py @@ -99,6 +99,19 @@ def is_schedule_free(self): self.PRODIGY_PLUS_SCHEDULE_FREE, ] + @property + def is_adv(self): + return self in [ + self.ADAMW_ADV, + self.ADOPT_ADV, + self.LION_ADV, + self.SINKSGD_ADV, + self.SIGNSGD_ADV, + self.PRODIGY_ADV, + self.MUON_ADV, + self.ADAMUON_ADV, + ] + def supports_fused_back_pass(self): return self in [ Optimizer.ADAFACTOR, diff --git a/modules/util/optimizer/tag_util.py b/modules/util/optimizer/tag_util.py new file mode 100644 index 000000000..c8ec09d42 --- /dev/null +++ b/modules/util/optimizer/tag_util.py @@ -0,0 +1,31 @@ +import torch +from modules.module.LoRAModule import LoRAModuleWrapper + +def tag_peft_parameters(model: torch.nn.Module | None): + """ + Tags PEFT parameters with attributes like `_is_lora_A`, `_is_lora_B`, + `_is_oft`, and `_is_dora_scale` based on their names. + This is required to apply correct scaling and optimizations for certain features in adv_optm library. + """ + def apply_tags(name: str, p: torch.nn.Parameter): + if name.endswith(("lora_down.weight", "lokr_w1_b", "lokr_w2_b")): + # Down projection + p._is_lora_A = True + elif name.endswith(("lora_up.weight", "lokr_w1_a", "lokr_w2_a")): + # Up projection + p._is_lora_B = True + elif name.endswith("dora_scale"): + # Vector in shape of >= 2D tensor + p._is_dora_scale = True + elif name.endswith("oft_R.weight"): + # Set of independent vectors (rank, n_elements) + p._is_oft = True + + for module_prefix, module in vars(model).items(): + if isinstance(module, LoRAModuleWrapper): + for lora_module in module.lora_modules.values(): + for param_name, p in lora_module.named_parameters(): + apply_tags(param_name, p) + elif isinstance(module, torch.nn.Module): + for param_name, p in module.named_parameters(): + apply_tags(param_name, p) From b3876ec29222f6b7dd0b20c7af84edc86ccbbc47 Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Sun, 31 May 2026 11:50:58 +0300 Subject: [PATCH 57/84] remove LoRAModule tagging --- modules/module/LoRAModule.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/modules/module/LoRAModule.py b/modules/module/LoRAModule.py index 91c2d456e..bb1410a23 100644 --- a/modules/module/LoRAModule.py +++ b/modules/module/LoRAModule.py @@ -314,8 +314,6 @@ def initialize_weights(self): self.lora_down, self.lora_up = self.create_layer() nn.init.kaiming_uniform_(self.lora_down.weight, a=math.sqrt(5)) nn.init.zeros_(self.lora_up.weight) - self.lora_down.weight._is_lora_A = True - self.lora_up.weight._is_lora_B = True def check_initialized(self): super().check_initialized() @@ -423,7 +421,6 @@ def initialize_weights(self): ) nn.init.zeros_(self.oft_R.weight) - self.oft_R.weight._is_oft = True def forward(self, x, *args, **kwargs): self.check_initialized() @@ -516,7 +513,6 @@ def initialize_weights(self): .to(device=self.orig_module.weight.device) ) - self.dora_scale._is_dora_scale = True del orig_weight def check_initialized(self): From 14985e78c3cb623596ad84709de1c3cf9434dd02 Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Sun, 31 May 2026 11:54:12 +0300 Subject: [PATCH 58/84] pre-commit --- modules/ui/MuonAdamWindow.py | 2 +- modules/util/create.py | 2 +- modules/util/optimizer/tag_util.py | 8 +++++--- 3 files changed, 7 insertions(+), 5 deletions(-) diff --git a/modules/ui/MuonAdamWindow.py b/modules/ui/MuonAdamWindow.py index 6c9a1b545..2fef43500 100644 --- a/modules/ui/MuonAdamWindow.py +++ b/modules/ui/MuonAdamWindow.py @@ -102,7 +102,7 @@ def create_adam_params_ui(self, master): if param_type != 'bool': components.entry(master, row, col + 1, self.adam_ui_state, key) - elif type == 'StatePrecision': + elif param_type == 'StatePrecision': components.options(master, row, col + 1, ["auto", "factored", "fp32", "bf16_sr", "int8_sr"], self.adam_ui_state, key) else: components.switch(master, row, col + 1, self.adam_ui_state, key) diff --git a/modules/util/create.py b/modules/util/create.py index 42ed28cda..6e2883bf3 100644 --- a/modules/util/create.py +++ b/modules/util/create.py @@ -33,8 +33,8 @@ from modules.util.optimizer.adafactor_extensions import patch_adafactor from modules.util.optimizer.adam_extensions import patch_adam from modules.util.optimizer.adamw_extensions import patch_adamw -from modules.util.optimizer.tag_util import tag_peft_parameters from modules.util.optimizer.muon_util import split_parameters_for_muon +from modules.util.optimizer.tag_util import tag_peft_parameters from modules.util.TrainProgress import TrainProgress from modules.zluda import ZLUDA diff --git a/modules/util/optimizer/tag_util.py b/modules/util/optimizer/tag_util.py index c8ec09d42..7859554e3 100644 --- a/modules/util/optimizer/tag_util.py +++ b/modules/util/optimizer/tag_util.py @@ -1,9 +1,11 @@ -import torch from modules.module.LoRAModule import LoRAModuleWrapper +import torch + + def tag_peft_parameters(model: torch.nn.Module | None): """ - Tags PEFT parameters with attributes like `_is_lora_A`, `_is_lora_B`, + Tags PEFT parameters with attributes like `_is_lora_A`, `_is_lora_B`, `_is_oft`, and `_is_dora_scale` based on their names. This is required to apply correct scaling and optimizations for certain features in adv_optm library. """ @@ -21,7 +23,7 @@ def apply_tags(name: str, p: torch.nn.Parameter): # Set of independent vectors (rank, n_elements) p._is_oft = True - for module_prefix, module in vars(model).items(): + for module in vars(model).values(): if isinstance(module, LoRAModuleWrapper): for lora_module in module.lora_modules.values(): for param_name, p in lora_module.named_parameters(): From 334881ccf2307d75bd93b1b4df41a4761fe34371 Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Sun, 31 May 2026 12:49:28 +0300 Subject: [PATCH 59/84] Change centered_wd_mode to full --- modules/util/config/TrainConfig.py | 2 +- modules/util/optimizer_util.py | 16 ++++++++-------- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/modules/util/config/TrainConfig.py b/modules/util/config/TrainConfig.py index d3243ecf9..55747715a 100644 --- a/modules/util/config/TrainConfig.py +++ b/modules/util/config/TrainConfig.py @@ -257,7 +257,7 @@ def default_values(): data.append(("spectral_normalization", False, bool, False)) data.append(("stochastic_sign", False, bool, False)) data.append(("centered_wd", 0.0, float, False)) - data.append(("centered_wd_mode", "float8", str, False)) + data.append(("centered_wd_mode", "full", str, False)) data.append(("factored_2nd", False, bool, False)) data.append(("fisher_wd", False, bool, False)) data.append(("state_precision", "float8", str, False)) diff --git a/modules/util/optimizer_util.py b/modules/util/optimizer_util.py index ee9acb108..2452282d3 100644 --- a/modules/util/optimizer_util.py +++ b/modules/util/optimizer_util.py @@ -459,7 +459,7 @@ def init_model_parameters( "orthogonal_gradient": False, "spectral_normalization": False, "centered_wd": 0.0, - "centered_wd_mode": "float8", + "centered_wd_mode": "full", "state_precision": "auto", }, Optimizer.ADAMW_ADV: { @@ -480,7 +480,7 @@ def init_model_parameters( "kourkoutas_beta": False, "spectral_normalization": False, "centered_wd": 0.0, - "centered_wd_mode": "float8", + "centered_wd_mode": "full", "state_precision": "auto", }, Optimizer.ADOPT_ADV: { @@ -501,7 +501,7 @@ def init_model_parameters( "kourkoutas_beta": False, "spectral_normalization": False, "centered_wd": 0.0, - "centered_wd_mode": "float8", + "centered_wd_mode": "full", "state_precision": "auto", }, Optimizer.PRODIGY_ADV: { @@ -529,7 +529,7 @@ def init_model_parameters( "orthogonal_gradient": False, "kourkoutas_beta": False, "centered_wd": 0.0, - "centered_wd_mode": "float8", + "centered_wd_mode": "full", }, Optimizer.SIGNSGD_ADV: { "stochastic_sign": False, @@ -547,7 +547,7 @@ def init_model_parameters( "orthogonal_gradient": False, "spectral_normalization": False, "centered_wd": 0.0, - "centered_wd_mode": "float8", + "centered_wd_mode": "full", "state_precision": "auto", }, Optimizer.LION_ADV: { @@ -564,7 +564,7 @@ def init_model_parameters( "orthogonal_gradient": False, "spectral_normalization": False, "centered_wd": 0.0, - "centered_wd_mode": "float8", + "centered_wd_mode": "full", }, Optimizer.MUON_ADV: { "beta1": 0.9, @@ -592,7 +592,7 @@ def init_model_parameters( "orthogonal_gradient": False, "approx_mars": False, "centered_wd": 0.0, - "centered_wd_mode": "float8", + "centered_wd_mode": "full", "state_precision": "auto", "muon_adam_config": {}, }, @@ -624,7 +624,7 @@ def init_model_parameters( "orthogonal_gradient": False, "approx_mars": False, "centered_wd": 0.0, - "centered_wd_mode": "float8", + "centered_wd_mode": "full", "state_precision": "auto", "muon_adam_config": {}, }, From 04940ed0a10fed6334a7495b2e7c2b76c1146f99 Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Mon, 1 Jun 2026 14:20:09 +0300 Subject: [PATCH 60/84] support for dora_multiplier --- modules/util/optimizer/tag_util.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modules/util/optimizer/tag_util.py b/modules/util/optimizer/tag_util.py index 7859554e3..a6c667729 100644 --- a/modules/util/optimizer/tag_util.py +++ b/modules/util/optimizer/tag_util.py @@ -16,7 +16,7 @@ def apply_tags(name: str, p: torch.nn.Parameter): elif name.endswith(("lora_up.weight", "lokr_w1_a", "lokr_w2_a")): # Up projection p._is_lora_B = True - elif name.endswith("dora_scale"): + elif name.endswith(("dora_scale", "dora_multiplier")): # Vector in shape of >= 2D tensor p._is_dora_scale = True elif name.endswith("oft_R.weight"): From fa91c0a1a7fbdf644efce7a121051dc15f9eb800 Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Thu, 4 Jun 2026 22:05:53 +0300 Subject: [PATCH 61/84] change to dora_log_multiplier --- modules/util/optimizer/tag_util.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modules/util/optimizer/tag_util.py b/modules/util/optimizer/tag_util.py index a6c667729..131ac1423 100644 --- a/modules/util/optimizer/tag_util.py +++ b/modules/util/optimizer/tag_util.py @@ -16,7 +16,7 @@ def apply_tags(name: str, p: torch.nn.Parameter): elif name.endswith(("lora_up.weight", "lokr_w1_a", "lokr_w2_a")): # Up projection p._is_lora_B = True - elif name.endswith(("dora_scale", "dora_multiplier")): + elif name.endswith(("dora_scale", "dora_log_multiplier")): # Vector in shape of >= 2D tensor p._is_dora_scale = True elif name.endswith("oft_R.weight"): From 091207ee570b11f1a47f0fa30117b056ff4ff734 Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Fri, 5 Jun 2026 03:02:17 +0300 Subject: [PATCH 62/84] remove leftover --- modules/util/create.py | 1 - modules/util/optimizer_util.py | 1 - 2 files changed, 2 deletions(-) diff --git a/modules/util/create.py b/modules/util/create.py index 6e2883bf3..4b5ded9d4 100644 --- a/modules/util/create.py +++ b/modules/util/create.py @@ -812,7 +812,6 @@ def create_optimizer( betas=(optimizer_config.beta1 if optimizer_config.beta1 is not None else 0.9, optimizer_config.beta2 if optimizer_config.beta2 is not None else 0.99), weight_decay=optimizer_config.weight_decay if optimizer_config.weight_decay is not None else 0.0, - clip_threshold=optimizer_config.clip_threshold if optimizer_config.clip_threshold is not None else 0.0, nnmf_factor=optimizer_config.nnmf_factor if optimizer_config.nnmf_factor is not None else False, cautious_wd=optimizer_config.cautious_wd if optimizer_config.cautious_wd is not None else False, stochastic_rounding=optimizer_config.stochastic_rounding, diff --git a/modules/util/optimizer_util.py b/modules/util/optimizer_util.py index 2452282d3..c064dd817 100644 --- a/modules/util/optimizer_util.py +++ b/modules/util/optimizer_util.py @@ -556,7 +556,6 @@ def init_model_parameters( "beta2": 0.99, "cautious_wd": False, "weight_decay": 0.0, - "clip_threshold": None, "nnmf_factor": False, "stochastic_rounding": True, "compile": False, From dfd34a8a799b714fbc7451dd599e56424e64917c Mon Sep 17 00:00:00 2001 From: Koratahiu Date: Fri, 5 Jun 2026 18:00:37 +0300 Subject: [PATCH 63/84] dev24: bugfixes, improved spectral scaling for OFT, fix normed momentum and nesterov --- requirements-global.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements-global.txt b/requirements-global.txt index 8388d2fd6..d1bb10dbb 100644 --- a/requirements-global.txt +++ b/requirements-global.txt @@ -41,7 +41,7 @@ prodigyopt==1.1.2 # prodigy optimizer schedulefree==1.4.1 # schedule-free optimizers pytorch_optimizer==3.6.0 # pytorch optimizers prodigy-plus-schedule-free==2.0.1 # Prodigy plus optimizer -adv_optm==2.4.dev23 # advanced optimizers +adv_optm==2.4.dev24 # advanced optimizers -e git+https://github.com/KellerJordan/Muon.git@f90a42b#egg=muon-optimizer # Profiling From a2374bfb14ea22684a421bd9052ef9c0114ea306 Mon Sep 17 00:00:00 2001 From: Koratahiu Date: Sat, 6 Jun 2026 03:59:30 +0300 Subject: [PATCH 64/84] dev25: Improved normed momentum and nesterov --- requirements-global.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements-global.txt b/requirements-global.txt index d1bb10dbb..05cab3411 100644 --- a/requirements-global.txt +++ b/requirements-global.txt @@ -41,7 +41,7 @@ prodigyopt==1.1.2 # prodigy optimizer schedulefree==1.4.1 # schedule-free optimizers pytorch_optimizer==3.6.0 # pytorch optimizers prodigy-plus-schedule-free==2.0.1 # Prodigy plus optimizer -adv_optm==2.4.dev24 # advanced optimizers +adv_optm==2.4.dev25 # advanced optimizers -e git+https://github.com/KellerJordan/Muon.git@f90a42b#egg=muon-optimizer # Profiling From 0deb5979f264f7ce69ae2d9aa7724a6e9c2e69e8 Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Sat, 6 Jun 2026 21:19:51 +0300 Subject: [PATCH 65/84] Rework orthogonal_gradient --- modules/ui/MuonAdamWindow.py | 4 +++- modules/ui/OptimizerParamsWindow.py | 5 ++++- modules/util/config/TrainConfig.py | 4 ++-- modules/util/create.py | 16 ++++++++-------- modules/util/optimizer_util.py | 16 ++++++++-------- 5 files changed, 25 insertions(+), 20 deletions(-) diff --git a/modules/ui/MuonAdamWindow.py b/modules/ui/MuonAdamWindow.py index 2fef43500..c0958e88b 100644 --- a/modules/ui/MuonAdamWindow.py +++ b/modules/ui/MuonAdamWindow.py @@ -73,7 +73,7 @@ def create_adam_params_ui(self, master): 'use_bias_correction': {'title': 'Bias Correction', 'tooltip': 'Turn on Adam\'s bias correction.', 'type': 'bool'}, 'weight_decay': {'title': 'Weight Decay', 'tooltip': 'Regularization to prevent overfitting.', 'type': 'float'}, 'use_orthograd': {'title': 'use_orthograd', 'tooltip': 'Use orthograd method', 'type': 'bool'}, - 'orthogonal_gradient': {'title': 'OrthoGrad', 'tooltip': 'Reduces overfitting by removing the gradient component parallel to the weight, thus improving generalization.', 'type': 'bool'}, + 'orthogonal_gradient': {'title': 'OrthoGrad', 'tooltip': 'Reduces overfitting by removing the gradient component parallel to the weight, thus improving generalization. This has two modes: 1. flattened: Standard vectorized OrthoGrad. Fastest, but loses the structural properties of matrices. 2. iterative: Matrix-wise OrthoGrad, preserves structure by iteratively projecting rows and columns.', 'type': 'OrthoGrad'}, 'use_atan2': {'title': 'Atan2 Scaling', 'tooltip': 'A robust replacement for eps, which also incorporates gradient clipping, bounding and stabilizing the optimizer updates.', 'type': 'bool'}, 'kourkoutas_beta': {'title': 'Kourkoutas Beta', 'tooltip': 'Enables a layer-wise dynamic β₂ adaptation. This feature makes the optimizer more responsive to "spiky" gradients by lowering β₂ during periods of high variance, and more stable during calm periods by raising β₂ towards its maximum. It can significantly improve training stability and final loss.', 'type': 'bool'}, 'nesterov_coef': {'title': 'Nesterov Coef', 'tooltip': 'Controls the mixing coefficient between momentum gradients and raw gradients in Nesterov momentum. For a factor of 0.8, the final update will be 80% of the momentum gradients and 20% raw gradient. Leaving it unset toggles the standard Nestrov behavior (where nesterov_coef = beta1 or momentum). Setting it to 0 cancels momentum contribution.', 'type': 'float'}, @@ -104,5 +104,7 @@ def create_adam_params_ui(self, master): components.entry(master, row, col + 1, self.adam_ui_state, key) elif param_type == 'StatePrecision': components.options(master, row, col + 1, ["auto", "factored", "fp32", "bf16_sr", "int8_sr"], self.adam_ui_state, key) + elif type == 'OrthoGrad': + components.options(master, row, col + 1, ["disabled", "flattened", "iterative"], self.adam_ui_state, key) else: components.switch(master, row, col + 1, self.adam_ui_state, key) diff --git a/modules/ui/OptimizerParamsWindow.py b/modules/ui/OptimizerParamsWindow.py index 7b6b3ef9f..7138d26c4 100644 --- a/modules/ui/OptimizerParamsWindow.py +++ b/modules/ui/OptimizerParamsWindow.py @@ -167,7 +167,7 @@ def create_dynamic_ui( 'use_schedulefree': {'title': 'use_schedulefree', 'tooltip': 'Use Schedulefree method', 'type': 'bool'}, 'use_orthograd': {'title': 'use_orthograd', 'tooltip': 'Use orthograd method', 'type': 'bool'}, 'nnmf_factor': {'title': 'Factored Optimizer', 'tooltip': 'Enables a memory-efficient mode by applying fast low-rank factorization to the optimizers states. It combines factorization for magnitudes with 1-bit compression for signs, drastically reducing VRAM usage and allowing for larger models or batch sizes. This is an approximation which may slightly alter training dynamics.', 'type': 'bool'}, - 'orthogonal_gradient': {'title': 'OrthoGrad', 'tooltip': 'Reduces overfitting by removing the gradient component parallel to the weight, thus improving generalization.', 'type': 'bool'}, + 'orthogonal_gradient': {'title': 'OrthoGrad', 'tooltip': 'Reduces overfitting by removing the gradient component parallel to the weight, thus improving generalization. This has two modes: 1. flattened: Standard vectorized OrthoGrad. Fastest, but loses the structural properties of matrices. 2. iterative: Matrix-wise OrthoGrad, preserves structure by iteratively projecting rows and columns.', 'type': 'OrthoGrad'}, 'use_atan2': {'title': 'Atan2 Scaling', 'tooltip': 'A robust replacement for eps, which also incorporates gradient clipping, bounding and stabilizing the optimizer updates.', 'type': 'bool'}, 'beta1_warmup': {'title': 'Beta1 Warmup Steps', 'tooltip': 'Number of warmup steps to gradually increase beta1 from Minimum Beta1 Value to its final value. During warmup, beta1 increases linearly. leave it empty to disable warmup and use constant beta1.', 'type': 'int'}, 'min_beta1': {'title': 'Minimum Beta1', 'tooltip': 'Starting beta1 value for warmup scheduling. Used only when beta1 warmup is enabled. Lower values allow faster initial adaptation, while higher values provide more smoothing. The final beta1 value is specified in the beta1 parameter.', 'type': 'float'}, @@ -244,6 +244,9 @@ def create_dynamic_ui( elif type == 'StatePrecision': components.options(master, row, col + 1, ["auto", "factored", "fp32", "bf16_sr", "int8_sr"], self.optimizer_ui_state, key, command=self.update_user_pref) + elif type == 'OrthoGrad': + components.options(master, row, col + 1, ["disabled", "flattened", "iterative"], self.optimizer_ui_state, key, + command=self.update_user_pref) elif type != 'bool': components.entry(master, row, col + 1, self.optimizer_ui_state, key, command=self.update_user_pref) diff --git a/modules/util/config/TrainConfig.py b/modules/util/config/TrainConfig.py index 55747715a..ceccc518c 100644 --- a/modules/util/config/TrainConfig.py +++ b/modules/util/config/TrainConfig.py @@ -110,7 +110,7 @@ class TrainOptimizerConfig(BaseConfig): use_schedulefree: True use_orthograd: False nnmf_factor: False - orthogonal_gradient: False + orthogonal_gradient: str use_atan2: False beta1_warmup: int min_beta1: float @@ -231,7 +231,7 @@ def default_values(): data.append(("use_schedulefree", True, bool, True)) data.append(("use_orthograd", False, bool, False)) data.append(("nnmf_factor", False, bool, False)) - data.append(("orthogonal_gradient", False, bool, False)) + data.append(("orthogonal_gradient", 'disabled', str, False)) data.append(("use_atan2", False, bool, False)) data.append(("beta1_warmup", None, int, True)) data.append(("min_beta1", None, float, True)) diff --git a/modules/util/create.py b/modules/util/create.py index 4b5ded9d4..ab23cc7bf 100644 --- a/modules/util/create.py +++ b/modules/util/create.py @@ -680,7 +680,7 @@ def create_optimizer( nesterov_coef=optimizer_config.nesterov_coef if optimizer_config.nesterov_coef is not None else None, cautious_wd=optimizer_config.cautious_wd if optimizer_config.cautious_wd is not None else False, stochastic_rounding=optimizer_config.stochastic_rounding, - orthogonal_gradient=optimizer_config.orthogonal_gradient if optimizer_config.orthogonal_gradient is not None else False, + orthogonal_gradient=optimizer_config.orthogonal_gradient if optimizer_config.orthogonal_gradient is not None else 'disabled', compiled_optimizer=optimizer_config.compile if optimizer_config.compile is not None else False, orthogonal_sinkhorn=optimizer_config.orthogonal_sinkhorn if optimizer_config.orthogonal_sinkhorn is not None else False, sinkhorn_iterations=optimizer_config.sinkhorn_iterations if optimizer_config.sinkhorn_iterations is not None else 5, @@ -707,7 +707,7 @@ def create_optimizer( cautious_wd=optimizer_config.cautious_wd if optimizer_config.cautious_wd is not None else False, stochastic_rounding=optimizer_config.stochastic_rounding, use_atan2=optimizer_config.use_atan2 if optimizer_config.use_atan2 is not None else False, - orthogonal_gradient=optimizer_config.orthogonal_gradient if optimizer_config.orthogonal_gradient is not None else False, + orthogonal_gradient=optimizer_config.orthogonal_gradient if optimizer_config.orthogonal_gradient is not None else 'disabled', kourkoutas_beta=optimizer_config.kourkoutas_beta if optimizer_config.kourkoutas_beta is not None else False, k_warmup_steps=(config.learning_rate_warmup_steps / config.gradient_accumulation_steps), compiled_optimizer=optimizer_config.compile if optimizer_config.compile is not None else False, @@ -734,7 +734,7 @@ def create_optimizer( cautious_wd=optimizer_config.cautious_wd if optimizer_config.cautious_wd is not None else False, stochastic_rounding=optimizer_config.stochastic_rounding, use_atan2=optimizer_config.use_atan2 if optimizer_config.use_atan2 is not None else False, - orthogonal_gradient=optimizer_config.orthogonal_gradient if optimizer_config.orthogonal_gradient is not None else False, + orthogonal_gradient=optimizer_config.orthogonal_gradient if optimizer_config.orthogonal_gradient is not None else 'disabled', kourkoutas_beta=optimizer_config.kourkoutas_beta if optimizer_config.kourkoutas_beta is not None else False, k_warmup_steps=(config.learning_rate_warmup_steps / config.gradient_accumulation_steps), compiled_optimizer=optimizer_config.compile if optimizer_config.compile is not None else False, @@ -768,7 +768,7 @@ def create_optimizer( prodigy_steps=optimizer_config.prodigy_steps if optimizer_config.prodigy_steps is not None else 0, d_limiter=optimizer_config.d_limiter if optimizer_config.d_limiter is not None else False, use_atan2=optimizer_config.use_atan2 if optimizer_config.use_atan2 is not None else False, - orthogonal_gradient=optimizer_config.orthogonal_gradient if optimizer_config.orthogonal_gradient is not None else False, + orthogonal_gradient=optimizer_config.orthogonal_gradient if optimizer_config.orthogonal_gradient is not None else 'disabled', kourkoutas_beta=optimizer_config.kourkoutas_beta if optimizer_config.kourkoutas_beta is not None else False, k_warmup_steps=(config.learning_rate_warmup_steps / config.gradient_accumulation_steps), compiled_optimizer=optimizer_config.compile if optimizer_config.compile is not None else False, @@ -791,7 +791,7 @@ def create_optimizer( weight_decay=optimizer_config.weight_decay if optimizer_config.weight_decay is not None else 0.0, cautious_wd=optimizer_config.cautious_wd if optimizer_config.cautious_wd is not None else False, stochastic_rounding=optimizer_config.stochastic_rounding, - orthogonal_gradient=optimizer_config.orthogonal_gradient if optimizer_config.orthogonal_gradient is not None else False, + orthogonal_gradient=optimizer_config.orthogonal_gradient if optimizer_config.orthogonal_gradient is not None else 'disabled', compiled_optimizer=optimizer_config.compile if optimizer_config.compile is not None else False, spectral_normalization=optimizer_config.spectral_normalization if optimizer_config.spectral_normalization is not None else False, centered_wd=optimizer_config.centered_wd if optimizer_config.centered_wd is not None else 0.0, @@ -815,7 +815,7 @@ def create_optimizer( nnmf_factor=optimizer_config.nnmf_factor if optimizer_config.nnmf_factor is not None else False, cautious_wd=optimizer_config.cautious_wd if optimizer_config.cautious_wd is not None else False, stochastic_rounding=optimizer_config.stochastic_rounding, - orthogonal_gradient=optimizer_config.orthogonal_gradient if optimizer_config.orthogonal_gradient is not None else False, + orthogonal_gradient=optimizer_config.orthogonal_gradient if optimizer_config.orthogonal_gradient is not None else 'disabled', compiled_optimizer=optimizer_config.compile if optimizer_config.compile is not None else False, stochastic_sign=optimizer_config.stochastic_sign if optimizer_config.stochastic_sign is not None else False, spectral_normalization=optimizer_config.spectral_normalization if optimizer_config.spectral_normalization is not None else False, @@ -867,7 +867,7 @@ def create_optimizer( low_rank_ortho=optimizer_config.low_rank_ortho if optimizer_config.low_rank_ortho is not None else False, ortho_rank=optimizer_config.ortho_rank if optimizer_config.ortho_rank is not None else 128, accelerated_ns=optimizer_config.accelerated_ns if optimizer_config.accelerated_ns is not None else False, - orthogonal_gradient=optimizer_config.orthogonal_gradient if optimizer_config.orthogonal_gradient is not None else False, + orthogonal_gradient=optimizer_config.orthogonal_gradient if optimizer_config.orthogonal_gradient is not None else 'disabled', approx_mars=optimizer_config.approx_mars if optimizer_config.approx_mars is not None else False, compiled_optimizer=optimizer_config.compile if optimizer_config.compile is not None else False, spectral_normalization=optimizer_config.spectral_normalization if optimizer_config.spectral_normalization is not None else False, @@ -924,7 +924,7 @@ def create_optimizer( ortho_rank=optimizer_config.ortho_rank if optimizer_config.ortho_rank is not None else 128, normuon_variant=optimizer_config.normuon_variant if optimizer_config.normuon_variant is not None else False, accelerated_ns=optimizer_config.accelerated_ns if optimizer_config.accelerated_ns is not None else False, - orthogonal_gradient=optimizer_config.orthogonal_gradient if optimizer_config.orthogonal_gradient is not None else False, + orthogonal_gradient=optimizer_config.orthogonal_gradient if optimizer_config.orthogonal_gradient is not None else 'disabled', approx_mars=optimizer_config.approx_mars if optimizer_config.approx_mars is not None else False, compiled_optimizer=optimizer_config.compile if optimizer_config.compile is not None else False, spectral_normalization=optimizer_config.spectral_normalization if optimizer_config.spectral_normalization is not None else False, diff --git a/modules/util/optimizer_util.py b/modules/util/optimizer_util.py index c064dd817..3483b1f50 100644 --- a/modules/util/optimizer_util.py +++ b/modules/util/optimizer_util.py @@ -456,7 +456,7 @@ def init_model_parameters( "compile": False, "fused_back_pass": False, "orthogonal_sinkhorn": False, - "orthogonal_gradient": False, + "orthogonal_gradient": 'disabled', "spectral_normalization": False, "centered_wd": 0.0, "centered_wd_mode": "full", @@ -476,7 +476,7 @@ def init_model_parameters( "compile": False, "fused_back_pass": False, "use_atan2": False, - "orthogonal_gradient": False, + "orthogonal_gradient": 'disabled', "kourkoutas_beta": False, "spectral_normalization": False, "centered_wd": 0.0, @@ -497,7 +497,7 @@ def init_model_parameters( "compile": False, "fused_back_pass": False, "use_atan2": True, - "orthogonal_gradient": False, + "orthogonal_gradient": 'disabled', "kourkoutas_beta": False, "spectral_normalization": False, "centered_wd": 0.0, @@ -526,7 +526,7 @@ def init_model_parameters( "prodigy_steps": 0, "d_limiter": False, "use_atan2": False, - "orthogonal_gradient": False, + "orthogonal_gradient": 'disabled', "kourkoutas_beta": False, "centered_wd": 0.0, "centered_wd_mode": "full", @@ -544,7 +544,7 @@ def init_model_parameters( "stochastic_rounding": True, "compile": False, "fused_back_pass": False, - "orthogonal_gradient": False, + "orthogonal_gradient": 'disabled', "spectral_normalization": False, "centered_wd": 0.0, "centered_wd_mode": "full", @@ -560,7 +560,7 @@ def init_model_parameters( "stochastic_rounding": True, "compile": False, "fused_back_pass": False, - "orthogonal_gradient": False, + "orthogonal_gradient": 'disabled', "spectral_normalization": False, "centered_wd": 0.0, "centered_wd_mode": "full", @@ -588,7 +588,7 @@ def init_model_parameters( "muon_te2_adam_lr": None, "normuon_variant": True, "beta2_normuon": 0.95, - "orthogonal_gradient": False, + "orthogonal_gradient": 'disabled', "approx_mars": False, "centered_wd": 0.0, "centered_wd_mode": "full", @@ -620,7 +620,7 @@ def init_model_parameters( "muon_te2_adam_lr": None, "use_atan2": False, "normuon_variant": True, - "orthogonal_gradient": False, + "orthogonal_gradient": 'disabled', "approx_mars": False, "centered_wd": 0.0, "centered_wd_mode": "full", From c31ae3f1784200374a5f99ee3f03056075525b71 Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Sat, 6 Jun 2026 22:06:18 +0300 Subject: [PATCH 66/84] Stable and final 2.5 version --- requirements-global.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements-global.txt b/requirements-global.txt index 8388d2fd6..f1fe8796b 100644 --- a/requirements-global.txt +++ b/requirements-global.txt @@ -41,7 +41,7 @@ prodigyopt==1.1.2 # prodigy optimizer schedulefree==1.4.1 # schedule-free optimizers pytorch_optimizer==3.6.0 # pytorch optimizers prodigy-plus-schedule-free==2.0.1 # Prodigy plus optimizer -adv_optm==2.4.dev23 # advanced optimizers +adv_optm==2.5 # advanced optimizers -e git+https://github.com/KellerJordan/Muon.git@f90a42b#egg=muon-optimizer # Profiling From a49c33340bb0bcdcfdb29fc69fe5d501fa33388c Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Sat, 6 Jun 2026 22:17:41 +0300 Subject: [PATCH 67/84] pre-commit --- modules/ui/MuonAdamWindow.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modules/ui/MuonAdamWindow.py b/modules/ui/MuonAdamWindow.py index c0958e88b..e895e7290 100644 --- a/modules/ui/MuonAdamWindow.py +++ b/modules/ui/MuonAdamWindow.py @@ -104,7 +104,7 @@ def create_adam_params_ui(self, master): components.entry(master, row, col + 1, self.adam_ui_state, key) elif param_type == 'StatePrecision': components.options(master, row, col + 1, ["auto", "factored", "fp32", "bf16_sr", "int8_sr"], self.adam_ui_state, key) - elif type == 'OrthoGrad': + elif param_type == 'OrthoGrad': components.options(master, row, col + 1, ["disabled", "flattened", "iterative"], self.adam_ui_state, key) else: components.switch(master, row, col + 1, self.adam_ui_state, key) From 0d06b4de0af849b4177cb139faf5cecc3f2d4329 Mon Sep 17 00:00:00 2001 From: Koratahiu Date: Sun, 7 Jun 2026 01:20:31 +0300 Subject: [PATCH 68/84] 2.5.1: Improve iterative OrthoGrad --- requirements-global.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements-global.txt b/requirements-global.txt index f1fe8796b..e10213770 100644 --- a/requirements-global.txt +++ b/requirements-global.txt @@ -41,7 +41,7 @@ prodigyopt==1.1.2 # prodigy optimizer schedulefree==1.4.1 # schedule-free optimizers pytorch_optimizer==3.6.0 # pytorch optimizers prodigy-plus-schedule-free==2.0.1 # Prodigy plus optimizer -adv_optm==2.5 # advanced optimizers +adv_optm==2.5.1 # advanced optimizers -e git+https://github.com/KellerJordan/Muon.git@f90a42b#egg=muon-optimizer # Profiling From 895e908348c60b79a3fc2bd77fac0efc17f64a21 Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Wed, 10 Jun 2026 13:20:06 +0300 Subject: [PATCH 69/84] 2.5.2: Improved Spectral scaling for OFT --- modules/util/create.py | 2 +- modules/util/optimizer/tag_util.py | 11 +++++++++-- requirements-global.txt | 2 +- 3 files changed, 11 insertions(+), 4 deletions(-) diff --git a/modules/util/create.py b/modules/util/create.py index ab23cc7bf..bc5640c4a 100644 --- a/modules/util/create.py +++ b/modules/util/create.py @@ -143,7 +143,7 @@ def create_optimizer( if config.optimizer.optimizer.is_adv: # Tag PEFT parameters based on parameter names. - tag_peft_parameters(model) + tag_peft_parameters(model, config) match config.optimizer.optimizer: diff --git a/modules/util/optimizer/tag_util.py b/modules/util/optimizer/tag_util.py index 131ac1423..60df705cf 100644 --- a/modules/util/optimizer/tag_util.py +++ b/modules/util/optimizer/tag_util.py @@ -1,9 +1,11 @@ from modules.module.LoRAModule import LoRAModuleWrapper +from modules.util.config.TrainConfig import TrainConfig import torch +import math -def tag_peft_parameters(model: torch.nn.Module | None): +def tag_peft_parameters(model: torch.nn.Module | None, config: TrainConfig): """ Tags PEFT parameters with attributes like `_is_lora_A`, `_is_lora_B`, `_is_oft`, and `_is_dora_scale` based on their names. @@ -20,8 +22,13 @@ def apply_tags(name: str, p: torch.nn.Parameter): # Vector in shape of >= 2D tensor p._is_dora_scale = True elif name.endswith("oft_R.weight"): - # Set of independent vectors (rank, n_elements) + # Set of independent flattened matrices (rank, n_elements) p._is_oft = True + if config.oft_scaled: + n_el = p.shape[-1] + b = (1 + math.sqrt(1 + 8 * n_el)) / 2 + scaling_factor = 2 * math.sqrt(b - 1) + p._oft_scale_factor = scaling_factor for module in vars(model).values(): if isinstance(module, LoRAModuleWrapper): diff --git a/requirements-global.txt b/requirements-global.txt index e10213770..723e31cbd 100644 --- a/requirements-global.txt +++ b/requirements-global.txt @@ -41,7 +41,7 @@ prodigyopt==1.1.2 # prodigy optimizer schedulefree==1.4.1 # schedule-free optimizers pytorch_optimizer==3.6.0 # pytorch optimizers prodigy-plus-schedule-free==2.0.1 # Prodigy plus optimizer -adv_optm==2.5.1 # advanced optimizers +adv_optm==2.5.2 # advanced optimizers -e git+https://github.com/KellerJordan/Muon.git@f90a42b#egg=muon-optimizer # Profiling From 4099562e09cd0dededdfeb7bcb238b935acd8dfd Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Wed, 10 Jun 2026 13:22:27 +0300 Subject: [PATCH 70/84] pre-commit --- modules/util/optimizer/tag_util.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/modules/util/optimizer/tag_util.py b/modules/util/optimizer/tag_util.py index 60df705cf..3c82bad58 100644 --- a/modules/util/optimizer/tag_util.py +++ b/modules/util/optimizer/tag_util.py @@ -1,8 +1,9 @@ +import math + from modules.module.LoRAModule import LoRAModuleWrapper from modules.util.config.TrainConfig import TrainConfig import torch -import math def tag_peft_parameters(model: torch.nn.Module | None, config: TrainConfig): From 6d4f25fa6900460c50dcadcb94bdd15cd9e2fe4d Mon Sep 17 00:00:00 2001 From: Koratahiu Date: Fri, 12 Jun 2026 18:39:01 +0300 Subject: [PATCH 71/84] 2.5.3 Improve OFT spectral scaling --- requirements-global.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements-global.txt b/requirements-global.txt index 723e31cbd..cc3376029 100644 --- a/requirements-global.txt +++ b/requirements-global.txt @@ -41,7 +41,7 @@ prodigyopt==1.1.2 # prodigy optimizer schedulefree==1.4.1 # schedule-free optimizers pytorch_optimizer==3.6.0 # pytorch optimizers prodigy-plus-schedule-free==2.0.1 # Prodigy plus optimizer -adv_optm==2.5.2 # advanced optimizers +adv_optm==2.5.3 # advanced optimizers -e git+https://github.com/KellerJordan/Muon.git@f90a42b#egg=muon-optimizer # Profiling From db4e48ec2cb339c0eafd518d62911b97433ae906 Mon Sep 17 00:00:00 2001 From: Koratahiu Date: Mon, 15 Jun 2026 03:19:17 +0300 Subject: [PATCH 72/84] 2.5.4: Small SignSGD bugfix and improvement --- requirements-global.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements-global.txt b/requirements-global.txt index cc3376029..db406e51a 100644 --- a/requirements-global.txt +++ b/requirements-global.txt @@ -41,7 +41,7 @@ prodigyopt==1.1.2 # prodigy optimizer schedulefree==1.4.1 # schedule-free optimizers pytorch_optimizer==3.6.0 # pytorch optimizers prodigy-plus-schedule-free==2.0.1 # Prodigy plus optimizer -adv_optm==2.5.3 # advanced optimizers +adv_optm==2.5.4 # advanced optimizers -e git+https://github.com/KellerJordan/Muon.git@f90a42b#egg=muon-optimizer # Profiling From 984dc35dca0df02f4652c997c467f259b8f36682 Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Mon, 15 Jun 2026 03:20:14 +0300 Subject: [PATCH 73/84] fix state_precision to auto --- modules/util/config/TrainConfig.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modules/util/config/TrainConfig.py b/modules/util/config/TrainConfig.py index 55747715a..eef351475 100644 --- a/modules/util/config/TrainConfig.py +++ b/modules/util/config/TrainConfig.py @@ -260,7 +260,7 @@ def default_values(): data.append(("centered_wd_mode", "full", str, False)) data.append(("factored_2nd", False, bool, False)) data.append(("fisher_wd", False, bool, False)) - data.append(("state_precision", "float8", str, False)) + data.append(("state_precision", "auto", str, False)) data.append(("orthogonal_sinkhorn", False, bool, False)) data.append(("sinkhorn_iterations", None, int, True)) data.append(("normed_momentum", False, bool, False)) From 4329c98752bc577818978c0725a41be1f4a78913 Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Mon, 15 Jun 2026 03:32:41 +0300 Subject: [PATCH 74/84] Fix invalid orthogonal_gradient and state_precision part of @BitcrushedHeart referenced commit --- modules/util/config/TrainConfig.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/modules/util/config/TrainConfig.py b/modules/util/config/TrainConfig.py index 5ab415392..291e4190b 100644 --- a/modules/util/config/TrainConfig.py +++ b/modules/util/config/TrainConfig.py @@ -150,6 +150,20 @@ class TrainOptimizerConfig(BaseConfig): def __init__(self, data: list[(str, Any, type, bool)]): super().__init__(data) + def from_dict(self, data: dict) -> "TrainOptimizerConfig": + sp = data.get("state_precision") + valid_sp = {"auto", "factored", "fp32", "fp16", "bf16_sr", "int8_sr"} + if sp is not None and sp not in valid_sp: + print(f"WARN: invalid optimizer state_precision '{sp}' in config, falling back to 'auto'.") + data = data.copy() + data["state_precision"] = "auto" + # orthogonal_gradient was a bool before adv_optm 2.5 made it a mode string + og_ortho = data.get("orthogonal_gradient") + if isinstance(og_ortho, bool): + data = data.copy() + data["orthogonal_gradient"] = "flattened" if og_ortho else "disabled" + return super().from_dict(data) + @staticmethod def default_values(): data = [] From 9d7e8588dcb9d19f65a03fd3c15e9172faa26813 Mon Sep 17 00:00:00 2001 From: Koratahiu Date: Mon, 15 Jun 2026 19:49:41 +0300 Subject: [PATCH 75/84] 2.5.5: Muon Variants Bugfixes --- requirements-global.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements-global.txt b/requirements-global.txt index db406e51a..b7a87b6bf 100644 --- a/requirements-global.txt +++ b/requirements-global.txt @@ -41,7 +41,7 @@ prodigyopt==1.1.2 # prodigy optimizer schedulefree==1.4.1 # schedule-free optimizers pytorch_optimizer==3.6.0 # pytorch optimizers prodigy-plus-schedule-free==2.0.1 # Prodigy plus optimizer -adv_optm==2.5.4 # advanced optimizers +adv_optm==2.5.5 # advanced optimizers -e git+https://github.com/KellerJordan/Muon.git@f90a42b#egg=muon-optimizer # Profiling From 39f7ee2d8f1b447d63e3d38873725167989eb77a Mon Sep 17 00:00:00 2001 From: Koratahiu Date: Tue, 16 Jun 2026 16:42:21 +0300 Subject: [PATCH 76/84] version bump bugfixes --- requirements-global.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements-global.txt b/requirements-global.txt index b7a87b6bf..d7f3e3ddc 100644 --- a/requirements-global.txt +++ b/requirements-global.txt @@ -41,7 +41,7 @@ prodigyopt==1.1.2 # prodigy optimizer schedulefree==1.4.1 # schedule-free optimizers pytorch_optimizer==3.6.0 # pytorch optimizers prodigy-plus-schedule-free==2.0.1 # Prodigy plus optimizer -adv_optm==2.5.5 # advanced optimizers +adv_optm==2.5.7 # advanced optimizers -e git+https://github.com/KellerJordan/Muon.git@f90a42b#egg=muon-optimizer # Profiling From 68410a82980ee548d56caf4a70d51fa490682e60 Mon Sep 17 00:00:00 2001 From: Koratahiu Date: Wed, 17 Jun 2026 00:07:01 +0300 Subject: [PATCH 77/84] version bump --- requirements-global.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements-global.txt b/requirements-global.txt index d7f3e3ddc..0fa8fb6d8 100644 --- a/requirements-global.txt +++ b/requirements-global.txt @@ -41,7 +41,7 @@ prodigyopt==1.1.2 # prodigy optimizer schedulefree==1.4.1 # schedule-free optimizers pytorch_optimizer==3.6.0 # pytorch optimizers prodigy-plus-schedule-free==2.0.1 # Prodigy plus optimizer -adv_optm==2.5.7 # advanced optimizers +adv_optm==2.5.8 # advanced optimizers -e git+https://github.com/KellerJordan/Muon.git@f90a42b#egg=muon-optimizer # Profiling From 266199a3e378d4a5798eed5ac42ce718ce6062fa Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Sat, 20 Jun 2026 09:38:15 +0300 Subject: [PATCH 78/84] Remove from_dict and add __migration_10 --- modules/util/config/TrainConfig.py | 41 +++++++++++++++++++----------- 1 file changed, 26 insertions(+), 15 deletions(-) diff --git a/modules/util/config/TrainConfig.py b/modules/util/config/TrainConfig.py index 291e4190b..be217385f 100644 --- a/modules/util/config/TrainConfig.py +++ b/modules/util/config/TrainConfig.py @@ -150,20 +150,6 @@ class TrainOptimizerConfig(BaseConfig): def __init__(self, data: list[(str, Any, type, bool)]): super().__init__(data) - def from_dict(self, data: dict) -> "TrainOptimizerConfig": - sp = data.get("state_precision") - valid_sp = {"auto", "factored", "fp32", "fp16", "bf16_sr", "int8_sr"} - if sp is not None and sp not in valid_sp: - print(f"WARN: invalid optimizer state_precision '{sp}' in config, falling back to 'auto'.") - data = data.copy() - data["state_precision"] = "auto" - # orthogonal_gradient was a bool before adv_optm 2.5 made it a mode string - og_ortho = data.get("orthogonal_gradient") - if isinstance(og_ortho, bool): - data = data.copy() - data["orthogonal_gradient"] = "flattened" if og_ortho else "disabled" - return super().from_dict(data) - @staticmethod def default_values(): data = [] @@ -599,7 +585,7 @@ class TrainConfig(BaseConfig): def __init__(self, data: list[(str, Any, type, bool)]): super().__init__( data, - config_version=10, + config_version=11, config_migrations={ 0: self.__migration_0, 1: self.__migration_1, @@ -611,6 +597,7 @@ def __init__(self, data: list[(str, Any, type, bool)]): 7: self.__migration_7, 8: self.__migration_8, 9: self.__migration_9, + 10: self.__migration_10, } ) @@ -830,6 +817,30 @@ def replace_dtype(part: str): return migrated_data + def __migration_10(self, data: dict) -> dict: + migrated_data = data.copy() + + def migrate_optimizer(opt_data: dict): + sp = opt_data.get("state_precision") + valid_sp = {"auto", "factored", "fp32", "fp16", "bf16_sr", "int8_sr"} + if sp is not None and sp not in valid_sp: + print(f"WARN: invalid optimizer state_precision '{sp}' in config, falling back to 'auto'.") + opt_data["state_precision"] = "auto" + + og_ortho = opt_data.get("orthogonal_gradient") + if isinstance(og_ortho, bool): + opt_data["orthogonal_gradient"] = "flattened" if og_ortho else "disabled" + + if "optimizer" in migrated_data and isinstance(migrated_data["optimizer"], dict): + migrate_optimizer(migrated_data["optimizer"]) + + if "optimizer_defaults" in migrated_data and isinstance(migrated_data["optimizer_defaults"], dict): + for opt_data in migrated_data["optimizer_defaults"].values(): + if isinstance(opt_data, dict): + migrate_optimizer(opt_data) + + return migrated_data + def weight_dtypes(self) -> ModelWeightDtypes: return ModelWeightDtypes( self.train_dtype, From 43c6560e2dbd9321fa6d1c40694c5fb3d9b75988 Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Sun, 21 Jun 2026 01:02:48 +0300 Subject: [PATCH 79/84] Revert and add sub-config migration --- modules/util/config/TrainConfig.py | 50 ++++++++++++++---------------- 1 file changed, 24 insertions(+), 26 deletions(-) diff --git a/modules/util/config/TrainConfig.py b/modules/util/config/TrainConfig.py index be217385f..1829f7dd9 100644 --- a/modules/util/config/TrainConfig.py +++ b/modules/util/config/TrainConfig.py @@ -148,7 +148,29 @@ class TrainOptimizerConfig(BaseConfig): geometric_wd: False def __init__(self, data: list[(str, Any, type, bool)]): - super().__init__(data) + super().__init__( + data, + config_version=1, + config_migrations={ + 0: self.__migration_0 + } + ) + + def __migration_0(self, data: dict) -> dict: + migrated_data = data.copy() + + sp = migrated_data.get("state_precision") + valid_sp = {"auto", "factored", "fp32", "fp16", "bf16_sr", "int8_sr"} + if sp is not None and sp not in valid_sp: + print(f"WARN: invalid optimizer state_precision '{sp}' in config, falling back to 'auto'.") + migrated_data["state_precision"] = "auto" + + # orthogonal_gradient was a bool before adv_optm 2.5 made it a mode string + og_ortho = migrated_data.get("orthogonal_gradient") + if isinstance(og_ortho, bool): + migrated_data["orthogonal_gradient"] = "flattened" if og_ortho else "disabled" + + return migrated_data @staticmethod def default_values(): @@ -585,7 +607,7 @@ class TrainConfig(BaseConfig): def __init__(self, data: list[(str, Any, type, bool)]): super().__init__( data, - config_version=11, + config_version=10, config_migrations={ 0: self.__migration_0, 1: self.__migration_1, @@ -597,7 +619,6 @@ def __init__(self, data: list[(str, Any, type, bool)]): 7: self.__migration_7, 8: self.__migration_8, 9: self.__migration_9, - 10: self.__migration_10, } ) @@ -817,29 +838,6 @@ def replace_dtype(part: str): return migrated_data - def __migration_10(self, data: dict) -> dict: - migrated_data = data.copy() - - def migrate_optimizer(opt_data: dict): - sp = opt_data.get("state_precision") - valid_sp = {"auto", "factored", "fp32", "fp16", "bf16_sr", "int8_sr"} - if sp is not None and sp not in valid_sp: - print(f"WARN: invalid optimizer state_precision '{sp}' in config, falling back to 'auto'.") - opt_data["state_precision"] = "auto" - - og_ortho = opt_data.get("orthogonal_gradient") - if isinstance(og_ortho, bool): - opt_data["orthogonal_gradient"] = "flattened" if og_ortho else "disabled" - - if "optimizer" in migrated_data and isinstance(migrated_data["optimizer"], dict): - migrate_optimizer(migrated_data["optimizer"]) - - if "optimizer_defaults" in migrated_data and isinstance(migrated_data["optimizer_defaults"], dict): - for opt_data in migrated_data["optimizer_defaults"].values(): - if isinstance(opt_data, dict): - migrate_optimizer(opt_data) - - return migrated_data def weight_dtypes(self) -> ModelWeightDtypes: return ModelWeightDtypes( From 885031c013e8a1dfc61e3f492cbab8e59db18000 Mon Sep 17 00:00:00 2001 From: Koratahiu Date: Sun, 21 Jun 2026 01:58:19 +0300 Subject: [PATCH 80/84] 2.5.9: Enhance compiled optimizer mode --- requirements-global.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements-global.txt b/requirements-global.txt index 0fa8fb6d8..7f8a20fbc 100644 --- a/requirements-global.txt +++ b/requirements-global.txt @@ -41,7 +41,7 @@ prodigyopt==1.1.2 # prodigy optimizer schedulefree==1.4.1 # schedule-free optimizers pytorch_optimizer==3.6.0 # pytorch optimizers prodigy-plus-schedule-free==2.0.1 # Prodigy plus optimizer -adv_optm==2.5.8 # advanced optimizers +adv_optm==2.5.9 # advanced optimizers -e git+https://github.com/KellerJordan/Muon.git@f90a42b#egg=muon-optimizer # Profiling From 767ec38eed44666f3870fa2d27a28c7d41c4c56f Mon Sep 17 00:00:00 2001 From: Koratahiu Date: Sat, 27 Jun 2026 01:46:03 +0300 Subject: [PATCH 81/84] Update requirements-global.txt --- requirements-global.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements-global.txt b/requirements-global.txt index 7f8a20fbc..2d2c28b70 100644 --- a/requirements-global.txt +++ b/requirements-global.txt @@ -41,7 +41,7 @@ prodigyopt==1.1.2 # prodigy optimizer schedulefree==1.4.1 # schedule-free optimizers pytorch_optimizer==3.6.0 # pytorch optimizers prodigy-plus-schedule-free==2.0.1 # Prodigy plus optimizer -adv_optm==2.5.9 # advanced optimizers +adv_optm==2.5.10 # advanced optimizers -e git+https://github.com/KellerJordan/Muon.git@f90a42b#egg=muon-optimizer # Profiling From a7c647bb9310084232f40343a487d73f410a3b42 Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Sat, 27 Jun 2026 11:31:40 +0300 Subject: [PATCH 82/84] initial --- modules/ui/OptimizerParamsWindow.py | 1 + modules/util/config/TrainConfig.py | 2 ++ modules/util/create.py | 8 ++++++++ modules/util/optimizer_util.py | 8 ++++++++ requirements-global.txt | 2 +- 5 files changed, 20 insertions(+), 1 deletion(-) diff --git a/modules/ui/OptimizerParamsWindow.py b/modules/ui/OptimizerParamsWindow.py index 9a08a1cb2..67ce94e03 100644 --- a/modules/ui/OptimizerParamsWindow.py +++ b/modules/ui/OptimizerParamsWindow.py @@ -202,6 +202,7 @@ def create_dynamic_ui( 'nesterov_coef': {'title': 'Nesterov Coef', 'tooltip': 'Controls the mixing coefficient between momentum gradients and raw gradients in Nesterov momentum. For a factor of 0.8, the final update will be 80% of the momentum gradients and 20% raw gradient. Leaving it unset toggles the standard Nestrov behavior (where nesterov_coef = beta1 or momentum). Setting it to 0 cancels momentum contribution.', 'type': 'float'}, 'snr_cond': {'title': 'SNR Preconditioning', 'tooltip': 'Applies a Signal-to-Noise Ratio (SNR) precondition to reshape the optimization curve. It prioritizes high-confidence signals and dampens noise. On-the-fly math with zero memory overhead. Requires Normalization-then-Momentum (NtM). ', 'type': 'bool'}, 'geometric_wd': {'title': 'Geometric Weight Decay', 'tooltip': 'Regularizes weights based on the geometric structure of the optimizer. Compatible with cautious weight decay.', 'type': 'bool'}, + 'scaled_wd': {'title': 'Scaled Weight Decay', 'tooltip': 'Decoupled, dimension-scaled weight decay. Recommended wd value: 0.1 for all ranks, widths and training methods.', 'type': 'bool'}, } # @formatter:on diff --git a/modules/util/config/TrainConfig.py b/modules/util/config/TrainConfig.py index 1829f7dd9..203cef832 100644 --- a/modules/util/config/TrainConfig.py +++ b/modules/util/config/TrainConfig.py @@ -146,6 +146,7 @@ class TrainOptimizerConfig(BaseConfig): nesterov_coef: float snr_cond: False geometric_wd: False + scaled_wd: False def __init__(self, data: list[(str, Any, type, bool)]): super().__init__( @@ -289,6 +290,7 @@ def default_values(): data.append(("nesterov_coef", None, float, True)) data.append(("snr_cond", False, bool, False)) data.append(("geometric_wd", False, bool, False)) + data.append(("scaled_wd", False, bool, False)) return TrainOptimizerConfig(data) diff --git a/modules/util/create.py b/modules/util/create.py index bc5640c4a..637eaa003 100644 --- a/modules/util/create.py +++ b/modules/util/create.py @@ -679,6 +679,7 @@ def create_optimizer( nesterov=optimizer_config.nesterov if optimizer_config.nesterov is not None else True, nesterov_coef=optimizer_config.nesterov_coef if optimizer_config.nesterov_coef is not None else None, cautious_wd=optimizer_config.cautious_wd if optimizer_config.cautious_wd is not None else False, + scaled_wd=optimizer_config.scaled_wd if optimizer_config.scaled_wd is not None else False, stochastic_rounding=optimizer_config.stochastic_rounding, orthogonal_gradient=optimizer_config.orthogonal_gradient if optimizer_config.orthogonal_gradient is not None else 'disabled', compiled_optimizer=optimizer_config.compile if optimizer_config.compile is not None else False, @@ -705,6 +706,7 @@ def create_optimizer( factored_2nd=optimizer_config.factored_2nd if optimizer_config.factored_2nd is not None else False, fisher_wd=optimizer_config.fisher_wd if optimizer_config.fisher_wd is not None else False, cautious_wd=optimizer_config.cautious_wd if optimizer_config.cautious_wd is not None else False, + scaled_wd=optimizer_config.scaled_wd if optimizer_config.scaled_wd is not None else False, stochastic_rounding=optimizer_config.stochastic_rounding, use_atan2=optimizer_config.use_atan2 if optimizer_config.use_atan2 is not None else False, orthogonal_gradient=optimizer_config.orthogonal_gradient if optimizer_config.orthogonal_gradient is not None else 'disabled', @@ -732,6 +734,7 @@ def create_optimizer( factored_2nd=optimizer_config.factored_2nd if optimizer_config.factored_2nd is not None else False, fisher_wd=optimizer_config.fisher_wd if optimizer_config.fisher_wd is not None else False, cautious_wd=optimizer_config.cautious_wd if optimizer_config.cautious_wd is not None else False, + scaled_wd=optimizer_config.scaled_wd if optimizer_config.scaled_wd is not None else False, stochastic_rounding=optimizer_config.stochastic_rounding, use_atan2=optimizer_config.use_atan2 if optimizer_config.use_atan2 is not None else False, orthogonal_gradient=optimizer_config.orthogonal_gradient if optimizer_config.orthogonal_gradient is not None else 'disabled', @@ -760,6 +763,7 @@ def create_optimizer( factored_2nd=optimizer_config.factored_2nd if optimizer_config.factored_2nd is not None else False, fisher_wd=optimizer_config.fisher_wd if optimizer_config.fisher_wd is not None else False, cautious_wd=optimizer_config.cautious_wd if optimizer_config.cautious_wd is not None else False, + scaled_wd=optimizer_config.scaled_wd if optimizer_config.scaled_wd is not None else False, stochastic_rounding=optimizer_config.stochastic_rounding, d0=optimizer_config.d0 if optimizer_config.d0 is not None else 1e-6, d_coef=optimizer_config.d_coef if optimizer_config.d_coef is not None else 1.0, @@ -790,6 +794,7 @@ def create_optimizer( momentum=optimizer_config.momentum if optimizer_config.momentum is not None else 0, weight_decay=optimizer_config.weight_decay if optimizer_config.weight_decay is not None else 0.0, cautious_wd=optimizer_config.cautious_wd if optimizer_config.cautious_wd is not None else False, + scaled_wd=optimizer_config.scaled_wd if optimizer_config.scaled_wd is not None else False, stochastic_rounding=optimizer_config.stochastic_rounding, orthogonal_gradient=optimizer_config.orthogonal_gradient if optimizer_config.orthogonal_gradient is not None else 'disabled', compiled_optimizer=optimizer_config.compile if optimizer_config.compile is not None else False, @@ -814,6 +819,7 @@ def create_optimizer( weight_decay=optimizer_config.weight_decay if optimizer_config.weight_decay is not None else 0.0, nnmf_factor=optimizer_config.nnmf_factor if optimizer_config.nnmf_factor is not None else False, cautious_wd=optimizer_config.cautious_wd if optimizer_config.cautious_wd is not None else False, + scaled_wd=optimizer_config.scaled_wd if optimizer_config.scaled_wd is not None else False, stochastic_rounding=optimizer_config.stochastic_rounding, orthogonal_gradient=optimizer_config.orthogonal_gradient if optimizer_config.orthogonal_gradient is not None else 'disabled', compiled_optimizer=optimizer_config.compile if optimizer_config.compile is not None else False, @@ -859,6 +865,7 @@ def create_optimizer( weight_decay=optimizer_config.weight_decay if optimizer_config.weight_decay is not None else 0.0, rms_rescaling=optimizer_config.rms_rescaling if optimizer_config.rms_rescaling is not None else True, cautious_wd=optimizer_config.cautious_wd if optimizer_config.cautious_wd is not None else False, + scaled_wd=optimizer_config.scaled_wd if optimizer_config.scaled_wd is not None else False, stochastic_rounding=optimizer_config.stochastic_rounding, nesterov=optimizer_config.nesterov if optimizer_config.nesterov is not None else True, nesterov_coef=optimizer_config.nesterov_coef if optimizer_config.nesterov_coef is not None else None, @@ -916,6 +923,7 @@ def create_optimizer( rms_rescaling=optimizer_config.rms_rescaling if optimizer_config.rms_rescaling is not None else True, weight_decay=optimizer_config.weight_decay if optimizer_config.weight_decay is not None else 0.0, cautious_wd=optimizer_config.cautious_wd if optimizer_config.cautious_wd is not None else False, + scaled_wd=optimizer_config.scaled_wd if optimizer_config.scaled_wd is not None else False, stochastic_rounding=optimizer_config.stochastic_rounding, nesterov=optimizer_config.nesterov if optimizer_config.nesterov is not None else True, nesterov_coef=optimizer_config.nesterov_coef if optimizer_config.nesterov_coef is not None else None, diff --git a/modules/util/optimizer_util.py b/modules/util/optimizer_util.py index 3483b1f50..3043b271b 100644 --- a/modules/util/optimizer_util.py +++ b/modules/util/optimizer_util.py @@ -449,6 +449,7 @@ def init_model_parameters( "nesterov_coef": None, "sinkhorn_iterations": 5, "cautious_wd": False, + "scaled_wd": False, "geometric_wd": False, "weight_decay": 0.0, "snr_cond": False, @@ -470,6 +471,7 @@ def init_model_parameters( "eps": 1e-8, "fisher_wd": False, "cautious_wd": False, + "scaled_wd": False, "weight_decay": 0.0, "factored_2nd": False, "stochastic_rounding": True, @@ -491,6 +493,7 @@ def init_model_parameters( "eps": 1e-6, "fisher_wd": False, "cautious_wd": False, + "scaled_wd": False, "weight_decay": 0.0, "factored_2nd": False, "stochastic_rounding": True, @@ -513,6 +516,7 @@ def init_model_parameters( "eps": 1e-8, "fisher_wd": False, "cautious_wd": False, + "scaled_wd": False, "weight_decay": 0.0, "nnmf_factor": False, "factored_2nd": False, @@ -539,6 +543,7 @@ def init_model_parameters( "nesterov_coef": None, "geometric_wd": False, "cautious_wd": False, + "scaled_wd": False, "weight_decay": 0.0, "snr_cond": False, "stochastic_rounding": True, @@ -555,6 +560,7 @@ def init_model_parameters( "beta1": 0.9, "beta2": 0.99, "cautious_wd": False, + "scaled_wd": False, "weight_decay": 0.0, "nnmf_factor": False, "stochastic_rounding": True, @@ -570,6 +576,7 @@ def init_model_parameters( "nesterov": True, "nesterov_coef": None, "cautious_wd": False, + "scaled_wd": False, "weight_decay": 0.0, "accelerated_ns": False, "ns_steps": 5, @@ -602,6 +609,7 @@ def init_model_parameters( "beta2": 0.95, "eps": 1e-8, "cautious_wd": False, + "scaled_wd": False, "weight_decay": 0.0, "accelerated_ns": False, "ns_steps": 5, diff --git a/requirements-global.txt b/requirements-global.txt index 2d2c28b70..506b39d6f 100644 --- a/requirements-global.txt +++ b/requirements-global.txt @@ -41,7 +41,7 @@ prodigyopt==1.1.2 # prodigy optimizer schedulefree==1.4.1 # schedule-free optimizers pytorch_optimizer==3.6.0 # pytorch optimizers prodigy-plus-schedule-free==2.0.1 # Prodigy plus optimizer -adv_optm==2.5.10 # advanced optimizers +adv_optm==2.6.1.dev1 # advanced optimizers -e git+https://github.com/KellerJordan/Muon.git@f90a42b#egg=muon-optimizer # Profiling From f86818c789cfc302277f77348845343949347446 Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Sat, 27 Jun 2026 11:38:20 +0300 Subject: [PATCH 83/84] dev2 --- requirements-global.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements-global.txt b/requirements-global.txt index 506b39d6f..6efa4dd90 100644 --- a/requirements-global.txt +++ b/requirements-global.txt @@ -41,7 +41,7 @@ prodigyopt==1.1.2 # prodigy optimizer schedulefree==1.4.1 # schedule-free optimizers pytorch_optimizer==3.6.0 # pytorch optimizers prodigy-plus-schedule-free==2.0.1 # Prodigy plus optimizer -adv_optm==2.6.1.dev1 # advanced optimizers +adv_optm==2.6.1.dev2 # advanced optimizers -e git+https://github.com/KellerJordan/Muon.git@f90a42b#egg=muon-optimizer # Profiling From 81c5a890caa704e4848dfa0a900c7ee03b0432d5 Mon Sep 17 00:00:00 2001 From: Koratahiu Date: Sun, 28 Jun 2026 13:39:38 +0300 Subject: [PATCH 84/84] Update requirements-global.txt --- requirements-global.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements-global.txt b/requirements-global.txt index 6efa4dd90..b64edfc3b 100644 --- a/requirements-global.txt +++ b/requirements-global.txt @@ -41,7 +41,7 @@ prodigyopt==1.1.2 # prodigy optimizer schedulefree==1.4.1 # schedule-free optimizers pytorch_optimizer==3.6.0 # pytorch optimizers prodigy-plus-schedule-free==2.0.1 # Prodigy plus optimizer -adv_optm==2.6.1.dev2 # advanced optimizers +adv_optm==2.6.1.dev3 # advanced optimizers -e git+https://github.com/KellerJordan/Muon.git@f90a42b#egg=muon-optimizer # Profiling