keras-nightly 3.12.0.dev2025100503__py3-none-any.whl → 3.14.0.dev2026011604__py3-none-any.whl

keras/src/optimizers/muon.py

@@ -20,7 +20,7 @@ class Muon(optimizer.Optimizer):
     The Muon optimizer can use both the Muon update step or the
     AdamW update step based on the following:
 
-    - For any variable that isn't 2D, 3D or 4D, the AdamW step
+    - For any variable that isn't 2D, the AdamW step
         will be used. This is not configurable.
     - If the argument `exclude_embeddings` (defaults to `True`) is set
     to `True`, the AdamW step will be used.
@@ -46,10 +46,12 @@ class Muon(optimizer.Optimizer):
             that takes no arguments and returns the actual value to use.
             The exponential decay rate for the 1st moment estimates. Defaults to
             `0.9`.
-        adam_beta_2: A float value or a constant float tensor, ora callable
+        adam_beta_2: A float value or a constant float tensor, or a callable
             that takes no arguments and returns the actual value to use.
             The exponential decay rate for the 2nd moment estimates. Defaults to
             `0.999`.
+        adam_weight_decay: Float. If set, weight decay is applied when using
+            the Adam optimizer.
         epsilon: A small constant for numerical stability. This is
             "epsilon hat" in the Kingma and Ba paper
             (in the formula just before Section 2.1),
@@ -67,11 +69,15 @@ class Muon(optimizer.Optimizer):
             It is recommended to use the default value
         adam_lr_ratio: Float, the ratio of the learning rate when
                 using Adam to the main learning rate.
-                it is recommended to set it to 0.1
+                It is recommended to set it to 1
         momentum: Float, momentum used by internal SGD.
         ns_steps: Integer, number of Newton-Schulz iterations to run.
         nesterov: Boolean, whether to use Nesterov-style momentum
         {{base_optimizer_keyword_args}}
