Deprecate aqt keep deps

.github/workflows/code_quality.yml

@@ -17,6 +17,9 @@
 name: Code Quality
 
 on:
+  push:
+    branches:
+      - deprecate-aqt-keep-deps
   pull_request:
 
 concurrency:
@@ -56,7 +59,9 @@ jobs:
 
     - name: Run pre-commit checks on just the files that have changed
       run: |
-        git fetch origin "$GITHUB_BASE_REF":"$GITHUB_BASE_REF"
-        git branch "$GITHUB_HEAD_REF"
+        BASE_REF="${GITHUB_BASE_REF:-main}"
+        HEAD_REF="${GITHUB_HEAD_REF:-$GITHUB_SHA}"
+        git fetch origin "$BASE_REF":"$BASE_REF"
+        git branch "$HEAD_REF" || true
         . "$GITHUB_WORKSPACE"/venv/bin/activate
-        pre-commit run --from-ref "$GITHUB_BASE_REF" --to-ref "$GITHUB_HEAD_REF" --show-diff-on-failure
+        pre-commit run --from-ref "$BASE_REF" --to-ref "$HEAD_REF" --show-diff-on-failure

.github/workflows/run_tests_against_package.yml

@@ -164,16 +164,19 @@ jobs:
             # Dynamically discover the 'nvidia' folder and prepend all its sub-library
             # directories (including nccl, cublas, cudnn) to LD_LIBRARY_PATH to prevent
             # JAX from partially loading incompatible system-level CUDA libraries.
