diff --git a/modules/module/LoRAModule.py b/modules/module/LoRAModule.py index 7b5e09f59..f180c6dc7 100644 --- a/modules/module/LoRAModule.py +++ b/modules/module/LoRAModule.py @@ -578,15 +578,19 @@ class OFTModule(PeftBase): oft_block_size: int block_share: bool oft_scaled: bool + oft_clipped_norm: float | None + oft_cans: bool dropout_probability: float adjustment_info: tuple[int, int] | None # for reporting - def __init__(self, prefix: str, orig_module: nn.Module | None, oft_block_size: int, block_share: bool, oft_scaled: bool, **kwargs): + def __init__(self, prefix: str, orig_module: nn.Module | None, oft_block_size: int, block_share: bool, oft_scaled: bool, oft_clipped_norm: float | None, oft_cans: bool, **kwargs): super().__init__(prefix, orig_module) self.oft_block_size = oft_block_size self.rank = 0 self.block_share = block_share self.oft_scaled = oft_scaled + self.oft_clipped_norm = oft_clipped_norm + self.oft_cans = oft_cans self.dropout_probability = kwargs.pop('dropout_probability', 0.0) self.oft_R = None self.adjustment_info = None @@ -656,6 +660,8 @@ def initialize_weights(self): use_cayley_neumann=True, num_cayley_neumann_terms=5, dropout_probability=self.dropout_probability, + oft_clipped_norm=self.oft_clipped_norm, + oft_cans=self.oft_cans, ) nn.init.zeros_(self.oft_R.weight) @@ -672,8 +678,8 @@ def forward(self, x, *args, **kwargs): effective_weight = self.oft_R.weight / scaling_factor # For Conv2d, we must rotate the weights, not the input, to preserve spatial information. - orth_rotate = self.oft_R._cayley_batch( - effective_weight, self.oft_R.block_size, self.oft_R.use_cayley_neumann, self.oft_R.num_cayley_neumann_terms + orth_rotate = self.oft_R._compute_orthogonal_matrix( + effective_weight, self.oft_R.block_size, self.oft_R.use_cayley_neumann, self.oft_R.num_cayley_neumann_terms, self.oft_R.oft_cans ) orth_rotate = self.oft_R.dropout(orth_rotate) @@ -864,6 +870,8 @@ def __init__( config.oft_block_size, config.oft_block_share, config.oft_scaled, + config.oft_clipped_norm, + config.oft_cans, ] self.additional_kwargs = { 'dropout_probability': config.dropout_probability, diff --git a/modules/module/oft_utils.py b/modules/module/oft_utils.py index 97de3199c..2a06029aa 100644 --- a/modules/module/oft_utils.py +++ b/modules/module/oft_utils.py @@ -47,7 +47,9 @@ def __init__( oft_scaled=False, use_cayley_neumann=True, num_cayley_neumann_terms=5, + oft_cans=False, dropout_probability=0.0, + oft_clipped_norm: float | None = 0.95, ): super().__init__() self.r = r @@ -62,14 +64,38 @@ def __init__( # allowing inference tools to automatically detect scaled oft. self.register_buffer("scaled_oft", torch.tensor(True)) self.oft_scaled = oft_scaled - self.use_cayley_neumann = use_cayley_neumann + self.use_cayley_neumann = use_cayley_neumann and not oft_cans self.num_cayley_neumann_terms = num_cayley_neumann_terms + self.oft_cans = oft_cans + if oft_cans: + self.register_buffer("cans_exp", torch.tensor([])) # Create indices for upper triangle (excluding diagonal) rows, cols = torch.triu_indices(block_size, block_size, 1) self.register_buffer("rows", rows, persistent=False) self.register_buffer("cols", cols, persistent=False) self.dropout = MultiplicativeDropoutLayer(p=dropout_probability) - + if not self.use_cayley_neumann: + id_mat = (torch.eye(block_size).unsqueeze(0).expand(r, block_size, block_size)) + self.register_buffer("id_mat", id_mat, persistent=False) + if oft_clipped_norm == -1: + if oft_cans: + # 0.95 * pi (~3.11) avoids the gradient ambiguity/singularity + # at exactly 180 degrees (pi). + self.oft_clipped_norm = 0.95 * math.pi + elif use_cayley_neumann: + # Neumann series diverges if norm >= 1.0, so 0.95 is the safe max. + self.oft_clipped_norm = 0.95 + else: + self.oft_clipped_norm = oft_clipped_norm + if oft_clipped_norm is not None: + self.register_buffer("clipped_oft", torch.tensor(self.oft_clipped_norm)) + # Initialize states for Spectral Normalization via Power Iteration + u = torch.randn(r, block_size) + u = u / u.norm(dim=1, keepdim=True).clamp_min(1e-12) + self.register_buffer("u_state", u, persistent=False) + v = torch.randn(r, block_size) + v = v / v.norm(dim=1, keepdim=True).clamp_min(1e-12) + self.register_buffer("v_state", v, persistent=False) def _pytorch_skew_symmetric(self, vec, block_size): batch_size = vec.shape[0] @@ -88,18 +114,131 @@ def _pytorch_skew_symmetric_inv(self, matrix, block_size): vec = matrix[:, self.rows, self.cols] return vec - def _cayley_batch( - self, Q: torch.Tensor, block_size: int, use_cayley_neumann: bool = True, num_neumann_terms: int = 5 + def _break_inductor_graph(self, x: torch.Tensor) -> torch.Tensor: + """ + Acts as an opaque boundary for TorchInductor. Forces materialization + of the tensor to sever deeply nested AST trees in torch.compile. + """ + return torch.bmm(x, torch.ones_like(x)) + + @torch.no_grad() + def _spectral_norm(self, Q_skew): + u = self.u_state.unsqueeze(-1).to(Q_skew.dtype) + v = self.v_state.unsqueeze(-1).to(Q_skew.dtype) + # Update v (Right Singular Vector) + v_raw = torch.bmm(Q_skew.mT, u) + v_norm = torch.linalg.vector_norm(v_raw, dim=1, keepdim=True) + candidate_v = v_raw / v_norm.clamp_min(1e-8) + next_v = torch.where(v_norm >= 1e-6, candidate_v, v) + # Update u (Left Singular Vector) + u_raw = torch.bmm(Q_skew, next_v) + u_norm = torch.linalg.vector_norm(u_raw, dim=1, keepdim=True) + candidate_u = u_raw / u_norm.clamp_min(1e-8) + next_u = torch.where(u_norm >= 1e-6, candidate_u, u) + if self.training: + self.v_state.copy_(next_v.squeeze(-1)) + self.u_state.copy_(next_u.squeeze(-1)) + return next_v, next_u + + def _cans_newton_schulz_iteration( + self, + G: torch.Tensor, + steps: int = 3, + eps: float = 1e-7, + ) -> torch.Tensor: + """ + Chebyshev-Optimized Newton-Schulz iteration with a dynamically computed Chebyshev lower bound. + """ + original_dtype = G.dtype + X = G + + # Max row sum is guaranteed to be >= the maximum singular value of X. + g_norm = X.abs().sum(dim=-1, keepdim=True).amax(dim=-2, keepdim=True).clamp_min(eps).detach() + X = X / g_norm + + # The 4th-order Taylor expansion of exp(Q) has a minimum singular value of exactly 0.5 + # (occurring at ||Q|| = sqrt(6) ~ 2.449). Therefore, the min_singular_value of normalized X + # is guaranteed to be >= 0.5 / g_norm. + lower_bound = 0.5 / g_norm + upper_bound = 1 + + for _ in range(steps): + lb, ub = lower_bound, upper_bound + lb_ub = lb * ub + e_sq = (lb**2 + lb_ub + ub**2) / 3.0 + K = 2.0 * e_sq**1.5 + L = lb_ub * (lb + ub) + denom = K + L + alpha = 6.0 / denom + c1 = alpha * e_sq + c3 = -alpha / 3.0 + + A = torch.bmm(X, X.mT) + X = c1 * X + c3 * torch.bmm(A, X) + + # Dynamically update bounds for the next step + eps_val = (K - L) / denom + eps_val = self._break_inductor_graph(eps_val) + + lower_bound = 1 - eps_val + upper_bound = 1 + eps_val + + return X.to(original_dtype) + + def _matrix_exp_cans(self, Q_skew: torch.Tensor) -> torch.Tensor: + """ + Approximates the Matrix Exponential using a 4th-order Taylor expansion, + Scaling & Squaring, and Chebyshev-Optimized Newton-Schulz (CANS). + """ + num_squarings = 2 + id_mat = self.id_mat + + # Scaling step + Q_scaled = Q_skew / (2 ** num_squarings) + Q_squared = torch.bmm(Q_scaled, Q_scaled) + + # 4th-order Taylor expansion: exp(Q) ≈ I + Q + Q^2/2 + Q^3/6 + Q^4/24 + # Factored to minimize matrix multiplications: (I + Q) + Q^2 * (0.5*I + 1/6*Q + 1/24*Q^2) + taylor_higher_order = 0.5 * id_mat + (1.0 / 6.0) * Q_scaled + (1.0 / 24.0) * Q_squared + G = torch.baddbmm(id_mat + Q_scaled, Q_squared, taylor_higher_order) + + # Orthogonalize the approximation (CANS) + # Empirically, CANS requires 3 steps to converge + R = self._cans_newton_schulz_iteration(G=G, steps=3) + + # Squaring step to recover full rotation + for _ in range(num_squarings): + R = torch.bmm(R, R) + + # Final standard Newton-Schulz step to correct drift caused by squaring in lower precision + # R_new = R + 0.5 * R * (I - R^T R) + residual = torch.baddbmm(id_mat, R.mT, R, beta=1.0, alpha=-1.0) + R = torch.baddbmm(R, R, residual, beta=1.0, alpha=0.5) + + return R + + def _compute_orthogonal_matrix( + self, Q: torch.Tensor, block_size: int, use_cayley_neumann: bool = True, num_neumann_terms: int = 5, oft_cans: bool = False, ) -> torch.Tensor: """ - Perform the Cayley parametrization on a batch of skew-symmetric matrices. + Converts learned weights into a batch of orthogonal matrices. """ b, _ = Q.shape previous_dtype = Q.dtype Q_skew = self._pytorch_skew_symmetric(Q, block_size) - if use_cayley_neumann: + # Spectral Normalization / Clipping + if self.oft_clipped_norm is not None: + v_norm, u_norm = self._spectral_norm(Q_skew) + u_raw_grad = torch.bmm(Q_skew, v_norm) + sigma = torch.sum(u_norm * u_raw_grad, dim=1, keepdim=True) + max_norm = self.oft_clipped_norm + Q_skew = Q_skew * (max_norm / torch.clamp(sigma, min=max_norm)) + + if oft_cans: + R = self._matrix_exp_cans(Q_skew) + elif use_cayley_neumann: R = torch.eye(block_size, device=Q.device, dtype=Q.dtype).repeat(b, 1, 1) if num_neumann_terms > 1: R.add_(Q_skew, alpha=2.0) @@ -114,12 +253,7 @@ def _cayley_batch( Q_power = torch.bmm(Q_power, Q_skew) R.add_(Q_power) else: - id_mat = ( - torch.eye(Q_skew.shape[-1], device=Q_skew.device) - .unsqueeze(0) - .expand(b, Q_skew.shape[-1], Q_skew.shape[-1]) - ) - R = torch.linalg.solve(id_mat + Q_skew, id_mat - Q_skew, left=False) + R = torch.linalg.solve(self.id_mat + Q_skew, self.id_mat - Q_skew, left=False) return R.to(previous_dtype) @@ -130,11 +264,18 @@ def forward(self, x): orig_shape = x.shape - scaling_factor = 2 * math.sqrt(self.block_size - 1) if self.oft_scaled else 1 - effective_weight = self.weight / scaling_factor + if self.oft_scaled: + # Cayley has a 2x gradient multiplier (I + 2Q). Exp has a 1x multiplier (I + Q). + # We drop the 2 for Exp/CANS to maintain consistent effective learning rates. + is_cayley = self.use_cayley_neumann and not self.oft_cans + multiplier = 2.0 if is_cayley else 1.0 + scaling_factor = multiplier * math.sqrt(self.block_size - 1) + effective_weight = self.weight / scaling_factor + else: + effective_weight = self.weight - orth_rotate = self._cayley_batch( - effective_weight, self.block_size, self.use_cayley_neumann, self.num_cayley_neumann_terms + orth_rotate = self._compute_orthogonal_matrix( + effective_weight, self.block_size, self.use_cayley_neumann, self.num_cayley_neumann_terms, self.oft_cans ) orth_rotate = self.dropout(orth_rotate) diff --git a/modules/ui/LoraTab.py b/modules/ui/LoraTab.py index 6e9975476..b614af813 100644 --- a/modules/ui/LoraTab.py +++ b/modules/ui/LoraTab.py @@ -134,6 +134,16 @@ def setup_lora(self, peft_type: PeftType): tooltip="Applies a scaling factor to the learned weights. This ensures that the effective learning rate remains consistent across different block sizes. Without this, different block sizes require significantly different learning rates.") components.switch(master, 2, 4, self.ui_state, "oft_scaled") + # CANS OFT + components.label(master, 4, 3, "Matrix Exponential CANS", + tooltip="Replaces Cayley-Neumann with Matrix Exponential with Chebyshev-Optimized Newton-Schulz (CANS) to improve orthogonalization stability.") + components.switch(master, 4, 4, self.ui_state, "oft_cans") + + # Clip OFT max norm + components.label(master, 5, 0, "Spectral Norm Clipping", + tooltip="Strictly clips the spectral norm of the OFT matrix to guarantee convergence of the Cayley parametrization (requires norm <= 1.0). Smaller values constrain the learned rotation to stay near the identity matrix, limiting adaptation. Default: 1.0 (e.g. 0.8 = 80% of maximum expressiveness). Leave empty to disable.") + components.entry(master, 5, 1, self.ui_state, "oft_clipped_norm") + # Dropout Percentage components.label(master, 2, 0, "Dropout Probability", tooltip="Dropout probability. This percentage of the rotated adapter nodes that will be randomly restored to the base model initial statue. Helps with overfitting. 0 disables, 1 maximum.") diff --git a/modules/ui/MuonAdamWindow.py b/modules/ui/MuonAdamWindow.py index 5879ab432..e895e7290 100644 --- a/modules/ui/MuonAdamWindow.py +++ b/modules/ui/MuonAdamWindow.py @@ -73,14 +73,13 @@ 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'}, + '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'}, - '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'}, + '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. '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 @@ -103,5 +102,9 @@ def create_adam_params_ui(self, master): if param_type != 'bool': 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 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) diff --git a/modules/ui/OptimizerParamsWindow.py b/modules/ui/OptimizerParamsWindow.py index 1d8cbe2b5..9a08a1cb2 100644 --- a/modules/ui/OptimizerParamsWindow.py +++ b/modules/ui/OptimizerParamsWindow.py @@ -167,14 +167,10 @@ 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'}, - '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'}, - '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'}, @@ -184,7 +180,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'}, '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'}, @@ -192,8 +188,20 @@ 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 based on layer dimensions and training method. 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'}, + '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'}, + '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'}, + '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'}, } # @formatter:on @@ -230,6 +238,15 @@ 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, ["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", "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 f52988502..e0ff55840 100644 --- a/modules/util/config/TrainConfig.py +++ b/modules/util/config/TrainConfig.py @@ -110,14 +110,10 @@ class TrainOptimizerConfig(BaseConfig): use_schedulefree: True use_orthograd: False nnmf_factor: False - orthogonal_gradient: False + orthogonal_gradient: str 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 @@ -136,11 +132,45 @@ class TrainOptimizerConfig(BaseConfig): accelerated_ns: False cautious_wd: False approx_mars: False - auto_kappa_p: False compile: False + spectral_normalization: False + stochastic_sign: False + centered_wd: float + centered_wd_mode: str + factored_2nd: False + fisher_wd: False + state_precision: str + orthogonal_sinkhorn: False + sinkhorn_iterations: int + normed_momentum: False + nesterov_coef: float + snr_cond: False + 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(): @@ -223,14 +253,10 @@ 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(("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)) @@ -249,8 +275,20 @@ 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)) + data.append(("centered_wd", 0.0, float, 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", "auto", str, False)) + 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)) + data.append(("snr_cond", False, bool, False)) + data.append(("geometric_wd", False, bool, False)) return TrainOptimizerConfig(data) @@ -522,6 +560,8 @@ class TrainConfig(BaseConfig): oft_block_size: int oft_block_share: bool oft_scaled: bool + oft_clipped_norm: float | None + oft_cans: bool # lokr lokr_dim: int @@ -800,6 +840,7 @@ def replace_dtype(part: str): return migrated_data + def weight_dtypes(self) -> ModelWeightDtypes: return ModelWeightDtypes( self.train_dtype, @@ -1161,6 +1202,8 @@ def default_values() -> 'TrainConfig': data.append(("oft_block_size", 32, int, False)) data.append(("oft_block_share", False, bool, False)) data.append(("oft_scaled", False, bool, False)) + data.append(("oft_clipped_norm", 0.95, float, True)) + data.append(("oft_cans", False, bool, False)) # lokr data.append(("lokr_dim", 16, int, False)) diff --git a/modules/util/create.py b/modules/util/create.py index 7c0194da8..bc5640c4a 100644 --- a/modules/util/create.py +++ b/modules/util/create.py @@ -34,6 +34,7 @@ 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.tag_util import tag_peft_parameters from modules.util.TrainProgress import TrainProgress from modules.zluda import ZLUDA @@ -125,7 +126,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 @@ -140,6 +141,10 @@ 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, config) + match config.optimizer.optimizer: # SGD Optimizer @@ -662,6 +667,31 @@ def create_optimizer( quant_block_size=optimizer_config.quant_block_size if optimizer_config.quant_block_size is not None else 2048 ) + # SINKSGD_ADV Optimizer + case Optimizer.SINKSGD_ADV: + from adv_optm import SinkSGD_adv + 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.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, + 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 '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, + 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", + 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, + ) + # ADAMW_ADV Optimizer case Optimizer.ADAMW_ADV: from adv_optm import AdamW_adv @@ -670,19 +700,23 @@ 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, - 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, - 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, + 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, + 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", + 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 @@ -693,21 +727,23 @@ 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, - 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, - 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, - 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, + 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, + 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", + 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 @@ -719,9 +755,10 @@ 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, - 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, @@ -731,15 +768,15 @@ 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, - 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, + 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, + 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 @@ -748,15 +785,22 @@ 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, - 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, - 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", + 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 @@ -768,13 +812,15 @@ 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, - 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, + 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, + 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", ) # MUON_ADV Optimizer @@ -783,7 +829,8 @@ 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) + # Prepare Adam-specific keyword arguments from the config adam_kwargs = {} @@ -811,20 +858,22 @@ 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, + 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, 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, - 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 ) @@ -834,7 +883,8 @@ 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) + # Prepare Adam-specific keyword arguments from the config adam_kwargs = {} @@ -865,20 +915,22 @@ 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, + 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, - 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, 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, + 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 ) @@ -887,7 +939,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/enum/Optimizer.py b/modules/util/enum/Optimizer.py index 945d73488..6326a1861 100644 --- a/modules/util/enum/Optimizer.py +++ b/modules/util/enum/Optimizer.py @@ -41,6 +41,7 @@ class Optimizer(Enum): # 32 bit is torch and not bnb SGD = 'SGD' SGD_8BIT = 'SGD_8BIT' + SINKSGD_ADV = 'SINKSGD_ADV' SIGNSGD_ADV = 'SIGNSGD_ADV' # Schedule-free optimizers @@ -98,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, @@ -113,6 +127,7 @@ def supports_fused_back_pass(self): Optimizer.MUON_ADV, Optimizer.ADAMUON_ADV, Optimizer.SIGNSGD_ADV, + Optimizer.SINKSGD_ADV, ] # Small helper for adjusting learning rates to adaptive optimizers. diff --git a/modules/util/optimizer/muon_util.py b/modules/util/optimizer/muon_util.py index b2630c16b..186019937 100644 --- a/modules/util/optimizer/muon_util.py +++ b/modules/util/optimizer/muon_util.py @@ -105,8 +105,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 +116,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/tag_util.py b/modules/util/optimizer/tag_util.py new file mode 100644 index 000000000..9646c62b8 --- /dev/null +++ b/modules/util/optimizer/tag_util.py @@ -0,0 +1,43 @@ +import math + +from modules.module.LoRAModule import LoRAModuleWrapper +from modules.util.config.TrainConfig import TrainConfig + +import torch + + +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. + 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", "dora_log_multiplier")): + # Vector in shape of >= 2D tensor + p._is_dora_scale = True + elif name.endswith("oft_R.weight"): + # 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) + else: + scaling_factor = 2 if config.oft_cans else 1 + p._oft_scale_factor = scaling_factor + + 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(): + apply_tags(param_name, p) + elif isinstance(module, torch.nn.Module): + for param_name, p in module.named_parameters(): + apply_tags(param_name, p) diff --git a/modules/util/optimizer_util.py b/modules/util/optimizer_util.py index 60fc98815..3483b1f50 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: @@ -448,50 +442,80 @@ def init_model_parameters( "min_8bit_size": 16384, "quant_block_size": 2048 }, + Optimizer.SINKSGD_ADV: { + "normed_momentum": False, + "momentum": 0.95, + "nesterov": True, + "nesterov_coef": None, + "sinkhorn_iterations": 5, + "cautious_wd": False, + "geometric_wd": False, + "weight_decay": 0.0, + "snr_cond": False, + "stochastic_rounding": True, + "compile": False, + "fused_back_pass": False, + "orthogonal_sinkhorn": False, + "orthogonal_gradient": 'disabled', + "spectral_normalization": False, + "centered_wd": 0.0, + "centered_wd_mode": "full", + "state_precision": "auto", + }, Optimizer.ADAMW_ADV: { "beta1": 0.9, + "nesterov": False, + "nesterov_coef": None, "beta2": 0.99, "eps": 1e-8, + "fisher_wd": False, "cautious_wd": False, "weight_decay": 0.0, - "nnmf_factor": False, + "factored_2nd": False, "stochastic_rounding": True, "compile": False, "fused_back_pass": False, "use_atan2": False, - "orthogonal_gradient": False, - "use_AdEMAMix": False, - "beta3_ema": 0.9999, - "alpha": 5, + "orthogonal_gradient": 'disabled', "kourkoutas_beta": False, + "spectral_normalization": False, + "centered_wd": 0.0, + "centered_wd_mode": "full", + "state_precision": "auto", }, Optimizer.ADOPT_ADV: { "beta1": 0.9, + "nesterov": False, + "nesterov_coef": None, "beta2": 0.9999, "eps": 1e-6, + "fisher_wd": False, "cautious_wd": False, "weight_decay": 0.0, - "nnmf_factor": False, + "factored_2nd": False, "stochastic_rounding": True, "compile": False, "fused_back_pass": False, "use_atan2": True, - "orthogonal_gradient": False, - "use_AdEMAMix": False, - "beta3_ema": 0.9999, - "alpha": 5, - "Simplified_AdEMAMix": False, - "alpha_grad": 100.0, + "orthogonal_gradient": 'disabled', "kourkoutas_beta": False, + "spectral_normalization": False, + "centered_wd": 0.0, + "centered_wd_mode": "full", + "state_precision": "auto", }, Optimizer.PRODIGY_ADV: { "beta1": 0.9, + "nesterov": False, + "nesterov_coef": None, "beta2": 0.99, "beta3": None, "eps": 1e-8, + "fisher_wd": False, "cautious_wd": False, "weight_decay": 0.0, "nnmf_factor": False, + "factored_2nd": False, "stochastic_rounding": True, "compile": False, "fused_back_pass": False, @@ -502,41 +526,49 @@ def init_model_parameters( "prodigy_steps": 0, "d_limiter": False, "use_atan2": False, - "orthogonal_gradient": False, - "use_AdEMAMix": False, - "beta3_ema": 0.9999, - "alpha": 5, - "Simplified_AdEMAMix": False, - "alpha_grad": 100.0, + "orthogonal_gradient": 'disabled', "kourkoutas_beta": False, + "centered_wd": 0.0, + "centered_wd_mode": "full", }, Optimizer.SIGNSGD_ADV: { + "stochastic_sign": False, + "normed_momentum": False, "momentum": 0.95, + "nesterov": False, + "nesterov_coef": None, + "geometric_wd": False, "cautious_wd": False, "weight_decay": 0.0, - "nnmf_factor": False, + "snr_cond": False, "stochastic_rounding": True, "compile": False, "fused_back_pass": False, - "orthogonal_gradient": False, - "Simplified_AdEMAMix": False, - "alpha_grad": 100.0, + "orthogonal_gradient": 'disabled', + "spectral_normalization": False, + "centered_wd": 0.0, + "centered_wd_mode": "full", + "state_precision": "auto", }, Optimizer.LION_ADV: { + "stochastic_sign": False, "beta1": 0.9, "beta2": 0.99, "cautious_wd": False, "weight_decay": 0.0, - "clip_threshold": None, "nnmf_factor": False, "stochastic_rounding": True, "compile": False, "fused_back_pass": False, - "orthogonal_gradient": False, - "auto_kappa_p": True, + "orthogonal_gradient": 'disabled', + "spectral_normalization": False, + "centered_wd": 0.0, + "centered_wd_mode": "full", }, Optimizer.MUON_ADV: { "beta1": 0.9, + "nesterov": True, + "nesterov_coef": None, "cautious_wd": False, "weight_decay": 0.0, "accelerated_ns": False, @@ -544,7 +576,7 @@ def init_model_parameters( "low_rank_ortho": False, "ortho_rank": 128, "rms_rescaling": True, - "nnmf_factor": False, + "spectral_normalization": False, "stochastic_rounding": True, "compile": False, "fused_back_pass": False, @@ -554,17 +586,19 @@ def init_model_parameters( "muon_adam_lr": 1e-6, "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, + "orthogonal_gradient": 'disabled', "approx_mars": False, + "centered_wd": 0.0, + "centered_wd_mode": "full", + "state_precision": "auto", "muon_adam_config": {}, }, Optimizer.ADAMUON_ADV: { "beta1": 0.95, + "nesterov": True, + "nesterov_coef": None, "beta2": 0.95, "eps": 1e-8, "cautious_wd": False, @@ -574,7 +608,7 @@ def init_model_parameters( "low_rank_ortho": False, "ortho_rank": 128, "rms_rescaling": True, - "nnmf_factor": False, + "spectral_normalization": False, "stochastic_rounding": True, "compile": False, "fused_back_pass": False, @@ -584,13 +618,13 @@ def init_model_parameters( "muon_adam_lr": 1e-6, "muon_te1_adam_lr": None, "muon_te2_adam_lr": None, - "nesterov": False, "use_atan2": False, - "Simplified_AdEMAMix": False, - "alpha_grad": 100.0, "normuon_variant": True, - "orthogonal_gradient": False, + "orthogonal_gradient": 'disabled', "approx_mars": False, + "centered_wd": 0.0, + "centered_wd_mode": "full", + "state_precision": "auto", "muon_adam_config": {}, }, Optimizer.ADABELIEF: { diff --git a/requirements-global.txt b/requirements-global.txt index 91014c6ff..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.2.3 # advanced optimizers +adv_optm==2.5.10 # advanced optimizers -e git+https://github.com/KellerJordan/Muon.git@f90a42b#egg=muon-optimizer # Profiling