+        rms_rate: Float. A parameter from https://arxiv.org/abs/2502.16982
+            that can enhance the stability of Muon, allowing it to use the
+            same learning rate and weight decay as Adam. Defaults to `0.2`.
+            Set to `None` to disable this feature.
     """
 
     def __init__(
@@ -79,8 +85,9 @@ class Muon(optimizer.Optimizer):
         learning_rate=0.001,
         adam_beta_1=0.9,
         adam_beta_2=0.999,
+        adam_weight_decay=0.004,
         epsilon=1e-7,
-        weight_decay=0.1,
+        weight_decay=0.004,
         clipnorm=None,
         clipvalue=None,
         global_clipnorm=None,
@@ -95,10 +102,11 @@ class Muon(optimizer.Optimizer):
         muon_a=3.4445,
         muon_b=-4.7750,
         muon_c=2.0315,
-        adam_lr_ratio=0.1,
+        adam_lr_ratio=1,
         momentum=0.95,
-        ns_steps=6,
+        ns_steps=5,
         nesterov=True,
+        rms_rate=0.2,
         **kwargs,
     ):
         super().__init__(
@@ -127,12 +135,13 @@ class Muon(optimizer.Optimizer):
         self.nesterov = nesterov
         self.exclude_embeddings = exclude_embeddings
         self.exclude_layers = exclude_layers or []
+        self.adam_weight_decay = adam_weight_decay
+        self.rms_rate = rms_rate
 
     def _should_use_adamw(self, variable):
-        # To use it with 4D convolutional filters,
         # it works well to just flatten their last 3 dimensions.
         # any {0,1}-D parameters should all be optimized by adam
-        if not 1 < len(variable.shape) < 4:
+        if len(variable.shape) != 2:
             return True
         if self.exclude_embeddings and "embedding" in variable.path.lower():
             return True
@@ -153,52 +162,50 @@ class Muon(optimizer.Optimizer):
         if self.built:
             return
         super().build(var_list)
-        self.adam_momentums = {}
-        self.adam_velocities = {}
-
-        self.muon_momentums = {}
-        self.muon_velocities = {}
+        # Momentums are for both Muon and Adam
+        self.momentums = [None] * len(var_list)
+        # Velocities are just for Adam
+        self.adam_velocities = [None] * len(var_list)
 
         for var in var_list:
             if not self._overwrite_variable_with_gradient(var):
-                self.adam_momentums[var.path] = (
+                self.momentums[self._get_variable_index(var)] = (
                     self.add_variable_from_reference(
                         reference_variable=var, name="momentum"
                     )
                 )
                 if self._should_use_adamw(var):
-                    self.adam_velocities[var.path] = (
+                    self.adam_velocities[self._get_variable_index(var)] = (
                         self.add_variable_from_reference(
                             reference_variable=var, name="velocity"
                         )
                     )
 
     def update_step(self, gradient, variable, learning_rate):
-        if self._should_use_adamw(variable):
+        variable_index = self._get_variable_index(variable)
+        m = self.momentums[variable_index]
+        v = self.adam_velocities[variable_index]
+
+        # The presence of the velocity tells us that this variable is for Adam
+        if v is not None:
             # It should be noted that lr is one-tenth when using adamw.
             self._adamw_update_step(
-                gradient, variable, learning_rate * self.adam_lr_ratio
+                gradient, variable, learning_rate * self.adam_lr_ratio, m, v
             )
         else:
-            self._muon_update_step(gradient, variable, learning_rate)
+            self._muon_update_step(gradient, variable, learning_rate, m)
 
-    def _muon_update_step(self, gradient, variable, lr):
-        m = self.adam_momentums[variable.path]
+    def _muon_update_step(self, gradient, variable, lr, m):
         self.assign_add(m, ops.add(gradient, m * (self.momentum - 1)))
-        shape = variable.shape
         if self.nesterov:
             g = ops.add(gradient, self.momentum * m)
         else:
             g = m
+        update = self.zeropower_via_newtonschulz5(g, self.ns_steps)
 
-        self.assign_sub(
-            variable,
-            lr
-            * self.zeropower_via_newtonschulz5(g, self.ns_steps)
-            * max(1, shape[0] / shape[1]) ** 0.5,
-        )
+        self.assign_sub(variable, self.lr_adjust(lr * update))
 
-    def _adamw_update_step(self, gradient, variable, learning_rate):
+    def _adamw_update_step(self, gradient, variable, learning_rate, m, v):
         """Update step given gradient and the associated model variable."""
         lr = ops.cast(learning_rate, variable.dtype)
         gradient = ops.cast(gradient, variable.dtype)
@@ -210,9 +217,6 @@ class Muon(optimizer.Optimizer):
             ops.cast(self.adam_beta_2, variable.dtype), local_step
         )
 
-        m = self.adam_momentums[variable.path]
-        v = self.adam_velocities[variable.path]
-
         alpha = lr * ops.sqrt(1 - adam_beta_2_power) / (1 - adam_beta_1_power)
 
         self.assign_add(
@@ -239,6 +243,20 @@ class Muon(optimizer.Optimizer):
         X = ops.transpose(X, temp_order)
         return X
 
+    def lr_adjust(self, x):
+        """Adjusts learning rate based on the Moonlight implementation.
+        This method enhances the stability of Muon, allowing it to use the same
+        learning rate and weight decay as Adam. For details, see
+        https://arxiv.org/abs/2502.16982.
+        For a 2D matrix, the update is scaled by `sqrt(max(n, m)) * rms_rate`,
+        where `n` and `m` are the dimensions of the matrix.
+        """
+        if self.rms_rate is None:
+            return x
+        # moonlight version
+        # https://github.com/MoonshotAI/Moonlight/blob/master/examples/toy_train.py
+        return x * ops.sqrt(ops.maximum(x.shape[0], x.shape[1])) * self.rms_rate
+
     def zeropower_via_newtonschulz5(self, x, steps: int):
         """We apply the Newton-Schulz iteration to compute matrix G.
 
@@ -268,6 +286,20 @@ class Muon(optimizer.Optimizer):
             x = self.transpose_last_axis(x)
         return x
 