-            if [ -d ".venv/lib" ]; then
-              NVIDIA_DIR=$(find .venv/lib/ -maxdepth 3 -name "nvidia" -type d 2>/dev/null | head -n 1)
-              if [ -n "${NVIDIA_DIR}" ]; then
-                for dir in "${NVIDIA_DIR}"/*; do
-                  if [ -d "$dir/lib" ]; then
-                    export LD_LIBRARY_PATH=$(pwd)/$dir/lib:${LD_LIBRARY_PATH}
-                  fi
-                done
-              fi
+            NVIDIA_DIR=$(find .venv/lib/ -maxdepth 3 -name "nvidia" -type d 2>/dev/null | head -n 1)
+            if [ -n "${NVIDIA_DIR}" ]; then
+              for dir in "${NVIDIA_DIR}"/*; do
+                if [ -d "$dir/lib" ]; then
+                  export LD_LIBRARY_PATH=$(pwd)/$dir/lib:${LD_LIBRARY_PATH}
+                fi
+              done
             fi
+
+            # Configure NCCL for GPU execution
+            # Removed NCCL_SOCKET_IFNAME=lo as it breaks multi-gpu collective ops in Docker
+            export NCCL_DEBUG=WARN
+            echo "Set NCCL_DEBUG=WARN for GPU execution."
           fi
           if [ "${INPUTS_TOTAL_WORKERS}" -gt 1 ]; then
             $PYTHON_EXE -m pip install --quiet pytest-split pytest-xdist

src/maxtext/configs/base.yml

@@ -112,7 +112,6 @@ dtype: "bfloat16"
 # used to configure quantization in the transformer layers, defaults to null implying bf16.
 # possible alternative settings are as follows:
 # 'int8' for dynamic range quantization using 8-bits
-# 'intmp' for mixed precision quantization for inference as described here: src/maxtext/configs/quantization/readme.md
 # 'fp8' for 8-bit floating-point gemms on nvidia gpus.
 # 'nanoo_fp8' for 8-bit floating-point gemms on amd mi300/mi325 gpus.
 # 'fp8_full' for fp8 quantization with static scaling.
@@ -123,10 +122,6 @@ constant_bound_config: ""
 # https://kolonist26-jax-kr.readthedocs.io/en/latest/jax.lax.html#jax.lax.precision
 matmul_precision: "default"
 activations_in_float32: false # sets activations to float32 before nonlinearity it true, else dtype
-# used to replicate the quantization scale to avoid the inefficient xla fusion for 2d sharding.
-replicate_quant_scale: false
-# path to file with quantization config for intmp.
-quant_cfg_path: ""
 quantize_kvcache: false # set to true to quantize kv cache values, defaults to false
 # valid kv_quant_axis values:
 #   - "" is valid only when quantize_kvcache is false
@@ -143,7 +138,7 @@ save_quantized_params_path: ""
 # when left as is, corresponds to training
 # accepted values are "inference"
 model_call_mode: ""
-use_qwix_quantization: false # [DEPRECATED: AQT will be removed in a future release. It is strongly recommended to set use_qwix_quantization to true] whether to use qwix for quantization. if set to true, the model will be quantized using qwix.
+use_qwix_quantization: true # whether to use qwix for quantization. if set to true, the model will be quantized using qwix.
 use_manual_quantization: false # a flag if to use manual quantization for batch split. Only used if use_batch_split_schedule is true.
 # quantization calibration method used for weights and activations. supported methods can be found in https://github.com/google/qwix/blob/dc2a0770351c740e5ab3cce7c0efe9f7beacce9e/qwix/qconfig.py#l70-l80
 weight_quantization_calibration_method: "absmax"

src/maxtext/configs/inference/multihost/disaggregation/llama3_405b_v6e-16-16.yml

@@ -5,8 +5,6 @@ sharding_strategy: "experimental"
 attention: 'dot_product'
 allow_split_physical_axes: true
 tokenizer_path: "assets/tokenizer_llama3.tiktoken"
-# Used to replicate the quantization scale to avoid the inefficient XLA fusion.
-replicate_quant_scale: true
 
 inference_server: "ExperimentalMaxtextDisaggregatedServer"
 

src/maxtext/configs/inference/multihost/interleaved/llama2_70b_v5e-16.yml

@@ -8,8 +8,6 @@ model_name: "llama2-70b"
 sharding_strategy: "experimental"
 attention: 'dot_product'
 allow_split_physical_axes: true
-# Used to replicate the quantization scale to avoid the inefficient XLA fusion.
-replicate_quant_scale: true
 
 logical_axis_rules: [
                       ['embed', []],

src/maxtext/configs/inference/multihost/interleaved/llama3_405b_v5e-64.yml

@@ -10,8 +10,6 @@ sharding_strategy: "experimental"
 attention: 'dot_product'
 allow_split_physical_axes: true
 tokenizer_path: "assets/tokenizer_llama3.tiktoken"
-# Used to replicate the quantization scale to avoid the inefficient XLA fusion.
-replicate_quant_scale: true
 
 logical_axis_rules: [
                       ['embed', []],

src/maxtext/configs/inference/multihost/interleaved/llama3_70b_v5e-16.yml

@@ -9,8 +9,6 @@ tokenizer_path: "assets/tokenizer_llama3.tiktoken"
 sharding_strategy: "experimental"
 attention: 'dot_product'
 allow_split_physical_axes: true
-# Used to replicate the quantization scale to avoid the inefficient XLA fusion.
-replicate_quant_scale: true
 
 logical_axis_rules: [
                       ['embed', []],

src/maxtext/configs/tpu/v5e/llama2_70b_v5e-16.yml

@@ -8,8 +8,6 @@ model_name: "llama2-70b"
 sharding_strategy: "experimental"
 attention: 'dot_product'
 allow_split_physical_axes: true
-# Used to replicate the quantization scale to avoid the inefficient XLA fusion.
-replicate_quant_scale: true
 
 logical_axis_rules: [
                       ['embed', []],

src/maxtext/configs/tpu/v5e/llama3_405b_v5e-64.yml

@@ -10,8 +10,6 @@ sharding_strategy: "experimental"
 attention: 'dot_product'
 allow_split_physical_axes: true
 tokenizer_path: "assets/tokenizer_llama3.tiktoken"
-# Used to replicate the quantization scale to avoid the inefficient XLA fusion.
-replicate_quant_scale: true
 
 logical_axis_rules: [
                       ['embed', []],

src/maxtext/configs/tpu/v5e/llama3_70b_v5e-16.yml

@@ -9,8 +9,6 @@ tokenizer_path: "assets/tokenizer_llama3.tiktoken"
 sharding_strategy: "experimental"
 attention: 'dot_product'
 allow_split_physical_axes: true
-# Used to replicate the quantization scale to avoid the inefficient XLA fusion.
-replicate_quant_scale: true
 
 logical_axis_rules: [
                       ['embed', []],

src/maxtext/configs/tpu/v6e/inference/llama4_maverick_v6e-64.yml

@@ -9,8 +9,6 @@ base_config: "inference/inference_jetstream.yml"
 sharding_strategy: "experimental"
 attention: 'dot_product'
 allow_split_physical_axes: true
-# Used to replicate the quantization scale to avoid the inefficient XLA fusion.
-replicate_quant_scale: true
 
 logical_axis_rules: [
                       ['embed', []],

src/maxtext/configs/types.py

@@ -86,7 +86,6 @@ class QuantizationType(str, Enum):
   NONE = ""
   INT4 = "int4"
   INT8 = "int8"
-  INTMP = "intmp"
   FP8_E5M2 = "fp8_e5m2"
   FP8_E4M3 = "fp8_e4m3"
   FP8 = "fp8"
@@ -341,7 +340,7 @@ class Checkpointing(BaseModel):
   )
   checkpoint_is_quantized: bool = Field(
       False,
-      description="Set to True if reading from a saved AQT quantized checkpoint.",
+      description="Set to True if reading from a saved quantized checkpoint.",
   )
   save_quantized_params_path: PathStr = Field("", description="Path to save params quantized on the fly.")
   enable_orbax_v1: bool = Field(False, description="Bool flag for enabling Orbax v1.")
@@ -429,16 +428,11 @@ class Quantization(BaseModel):
       QuantizationType.NONE,
       description="Activates quantization for transformer layers.",
   )
-  replicate_quant_scale: bool = Field(
-      False,
-      description="Replicates quantization scale to avoid inefficient XLA fusion.",
-  )
-  quant_cfg_path: PathStr = Field("", description="Path to the configuration file for 'intmp' quantization.")
   quantize_kvcache: bool = Field(False, description="If True, quantizes the Key-Value cache.")
   kv_quant_axis: KvQuantAxis = Field(KvQuantAxis.HEADS_AND_DKV, description="Axes to quantize over for the KV cache.")
   kv_quant_dtype: Literal["int8", "int4"] = Field("int8", description="Data type for KV cache quantization.")
   quantization_local_shard_count: int = Field(-1, description="Shards the range finding operation for quantization.")
-  use_qwix_quantization: bool = Field(False, description="Whether to use qwix for quantization.")
+  use_qwix_quantization: bool = Field(True, description="Whether to use qwix for quantization.")
   use_manual_quantization: bool = Field(
       False,
       description="Whether to use manual quantization for batch split. Only used if use_batch_split_schedule is True.",
@@ -2512,7 +2506,7 @@ def _validate_check_vma_is_supported(self):
     if self.use_ring_of_experts:
       raise ValueError("check_vma is not yet supported with ring of experts.")
     _allowed = {"ici_expert_parallelism", "ici_fsdp_parallelism"}
-    active = [name for name in IciParallelism.model_fields if name not in _allowed and getattr(self, name) != 1]
+    active = [name for name in IciParallelism.model_fields if name not in _allowed and getattr(self, name) != 1]  # pylint: disable=not-an-iterable
     if active:
       raise ValueError(
           f"check_vma=True only supports ici_expert_parallelism and ici_fsdp_parallelism. "
@@ -2756,12 +2750,19 @@ def get_num_target_devices():
     }
     self.num_slices = max_utils.get_num_slices(raw_keys_for_num_slices)
 
-    # Check for AQT deprecation warning
+    # Enforce that Qwix is required for non-native/non-TE quantization
     if self.quantization and not self.use_qwix_quantization:
-      if self.quantization not in ("fp8", "nanoo_fp8") and not self.quantization.startswith("te_"):
-        logger.warning(
-            "WARNING: AQT quantization is deprecated and will be removed in a future release. "
-            "Please migrate to Qwix by setting use_qwix_quantization=True."
+      is_native_or_te = self.quantization in (
+          QuantizationType.FP8,
+          QuantizationType.NANOO_FP8,
+          QuantizationType.FP8_NANO_V2,
+          QuantizationType.FP8_GPU,
+      ) or self.quantization.startswith("te_")
+      if not is_native_or_te:
+        raise ValueError(
+            f"Quantization type '{self.quantization}' without Qwix (use_qwix_quantization=False) "
+            f"is unsupported because legacy AQT has been completely removed. "
+            f"Please migrate to Qwix by setting use_qwix_quantization=True."
         )
 
     # Default quantization sharding count to number of local devices if not set.
@@ -2967,15 +2968,6 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
         self.data_sharding[0].remove("stage")
         self.data_sharding[0].insert(0, "stage")
 
-      # Add sharding for FP8 amax history when using pipeline parallelism.
-      if self.quantization and self.quantization in (
-          "fp8",
-          "nanoo_fp8",
-          "fp8_gpu",
-          "te_fp8_delayedscaling",
-      ):
-        self.logical_axis_rules.append(["aqt_amax_history", ("stage",)])
-
     # H. RESOLVE local_sa_* FLAGS: inherit from global sa_* if not explicitly set.
     if self.local_sa_block_q is None:
       self.local_sa_block_q = self.sa_block_q
@@ -3290,7 +3282,7 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
       self.use_grpo = False
 
     if self.use_batch_split_schedule:
-      if self.quantization and not self.quantization == "fp8_full":
+      if self.quantization and self.quantization != "fp8_full":
         raise ValueError("Batch split quantization only supports `quantization=fp8_full`")
 
     if self.opt_type == "muon" and self.decoder_block not in [

src/maxtext/inference/kvcache.py

@@ -22,9 +22,12 @@
 from flax import linen as nn
 from flax import nnx
 
-from aqt.jax.v2 import config as aqt_config
-from aqt.jax.v2.aqt_tensor import QTensor as KVTensor
-from aqt.jax.v2.flax import aqt_flax
+
+class KVTensor:
+
+  def __init__(self, *args, **kwargs):
+    raise NotImplementedError("KV Cache quantization is not supported because AQT is deprecated.")
+
 
 from maxtext.layers import nnx_wrappers
 from maxtext.layers.initializers import variable_to_logically_partitioned
@@ -96,56 +99,10 @@ def einsum_fn_with_rhs_qtensor(
       lhs_dequant_mode=None,
       lhs_calibration_mode=None,
   ):
-    """einsum function where QTensor is the right-hand-side"""
-    # Assumes kv is already quantized.
-    einsum = jnp.einsum
-    if self.dtype != jnp.float8_e4m3fn:
-      num_bits = 4 if self.dtype == jnp.int4 else 8
-      kv_cfg = aqt_config.dot_general_make(
-          lhs_bits=None,
-          rhs_bits=num_bits,
-          bwd_bits=None,
-          use_fwd_quant=False,
-      )
-    else:
-      kv_cfg = aqt_config.config_fwd_fp8()
-
-    if rhs_dequant_mode:
-      aqt_config.set_fwd_dequant_mode(kv_cfg, rhs_dequant_mode=rhs_dequant_mode)
-    if rhs_calibration_mode:
-      aqt_config.set_fwd_calibration_mode(
-          kv_cfg,
-          rhs_calibration_mode=rhs_calibration_mode,
-      )
-    if lhs_dequant_mode:
-      aqt_config.set_fwd_dequant_mode(kv_cfg, lhs_dequant_mode=lhs_dequant_mode)
-    if lhs_calibration_mode:
-      aqt_config.set_fwd_calibration_mode(
-          kv_cfg,
-          lhs_calibration_mode=lhs_calibration_mode,
-      )
-    einsum = aqt_flax.AqtEinsum(
-        rhs_quant_mode=aqt_flax.QuantMode.TRAIN,
-        lhs_freeze_mode=aqt_flax.FreezerMode.NONE,
-        rhs_freeze_mode=aqt_flax.FreezerMode.NONE,
-        cfg=kv_cfg,
-    )
-    return einsum
+    raise NotImplementedError("KV Cache quantization is not supported because AQT is deprecated.")
 
   def einsum_fn_with_rhs_qtensor_and_dequant(self):
-    """Get einstein summation for different dequant modes."""
-    if self.dtype == jnp.float8_e4m3fn:
-      return self.einsum_fn_with_rhs_qtensor(
-          lhs_dequant_mode=aqt_config.DequantMode.THIS_INPUT,
-          lhs_calibration_mode=aqt_config.CalibrationMode.REMAINING_AXIS,
-          rhs_dequant_mode=aqt_config.DequantMode.OTHER_INPUT,
-          rhs_calibration_mode=aqt_config.CalibrationMode.REMAINING_AXIS,
-      )
-    else:
-      return self.einsum_fn_with_rhs_qtensor(
-          rhs_dequant_mode=aqt_config.DequantMode.OTHER_INPUT,
-          rhs_calibration_mode=aqt_config.CalibrationMode.REMAINING_AXIS,
-      )
+    raise NotImplementedError("KV Cache quantization is not supported because AQT is deprecated.")
 
 
 def kv_cache_as_linen(