diff --git a/modules/ui/MuonAdamWindow.py b/modules/ui/MuonAdamWindow.py index 2f1682f36..5879ab432 100644 --- a/modules/ui/MuonAdamWindow.py +++ b/modules/ui/MuonAdamWindow.py @@ -76,14 +76,11 @@ def create_adam_params_ui(self, master): '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'}, - 'grams_moment': {'title': 'GRAMS Variant', 'tooltip': 'Aligns the momentum direction with the current gradient direction while preserving its accumulated magnitude.', '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'}, - 'k_warmup_steps': {'title': 'K-β Warmup Steps ', 'tooltip': 'When using Kourkoutas Beta, the number of initial training steps during which the dynamic β₂ logic is held off. In this period, β₂ is set to its fixed value to allow for initial training stability before the adaptive mechanism activates.', 'type': 'int'}, } # @formatter:on diff --git a/modules/ui/OptimizerParamsWindow.py b/modules/ui/OptimizerParamsWindow.py index ed50e9f57..16063c26c 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'}, - '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'}, - 'grams_moment': {'title': 'GRAMS Variant', 'tooltip': 'Aligns the momentum direction with the current gradient direction while preserving its accumulated magnitude.', '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'}, @@ -178,7 +176,6 @@ def create_dynamic_ui( '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'}, - 'k_warmup_steps': {'title': 'K-β Warmup Steps ', 'tooltip': 'When using Kourkoutas Beta, the number of initial training steps during which the dynamic β₂ logic is held off. In this period, β₂ is set to its fixed value to allow for initial training stability before the adaptive mechanism activates.', 'type': 'int'}, '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'}, 'MuonWithAuxAdam': {'title': 'MuonWithAuxAdam', 'tooltip': 'Whether to use the standard way of Muon. Non-hidden layers fallback to ADAMW, and MUON takes the rest. Note: The auxiliary Adam (ADAMW) is typically only relevant for training "full" LoRA (LoRA for all layers) or full finetune and is irrelevant for most common LoRA use cases.', 'type': 'bool'}, @@ -190,13 +187,11 @@ def create_dynamic_ui( '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'}, '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'}, 'low_rank_ortho': {'title': 'Low-rank Orthogonalization', 'tooltip': 'Use low-rank orthogonalization to accelerate Muon by orthogonalizing only in a low-dimensional subspace, improving speed and noise robustness.', 'type': 'bool'}, 'ortho_rank': {'title': 'Ortho Rank', 'tooltip': 'Target rank for low-rank orthogonalization. Controls the dimensionality of the subspace used for efficient and noise-robust orthogonalization.', 'type': 'int'}, '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'}, - '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'}, } diff --git a/modules/util/config/TrainConfig.py b/modules/util/config/TrainConfig.py index ad9fbfd6c..1c940da66 100644 --- a/modules/util/config/TrainConfig.py +++ b/modules/util/config/TrainConfig.py @@ -118,10 +118,7 @@ class TrainOptimizerConfig(BaseConfig): beta1_warmup: int min_beta1: float Simplified_AdEMAMix: False - cautious_mask: False - grams_moment: False kourkoutas_beta: False - k_warmup_steps: int schedulefree_c: float ns_steps: int MuonWithAuxAdam: False @@ -134,13 +131,11 @@ class TrainOptimizerConfig(BaseConfig): rms_rescaling: True normuon_variant: False beta2_normuon: float - normuon_eps: float low_rank_ortho: False ortho_rank: int accelerated_ns: False cautious_wd: False approx_mars: False - kappa_p: float auto_kappa_p: False compile: False @@ -236,10 +231,7 @@ def default_values(): data.append(("beta1_warmup", None, int, True)) data.append(("min_beta1", None, float, True)) data.append(("Simplified_AdEMAMix", False, bool, False)) - data.append(("cautious_mask", False, bool, False)) - data.append(("grams_moment", False, bool, False)) data.append(("kourkoutas_beta", False, bool, False)) - data.append(("k_warmup_steps", None, int, True)) data.append(("schedulefree_c", None, float, True)) data.append(("ns_steps", None, int, True)) data.append(("MuonWithAuxAdam", False, bool, False)) @@ -252,13 +244,11 @@ def default_values(): data.append(("rms_rescaling", True, bool, True)) data.append(("normuon_variant", False, bool, False)) data.append(("beta2_normuon", None, float, True)) - data.append(("normuon_eps", None, float, True)) data.append(("low_rank_ortho", False, bool, False)) data.append(("ortho_rank", None, int, True)) data.append(("accelerated_ns", False, bool, False)) data.append(("cautious_wd", False, bool, False)) data.append(("approx_mars", False, bool, False)) - data.append(("kappa_p", None, float, True)) data.append(("auto_kappa_p", False, bool, False)) data.append(("compile", False, bool, False)) diff --git a/modules/util/create.py b/modules/util/create.py index 639f86beb..210c0980d 100644 --- a/modules/util/create.py +++ b/modules/util/create.py @@ -672,19 +672,16 @@ def create_optimizer( 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, 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, 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, use_atan2=optimizer_config.use_atan2 if optimizer_config.use_atan2 is not None else False, - cautious_mask=optimizer_config.cautious_mask if optimizer_config.cautious_mask is not None else False, - grams_moment=optimizer_config.grams_moment if optimizer_config.grams_moment 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=optimizer_config.k_warmup_steps if optimizer_config.k_warmup_steps is not None else 0, + 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, ) @@ -702,8 +699,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, use_atan2=optimizer_config.use_atan2 if optimizer_config.use_atan2 is not None else False, - cautious_mask=optimizer_config.cautious_mask if optimizer_config.cautious_mask is not None else False, - grams_moment=optimizer_config.grams_moment if optimizer_config.grams_moment 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, @@ -711,7 +706,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, 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, + 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, ) @@ -736,8 +731,6 @@ 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, - cautious_mask=optimizer_config.cautious_mask if optimizer_config.cautious_mask is not None else False, - grams_moment=optimizer_config.grams_moment if optimizer_config.grams_moment 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, @@ -745,30 +738,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, 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, - ) - - # SIMPLIFIED_AdEMAMix Optimizer - case Optimizer.SIMPLIFIED_AdEMAMix: - from adv_optm import Simplified_AdEMAMix - optimizer = Simplified_AdEMAMix( - params=parameters, - lr=config.learning_rate, - betas=(optimizer_config.beta1 if optimizer_config.beta1 is not None else 0.99, - optimizer_config.beta2 if optimizer_config.beta2 is not None else 0.999), - 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, - alpha_grad=optimizer_config.alpha_grad if optimizer_config.alpha_grad is not None else 100, - 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, - 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, - 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, + 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, ) @@ -804,34 +774,6 @@ def create_optimizer( stochastic_rounding=optimizer_config.stochastic_rounding, cautious_mask=optimizer_config.cautious_mask if optimizer_config.cautious_mask is not None else False, orthogonal_gradient=optimizer_config.orthogonal_gradient if optimizer_config.orthogonal_gradient is not None else False, - 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, - ) - - # LION_PRODIGY_ADV Optimizer - case Optimizer.LION_PRODIGY_ADV: - from adv_optm import Lion_Prodigy_adv - optimizer = Lion_Prodigy_adv( - params=parameters, - lr=config.learning_rate, - 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), - beta3=optimizer_config.beta3 if optimizer_config.beta3 is not None else None, - 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, - 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, - growth_rate=optimizer_config.growth_rate if optimizer_config.growth_rate is not None else float('inf'), - slice_p=optimizer_config.slice_p if optimizer_config.slice_p is not None else 1, - 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, - cautious_mask=optimizer_config.cautious_mask if optimizer_config.cautious_mask is not None else False, - orthogonal_gradient=optimizer_config.orthogonal_gradient if optimizer_config.orthogonal_gradient is not None else False, - 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, ) @@ -876,7 +818,6 @@ def create_optimizer( nesterov=optimizer_config.nesterov if optimizer_config.nesterov is not None else True, 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, - normuon_eps=optimizer_config.normuon_eps if optimizer_config.normuon_eps is not None else 1e-8, 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, diff --git a/modules/util/enum/Optimizer.py b/modules/util/enum/Optimizer.py index 58edf419c..945d73488 100644 --- a/modules/util/enum/Optimizer.py +++ b/modules/util/enum/Optimizer.py @@ -21,7 +21,6 @@ class Optimizer(Enum): AdEMAMix = 'AdEMAMix' AdEMAMix_8BIT = "AdEMAMix_8BIT" - SIMPLIFIED_AdEMAMix = "SIMPLIFIED_AdEMAMix" ADOPT = 'ADOPT' ADOPT_ADV = 'ADOPT_ADV' @@ -59,7 +58,6 @@ class Optimizer(Enum): PRODIGY = 'PRODIGY' PRODIGY_PLUS_SCHEDULE_FREE = 'PRODIGY_PLUS_SCHEDULE_FREE' PRODIGY_ADV = 'PRODIGY_ADV' - LION_PRODIGY_ADV = 'LION_PRODIGY_ADV' # ADAFACTOR ADAFACTOR = 'ADAFACTOR' @@ -90,7 +88,6 @@ def is_adaptive(self): self.PRODIGY, self.PRODIGY_PLUS_SCHEDULE_FREE, self.PRODIGY_ADV, - self.LION_PRODIGY_ADV, ] @property @@ -110,11 +107,9 @@ def supports_fused_back_pass(self): Optimizer.ADAMW, Optimizer.ADAMW_ADV, Optimizer.ADOPT_ADV, - Optimizer.SIMPLIFIED_AdEMAMix, Optimizer.PRODIGY_PLUS_SCHEDULE_FREE, Optimizer.PRODIGY_ADV, Optimizer.LION_ADV, - Optimizer.LION_PRODIGY_ADV, Optimizer.MUON_ADV, Optimizer.ADAMUON_ADV, Optimizer.SIGNSGD_ADV, diff --git a/modules/util/optimizer_util.py b/modules/util/optimizer_util.py index 448154038..bec700598 100644 --- a/modules/util/optimizer_util.py +++ b/modules/util/optimizer_util.py @@ -454,20 +454,16 @@ def init_model_parameters( "eps": 1e-8, "cautious_wd": False, "weight_decay": 0.0, - "use_bias_correction": True, "nnmf_factor": False, "stochastic_rounding": True, "compile": False, "fused_back_pass": False, "use_atan2": False, - "cautious_mask": False, - "grams_moment": False, "orthogonal_gradient": False, "use_AdEMAMix": False, "beta3_ema": 0.9999, "alpha": 5, "kourkoutas_beta": False, - "k_warmup_steps": None, }, Optimizer.ADOPT_ADV: { "beta1": 0.9, @@ -479,9 +475,7 @@ def init_model_parameters( "stochastic_rounding": True, "compile": False, "fused_back_pass": False, - "use_atan2": False, - "cautious_mask": False, - "grams_moment": False, + "use_atan2": True, "orthogonal_gradient": False, "use_AdEMAMix": False, "beta3_ema": 0.9999, @@ -489,7 +483,6 @@ def init_model_parameters( "Simplified_AdEMAMix": False, "alpha_grad": 100.0, "kourkoutas_beta": False, - "k_warmup_steps": None, }, Optimizer.PRODIGY_ADV: { "beta1": 0.9, @@ -509,8 +502,6 @@ def init_model_parameters( "prodigy_steps": 0, "d_limiter": False, "use_atan2": False, - "cautious_mask": False, - "grams_moment": False, "orthogonal_gradient": False, "use_AdEMAMix": False, "beta3_ema": 0.9999, @@ -518,25 +509,6 @@ def init_model_parameters( "Simplified_AdEMAMix": False, "alpha_grad": 100.0, "kourkoutas_beta": False, - "k_warmup_steps": None, - }, - Optimizer.SIMPLIFIED_AdEMAMix: { - "beta1": 0.99, - "beta2": 0.99, - "eps": 1e-8, - "cautious_wd": False, - "weight_decay": 0.0, - "alpha_grad": 100.0, - "beta1_warmup": None, - "min_beta1": 0.9, - "use_bias_correction": True, - "nnmf_factor": False, - "stochastic_rounding": True, - "compile": False, - "fused_back_pass": False, - "orthogonal_gradient": False, - "kourkoutas_beta": False, - "k_warmup_steps": None, }, Optimizer.SIGNSGD_ADV: { "momentum": 0.95, @@ -562,29 +534,6 @@ def init_model_parameters( "fused_back_pass": False, "cautious_mask": False, "orthogonal_gradient": False, - "kappa_p": 1.0, - "auto_kappa_p": True, - }, - Optimizer.LION_PRODIGY_ADV: { - "beta1": 0.9, - "beta2": 0.99, - "beta3": None, - "cautious_wd": False, - "weight_decay": 0.0, - "clip_threshold": None, - "nnmf_factor": False, - "stochastic_rounding": True, - "compile": False, - "fused_back_pass": False, - "d0": 1e-6, - "d_coef": 1.0, - "growth_rate": float('inf'), - "slice_p": 11, - "prodigy_steps": 0, - "d_limiter": True, - "cautious_mask": False, - "orthogonal_gradient": False, - "kappa_p": 1.0, "auto_kappa_p": True, }, Optimizer.MUON_ADV: { @@ -611,7 +560,6 @@ def init_model_parameters( "alpha_grad": 100.0, "normuon_variant": True, "beta2_normuon": 0.95, - "normuon_eps": 1e-8, "orthogonal_gradient": False, "approx_mars": False, "muon_adam_config": {},