+    def _apply_weight_decay(self, variables):
+        for variable in variables:
+            if not self._use_weight_decay(variable):
+                continue
+            if self._should_use_adamw(variable):
+                weight_decay_value = self.adam_weight_decay
+            else:
+                weight_decay_value = self.weight_decay
+            if weight_decay_value is None:
+                continue
+            wd = ops.cast(weight_decay_value, variable.dtype)
+            lr = ops.cast(self.learning_rate, variable.dtype)
+            variable.assign(variable - variable * wd * lr)
+
     def get_config(self):
         config = super().get_config()
         config.update(
@@ -284,6 +316,8 @@ class Muon(optimizer.Optimizer):
                 "ns_steps": self.ns_steps,
                 "nesterov": self.nesterov,
                 "exclude_embeddings": self.exclude_embeddings,
+                "adam_weight_decay": self.adam_weight_decay,
+                "rms_rate": self.rms_rate,
             }
         )
         return config

keras/src/optimizers/schedules/learning_rate_schedule.py

@@ -584,9 +584,10 @@ class CosineDecay(LearningRateSchedule):
     schedule applies a linear increase per optimizer step to our learning rate
     from `initial_learning_rate` to `warmup_target` for a duration of
     `warmup_steps`. Afterwards, it applies a cosine decay function taking our
-    learning rate from `warmup_target` to `alpha` for a duration of
-    `decay_steps`. If `warmup_target` is None we skip warmup and our decay
-    will take our learning rate from `initial_learning_rate` to `alpha`.
+    learning rate from `warmup_target` to `warmup_target * alpha` for a
+    duration of `decay_steps`. If `warmup_target` is None we skip warmup and
+    our decay will take our learning rate from `initial_learning_rate` to
+    `initial_learning_rate * alpha`.
     It requires a `step` value to  compute the learning rate. You can
     just pass a backend variable that you increment at each training step.
 

keras/src/quantizers/__init__.py

@@ -1,6 +1,11 @@
 import inspect
 
 from keras.src.api_export import keras_export
+from keras.src.quantizers.awq_config import AWQConfig
+from keras.src.quantizers.quantization_config import Float8QuantizationConfig
+from keras.src.quantizers.quantization_config import Int4QuantizationConfig
+from keras.src.quantizers.quantization_config import Int8QuantizationConfig
+from keras.src.quantizers.quantization_config import QuantizationConfig
 from keras.src.quantizers.quantizers import AbsMaxQuantizer
 from keras.src.quantizers.quantizers import Quantizer
 from keras.src.quantizers.quantizers import abs_max_quantize
@@ -13,7 +18,15 @@ from keras.src.quantizers.quantizers import unpack_int4
 from keras.src.saving import serialization_lib
 from keras.src.utils.naming import to_snake_case
 
