diff --git a/modules/module/LoRAModule.py b/modules/module/LoRAModule.py index 7b5e09f59..fceae77f6 100644 --- a/modules/module/LoRAModule.py +++ b/modules/module/LoRAModule.py @@ -578,15 +578,17 @@ class OFTModule(PeftBase): oft_block_size: int block_share: bool oft_scaled: bool + 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_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_cans = oft_cans self.dropout_probability = kwargs.pop('dropout_probability', 0.0) self.oft_R = None self.adjustment_info = None @@ -656,6 +658,7 @@ def initialize_weights(self): use_cayley_neumann=True, num_cayley_neumann_terms=5, dropout_probability=self.dropout_probability, + oft_cans=self.oft_cans, ) nn.init.zeros_(self.oft_R.weight) @@ -673,7 +676,7 @@ def forward(self, x, *args, **kwargs): # 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 + 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 +867,7 @@ def __init__( config.oft_block_size, config.oft_block_share, config.oft_scaled, + 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..4bdbb68f1 100644 --- a/modules/module/oft_utils.py +++ b/modules/module/oft_utils.py @@ -47,6 +47,7 @@ def __init__( oft_scaled=False, use_cayley_neumann=True, num_cayley_neumann_terms=5, + oft_cans=False, dropout_probability=0.0, ): super().__init__() @@ -62,14 +63,19 @@ 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_oft", torch.tensor(True)) # 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) def _pytorch_skew_symmetric(self, vec, block_size): batch_size = vec.shape[0] @@ -88,8 +94,57 @@ def _pytorch_skew_symmetric_inv(self, matrix, block_size): vec = matrix[:, self.rows, self.cols] return vec + def _cans_newton_schulz_iteration( + self, + G: torch.Tensor, + steps: int = 7, + eps: float = 1e-7, + ) -> torch.Tensor: + """ + Chebyshev-Optimized Newton-Schulz iteration with a dynamically computed Chebyshev lower bound. + Optimized for G = I + Q (where Q is skew-symmetric). + """ + 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 + + # Since min_singular_value(I + Q) >= 1, the min_singular_value of normalized X + # is guaranteed to be >= 1 / ||G||_F. + lower_bound = (1.0 / 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 + + # bmm acts as an opaque boundary, forcing Inductor to + # materialize eps_val and severing the exponential AST tree. + # Shape is (B, 1, 1), so ones_like acts as an identity. + eps_val = torch.bmm(eps_val, torch.ones_like(eps_val)) + + lower_bound = 1 - eps_val + upper_bound = 1 + eps_val + + return X.to(original_dtype) + def _cayley_batch( - self, Q: torch.Tensor, block_size: int, use_cayley_neumann: bool = True, num_neumann_terms: int = 5 + 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. @@ -113,13 +168,17 @@ def _cayley_batch( R.add_(Q_power, alpha=2.0) Q_power = torch.bmm(Q_power, Q_skew) R.add_(Q_power) + elif oft_cans: + # Compute G = (I + Q)^2 = I + 2Q + Q^2 + # Squaring the matrix doubles the rotation range and matches Cayley (I + 2Q). + Q_squared = torch.bmm(Q_skew, Q_skew) + G = self.id_mat + 2 * Q_skew + Q_squared + # Empirically, BF16 requires 5 steps to converge to ortho error ~1e-2 (its limit) + # While FP32 takes 7 steps to converge to ortho error ~1e-6 + steps = 5 if G.dtype == torch.bfloat16 else 7 + R = self._cans_newton_schulz_iteration(G=G, steps=steps) 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) @@ -134,7 +193,7 @@ def forward(self, x): effective_weight = self.weight / scaling_factor orth_rotate = self._cayley_batch( - effective_weight, self.block_size, self.use_cayley_neumann, self.num_cayley_neumann_terms + 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..6614fe2f1 100644 --- a/modules/ui/LoraTab.py +++ b/modules/ui/LoraTab.py @@ -134,6 +134,11 @@ 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, "Accelerated Newton-Schulz", + tooltip="Replaces Cayley-Neumann with Chebyshev-Optimized Newton-Schulz (CANS) to improve orthogonalization stability and reduce error without the high computational cost of the exact solver.") + components.switch(master, 4, 4, self.ui_state, "oft_cans") + # 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/util/config/TrainConfig.py b/modules/util/config/TrainConfig.py index f52988502..04a66f054 100644 --- a/modules/util/config/TrainConfig.py +++ b/modules/util/config/TrainConfig.py @@ -522,6 +522,7 @@ class TrainConfig(BaseConfig): oft_block_size: int oft_block_share: bool oft_scaled: bool + oft_cans: bool # lokr lokr_dim: int @@ -1161,6 +1162,7 @@ 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_cans", False, bool, False)) # lokr data.append(("lokr_dim", 16, int, False))