-ALL_OBJECTS = {Quantizer, AbsMaxQuantizer}
+ALL_OBJECTS = {
+    Quantizer,
+    AbsMaxQuantizer,
+    QuantizationConfig,
+    Int8QuantizationConfig,
+    Int4QuantizationConfig,
+    Float8QuantizationConfig,
+    AWQConfig,
+}
 ALL_OBJECTS_DICT = {cls.__name__: cls for cls in ALL_OBJECTS}
 ALL_OBJECTS_DICT.update(
     {to_snake_case(cls.__name__): cls for cls in ALL_OBJECTS}

keras/src/quantizers/awq.py

@@ -0,0 +1,361 @@
+"""AWQ (Activation-aware Weight Quantization) algorithm implementation.
+
+AWQ protects salient weights by finding optimal per-channel scales based on
+activation magnitudes, then applies those scales before quantization.
+
+Reference: https://arxiv.org/abs/2306.00978
+"""
+
+import types
+
+from keras.src import ops
+from keras.src.layers import Dense
+from keras.src.layers import EinsumDense
+from keras.src.quantizers.quantizers import compute_quantization_parameters
+from keras.src.quantizers.quantizers import dequantize_with_sz_map
+from keras.src.quantizers.quantizers import dequantize_with_zero_point
+from keras.src.quantizers.quantizers import quantize_with_sz_map
+from keras.src.quantizers.quantizers import quantize_with_zero_point
+
+
+def awq_search_optimal_scales(
+    weights,
+    activation_magnitudes,
+    *,
+    num_grid_points=20,
+    group_size=-1,
+):
+    """Search for optimal AWQ scales using grid search.
+
+    The AWQ algorithm finds scaling factors that protect salient weights.
+    For each channel, we search for an optimal ratio in [0, 1] that minimizes
+    the activation-weighted quantization error.
+
+    The key insight: we MULTIPLY weights by scales before quantization to
+    expand salient weights. This ensures quantization noise is small relative
+    to the expanded weight magnitude. During inference, we divide by scales
+    to restore the original magnitude.
+
+    Scale formula: scales = x_max.pow(ratio).clamp(min=1e-4)
+    Loss function: Activation-weighted MSE (approximates output error)
+
+    Args:
+        weights: Weight tensor [out_features, in_features] (transposed kernel).
+        activation_magnitudes: Per-channel activation magnitudes [in_features].
+        num_grid_points: Number of grid search points. Defaults to 20.
+        group_size: Group size for quantization (-1 for per-channel).
+
+    Returns:
+        best_scales: Optimal per-channel scales [in_features].
+    """
+    in_features = ops.shape(weights)[1]
+
+    # Compute per-channel activation magnitudes (x_max)
+    # activations should already be per-channel max magnitudes
+    x_max = ops.cast(activation_magnitudes, "float32")
+    # Avoid zero or very small values
+    x_max = ops.where(ops.less(x_max, 1e-8), ops.ones_like(x_max), x_max)
+
+    best_loss = None
+    best_scales = ops.ones((in_features,), dtype="float32")
+
+    # Grid search over ratio values from 0 to 1
+    for i in range(num_grid_points + 1):
+        ratio = i / num_grid_points
+
+        # Compute scales: x_max^ratio (clipped to avoid numerical issues)
+        if ratio == 0:
+            scales = ops.ones_like(x_max)
+        else:
+            scales = ops.power(x_max, ratio)
+        scales = ops.maximum(scales, 1e-4)
+
+        # Normalize scales to avoid extreme values
+        scale_mean = ops.sqrt(ops.multiply(ops.max(scales), ops.min(scales)))
+        scale_mean = ops.maximum(scale_mean, 1e-8)
+        scales = ops.divide(scales, scale_mean)
+
+        # Apply scales to weights by MULTIPLYING (expand salient weights)
+        # weights_scaled: [out_features, in_features]
+        weights_scaled = ops.multiply(weights, scales)
+
+        if group_size == -1:
+            # Per-channel quantization (no grouping)
+            scale_q, zero_q, maxq = compute_quantization_parameters(
+                weights_scaled,
+                bits=4,
+                symmetric=False,
+                per_channel=True,
+                group_size=-1,
+                compute_dtype="float32",
+            )
+
+            # Quantize and dequantize
+            quantized = quantize_with_zero_point(
+                weights_scaled, scale_q, zero_q, maxq
+            )
+            dequantized = dequantize_with_zero_point(quantized, scale_q, zero_q)
+        else:
+            # Grouped quantization - use proper per-row grouping
+            scale_q, zero_q, maxq = compute_quantization_parameters(
+                weights_scaled,
+                bits=4,
+                symmetric=False,
+                per_channel=True,
+                group_size=group_size,
+                compute_dtype="float32",
+            )
+
+            # Compute group indices: maps each input feature to its group
+            g_idx = ops.cast(ops.arange(0, in_features) // group_size, "int32")
+
+            # Quantize and dequantize using group index mapping
+            quantized = quantize_with_sz_map(
+                weights_scaled, scale_q, zero_q, g_idx, maxq
+            )
+            dequantized = dequantize_with_sz_map(
+                quantized, scale_q, zero_q, g_idx
+            )
+
+        # Scale back down by DIVIDING to restore original magnitude
+        reconstructed = ops.divide(dequantized, scales)
+
+        # Compute activation-weighted MSE loss
+        # This approximates the output error: ||W*X - W_hat*X||^2
+        # by weighting each channel's error by x_max^2
+        weight_error = ops.square(ops.subtract(weights, reconstructed))
+        # Weight by activation magnitudes squared (broadcast over out_features)
+        weighted_error = ops.multiply(weight_error, ops.square(x_max))
+        loss = ops.mean(weighted_error)
+
+        # Track best
+        if best_loss is None:
+            best_loss = loss
+            best_scales = scales
+        else:
+            is_better = ops.less(loss, best_loss)
+            if is_better:
+                best_loss = loss
+                best_scales = scales
+
+    return best_scales
+
+
+def awq_quantize_matrix(
+    weights_transpose,
+    activation_magnitudes,
+    *,
+    num_grid_points=20,
+    group_size=-1,
+):
+    """Quantize a weight matrix using AWQ.
+
+    This function performs the complete AWQ quantization process:
+    1. Find optimal per-channel scales via grid search
+    2. Apply scales to weights
+    3. Compute quantization parameters
+    4. Quantize weights
+
+    Args:
+        weights_transpose: Weight matrix [out_features, in_features].
+        activation_magnitudes: Per-channel activation magnitudes [in_features].
+        num_grid_points: Number of grid search points.
+        group_size: Group size for quantization.
+
+    Returns:
+        quantized_weights: Quantized weights [out_features, in_features].
+        scales: Quantization scales [out_features, num_groups].
+        zeros: Zero points [out_features, num_groups].
+        awq_scales: AWQ per-channel scales [in_features].
+        g_idx: Group indices [in_features].
+    """
+    in_features = ops.shape(weights_transpose)[1]
+
+    # Step 1: Find optimal AWQ scales via grid search
+    awq_scales = awq_search_optimal_scales(
+        weights_transpose,
+        activation_magnitudes,
+        num_grid_points=num_grid_points,
+        group_size=group_size,
+    )
+
+    # Step 2: Apply AWQ scales by MULTIPLYING (expand salient weights)
+    # weights_scaled: [out_features, in_features]
+    weights_scaled = ops.multiply(weights_transpose, awq_scales)
+
+    if group_size == -1:
+        # Per-channel quantization (no grouping)
+        scale_q, zero_q, maxq = compute_quantization_parameters(
+            weights_scaled,
+            bits=4,
+            symmetric=False,
+            per_channel=True,
+            group_size=-1,
+            compute_dtype="float32",
+        )
+
+        # Quantize
+        quantized = quantize_with_zero_point(
+            weights_scaled, scale_q, zero_q, maxq
+        )
+
+        # Build group indices (all 0s for per-channel)
+        g_idx = ops.zeros((in_features,), dtype="float32")
+    else:
+        # Grouped quantization - use proper per-row grouping
+        scale_q, zero_q, maxq = compute_quantization_parameters(
+            weights_scaled,
+            bits=4,
+            symmetric=False,
+            per_channel=True,
+            group_size=group_size,
+            compute_dtype="float32",
+        )
+
+        # Compute group indices: maps each input feature to its group
+        g_idx = ops.cast(ops.arange(0, in_features) // group_size, "int32")
+
+        # Quantize using group index mapping
+        quantized = quantize_with_sz_map(
+            weights_scaled, scale_q, zero_q, g_idx, maxq
+        )
+
+        # Convert g_idx to float for storage
+        g_idx = ops.cast(g_idx, "float32")
+
+    return quantized, scale_q, zero_q, awq_scales, g_idx
+
+
+class AWQ:
+    """AWQ quantizer for a single layer.
+
+    This class accumulates activation statistics during calibration and
+    performs AWQ quantization on layer weights.
+
+    The AWQ algorithm works by:
+    1. Collecting per-channel maximum activation magnitudes
+    2. Using activation magnitudes to determine weight saliency
+    3. Finding optimal per-channel scales via grid search
+    4. Applying scales before quantization to protect salient weights
+
+    Args:
+        layer: The layer to quantize (Dense or EinsumDense).
+        config: AWQConfig instance with quantization parameters.
+    """
+
+    def __init__(self, layer, config=None):
+        from keras.src.quantizers.awq_config import AWQConfig
+
+        self.original_layer = layer
+        self.config = config or AWQConfig(dataset=None, tokenizer=None)
+        self.num_samples = 0
+
+        # Handle Dense and EinsumDense layers
+        if isinstance(layer, Dense) or (
+            isinstance(layer, EinsumDense) and layer.kernel.ndim == 2
+        ):
+            self.kernel_shape = layer.kernel.shape
+            self.rows = self.kernel_shape[0]  # in_features
+            self.columns = self.kernel_shape[1]  # out_features
+            self.layer = layer
+        elif isinstance(layer, EinsumDense) and layer.kernel.ndim == 3:
+            # Handle 3D EinsumDense layers (typically from attention blocks)
+            self.kernel_shape = layer.kernel.shape
+            shape = list(self.kernel_shape)
+            d_model_dim_index = shape.index(max(shape))
+
+            if d_model_dim_index == 0:  # QKV projection case
+                in_features, heads, head_dim = shape
+                self.rows = in_features
+                self.columns = heads * head_dim
+            elif d_model_dim_index in [1, 2]:  # Attention Output case
+                heads, head_dim, out_features = shape
+                self.rows = heads * head_dim
+                self.columns = out_features
+            else:
+                raise ValueError(
+                    f"Cannot determine dimensions for EinsumDense kernel "
+                    f"shape {shape}"
+                )
+
+            # Create a temporary object that holds a reshaped 2D version
+            self.layer = types.SimpleNamespace(
+                kernel=ops.reshape(layer.kernel, (self.rows, self.columns)),
+            )
+        else:
+            raise TypeError(f"Unsupported layer type for AWQ: {type(layer)}")
+
+        # Initialize activation magnitude accumulator (per-channel max)
+        self.activation_magnitudes = ops.zeros((self.rows,), dtype="float32")
+
+    def update_activation_magnitudes(self, input_batch):
+        """Update per-channel activation magnitude statistics.
+
+        This method tracks the maximum absolute activation value for each
+        input channel across all calibration batches.
+
+        Args:
+            input_batch: Input activations tensor [batch, ..., in_features].
+        """
+        if input_batch is None:
+            raise ValueError("Input tensor cannot be None.")
+        if ops.size(input_batch) == 0:
+            raise ValueError("Input tensor cannot be empty.")
+
+        # Flatten to [batch_samples, in_features]
+        if len(input_batch.shape) > 2:
+            input_batch = ops.reshape(input_batch, (-1, input_batch.shape[-1]))
+
+        x = ops.cast(input_batch, "float32")
+
+        # Compute per-channel max absolute value for this batch
+        batch_max = ops.max(ops.abs(x), axis=0)
+
+        # Update running max
+        self.activation_magnitudes = ops.maximum(
+            self.activation_magnitudes, batch_max
+        )
+        self.num_samples = self.num_samples + int(ops.shape(x)[0])
+
+    def quantize_layer(self):
+        """Perform AWQ quantization on the layer.
+
+        This method:
+        1. Runs the AWQ grid search to find optimal scales
+        2. Quantizes the layer weights
+        3. Updates the layer's quantized variables
+        """
+        from keras.src import quantizers
+
+        weights_matrix = ops.transpose(self.layer.kernel)
+
+        # Perform AWQ quantization
+        quantized, scale, zero, awq_scales, g_idx = awq_quantize_matrix(
+            weights_matrix,
+            self.activation_magnitudes,
+            num_grid_points=self.config.num_grid_points,
+            group_size=self.config.group_size,
+        )
+
+        # Cast to uint8 for storage
+        # quantized is already [out_features, in_features]
+        quantized = ops.cast(quantized, "uint8")
+
+        # Pack to 4-bit along axis 0 (output features)
+        quantized_packed, _, _ = quantizers.pack_int4(
+            quantized, axis=0, dtype="uint8"
+        )
+
+        # Assign to layer variables
+        del self.original_layer._kernel
+        self.original_layer.quantized_kernel.assign(quantized_packed)
+        self.original_layer.kernel_scale.assign(scale)
+        self.original_layer.kernel_zero.assign(zero)
+        self.original_layer.awq_scales.assign(awq_scales)
+        self.original_layer.g_idx.assign(g_idx)
+        self.original_layer.is_awq_calibrated = True
+
+    def free(self):
+        """Free memory used by the quantizer."""
+        del self.activation_magnitudes
+        del self.layer