diff --git a/.github/workflows/code_quality.yml b/.github/workflows/code_quality.yml index 6f5c2a277a..23a5d8a7ce 100644 --- a/.github/workflows/code_quality.yml +++ b/.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 diff --git a/.github/workflows/run_tests_against_package.yml b/.github/workflows/run_tests_against_package.yml index 34a05d8b21..bbaa714040 100644 --- a/.github/workflows/run_tests_against_package.yml +++ b/.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 diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index 0cc69fc2b0..1407fc4972 100644 --- a/src/maxtext/configs/base.yml +++ b/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" diff --git a/src/maxtext/configs/inference/multihost/disaggregation/llama3_405b_v6e-16-16.yml b/src/maxtext/configs/inference/multihost/disaggregation/llama3_405b_v6e-16-16.yml index 2b83fadc56..02402e07dd 100644 --- a/src/maxtext/configs/inference/multihost/disaggregation/llama3_405b_v6e-16-16.yml +++ b/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" diff --git a/src/maxtext/configs/inference/multihost/interleaved/llama2_70b_v5e-16.yml b/src/maxtext/configs/inference/multihost/interleaved/llama2_70b_v5e-16.yml index dba4bc03ce..eb366a7ee5 100644 --- a/src/maxtext/configs/inference/multihost/interleaved/llama2_70b_v5e-16.yml +++ b/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', []], diff --git a/src/maxtext/configs/inference/multihost/interleaved/llama3_405b_v5e-64.yml b/src/maxtext/configs/inference/multihost/interleaved/llama3_405b_v5e-64.yml index b71b7990f1..220c496732 100644 --- a/src/maxtext/configs/inference/multihost/interleaved/llama3_405b_v5e-64.yml +++ b/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', []], diff --git a/src/maxtext/configs/inference/multihost/interleaved/llama3_70b_v5e-16.yml b/src/maxtext/configs/inference/multihost/interleaved/llama3_70b_v5e-16.yml index 525d30e30c..1b1c24dd86 100644 --- a/src/maxtext/configs/inference/multihost/interleaved/llama3_70b_v5e-16.yml +++ b/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', []], diff --git a/src/maxtext/configs/tpu/v5e/llama2_70b_v5e-16.yml b/src/maxtext/configs/tpu/v5e/llama2_70b_v5e-16.yml index dba4bc03ce..eb366a7ee5 100644 --- a/src/maxtext/configs/tpu/v5e/llama2_70b_v5e-16.yml +++ b/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', []], diff --git a/src/maxtext/configs/tpu/v5e/llama3_405b_v5e-64.yml b/src/maxtext/configs/tpu/v5e/llama3_405b_v5e-64.yml index b71b7990f1..220c496732 100644 --- a/src/maxtext/configs/tpu/v5e/llama3_405b_v5e-64.yml +++ b/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', []], diff --git a/src/maxtext/configs/tpu/v5e/llama3_70b_v5e-16.yml b/src/maxtext/configs/tpu/v5e/llama3_70b_v5e-16.yml index 525d30e30c..1b1c24dd86 100644 --- a/src/maxtext/configs/tpu/v5e/llama3_70b_v5e-16.yml +++ b/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', []], diff --git a/src/maxtext/configs/tpu/v6e/inference/llama4_maverick_v6e-64.yml b/src/maxtext/configs/tpu/v6e/inference/llama4_maverick_v6e-64.yml index 68f6c839e8..1e2bbc635a 100644 --- a/src/maxtext/configs/tpu/v6e/inference/llama4_maverick_v6e-64.yml +++ b/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', []], diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index e270dfeec2..e806087e4e 100644 --- a/src/maxtext/configs/types.py +++ b/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 [ diff --git a/src/maxtext/inference/kvcache.py b/src/maxtext/inference/kvcache.py index ba2266060f..e14692de84 100644 --- a/src/maxtext/inference/kvcache.py +++ b/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( diff --git a/src/maxtext/inference/maxengine/maxengine.py b/src/maxtext/inference/maxengine/maxengine.py index f5232e0d89..b1d3e6bb15 100644 --- a/src/maxtext/inference/maxengine/maxengine.py +++ b/src/maxtext/inference/maxengine/maxengine.py @@ -29,7 +29,7 @@ if jax.__version_info__ >= (0, 6, 3): from jax.experimental.layout import Layout as DLL # type: ignore else: - from jax.experimental.layout import DeviceLocalLayout as DLL # type: ignore + from jax.experimental.layout import DeviceLocalLayout as DLL # type: ignore # pylint: disable=no-name-in-module from flax import linen as nn from flax import nnx @@ -123,7 +123,7 @@ def __init__(self, config: Any, devices: Any | None = None): # quant enabled we stay in `train` mode and let AQT quantize per-forward # against the full-precision kernel — same numerical result as `serve` # for absmax calibration, just slower. - nnx_quant_mode_str = "serve" if (quant is not None and config.checkpoint_is_quantized) else "train" + nnx_quant_mode_str = "serve" if (config.quantization and config.checkpoint_is_quantized) else "train" # We need both PREFILL and AR abstract models because the cache vars inherit # CACHE_BATCH_PREFILL vs CACHE_BATCH from the construction model_mode, and # bulk_insert searches for the substring "cache_batch" in the AR-mode names. @@ -326,9 +326,8 @@ def load_params(self, *args, params=None, rng: PRNGKeyType | None = None, **kwar if self.model.quant and self.config.checkpoint_is_quantized: print("Loading from the quantized checkpoint...") - self.model.quant.quant_mode = quantizations.get_quant_mode("serve") - rng1, rng2, rng3 = jax.random.split(rng, 3) + rng1, rng2 = jax.random.split(rng, 2) if params: print("Resharding given params") init_state_fn = functools.partial(maxtext_utils.init_initial_state, self.model, None, self.config, False, rng) @@ -373,7 +372,10 @@ def load_params(self, *args, params=None, rng: PRNGKeyType | None = None, **kwar ) if self.model.quant and not self.config.checkpoint_is_quantized: - params = self.quantize_params(state, rng3) + raise ValueError( + "On-the-fly parameter quantization is not supported for the modern backends. " + "Please load a pre-quantized checkpoint." + ) else: params = state.params @@ -382,15 +384,7 @@ def load_params(self, *args, params=None, rng: PRNGKeyType | None = None, **kwar return params def _load_params_nnx(self, params, rng): - """NNX equivalent of load_params: returns an nnx.Param state and populates KV cache shardings. - - Quantization handling: - * `checkpoint_is_quantized=True`: model built in `serve` mode (no full - kernel), `from_pretrained` reads `qrhs.frozen` from disk. - * `checkpoint_is_quantized=False` + `quantization=...`: model built in - `train` mode, full-precision kernel loaded; AQT layers quantize per - forward. Same output as serve mode (absmax calibration), slower. - """ + """NNX equivalent of load_params: returns an nnx.Param state and populates KV cache shardings.""" if params: print("Resharding given NNX params") @@ -399,9 +393,7 @@ def _load_params_nnx(self, params, rng): # axis metadata but no physical .sharding. Resolve logical to physical here so # device_put actually reshards instead of being a no-op. with nn_partitioning.axis_rules(self.config.logical_axis_rules): - target_shardings = sharding.nnx_construct_named_sharding( - params_abs, self._mesh - ) + target_shardings = sharding.nnx_construct_named_sharding(params_abs, self._mesh) params_state = jax.device_put(params, target_shardings) # We only need a concrete `rest` (RNG vars) for nnx.merge. create_nnx_sharded_model # builds the model with a jitted out_shardings so params are produced already @@ -409,9 +401,7 @@ def _load_params_nnx(self, params, rng): # large models). self.model is abstract with no .sharding, so pass an explicit one. _, full_abs = nnx.split(self.model) with nn_partitioning.axis_rules(self.config.logical_axis_rules): - full_sharding = sharding.nnx_construct_named_sharding( - full_abs, self._mesh - ) + full_sharding = sharding.nnx_construct_named_sharding(full_abs, self._mesh) concrete_model = maxtext_utils_nnx.create_nnx_sharded_model( self.model, self._create_model_fn, mesh=self._mesh, named_sharding=full_sharding ) @@ -549,61 +539,6 @@ def unapply_adapter(self, base_params, adapter_config, adapter_params): else: lora_utils.unapply_lora_from_base_params(base_params, adapter_params, lora_scale_factor) - def quantize_params(self, state, rng: PRNGKeyType | None = None): - """Forward pass to quantize decode params.""" - if rng is None: - rng = jax.random.PRNGKey(0) - if self.config.pure_nnx: - # NNX takes a different code path: convert-on-load lives in `_load_params_nnx` - # via `_convert_and_quantize_nnx`, which runs the dummy forward against a - # CONVERT-mode model and transfers `qrhs.frozen` into the SERVE model. - # The standalone `quantize_params(state, rng)` API expects a Linen-shape - # `state.params` dict and isn't reachable on the NNX pathway in maxengine - # (load_params already dispatched to _load_params_nnx). - raise NotImplementedError( - "Use load_params() on NNX — the convert step runs inside _load_params_nnx via " - "_convert_and_quantize_nnx. quantize_params(state, rng) is the Linen API." - ) - - self.model.quant.quant_mode = quantizations.get_quant_mode("convert") - - @jax.jit - def model_apply(_p, _rng): - image_shape = mm_processor.get_dummy_image_shape_for_init( - model_name=self.config.model_name, - batch_size=self.config.micro_batch_size_to_train_on, - ) - audio_shape = mm_processor.get_dummy_audio_shape_for_init(self.config) - return self.model.apply( - _p | {"aqt": {}}, - jnp.ones((1, self.config.max_prefill_predict_length), dtype=jnp.int32), - jnp.ones((1, self.config.max_prefill_predict_length), dtype=jnp.int32), - encoder_images=jnp.ones(image_shape, dtype=jnp.float32) if self.config.use_multimodal else None, - # encoder_image_masks indicates valid tiles if image tiling + padding is used in vision encoder input. - encoder_image_masks=jnp.ones(image_shape[:2], dtype=jnp.int32) - if self.config.use_multimodal and "llama4" in self.config.model_name - else None, - encoder_audios=jnp.ones(audio_shape, dtype=jnp.float32) if self.config.use_audio else None, - decoder_segment_ids=jnp.zeros((1, self.config.max_prefill_predict_length), dtype=jnp.int32), - enable_dropout=False, - model_mode=MODEL_MODE_PREFILL, - rngs={"params": _rng}, - mutable=True, - ) - - _, new_vars = model_apply(state.params, rng) - # Remove param values which have corresponding qtensors in aqt to save memory. - params = {} - params["aqt"] = new_vars["aqt"] - params["params"] = quantizations.remove_quantized_params(state.params["params"], new_vars["aqt"]) - self.abstract_params = jax.tree_util.tree_map( - lambda x: jax.ShapeDtypeStruct(shape=x.shape, dtype=x.dtype, sharding=x.sharding), - params, - ) - maxtext_utils.save_quantized_checkpoint_if_configured(self.config, params) - self.model.quant.quant_mode = quantizations.get_quant_mode("serve") - return params - def _maybe_stack_prefill_result_cache(self, cache): """Stack the caches across the layers.""" if not self.config.stack_prefill_result_cache: @@ -2047,8 +1982,7 @@ def set_engine_vars_from_base_engine( """Set internal vars from base_engine, which has already loaded the checkpoint and has sharding, mesh, and kv cache related vars set. """ - if base_engine.model.quant: - engine.model.quant.quant_mode = base_engine.model.quant.quant_mode + engine.state_mesh_annotations = base_engine.state_mesh_annotations engine.abstract_params = base_engine.abstract_params engine.kv_cache_annotations = maxtext_utils.get_kv_cache_annotations(engine.model, engine.config, rng, engine.mesh) # pylint: disable=protected-access diff --git a/src/maxtext/layers/attention_compressed.py b/src/maxtext/layers/attention_compressed.py index 65869fb186..337b1db4a9 100644 --- a/src/maxtext/layers/attention_compressed.py +++ b/src/maxtext/layers/attention_compressed.py @@ -38,7 +38,7 @@ from maxtext.layers.initializers import nd_dense_init, NdInitializer, variable_to_logically_partitioned from maxtext.layers.linears import DenseGeneral, DeepSeekV4GroupedLinear from maxtext.layers.normalizations import RMSNorm -from maxtext.layers.quantizations import AqtQuantization as Quant +from maxtext.layers.quantizations import Quantization as Quant from maxtext.inference.kvcache import KVQuant diff --git a/src/maxtext/layers/attention_mla.py b/src/maxtext/layers/attention_mla.py index df7fd16ea2..11449e8f72 100644 --- a/src/maxtext/layers/attention_mla.py +++ b/src/maxtext/layers/attention_mla.py @@ -68,7 +68,7 @@ from maxtext.layers.initializers import nd_dense_init, NdInitializer, variable_to_logically_partitioned from maxtext.layers.linears import DenseGeneral from maxtext.layers.normalizations import RMSNorm -from maxtext.layers.quantizations import AqtQuantization as Quant +from maxtext.layers.quantizations import Quantization as Quant from maxtext.inference import kvcache from maxtext.inference.kvcache import KVQuant from maxtext.utils.sharding import create_sharding diff --git a/src/maxtext/layers/attention_op.py b/src/maxtext/layers/attention_op.py index f3eb515547..2d81fbd5f0 100644 --- a/src/maxtext/layers/attention_op.py +++ b/src/maxtext/layers/attention_op.py @@ -70,7 +70,7 @@ from maxtext.kernels.attention.ragged_attention import ragged_mha from maxtext.layers import nnx_wrappers from maxtext.layers.initializers import variable_to_logically_partitioned -from maxtext.layers.quantizations import AqtQuantization as Quant +from maxtext.layers.quantizations import Quantization as Quant from maxtext.utils import max_utils from maxtext.utils.sharding import logical_to_mesh_axes, maybe_shard_with_pspec import numpy as np @@ -130,7 +130,7 @@ def validate_gpu_flash_attention(sinks: Array | None, record_max_logits: bool) - # TODO(agagik): change splash_attention_mask._ComputableMask to be non protected -class ChunkedCausalMask(splash_attention_mask._ComputableMask): # pylint: disable=protected-access +class ChunkedCausalMask(splash_attention_mask._ComputableMask): # pylint: disable=protected-access,abstract-method """Lazy chunked causal mask. Attention is causal within each chunk (0, K), (K, 2K), (2K, 3K), ... tokens @@ -2158,7 +2158,7 @@ def __call__( # pylint: disable=protected-access -class LoadBalancedCausalMask(splash_attention_mask._ComputableMask): +class LoadBalancedCausalMask(splash_attention_mask._ComputableMask): # pylint: disable=abstract-method """Lazy causal mask, prevents the model from attending to future tokens. Attributes: diff --git a/src/maxtext/layers/attentions.py b/src/maxtext/layers/attentions.py index ab7673d1d4..242e33d024 100644 --- a/src/maxtext/layers/attentions.py +++ b/src/maxtext/layers/attentions.py @@ -64,7 +64,7 @@ from maxtext.layers.initializers import nd_dense_init, NdInitializer, variable_to_logically_partitioned, default_bias_init from maxtext.layers.linears import DenseGeneral, canonicalize_tuple, normalize_axes from maxtext.layers.normalizations import RMSNorm, Qwen3NextRMSNorm, GlobalRMSNorm -from maxtext.layers.quantizations import AqtQuantization as Quant +from maxtext.layers.quantizations import Quantization as Quant from maxtext.inference import kvcache from maxtext.inference.kvcache import KVQuant from maxtext.utils.sharding import maybe_shard_with_logical, create_sharding diff --git a/src/maxtext/layers/decoders.py b/src/maxtext/layers/decoders.py index 6bca4b98c7..bbbc71d16e 100644 --- a/src/maxtext/layers/decoders.py +++ b/src/maxtext/layers/decoders.py @@ -38,7 +38,7 @@ from maxtext.layers.attentions import attention_as_linen from maxtext.layers.embeddings import attend_on_embedding, embed_as_linen, positional_embedding_as_linen from maxtext.layers.normalizations import rms_norm -from maxtext.layers.quantizations import AqtQuantization as Quant +from maxtext.layers.quantizations import Quantization as Quant from maxtext.models import ( deepseek, deepseek4, diff --git a/src/maxtext/layers/engram.py b/src/maxtext/layers/engram.py index 3b2eb4e2b5..855ba2ee79 100644 --- a/src/maxtext/layers/engram.py +++ b/src/maxtext/layers/engram.py @@ -30,7 +30,7 @@ from maxtext.layers.initializers import NdInitializer, nd_dense_init from maxtext.layers.linears import DenseGeneral from maxtext.layers.normalizations import RMSNorm -from maxtext.layers.quantizations import AqtQuantization as Quant +from maxtext.layers.quantizations import Quantization as Quant import numpy as np import sympy import tokenizers diff --git a/src/maxtext/layers/initializers.py b/src/maxtext/layers/initializers.py index bbc6605057..97543eb87d 100644 --- a/src/maxtext/layers/initializers.py +++ b/src/maxtext/layers/initializers.py @@ -20,7 +20,6 @@ from flax import linen as nn from flax import nnx -from aqt.jax.v2 import aqt_tensor from maxtext.common.common_types import Array, DType, Shape, PRNGKey @@ -68,7 +67,7 @@ def variable_to_logically_partitioned(variable: nnx.Variable): present, it wraps the variable's value in `nn.LogicallyPartitioned` to apply the specified sharding constraints. - It handles special cases for `aqt_tensor.QTensor` and variables of type + It handles special cases for variables of type `_overwrite_with_gradient` by returning their values directly without wrapping. @@ -79,8 +78,6 @@ def variable_to_logically_partitioned(variable: nnx.Variable): The variable's value, potentially wrapped in `nn.LogicallyPartitioned`. """ val = variable.get_value() - if isinstance(val, aqt_tensor.QTensor): - return val if variable.type.__name__ == "_overwrite_with_gradient": return val diff --git a/src/maxtext/layers/learn_to_init_layer.py b/src/maxtext/layers/learn_to_init_layer.py index 2530c17336..368f69674e 100644 --- a/src/maxtext/layers/learn_to_init_layer.py +++ b/src/maxtext/layers/learn_to_init_layer.py @@ -24,7 +24,7 @@ from maxtext.common.common_types import DType, ShardMode, Array from maxtext.layers.nnx_wrappers import ToNNX -from maxtext.layers.quantizations import AqtQuantization as Quant +from maxtext.layers.quantizations import Quantization as Quant from maxtext.layers.initializers import NdInitializer, nd_dense_init from maxtext.utils import max_logging, max_utils diff --git a/src/maxtext/layers/linears.py b/src/maxtext/layers/linears.py index ca99776f38..e5183506eb 100644 --- a/src/maxtext/layers/linears.py +++ b/src/maxtext/layers/linears.py @@ -31,10 +31,10 @@ from maxtext.common.common_types import DecoderBlockType, ShardMode, DType, Array, Config from maxtext.common.common_types import MODEL_MODE_PREFILL -from maxtext.layers import nnx_wrappers, quantizations +from maxtext.layers import nnx_wrappers from maxtext.layers import normalizations from maxtext.layers.initializers import NdInitializer, nd_dense_init, default_bias_init, variable_to_logically_partitioned -from maxtext.layers.quantizations import AqtQuantization as Quant +from maxtext.layers.quantizations import Quantization as Quant from maxtext.utils import max_logging from maxtext.utils import max_utils from maxtext.utils.sharding import maybe_shard_with_logical @@ -160,17 +160,16 @@ def __init__( kernel_in_axis = np.arange(len(self.axis)) kernel_out_axis = np.arange(len(self.axis), len(self.axis) + len(self.out_features_shape)) - if not quantizations.in_serve_mode(self.quant): - self.kernel = nnx.Param( - self.kernel_init( - rngs.params(), - kernel_shape, - self.weight_dtype, - kernel_in_axis, - kernel_out_axis, - ), - sharding=self.kernel_axes, - ) + self.kernel = nnx.Param( + self.kernel_init( + rngs.params(), + kernel_shape, + self.weight_dtype, + kernel_in_axis, + kernel_out_axis, + ), + sharding=self.kernel_axes, + ) if self.use_bias: bias_axes = self.kernel_axes[-len(self.out_features_shape) :] @@ -219,16 +218,12 @@ def __call__(self, inputs: Array, _initializing: bool = False, out_sharding: Nam f"does not match expected input feature size {self.in_features_shape[i]}" ) - if quantizations.in_serve_mode(self.quant): - kernel_shape = self.in_features_shape + self.out_features_shape - kernel = jnp.zeros(kernel_shape, dtype=self.dtype) - else: - kernel = self.kernel[...] - # Move logit_dense kernel to device if parameter offloading is enabled - if self.parameter_memory_host_offload: - max_logging.log("linear.py: Moving parameter logits_dense kernel to device") - kernel = jax.device_put(kernel, max_utils.device_space()) - kernel = jnp.asarray(kernel, self.dtype) + kernel = self.kernel[...] + # Move logit_dense kernel to device if parameter offloading is enabled + if self.parameter_memory_host_offload: + max_logging.log("linear.py: Moving parameter logits_dense kernel to device") + kernel = jax.device_put(kernel, max_utils.device_space()) + kernel = jnp.asarray(kernel, self.dtype) # out_sharding should be None for auto mesh axis if self.shard_mode != ShardMode.EXPLICIT: diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py index 18d162e273..52ab9accf9 100644 --- a/src/maxtext/layers/moe.py +++ b/src/maxtext/layers/moe.py @@ -13,6 +13,8 @@ # limitations under the License. +# pylint: disable=assignment-from-none + """MoE related Layers.""" import enum @@ -21,7 +23,7 @@ import random from typing import Iterable, Optional, Tuple, Union -from aqt.jax.v2 import aqt_tensor as aqt + from flax import nnx from flax import struct import jax @@ -229,7 +231,7 @@ def __init__( kernel_axes: Tuple[Optional[str], ...] = (), use_bias: bool = False, score_func: str = "", - quant: Optional[quantizations.AqtQuantization] = None, + quant: Optional[quantizations.Quantization] = None, shard_mode: ShardMode = ShardMode.AUTO, matmul_precision: str = "default", ): @@ -272,17 +274,16 @@ def __init__( kernel_in_axis = np.arange(len(self.axis)) kernel_out_axis = np.arange(len(self.axis), len(self.axis) + len(self.out_features_shape)) - if not quantizations.in_serve_mode(self.quant): - self.kernel = nnx.Param( - self.kernel_init( - rngs.params(), - kernel_shape, - self.weight_dtype, - kernel_in_axis, - kernel_out_axis, - ), - out_sharding=self.kernel_axes, - ) + self.kernel = nnx.Param( + self.kernel_init( + rngs.params(), + kernel_shape, + self.weight_dtype, + kernel_in_axis, + kernel_out_axis, + ), + out_sharding=self.kernel_axes, + ) if self.use_bias: bias_axes = self.kernel_axes[-len(self.out_features_shape) :] @@ -315,11 +316,7 @@ def __call__(self, inputs: jax.Array, _initializing: bool = False) -> Tuple[jax. inputs = jnp.asarray(inputs, self.dtype) norm_axis = linears.normalize_axes(self.axis, inputs.ndim) - if quantizations.in_serve_mode(self.quant): - kernel_shape = self.in_features_shape + self.out_features_shape - kernel = jnp.zeros(kernel_shape, dtype=self.dtype) - else: - kernel = self.kernel[...] + kernel = self.kernel[...] kernel = jnp.asarray(kernel, self.dtype) contract_ind = tuple(range(0, len(norm_axis))) @@ -368,7 +365,7 @@ def __init__( intermediate_dim: int = 2048, weight_dtype: ctypes.DType = jnp.float32, dtype: ctypes.DType = jnp.float32, - quant: Optional[quantizations.AqtQuantization] = None, + quant: Optional[quantizations.Quantization] = None, is_hash_routing: bool = False, ): """Initializes the RoutedMoE module. @@ -500,14 +497,7 @@ def __init__( kernel_in_axis = np.arange(1) kernel_out_axis = np.arange(1, 2) - if quantizations.in_serve_mode(self.quant): - # During aqt convert state we delete kernel weight from params to save - # memory. Instead they are retrieved from the tensors stored in the 'aqt' - # collection. - self.wi_0 = jnp.zeros((num_experts, self.moe_expert_input_dim, intermediate_dim)) - self.wi_1 = jnp.zeros((num_experts, self.moe_expert_input_dim, intermediate_dim)) - self.wo = jnp.zeros((num_experts, intermediate_dim, self.moe_expert_input_dim)) - elif self.config.prefuse_moe_weights and self.config.attention == "vllm_rpa": + if self.config.prefuse_moe_weights and self.config.attention == "vllm_rpa": # Pad model dimension in Fused MoE weight kernels for GMM_v2 execution. moe_intermediate_dim = ( self.config.padded_base_moe_mlp_dim @@ -522,7 +512,7 @@ def __init__( kernel_in_axis, kernel_out_axis, ), - out_sharding=self.wi_kernel_axes, + sharding=self.wi_kernel_axes, ) self.wo = nnx.Param( self.kernel_init( @@ -532,7 +522,7 @@ def __init__( kernel_in_axis, kernel_out_axis, ), - out_sharding=self.wo_kernel_axes, + sharding=self.wo_kernel_axes, ) else: # Pad model dimension in Unfused MoE weight kernels for GMM_v2 execution. @@ -549,7 +539,7 @@ def __init__( kernel_in_axis, kernel_out_axis, ), - out_sharding=self.wi_kernel_axes, + sharding=self.wi_kernel_axes, ) self.wi_1 = nnx.Param( self.kernel_init( @@ -559,7 +549,7 @@ def __init__( kernel_in_axis, kernel_out_axis, ), - out_sharding=self.wi_kernel_axes, + sharding=self.wi_kernel_axes, ) self.wo = nnx.Param( self.kernel_init( @@ -569,7 +559,7 @@ def __init__( kernel_in_axis, kernel_out_axis, ), - out_sharding=self.wo_kernel_axes, + sharding=self.wo_kernel_axes, ) if self.config.mlp_bias: @@ -579,15 +569,15 @@ def __init__( wo_bias_shape = (self.num_experts, self.moe_expert_input_dim) self.wi_0_bias = nnx.Param( default_bias_init(self.rngs.params(), wi_bias_shape, self.weight_dtype), - out_sharding=wi_bias_axes, + sharding=wi_bias_axes, ) self.wi_1_bias = nnx.Param( default_bias_init(self.rngs.params(), wi_bias_shape, self.weight_dtype), - out_sharding=wi_bias_axes, + sharding=wi_bias_axes, ) self.wo_bias = nnx.Param( default_bias_init(self.rngs.params(), wo_bias_shape, self.weight_dtype), - out_sharding=wo_bias_axes, + sharding=wo_bias_axes, ) else: self.wi_0_bias = None @@ -597,7 +587,7 @@ def __init__( if self.config.decoder_block == ctypes.DecoderBlockType.GEMMA4: self.per_expert_scale = nnx.Param( jnp.ones((self.num_experts,), dtype=self.weight_dtype), - out_sharding=("exp",), + sharding=("exp",), ) else: self.per_expert_scale = None @@ -1252,11 +1242,7 @@ def jax_ragged_dot_gmm(inputs, kernel, tiling, group_sizes, expert_assignments, min(tiling[2], n), ) rhs_inputs = kernel - if isinstance(kernel, aqt.QTensor): - if kernel.bias or kernel.sparsity_mask or len(kernel.scale) > 1: - raise ValueError("Unsupported usecase for ragged_dot with quantized kernel.") - rhs_inputs = kernel.qvalue - if self.config.use_qwix_quantization: + if self.config.use_qwix_quantization and self.config.quantization: # Use full contraction for QWIX quantization to allow quantization # fusion (max reduce over contracting dimension). tiling = (tiling[0], k, tiling[2]) @@ -1274,21 +1260,11 @@ def jax_ragged_dot_gmm(inputs, kernel, tiling, group_sizes, expert_assignments, group_sizes=group_sizes, preferred_element_type=self.dtype, ) - if isinstance(kernel, aqt.QTensor): - # Multiply outputs by the kernely scale - scales = jnp.take(kernel.scale[0].squeeze(), indices=expert_assignments, axis=0) - if padding_amount > 0: - scales = jax.lax.pad( - scales, - jnp.array(0.0, dtype=scales.dtype), - [(0, padding_amount, 0), (0, 0, 0)], - ) - output *= scales return output def get_tokamax_group_sizes(group_sizes, inputs, kernel): # TODO (b/491979205) pipeline fsdp ag per repeat fails tokamax gmm - if self.config.use_qwix_quantization or ( + if (self.config.use_qwix_quantization and self.config.quantization) or ( self.config.using_pipeline_parallelism and self.config.pipeline_fsdp_ag_per_repeat ): return group_sizes @@ -1401,15 +1377,6 @@ def explicitly_weight_ag(shard_exp_on_fsdp): return True return False - def maybe_aqt_partition(w0_kernel, w0_pspec, w1_kernel, w1_pspec, wo_kernel, wo_pspec): - if isinstance(w0_kernel, aqt.QTensor): - w0_pspec = aqt.partition_spec(w0_pspec, (1,), w0_kernel.dtype, use_bias=False) - if isinstance(w1_kernel, aqt.QTensor): - w1_pspec = aqt.partition_spec(w1_pspec, (1,), w1_kernel.dtype, use_bias=False) - if isinstance(wo_kernel, aqt.QTensor): - wo_pspec = aqt.partition_spec(wo_pspec, (1,), wo_kernel.dtype, use_bias=False) - return w0_pspec, w1_pspec, wo_pspec - def get_routed_moe_shardings(is_batch_sharded_by_expert, has_input_ids): if is_batch_sharded_by_expert: batch_logical_axis = "activation_batch" @@ -1494,7 +1461,6 @@ def get_routed_moe_shardings(is_batch_sharded_by_expert, has_input_ids): wo_bias_pspec, decoder_tokens_pspec, ) = get_routed_moe_shardings(is_batch_sharded_by_expert, input_ids is not None) - w0_pspec, w1_pspec, wo_pspec = maybe_aqt_partition(w0_kernel, w0_pspec, w1_kernel, w1_pspec, wo_kernel, wo_pspec) def route(x, logits, pre_bias_logits, rngs, input_ids=None): """Performs both across device and within device token routing/sorting""" @@ -2146,14 +2112,11 @@ def get_einsum( if self.quant: - def aqt_einsum(*args, **kwargs): # pylint: disable=unused-argument - # simply skip kwargs, since aqt einsum doesn't support any kwargs - # like precision - is_aqt = not isinstance(self.quant, quantizations.Fp8Quantization) - kw = {"mesh_axes": rhs_mesh_axes} if is_aqt else {"dtype": self.dtype} - return self.quant.einsum(**kw)(*args) # pytype: disable=attribute-error + def quant_einsum(*args, **kwargs): + kw = {"dtype": self.dtype} + return self.quant.einsum(**kw)(*args, **kwargs) # pytype: disable=attribute-error - einsum_op = aqt_einsum + einsum_op = quant_einsum else: einsum_op = jnp.einsum return einsum_op @@ -2533,35 +2496,6 @@ def fused_moe_matmul( output = jnp.reshape(output_2d, (batch_size, seq_len, emb_dim)) return output, None, None - def retrieve_quantized_weight( - self, - inputs, - gate_logits, - pre_bias_logits, - w0_kernel, - w1_kernel, - wo_kernel, - w0_bias, - w1_bias, - wo_bias, - ) -> tuple[aqt.QTensor, aqt.QTensor, aqt.QTensor]: - """Retrieve quantized weights.""" - # This is called only during tracing. This is to invoke creation of - # quantized tensor inside AqtEinsum. After jit, this will become no-op and - # will not affect performance. - _ = self.dense_matmul( - inputs, gate_logits, pre_bias_logits, w0_kernel, w1_kernel, wo_kernel, w0_bias, w1_bias, wo_bias - ) - - w0_kernel = self.variables["aqt"]["AqtEinsum_0"]["AqtDotGeneral_0"]["qrhs"]["frozen"] - w1_kernel = self.variables["aqt"]["AqtEinsum_1"]["AqtDotGeneral_0"]["qrhs"]["frozen"] - wo_kernel = self.variables["aqt"]["AqtEinsum_2"]["AqtDotGeneral_0"]["qrhs"]["frozen"] - - w0_kernel = max_utils.unbox_logicallypartioned(w0_kernel) - w1_kernel = max_utils.unbox_logicallypartioned(w1_kernel) - wo_kernel = max_utils.unbox_logicallypartioned(wo_kernel) - return w0_kernel, w1_kernel, wo_kernel - def __call__( self, inputs: jax.Array, @@ -2623,18 +2557,6 @@ def __call__( inputs, gate_logits, wo_kernel, w0_kernel=w0_kernel, w1_kernel=w1_kernel, fused_kernel=fused_kernel ) elif cfg.sparse_matmul: - if quantizations.in_serve_mode(self.quant): - w0_kernel, w1_kernel, wo_kernel = self.retrieve_quantized_weight( - inputs, - gate_logits, - pre_bias_logits, - w0_kernel, - w1_kernel, - wo_kernel, - w0_bias, - w1_bias, - wo_bias, - ) output, lb_loss, bias_updates = self.sparse_matmul( inputs, gate_logits, pre_bias_logits, w0_kernel, w1_kernel, wo_kernel, w0_bias, w1_bias, wo_bias, input_ids ) @@ -2657,7 +2579,7 @@ def __init__( rngs: nnx.Rngs, weight_dtype: ctypes.DType = jnp.float32, dtype: ctypes.DType = jnp.float32, - quant: Optional[quantizations.AqtQuantization] = None, + quant: Optional[quantizations.Quantization] = None, is_hash_routing: bool = False, ): """Initializes the RoutedAndSharedMoE module. @@ -2764,7 +2686,7 @@ def get_gate_logit( kernel_axes: Tuple[Optional[str], ...] = (), use_bias: bool = False, score_func: str = "", - quant: Optional[quantizations.AqtQuantization] = None, + quant: Optional[quantizations.Quantization] = None, matmul_precision: str = "default", name: Optional[str] = None, ): @@ -2804,7 +2726,7 @@ def get_routed_moe( intermediate_dim: int = 2048, weight_dtype: ctypes.DType = jnp.float32, dtype: ctypes.DType = jnp.float32, - quant: Optional[quantizations.AqtQuantization] = None, + quant: Optional[quantizations.Quantization] = None, name: Optional[str] = None, ): """Creates a RoutedMoE Linen module.""" @@ -2835,7 +2757,7 @@ def get_routed_and_shared_moe( kernel_axes: Tuple[Optional[str], ...], weight_dtype: ctypes.DType = jnp.float32, dtype: ctypes.DType = jnp.float32, - quant: Optional[quantizations.AqtQuantization] = None, + quant: Optional[quantizations.Quantization] = None, name: Optional[str] = None, is_hash_routing: bool = False, ): diff --git a/src/maxtext/layers/nnx_decoders.py b/src/maxtext/layers/nnx_decoders.py index 5d7262e5a4..0e580101ec 100644 --- a/src/maxtext/layers/nnx_decoders.py +++ b/src/maxtext/layers/nnx_decoders.py @@ -42,7 +42,7 @@ from maxtext.layers.attentions import Attention from maxtext.layers.embeddings import Embed, PositionalEmbedding, attend_on_embedding from maxtext.layers.normalizations import RMSNorm -from maxtext.layers.quantizations import AqtQuantization as Quant +from maxtext.layers.quantizations import Quantization as Quant from maxtext.models import ( deepseek, deepseek_batchsplit, diff --git a/src/maxtext/layers/nnx_wrappers.py b/src/maxtext/layers/nnx_wrappers.py index 07faeefda2..c08011c062 100644 --- a/src/maxtext/layers/nnx_wrappers.py +++ b/src/maxtext/layers/nnx_wrappers.py @@ -14,6 +14,7 @@ """NNX <> Linen interoperability.""" +import types from functools import partial import typing as tp from typing import Any @@ -554,9 +555,11 @@ def __getattr__(self, name: str): return self.kwargs[name] maybe_method = getattr(self.nnx_class, name, None) if callable(maybe_method): - method = partial(self.__call__, nnx_method=maybe_method) - method.__self__ = self - return method + + def unbound_func(instance, *args, **kwargs): + return instance.__call__(*args, nnx_method=maybe_method, **kwargs) + + return types.MethodType(unbound_func, self) return super().__getattribute__(name) def _update_variables(self, module): diff --git a/src/maxtext/layers/quantizations.py b/src/maxtext/layers/quantizations.py index 2c5f638f8c..e3c0e6e3f8 100644 --- a/src/maxtext/layers/quantizations.py +++ b/src/maxtext/layers/quantizations.py @@ -15,17 +15,10 @@ """Quantization library.""" import functools -import json import qwix.pallas as qpl -import re -from typing import Tuple, Sequence, Callable +from typing import Tuple, Callable from dataclasses import dataclass -from aqt.jax.v2 import config as aqt_config -from aqt.jax.v2 import aqt_tensor -from aqt.jax.v2.flax import aqt_flax -from aqt.jax.v2 import tiled_dot_general -from aqt.jax.v2 import calibration import qwix from qwix._src.core import dot_general_qt @@ -33,7 +26,6 @@ import jax import jax.numpy as jnp -from jax.tree_util import tree_flatten_with_path, tree_unflatten from flax.linen import fp8_ops from flax.linen import initializers as flax_initializers @@ -48,7 +40,7 @@ from maxtext.layers import nnx_wrappers from maxtext.common.common_types import DType, Config -from maxtext.inference.kvcache import KVQuant + # Params used to define mixed precision quantization configs DEFAULT = "__default__" # default config @@ -70,141 +62,6 @@ def einsum(self, dtype: DType = jnp.float32): """Placeholder for einsum implementation in subclasses.""" -def _tiling_fn(lhs, rhs, dimension_numbers, tile_size): - """apply tiling function""" - del lhs, rhs - - (lhs_ca, rhs_ca), _ = dimension_numbers - ret = tiled_dot_general.Cfg( - lhs=tiled_dot_general.TensorTiling(contraction_axes=[], remaining_axes=[]), - rhs=tiled_dot_general.TensorTiling(contraction_axes=[], remaining_axes=[]), - ) - - for lhs_idx, rhs_idx in zip(lhs_ca, rhs_ca): - ret.lhs.contraction_axes.append(tiled_dot_general.AxisTiling(axis=lhs_idx, tile_size=tile_size, tile_count=None)) - ret.rhs.contraction_axes.append(tiled_dot_general.AxisTiling(axis=rhs_idx, tile_size=tile_size, tile_count=None)) - - return ret - - -def _rhs_axis_metadata_wrapper( - x: jnp.ndarray, - tile_map, - no_sharding_axis: Sequence[int], - mesh_axes: Tuple[str, ...], - is_tiled: bool, - replicate_scale: bool = False, -): - """right-hand-side axis metadata wrapper""" - if replicate_scale: - # Temporarily using the shape to identify the scale. - # TODO: remove the replication once the 2d sharding quantization - # works as expected. - if len(x.shape) == 1: - return nn.with_logical_partitioning((lambda: x), tuple(None for _ in mesh_axes))() - - mesh_axes = list(mesh_axes) - if is_tiled: - # tile_map is a mapping between original rank and a list of new, tiled rank. - if len(mesh_axes) < len(tile_map): - mesh_axes = [None] * (len(tile_map) - len(mesh_axes)) + mesh_axes - new_mesh_axes = [None] * len(x.shape) - for orig_rank, new_rank in tile_map.items(): - assert new_rank - assert len(new_rank) <= 2 - new_mesh_axes[new_rank[-1]] = mesh_axes[orig_rank] - mesh_axes = new_mesh_axes - - if mesh_axes is not None and len(mesh_axes) > 0: - for no_shard_idx in no_sharding_axis: - if no_shard_idx < len(mesh_axes): - mesh_axes[no_shard_idx] = None - - return nn.with_logical_partitioning((lambda: x), mesh_axes)() - - -@dataclass -class AqtQuantization: - """Configures AQT quantization github.com/google/aqt.""" - - quant_dg: aqt_config.DotGeneral - quant_mode: aqt_flax.QuantMode = aqt_flax.QuantMode.TRAIN - replicate_scale: bool = False - - def _get_mixed_precision_cfg(self): - """get configuration for mixed precision""" - quant_dg = None - is_tiled = False - tiling_fn = None - # pylint: disable=protected-access - module_path = "/".join(nn.module._context.module_stack[-1].path) - tile_size = -1 - for layer_name_re, layer_quant_dg in self.quant_dg.items(): - if re.fullmatch(layer_name_re, module_path): - quant_dg, tile_size = layer_quant_dg - if quant_dg is None: - quant_dg, tile_size = self.quant_dg[DEFAULT] - if tile_size != -1: - is_tiled = True - tiling_fn = functools.partial(_tiling_fn, tile_size=tile_size) - return quant_dg, is_tiled, tiling_fn - - def _get_rhs_axis_metadata_wrapper( - self, mesh_axes: Tuple[str, ...] = (), is_tiled: bool = False, replicate_scale: bool = False - ): - if self.quant_mode == aqt_flax.QuantMode.CONVERT: - return None - return functools.partial( - _rhs_axis_metadata_wrapper, mesh_axes=mesh_axes, is_tiled=is_tiled, replicate_scale=replicate_scale - ) - - def dot_general_cls(self, mesh_axes: Tuple[str, ...] = ()): - """Returns dot_general configured with aqt params.""" - if isinstance(self.quant_dg, dict): - quant_dg, is_tiled, tiling_fn = self._get_mixed_precision_cfg() - else: - quant_dg, is_tiled, tiling_fn = self.quant_dg, False, None - rhs_axis_metadata_wrapper = self._get_rhs_axis_metadata_wrapper( - mesh_axes, is_tiled, replicate_scale=self.replicate_scale - ) - # module_path = "/".join(nn.module._context.module_stack[-1].path) - # print(f"quant_dg: {quant_dg}, is_tiled: {is_tiled}, module_path: {module_path}") - aqt_dg_cls = functools.partial( - aqt_flax.AqtDotGeneral, - quant_dg, - rhs_quant_mode=self.quant_mode, - lhs_freeze_mode=aqt_flax.FreezerMode.NONE, - rhs_freeze_mode=aqt_flax.FreezerMode.CALIBRATION_AND_VALUE, - rhs_axis_metadata_wrapper=rhs_axis_metadata_wrapper, - use_legacy_freezer=False, - tiling_fn=tiling_fn, - ) - return aqt_dg_cls - - def einsum(self, mesh_axes: Tuple[str, ...] = ()): - """Returns einsum configured with aqt params.""" - if isinstance(self.quant_dg, dict): - quant_dg, is_tiled, tiling_fn = self._get_mixed_precision_cfg() - else: - quant_dg, is_tiled, tiling_fn = self.quant_dg, False, None - - rhs_axis_metadata_wrapper = self._get_rhs_axis_metadata_wrapper( - mesh_axes, is_tiled, replicate_scale=self.replicate_scale - ) - aqt_einsum = functools.partial( - aqt_flax.AqtEinsum( - cfg=quant_dg, - rhs_quant_mode=self.quant_mode, - lhs_freeze_mode=aqt_flax.FreezerMode.NONE, - rhs_freeze_mode=aqt_flax.FreezerMode.CALIBRATION_AND_VALUE, - rhs_axis_metadata_wrapper=rhs_axis_metadata_wrapper, - use_legacy_freezer=False, - tiling_fn=tiling_fn, - ) - ) - return aqt_einsum - - @dataclass class QwixQuantization: """Configures Qwix quantization github.com/google/qwix, for training only.""" @@ -232,7 +89,7 @@ def dot_general_cls(self, mesh_axes: Tuple[str, ...] = ()): """Returns Qwix dot_general.""" return functools.partial(QwixDotGeneral, config=self._get_fp8_full_qwix_config()) - def einsum(self, mesh_axes: Tuple[str, ...] = ()): + def einsum(self, mesh_axes: Tuple[str, ...] = (), **kwargs): """Returns Qwix einsum.""" return QwixEinsum(config=self._get_fp8_full_qwix_config()) @@ -385,261 +242,22 @@ def dot_general_cls(self, mesh_axes: Tuple[str, ...] = ()): return nn.NANOOFp8DotGeneralOp -def _get_int8_quant_config(config): - drhs_bits = None - drhs_accumulator_dtype = None - drhs_local_aqt = None - if config.quantization_local_shard_count != 0: - drhs_bits = 8 - drhs_accumulator_dtype = jnp.int32 - drhs_local_aqt = aqt_config.LocalAqt(contraction_axis_shard_count=config.quantization_local_shard_count) - return aqt_config.config_v3( - fwd_bits=8, - dlhs_bits=8, - drhs_bits=drhs_bits, - rng_type="jax.uniform", - dlhs_local_aqt=None, - drhs_local_aqt=drhs_local_aqt, - fwd_accumulator_dtype=jnp.int32, - dlhs_accumulator_dtype=jnp.int32, - drhs_accumulator_dtype=drhs_accumulator_dtype, - ) - - -@dataclass(frozen=True) -class ConstantBoundConfig: - fwd_lhs_bound: float | None = None - fwd_rhs_bound: float | None = None - dlhs_lhs_bound: float | None = None - dlhs_rhs_bound: float | None = None - drhs_lhs_bound: float | None = None - drhs_rhs_bound: float | None = None - - -def _build_const_scale_config( - aqt_dg: aqt_config.DotGeneral, - cst_bound_config: ConstantBoundConfig, -) -> aqt_config.DotGeneral: - """Build a constant scale config for AQT dot general. - - Args: - aqt_dg: The AQT dot general config. - cst_bound_config: The constant bound config. - - Returns: - The AQT dot general config with constant scale config. - """ - if cst_bound_config.fwd_lhs_bound is not None: - aqt_dg.fwd.dg_quantizer.lhs.calibration = functools.partial( - calibration.ConstantCalibration, bound=cst_bound_config.fwd_lhs_bound - ) - if cst_bound_config.fwd_rhs_bound is not None: - aqt_dg.fwd.dg_quantizer.rhs.calibration = functools.partial( - calibration.ConstantCalibration, bound=cst_bound_config.fwd_rhs_bound - ) - if cst_bound_config.dlhs_lhs_bound: - aqt_dg.dlhs.dg_quantizer.lhs.calibration = functools.partial( - calibration.ConstantCalibration, bound=cst_bound_config.dlhs_lhs_bound - ) - - if cst_bound_config.dlhs_rhs_bound is not None: - aqt_dg.dlhs.dg_quantizer.rhs.calibration = functools.partial( - calibration.ConstantCalibration, bound=cst_bound_config.dlhs_rhs_bound - ) - - if cst_bound_config.drhs_lhs_bound is not None: - aqt_dg.drhs.dg_quantizer.lhs.calibration = functools.partial( - calibration.ConstantCalibration, bound=cst_bound_config.drhs_lhs_bound - ) - - if cst_bound_config.drhs_rhs_bound is not None: - aqt_dg.drhs.dg_quantizer.rhs.calibration = functools.partial( - calibration.ConstantCalibration, bound=cst_bound_config.drhs_rhs_bound - ) - - return aqt_dg - - -@dataclass -class PerTensorScales: - fwd_lhs: bool = False - fwd_rhs: bool = False - dlhs_lhs: bool = False - dlhs_rhs: bool = False - drhs_lhs: bool = False - drhs_rhs: bool = False - - -def _build_per_tensor_config( - aqt_dg: aqt_config.DotGeneral, - per_tensor_scales: PerTensorScales, -) -> aqt_config.DotGeneral: - """Build a per tensor config for AQT dot general. - - Args: - aqt_dg: The AQT dot general config. - per_tensor_scales: The per tensor scales config. - - Returns: - The AQT dot general config with per tensor config. - """ - if per_tensor_scales.fwd_lhs: - aqt_dg.fwd.dg_quantizer.lhs.calib_shared_axes = "per_tensor" - if per_tensor_scales.fwd_rhs: - aqt_dg.fwd.dg_quantizer.rhs.calib_shared_axes = "per_tensor" - if per_tensor_scales.dlhs_lhs: - aqt_dg.dlhs.dg_quantizer.lhs.calib_shared_axes = "per_tensor" - if per_tensor_scales.dlhs_rhs: - aqt_dg.dlhs.dg_quantizer.rhs.calib_shared_axes = "per_tensor" - if per_tensor_scales.drhs_lhs: - aqt_dg.drhs.dg_quantizer.lhs.calib_shared_axes = "per_tensor" - if per_tensor_scales.drhs_rhs: - aqt_dg.drhs.dg_quantizer.rhs.calib_shared_axes = "per_tensor" - return aqt_dg - - -# fp8 training recipe of dynamic scaling with configurable constant_bound_config for static scaling option -def _get_aqt_fp8_default_config(config): - """Get aqt for 8-bit floating point quantization configuration.""" - aqt_dg = aqt_config.config_v4( - fwd_bits="e4m3", - dlhs_bits="e5m2", - drhs_bits="e5m2", - use_dummy_static_bound=False, - fwd_accumulator_dtype=jnp.bfloat16, - dlhs_accumulator_dtype=jnp.bfloat16, - drhs_accumulator_dtype=jnp.bfloat16, - dlhs_use_fwd_quant=False, - drhs_use_fwd_quant=False, - ) - constant_bound_config = None - - if len(config.constant_bound_config) == 6: - ( - fwd_lhs_bound, - fwd_rhs_bound, - dlhs_lhs_bound, - dlhs_rhs_bound, - drhs_lhs_bound, - drhs_rhs_bound, - ) = config.constant_bound_config - constant_bound_config = ConstantBoundConfig( - fwd_lhs_bound=fwd_lhs_bound, - fwd_rhs_bound=fwd_rhs_bound, - dlhs_lhs_bound=dlhs_lhs_bound, - dlhs_rhs_bound=dlhs_rhs_bound, - drhs_lhs_bound=drhs_lhs_bound, - drhs_rhs_bound=drhs_rhs_bound, - ) - aqt_dg = _build_const_scale_config(aqt_dg, constant_bound_config) - - aqt_config.set_stochastic_rounding( - aqt_dg, - vjp_lhs_stochastic_rounding=False, - vjp_rhs_stochastic_rounding=False, - implementation="jax.uniform", - ) - - per_tensor_scales = PerTensorScales( - fwd_lhs=True, - fwd_rhs=True, - dlhs_lhs=True, - dlhs_rhs=True, - drhs_lhs=True, - drhs_rhs=True, - ) - return _build_per_tensor_config(aqt_dg, per_tensor_scales) - - -def _get_aqt_fp8_quant_config(config): - """get aqt for 8-bit floating point quantization configuration""" - cfg = aqt_config.config_v4(fwd_bits="e4m3", dlhs_bits=None, drhs_bits=None, fwd_accumulator_dtype=jnp.bfloat16) - return cfg - - -def _dot_general_make(quant_cfg): - """Create quantization configs for input matrices to a matmul""" - lhs_bits = quant_cfg[_A_BITS] - lhs_scale = quant_cfg[_A_SCALE] - rhs_bits = quant_cfg[_W_BITS] - rhs_scale = quant_cfg[_W_SCALE] - aqt_dg = aqt_config.dot_general_make(lhs_bits=lhs_bits, rhs_bits=rhs_bits) - if lhs_scale < 1.0: - aqt_dg.fwd.dg_quantizer.lhs.calibration = functools.partial(calibration.AbsMaxCalibration, scale=lhs_scale) - if rhs_scale < 1.0: - aqt_dg.fwd.dg_quantizer.rhs.calibration = functools.partial(calibration.AbsMaxCalibration, scale=rhs_scale) - return aqt_dg - - -def _get_default_mp_config(default=None): - default_config = {_W_BITS: None, _A_BITS: None, _W_SCALE: 1.0, _A_SCALE: 1.0, _TILE_SIZE: -1} - if default: - default_config.update(default) - return default_config - - -def _get_mixed_precision_quant_config(mixed_precision_config): - """Set quantization params based on user configuration.""" - ret_config = {} - default_mp_config = _get_default_mp_config(default=mixed_precision_config.get(DEFAULT, None)) - for layer_name_re, layer_quantization_config in mixed_precision_config.items(): - # Make a copy of default_mp_config to avoid updating original dict - quant_config = default_mp_config.copy() - # print(f"Mixed precision config: processing - # {layer_name_re} - {layer_quantization_config}, default config - {quant_config}") - if layer_name_re != DEFAULT: - for k in quant_config: - quant_config[k] = layer_quantization_config.get(k, default_mp_config[k]) - ret_config[layer_name_re] = [_dot_general_make(quant_config), quant_config["tile_size"]] - return ret_config - - def _get_quant_config(config): """Set quantization params based on user configuration.""" if not config.quantization or config.quantization == "": return None - if config.quantization == "int8": - return _get_int8_quant_config(config) - if config.quantization == "intmp": - assert config.quant_cfg_path, "Must specify quant_cfg for mixed precision quantization" - with open(config.quant_cfg_path, "rt", encoding="utf8") as config_file: - mixed_precision_config = json.load(config_file) - return _get_mixed_precision_quant_config(mixed_precision_config) if config.quantization == "fp8": return "fp8" if config.quantization == "nanoo_fp8": return "nanoo_fp8" - if config.quantization == "aqt_fp8": - return _get_aqt_fp8_quant_config(config) - if config.quantization == "aqt_fp8_full": - return _get_aqt_fp8_default_config(config) if config.quantization.startswith("te_"): return config.quantization - - raise ValueError(f"Invalid value configured for quantization {config.quantization}.") - - -def in_convert_mode(quant): - return quant and (quant.quant_mode == aqt_flax.QuantMode.CONVERT) - - -def in_serve_mode(quant): - return quant and (quant.quant_mode == aqt_flax.QuantMode.SERVE) - - -def get_quant_mode(quant_mode_str: str = "train"): - """Set quant mode.""" - if quant_mode_str == "train": - return aqt_flax.QuantMode.TRAIN - elif quant_mode_str == "serve": - return aqt_flax.QuantMode.SERVE - elif quant_mode_str == "convert": - return aqt_flax.QuantMode.CONVERT - raise ValueError(f"Invalid quantization mode {quant_mode_str}.") + return None def configure_quantization(config: Config, quant_mode_str: str = "train"): """Configure quantization based on user config and quant mode.""" + del quant_mode_str # Unused since AQT is removed if config.use_batch_split_schedule and config.quantization: # The older version of batch-split that fully uses qwix quantization. if config.quantization == "fp8_full" and not config.use_manual_quantization: @@ -661,61 +279,15 @@ def configure_quantization(config: Config, quant_mode_str: str = "train"): return NANOOFp8Quantization() elif isinstance(quant_cfg, str) and quant_cfg.startswith("te_"): return TransformerEngineQuantization(config) - quant_mode = get_quant_mode(quant_mode_str) - replicate_scale = config.replicate_quant_scale if config.replicate_quant_scale else False - return AqtQuantization(quant_dg=quant_cfg, quant_mode=quant_mode, replicate_scale=replicate_scale) return None -def match_aqt_and_unquantized_param(aqt_params, params): - """match aqt and unquantized params""" - aqt_param_flat, aqt_tree_def = jax.tree_util.tree_flatten_with_path( - aqt_params, is_leaf=lambda x: isinstance(x, aqt_tensor.QTensor) - ) - param_tree_flat, _ = jax.tree_util.tree_flatten_with_path(params) - aqt_paths = [] - # Original path of quantized AQT param path. - param_paths = [] - - for aqt_k, _ in aqt_param_flat: - index = None - for index, (k, _) in enumerate(param_tree_flat): - path_depth = len(k) - # every quantized parameter has AQT.. as the leaf node - # AqtDotGeneral and AqtEinsum replace leaf node. - # Therefore, leaf node should be ignored for path matching - # Note: Aqt only operates on kernels so don't pop bias parameters. - # Ref: https://github.com/AI-Hypercomputer/maxtext/compare/main...quantize_r1 - if k[: path_depth - 1] == aqt_k[: path_depth - 1] and k[-1].key != "bias": - aqt_paths.append(aqt_k) - param_paths.append(k) - break - assert index is not None - # since the parameter is already added, we can delete it. - param_tree_flat.pop(index) - return jax.tree_util.tree_unflatten(aqt_tree_def, param_paths) - - -def _get_aqt_key_paths(aqt_vars, params): - """Generate a list of paths which have aqt state""" - aqt_to_unquantized_key_path = match_aqt_and_unquantized_param(aqt_vars, params) - aqt_key_paths, _ = jax.tree_util.tree_flatten(aqt_to_unquantized_key_path, is_leaf=lambda x: isinstance(x, tuple)) - return list(aqt_key_paths) - - -def remove_quantized_params(params, aqt_vars): - """Remove param values with aqt tensors to Null to optimize memory.""" - quantized_param_paths = _get_aqt_key_paths(aqt_vars, params) - tree_flat, tree_struct = tree_flatten_with_path(params) - for i, (k, v) in enumerate(tree_flat): - if k in quantized_param_paths: - v = {} - tree_flat[i] = v - return tree_unflatten(tree_struct, tree_flat) - - def configure_kv_quant(config): - return None if not config.quantize_kvcache else KVQuant(config) + if config.quantize_kvcache: + raise ValueError( + "KV cache quantization (quantize_kvcache=True) is no longer supported " + "because Accurate Quantized Training (AQT) has been deprecated and removed from MaxText." + ) def _apply_linen_module_in_nnx(linen_module_cls, op_id, *args, **kwargs): @@ -846,7 +418,7 @@ def get_qt_provider(config): def maybe_quantize_model(model, config): """Quantize the model if quantization is enabled.""" # Batch split is not using Qwix's interception feature but manual plumbing - if config.use_qwix_quantization and not config.use_batch_split_schedule: + if config.use_qwix_quantization and not config.use_batch_split_schedule and not config.pure_nnx: quantization_provider = get_qt_provider(config) if quantization_provider: if config.pure_nnx: diff --git a/src/maxtext/models/deepseek.py b/src/maxtext/models/deepseek.py index d3a72b31bf..775865ccc9 100644 --- a/src/maxtext/models/deepseek.py +++ b/src/maxtext/models/deepseek.py @@ -63,7 +63,7 @@ def __init__( model_mode: str, mesh: Mesh, rngs: nnx.Rngs, - quant: Optional[quantizations.AqtQuantization] = None, + quant: Optional[quantizations.Quantization] = None, layer_idx: int = -1, ) -> None: self.config = config @@ -317,7 +317,7 @@ def __init__( model_mode: str, mesh: Mesh, rngs: nnx.Rngs, - quant: Optional[quantizations.AqtQuantization] = None, + quant: Optional[quantizations.Quantization] = None, layer_idx: int = -1, ) -> None: super().__init__(config, model_mode, mesh, rngs, quant, layer_idx) @@ -406,7 +406,7 @@ def __init__( model_mode: str, mesh: Mesh, rngs: nnx.Rngs, - quant: Optional[quantizations.AqtQuantization] = None, + quant: Optional[quantizations.Quantization] = None, layer_idx: int = -1, ) -> None: super().__init__(config, model_mode, mesh, rngs, quant, layer_idx) diff --git a/src/maxtext/models/deepseek4.py b/src/maxtext/models/deepseek4.py index c259db8c66..25eca866f9 100644 --- a/src/maxtext/models/deepseek4.py +++ b/src/maxtext/models/deepseek4.py @@ -56,7 +56,7 @@ def __init__( model_mode: str, mesh: Mesh, rngs: nnx.Rngs, - quant: Optional[quantizations.AqtQuantization] = None, + quant: Optional[quantizations.Quantization] = None, layer_idx: int = -1, compress_ratio: Optional[int] = None, is_hash_routing: Optional[bool] = None, @@ -184,7 +184,7 @@ def __init__( mesh: Mesh, model_mode: str, rngs: nnx.Rngs, - quant: None | quantizations.AqtQuantization = None, + quant: None | quantizations.Quantization = None, ): self.config = config self.mesh = mesh diff --git a/src/maxtext/models/gemma.py b/src/maxtext/models/gemma.py index 84f4f6817d..3f9ff0682e 100644 --- a/src/maxtext/models/gemma.py +++ b/src/maxtext/models/gemma.py @@ -29,7 +29,7 @@ from maxtext.layers.attentions import Attention from maxtext.layers.linears import Dropout, MlpBlock from maxtext.layers.normalizations import RMSNorm -from maxtext.layers.quantizations import AqtQuantization as Quant +from maxtext.layers.quantizations import Quantization as Quant from maxtext.utils import max_utils diff --git a/src/maxtext/models/gemma2.py b/src/maxtext/models/gemma2.py index a7315763eb..e3ff18fc32 100644 --- a/src/maxtext/models/gemma2.py +++ b/src/maxtext/models/gemma2.py @@ -30,7 +30,7 @@ from maxtext.layers.attentions import Attention from maxtext.layers.linears import Dropout, MlpBlock from maxtext.layers.normalizations import RMSNorm -from maxtext.layers.quantizations import AqtQuantization as Quant +from maxtext.layers.quantizations import Quantization as Quant from maxtext.utils import max_utils diff --git a/src/maxtext/models/gemma3.py b/src/maxtext/models/gemma3.py index 344d98ed88..7e98f6c3ae 100644 --- a/src/maxtext/models/gemma3.py +++ b/src/maxtext/models/gemma3.py @@ -29,7 +29,7 @@ from maxtext.layers.attentions import Attention from maxtext.layers.linears import DenseGeneral, MlpBlock, Dropout from maxtext.layers.normalizations import RMSNorm -from maxtext.layers.quantizations import AqtQuantization as Quant +from maxtext.layers.quantizations import Quantization as Quant from maxtext.layers.initializers import variable_to_logically_partitioned from maxtext.utils import max_utils diff --git a/src/maxtext/models/gemma4.py b/src/maxtext/models/gemma4.py index 626d2ff54c..3ac037a983 100644 --- a/src/maxtext/models/gemma4.py +++ b/src/maxtext/models/gemma4.py @@ -33,7 +33,7 @@ import jax.sharding from maxtext.layers.normalizations import RMSNorm -from maxtext.layers.quantizations import AqtQuantization as Quant +from maxtext.layers.quantizations import Quantization as Quant from maxtext.utils import max_utils diff --git a/src/maxtext/models/gemma4_small.py b/src/maxtext/models/gemma4_small.py index ca33470bf2..1ee945420e 100644 --- a/src/maxtext/models/gemma4_small.py +++ b/src/maxtext/models/gemma4_small.py @@ -29,7 +29,7 @@ from maxtext.layers.attentions import Attention from maxtext.layers.linears import DenseGeneral, MlpBlock from maxtext.layers.normalizations import RMSNorm -from maxtext.layers.quantizations import AqtQuantization as Quant +from maxtext.layers.quantizations import Quantization as Quant from maxtext.utils import max_utils diff --git a/src/maxtext/models/gpt3.py b/src/maxtext/models/gpt3.py index a6b08d8b24..df16ad1ecc 100644 --- a/src/maxtext/models/gpt3.py +++ b/src/maxtext/models/gpt3.py @@ -35,7 +35,7 @@ from maxtext.layers import linears from maxtext.layers.attentions import AttentionOp, KVQuant from maxtext.layers.initializers import Initializer, NdInitializer, nd_dense_init -from maxtext.layers.quantizations import AqtQuantization as Quant +from maxtext.layers.quantizations import Quantization as Quant from maxtext.utils import max_logging from maxtext.utils import max_utils diff --git a/src/maxtext/models/gpt_oss.py b/src/maxtext/models/gpt_oss.py index e854a75556..7b03cba502 100644 --- a/src/maxtext/models/gpt_oss.py +++ b/src/maxtext/models/gpt_oss.py @@ -35,7 +35,7 @@ from maxtext.layers import quantizations from maxtext.layers.attentions import Attention from maxtext.layers.normalizations import RMSNorm -from maxtext.layers.quantizations import AqtQuantization as Quant +from maxtext.layers.quantizations import Quantization as Quant from maxtext.utils import max_utils # ----------------------------------------- diff --git a/src/maxtext/models/llama2.py b/src/maxtext/models/llama2.py index 0c3e0cca7c..260934acc1 100644 --- a/src/maxtext/models/llama2.py +++ b/src/maxtext/models/llama2.py @@ -30,7 +30,7 @@ from maxtext.layers.attentions import Attention from maxtext.layers.linears import Dropout, MlpBlock from maxtext.layers.normalizations import RMSNorm -from maxtext.layers.quantizations import AqtQuantization as Quant +from maxtext.layers.quantizations import Quantization as Quant from maxtext.utils import max_utils from maxtext.utils.sharding import create_sharding, maybe_shard_with_logical from maxtext.layers.learn_to_init_layer import apply_lti_modification diff --git a/src/maxtext/models/llama4.py b/src/maxtext/models/llama4.py index 66dea4295c..806dc11344 100644 --- a/src/maxtext/models/llama4.py +++ b/src/maxtext/models/llama4.py @@ -35,7 +35,7 @@ from maxtext.layers.linears import MlpBlock from maxtext.layers.moe import RoutedAndSharedMoE from maxtext.layers.normalizations import RMSNorm -from maxtext.layers.quantizations import AqtQuantization as Quant +from maxtext.layers.quantizations import Quantization as Quant from maxtext.utils import max_utils #### Multi modal model implementation diff --git a/src/maxtext/models/mistral.py b/src/maxtext/models/mistral.py index 49a74c95db..8c4747e166 100644 --- a/src/maxtext/models/mistral.py +++ b/src/maxtext/models/mistral.py @@ -28,7 +28,7 @@ from maxtext.layers.attentions import Attention from maxtext.layers.linears import Dropout, MlpBlock from maxtext.layers.normalizations import RMSNorm -from maxtext.layers.quantizations import AqtQuantization as Quant +from maxtext.layers.quantizations import Quantization as Quant from maxtext.utils import max_utils # ----------------------------------------- diff --git a/src/maxtext/models/mixtral.py b/src/maxtext/models/mixtral.py index faf69273c6..38a9d5df41 100644 --- a/src/maxtext/models/mixtral.py +++ b/src/maxtext/models/mixtral.py @@ -29,7 +29,7 @@ from maxtext.layers.attentions import Attention from maxtext.layers.linears import Dropout from maxtext.layers.normalizations import RMSNorm -from maxtext.layers.quantizations import AqtQuantization as Quant +from maxtext.layers.quantizations import Quantization as Quant from maxtext.utils import max_utils # ----------------------------------------- diff --git a/src/maxtext/models/models.py b/src/maxtext/models/models.py index ac908c0f96..4785fb1ee6 100644 --- a/src/maxtext/models/models.py +++ b/src/maxtext/models/models.py @@ -33,7 +33,7 @@ from maxtext.layers.embeddings import Embed, embed_as_linen from maxtext.layers.encoders import AudioEncoder, VisionEncoder, audio_encoder_as_linen, vision_encoder_as_linen from maxtext.layers.multi_token_prediction import MultiTokenPredictionBlock, multi_token_prediction_block_as_linen -from maxtext.layers.quantizations import AqtQuantization as Quant +from maxtext.layers.quantizations import Quantization as Quant from maxtext.multimodal import processor as mm_processor from maxtext.utils import max_utils @@ -533,6 +533,8 @@ def __call__( mutable_collections.append("intermediates") if self.config.load_balance_loss_weight > 0.0 and "intermediates" not in mutable_collections: mutable_collections.append("intermediates") + if self.config.num_vocab_tiling > 1 and "intermediates" not in mutable_collections: + mutable_collections.append("intermediates") if self.config.pure_nnx_decoder: logits, hidden_state, kv_caches = self.decoder( diff --git a/src/maxtext/models/olmo3.py b/src/maxtext/models/olmo3.py index fe8a4e489e..3cebd6170f 100644 --- a/src/maxtext/models/olmo3.py +++ b/src/maxtext/models/olmo3.py @@ -36,7 +36,7 @@ from maxtext.layers.attentions import Attention from maxtext.layers.linears import MlpBlock from maxtext.layers.normalizations import RMSNorm -from maxtext.layers.quantizations import AqtQuantization as Quant +from maxtext.layers.quantizations import Quantization as Quant from maxtext.utils import max_utils diff --git a/src/maxtext/models/qwen2.py b/src/maxtext/models/qwen2.py index 69555c176d..995e7d4c54 100644 --- a/src/maxtext/models/qwen2.py +++ b/src/maxtext/models/qwen2.py @@ -30,7 +30,7 @@ from maxtext.layers import nnx_wrappers from maxtext.layers import quantizations from maxtext.layers.normalizations import RMSNorm -from maxtext.layers.quantizations import AqtQuantization as Quant +from maxtext.layers.quantizations import Quantization as Quant from maxtext.layers.attentions import Attention from maxtext.layers.linears import MlpBlock from maxtext.utils import max_utils diff --git a/src/maxtext/models/qwen3.py b/src/maxtext/models/qwen3.py index 3b5f08f0c1..7eae6cd6b5 100644 --- a/src/maxtext/models/qwen3.py +++ b/src/maxtext/models/qwen3.py @@ -41,7 +41,7 @@ from maxtext.layers import quantizations from maxtext.layers.embeddings import Qwen3OmniMoeVisionPosEmbedInterpolate, PositionalEmbedding from maxtext.layers.normalizations import RMSNorm, l2norm, Qwen3NextRMSNorm, Qwen3NextRMSNormGated -from maxtext.layers.quantizations import AqtQuantization as Quant +from maxtext.layers.quantizations import Quantization as Quant from maxtext.layers.attentions import Attention from maxtext.layers.linears import DenseGeneral, MlpBlock from maxtext.layers.moe import RoutedMoE diff --git a/src/maxtext/models/qwen3_5.py b/src/maxtext/models/qwen3_5.py index de8bd1cf18..549d321713 100644 --- a/src/maxtext/models/qwen3_5.py +++ b/src/maxtext/models/qwen3_5.py @@ -28,7 +28,7 @@ from maxtext.layers import initializers as max_initializers from maxtext.layers import nnx_wrappers from maxtext.layers.normalizations import Qwen3NextRMSNorm -from maxtext.layers.quantizations import AqtQuantization as Quant +from maxtext.layers.quantizations import Quantization as Quant from maxtext.utils import max_utils from maxtext.models.qwen3 import ( diff --git a/src/maxtext/models/qwen3_custom.py b/src/maxtext/models/qwen3_custom.py index e0ca4bb512..a78beb6b44 100644 --- a/src/maxtext/models/qwen3_custom.py +++ b/src/maxtext/models/qwen3_custom.py @@ -25,7 +25,7 @@ from maxtext.layers import moe from maxtext.layers import nnx_wrappers from maxtext.layers import quantizations -from maxtext.layers.quantizations import AqtQuantization as Quant +from maxtext.layers.quantizations import Quantization as Quant from maxtext.layers.attentions import Attention from maxtext.layers.linears import DenseGeneral from maxtext.utils import max_utils diff --git a/src/maxtext/models/simple_layer.py b/src/maxtext/models/simple_layer.py index ac4eb915a8..1e72b38143 100644 --- a/src/maxtext/models/simple_layer.py +++ b/src/maxtext/models/simple_layer.py @@ -37,7 +37,7 @@ def __init__( mesh: Mesh, model_mode: str, rngs: nnx.Rngs, - quant: Optional[quantizations.AqtQuantization] = None, + quant: Optional[quantizations.Quantization] = None, ) -> None: self.config = config @@ -93,7 +93,7 @@ def __init__( mesh: Mesh, model_mode: str, rngs: nnx.Rngs, - quant: Optional[quantizations.AqtQuantization] = None, + quant: Optional[quantizations.Quantization] = None, ) -> None: self.config = config diff --git a/src/maxtext/utils/layerwise_quantization.py b/src/maxtext/utils/layerwise_quantization.py deleted file mode 100644 index 3279b59a6f..0000000000 --- a/src/maxtext/utils/layerwise_quantization.py +++ /dev/null @@ -1,470 +0,0 @@ -# Copyright 2023–2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# https://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - - -r"""Layerwise quantization for large models - -Provides a utility to load and quantize a checkpoint layer by layer. Currently, it supports DeepSeek-family models only. - -Example cmd: - -python3 -m maxtext.utils.layerwise_quantization src/maxtext/configs/base.yml \ - tokenizer_path=${TOKENIZER_PATH?} load_parameters_path=${LOAD_PARAMS_PATH?} \ - model_name=deepseek2-16b ici_fsdp_parallelism=1 ici_autoregressive_parallelism=1 \ - ici_tensor_parallelism=-1 scan_layers=false weight_dtype=bfloat16 per_device_batch_size=1 \ - attention=dot_product quantization=int8 async_checkpointing=false enable_single_controller=true \ - tokenizer_type=huggingface megablox=false sparse_matmul=false \ - save_quantized_params_path=${SAVE_PARAMS_PATH?} checkpoint_storage_use_ocdbt=False \ - checkpoint_storage_use_zarr3=False - -""" - -import functools -import os -from typing import Any, Sequence - -from absl import app -from aqt.jax.v2 import aqt_tensor -from flax import nnx -from flax.linen import partitioning as nn_partitioning -import jax -import jax.numpy as jnp -from maxtext.common import common_types -from maxtext.common import checkpointing -from maxtext.layers import quantizations -from maxtext.models import deepseek, models -from maxtext.utils import max_logging -from maxtext.utils import max_utils -from maxtext.utils import maxtext_utils -from maxtext.utils import maxtext_utils_nnx -from maxtext.utils import model_creation_utils -import orbax.checkpoint as ocp -from tqdm import tqdm -from maxtext.configs import pyconfig - -IGNORE = ocp.PLACEHOLDER -PRNGKeyType = Any -DictKey = jax.tree_util.DictKey - - -def get_original_path_key(aqt_k_tuple: tuple[DictKey, ...]) -> tuple[DictKey, ...] | None: - """ - Maps an AQT PyTree path (tuple of keys) to its corresponding original parameter path. - Only returns a path if it corresponds to a parameter to be removed. - """ - aqt_k = list(aqt_k_tuple) - str_path = jax.tree_util.keystr(aqt_k_tuple) - if "AqtEinsum_" in str_path: - return None - if "AqtDotGeneral_" not in str_path: - return None - aqt_module_index = -1 - for i, key in enumerate(aqt_k): - if isinstance(key, DictKey) and key.key.startswith("AqtDotGeneral_"): - aqt_module_index = i - break - if aqt_module_index == -1: - return None - if ( - len(aqt_k) > aqt_module_index + 2 - and isinstance(aqt_k[aqt_module_index + 1], DictKey) - and aqt_k[aqt_module_index + 1].key == "qrhs" - and isinstance(aqt_k[aqt_module_index + 2], DictKey) - and aqt_k[aqt_module_index + 2].key == "frozen" - ): - parent_path = tuple(aqt_k[:aqt_module_index]) - return parent_path + (DictKey("kernel"),) - return None - - -def get_quantized_param_paths(aqt_params: Any, params: Any) -> set[tuple[DictKey, ...]]: - """ - Identifies the set of paths in the original params tree that have been quantized. - """ - - def is_qtensor(x): - return isinstance(x, aqt_tensor.QTensor) - - aqt_param_flat, _ = jax.tree_util.tree_flatten_with_path(aqt_params, is_leaf=is_qtensor) - if not aqt_param_flat: - return set() - param_tree_flat_with_path, _ = jax.tree_util.tree_flatten_with_path(params) - params_path_set: set[tuple[DictKey, ...]] = {tuple(k) for k, _ in param_tree_flat_with_path} - original_param_paths_to_remove: set[tuple[DictKey, ...]] = set() - for aqt_k_tuple, _ in aqt_param_flat: - original_k_tuple = get_original_path_key(aqt_k_tuple) - if original_k_tuple is None: - continue - if original_k_tuple in params_path_set: - original_param_paths_to_remove.add(original_k_tuple) - continue - params_keys_str = {jax.tree_util.keystr(k) for k in params_path_set} - raise ValueError( - f"Mapped AQT path {jax.tree_util.keystr(aqt_k_tuple)} to {jax.tree_util.keystr(original_k_tuple)}," - f" but not found in params. Available: {params_keys_str}" - ) - return original_param_paths_to_remove - - -def remove_quantized_params(params: Any, aqt_vars: Any) -> Any: - """Replaces the values in the original params tree that are now quantized with empty dicts.""" - quantized_param_path_set = get_quantized_param_paths(aqt_vars, params) - if not quantized_param_path_set: - return params - - def _map_fn(path, value): - return {} if tuple(path) in quantized_param_path_set else value - - return jax.tree_util.tree_map_with_path(_map_fn, params) - - -# --- Function to restructure NNX-Run AQT tree to match PURE LINEN saved format --- -def insert_deepseekmoeblock_scope(aqt_layer_tree: dict[str, Any]) -> dict[str, Any]: - """ - Moves top-level AqtEinsum_* entries into the existing 'DeepSeekMoeBlock_0' - dict to match the pure Linen AQT structure. - """ - if not isinstance(aqt_layer_tree, dict): - return aqt_layer_tree - - new_tree = dict(aqt_layer_tree) # Start with a copy - - einsum_items = {key: new_tree.pop(key) for key in list(new_tree.keys()) if key.startswith("AqtEinsum_")} - - if einsum_items: - if "DeepSeekMoeBlock_0" not in new_tree: - # This case indicates the MoE block itself was missing, which is unexpected - max_logging.log("Error: 'DeepSeekMoeBlock_0' not found in AQT vars for MoE layer.") - new_tree["DeepSeekMoeBlock_0"] = {} - elif not isinstance(new_tree["DeepSeekMoeBlock_0"], dict): - max_logging.log(f"Error: 'DeepSeekMoeBlock_0' is not a dict, type: {type(new_tree['DeepSeekMoeBlock_0'])}") - new_tree["DeepSeekMoeBlock_0"] = {} - - # Merge einsum_items into the DeepSeekMoeBlock_0 dict - new_tree["DeepSeekMoeBlock_0"].update(einsum_items) - - return new_tree - - -class LayerwiseQuantization: - """ - Layerwise quantization for large models. - """ - - def __init__(self, config: Any, rng: PRNGKeyType): - self.config = config - self.rng = rng - - # The Linen path runs layer-by-layer (memory-efficient for big DeepSeek - # models) and is DeepSeek-specific because it relies on the per-layer - # `DeepSeek*ToLinen` wrappers. The NNX path runs whole-model convert - # forward and is model-agnostic — see `_load_and_quantize_nnx`. - if not config.pure_nnx: - assert config.decoder_block == common_types.DecoderBlockType.DEEPSEEK, ( - f"Linen layerwise quantization only supports {common_types.DecoderBlockType.DEEPSEEK}, " - f"got {config.decoder_block}." - ) - # Mesh definition - devices_array = maxtext_utils.create_device_mesh(config=config) - self._mesh = jax.sharding.Mesh(devices_array, config.mesh_axes) - - self.quant = quantizations.configure_quantization(config) - if config.pure_nnx: - # NNX takes a separate code path that builds the model via from_pretrained; - # no Linen abstract-state bookkeeping is needed here. - self.unboxed_abstract_state = None - return - model = models.transformer_as_linen( - config, mesh=self._mesh, quant=self.quant, model_mode=common_types.MODEL_MODE_TRAIN - ) - init_state_fn = functools.partial(maxtext_utils.init_initial_state, model, None, self.config, False, self.rng) - - self.unboxed_abstract_state, _, _ = maxtext_utils.get_abstract_state(self.config, self._mesh, init_state_fn, False) - - def load_and_quantize(self) -> None: - """ - Load parameters layer by layer and quantize them. - """ - if self.config.pure_nnx: - self._load_and_quantize_nnx() - return - quantized_params = {} - quantized_params["params"] = {"decoder": {}} - quantized_params["aqt"] = {"decoder": {}} - config = self.config - self.quant.quant_mode = quantizations.get_quant_mode("convert") - model_mode = common_types.MODEL_MODE_PREFILL - _, rng_quant_params = jax.random.split(self.rng) - - layers = [ - deepseek.DeepSeekDenseLayerToLinen( - config=config, mesh=self._mesh, quant=self.quant, model_mode=model_mode, rngs=nnx.Rngs(self.rng) - ), - deepseek.DeepSeekMoELayerToLinen( - config=config, mesh=self._mesh, quant=self.quant, model_mode=model_mode, rngs=nnx.Rngs(self.rng) - ), - ] - layer_prefixes = [ - "dense_layers", - "moe_layers", - ] - num_moe_layers = config.num_decoder_layers - config.first_num_dense_layers - num_layers_list = [ - config.first_num_dense_layers, - num_moe_layers, - ] - - def model_apply(_p, _rng, layer): - return layer.apply( - _p | {"aqt": {}}, - jnp.ones((1, self.config.max_prefill_predict_length, self.config.base_emb_dim), dtype=jnp.int32), - None, - jnp.zeros((1, self.config.max_prefill_predict_length), dtype=jnp.int32), - True, - model_mode=model_mode, - rngs={"params": _rng}, - mutable=True, - ) - - for layer, num_layers, layer_prefix in zip(layers, num_layers_list, layer_prefixes): - for index in tqdm(range(num_layers)): - layer_name = f"{layer_prefix}_{index}" - max_logging.log(f"Processing layer: {layer_name}") - - params = self._load_layer(layer_name) - params["params"] = params["params"]["decoder"][layer_name] - - _, new_vars = model_apply(params, rng_quant_params, layer) - - if "aqt" not in new_vars: - max_logging.log( - f"Warning: 'aqt' not found in new_vars for {layer_name}. Skipping AQT processing for this layer." - ) - quantized_params["params"]["decoder"][layer_name] = params["params"] # Keep original params - continue - - aqt_vars = new_vars["aqt"] - - try: - removed_params = remove_quantized_params(params["params"], aqt_vars) - quantized_params["params"]["decoder"][layer_name] = removed_params - except Exception as e: - max_logging.log(f"ERROR: Failed to remove quantized params for {layer_name}: {e}") - max_logging.log(f"Dumping params['params'] keys for {layer_name}:") - jax.tree_util.tree_map_with_path( - lambda path, _: max_logging.log(f" {jax.tree_util.keystr(path)}"), params["params"] - ) - max_logging.log(f"Dumping new_vars['aqt'] keys for {layer_name}:") - jax.tree_util.tree_map_with_path(lambda path, _: max_logging.log(f" {jax.tree_util.keystr(path)}"), aqt_vars) - raise - - # Restructure the aqt_vars for this layer to match pure Linen format for saving - if layer_prefix == "moe_layers": - structured_aqt = insert_deepseekmoeblock_scope(aqt_vars) - else: - structured_aqt = aqt_vars - quantized_params["aqt"]["decoder"][layer_name] = structured_aqt - - unquantized_layers = ["decoder_norm", "logits_dense"] - for unquantized_layer in unquantized_layers: - params = self._load_layer(unquantized_layer) - quantized_params["params"]["decoder"][unquantized_layer] = params["params"]["decoder"][unquantized_layer] - quantized_params["params"]["token_embedder"] = self._load_layer("token_embedder")["params"]["token_embedder"] - - maxtext_utils.save_quantized_checkpoint_if_configured(self.config, quantized_params) - - def _load_and_quantize_nnx(self) -> None: - """Whole-model NNX convert: load full-precision via TRAIN-mode `from_pretrained`, - transfer kernels into a fresh CONVERT-mode model, run a forward (the - `ToNNX(AqtDotGeneral)` bridge auto-captures `qrhs.frozen`), strip kernels at - quantized paths, and save the serve-mode-shaped state. - - Two-step load: input checkpoints are typically full-precision (no AQT state - on disk), so we can't `from_pretrained(quant_mode_str="convert")` directly — - orbax would fail to find the missing `qrhs.frozen` leaves. Instead we load - in TRAIN mode (which has only kernels), then copy them into a randomly - initialized CONVERT model that already has the AQT variables provisioned. - """ - config = self.config - # MODEL_MODE_TRAIN avoids the PREFILL/AUTOREGRESSIVE cache plumbing — AQT - # layers populate `qrhs.frozen` regardless of model_mode, so train mode is - # simpler and faster. - max_logging.log("Loading full-precision NNX checkpoint in TRAIN mode...") - with self._mesh: - train_model = model_creation_utils.from_pretrained( - config, - mesh=self._mesh, - model_mode=common_types.MODEL_MODE_TRAIN, - quant_mode_str="train", - ) - - max_logging.log("Building CONVERT-mode model (random init) and copying kernels in...") - rngs = maxtext_utils_nnx.create_nnx_rngs(config, rng_key=self.rng) - with nn_partitioning.axis_rules(config.logical_axis_rules): - convert_model = model_creation_utils.from_config( - config, - mesh=self._mesh, - rngs=rngs, - model_mode=common_types.MODEL_MODE_TRAIN, - quant_mode_str="convert", - ) - self._copy_kernel_leaves_(convert_model, train_model) - del train_model - - # Forward populates AqtDotGeneral_0.qrhs.frozen on every quantized layer. - L = config.max_target_length - decoder_input_tokens = jnp.zeros((1, L), dtype=jnp.int32) - decoder_positions = jnp.arange(L, dtype=jnp.int32)[None, :] - decoder_segment_ids = jnp.ones((1, L), dtype=jnp.int32) - max_logging.log("Running CONVERT-mode forward to populate AQT scale factors...") - with nn_partitioning.axis_rules(config.logical_axis_rules): - _ = convert_model( - decoder_input_tokens, - decoder_positions, - decoder_segment_ids=decoder_segment_ids, - enable_dropout=False, - model_mode=common_types.MODEL_MODE_TRAIN, - ) - - # Convert-mode state has both `kernel` (full precision) and `AqtDotGeneral_0.qrhs.frozen` - # at every quantized DenseGeneral; the serve-mode reader expects only the latter. - convert_state = nnx.state(convert_model).to_pure_dict() - serve_state = self._strip_kernels_at_quantized_paths(convert_state) - - if config.save_quantized_params_path: - max_logging.log(f"Saving NNX-format quantized checkpoint to {config.save_quantized_params_path}") - - # Wrap each leaf in `{"value": }` so the on-disk shape matches what - # `from_pretrained`'s NNX-detection branch reads back (it later does - # `tree.map(lambda v: v["value"], ...)` on each leaf). Save directly via - # orbax — `save_params_to_path` would add an outer `{"params": ...}` wrap - # that the NNX path doesn't expect. - def _wrap_value(node): - if isinstance(node, dict): - return {k: _wrap_value(v) for k, v in node.items()} - return {"value": node} - - wrapped = _wrap_value(serve_state) - orbax_checkpointer = ocp.PyTreeCheckpointer( - use_ocdbt=config.checkpoint_storage_use_ocdbt, - use_zarr3=config.checkpoint_storage_use_zarr3, - ) - orbax_checkpointer.save(config.save_quantized_params_path, wrapped, force=True) - max_logging.log(f"Saved NNX-format quantized checkpoint at: {config.save_quantized_params_path}") - else: - max_logging.log("Skipping save: save_quantized_params_path is null.") - - @staticmethod - def _copy_kernel_leaves_(dst_model, src_model): - """Copy the full-precision parameter leaves (kernel/embedding/scale/bias) - from src into dst, leaving dst's AQT and RNG variables untouched. - """ - src_dict = nnx.state(src_model).to_pure_dict() - dst_state = nnx.state(dst_model) - dst_dict = dst_state.to_pure_dict() - - def walk(d_node, s_node): - if not (isinstance(d_node, dict) and isinstance(s_node, dict)): - return - for key, d_child in d_node.items(): - if key not in s_node: - continue - s_child = s_node[key] - if key in ("kernel", "embedding", "scale", "bias") and not isinstance(d_child, dict): - d_node[key] = s_child - elif isinstance(d_child, dict): - walk(d_child, s_child) - - walk(dst_dict, src_dict) - nnx.replace_by_pure_dict(dst_state, dst_dict) - nnx.update(dst_model, dst_state) - - @staticmethod - def _strip_kernels_at_quantized_paths(state_dict): - """Drop `kernel` keys at any node that has a sibling `AqtDotGeneral_0`. - - In convert mode each quantized DenseGeneral keeps both the full-precision - `kernel` (an nnx.Param) and the AQT-quantized `AqtDotGeneral_0.qrhs.frozen` - side-by-side. Serve mode (the on-disk shape `from_pretrained` reads back) - only carries the latter; the kernel is recreated as a dummy zero in - `linears.DenseGeneral.__call__`. - """ - if not isinstance(state_dict, dict): - return state_dict - has_aqt = "AqtDotGeneral_0" in state_dict - out = {} - for k, v in state_dict.items(): - if k == "kernel" and has_aqt: - continue - out[k] = LayerwiseQuantization._strip_kernels_at_quantized_paths(v) if isinstance(v, dict) else v - return out - - def _load_layer(self, layer_name): - """Loads a specific layer's parameters from the checkpoint.""" - - config = self.config - with nn_partitioning.axis_rules(config.logical_axis_rules): - - params = checkpointing.load_params_from_path( - config.load_parameters_path, - self._create_partial_abstract_params(self.unboxed_abstract_state.params, layer_name), - config.checkpoint_storage_concurrent_gb, - config.checkpoint_storage_use_ocdbt, - config.checkpoint_storage_use_zarr3, - ) - return params - - def _create_partial_abstract_params(self, abstract_unboxed_params, layer): - """Creates a partial abstract params structure using ocp.PLACEHOLDER.""" - - def _should_keep(path, _): - # True if the layer name is part of the path - return any(isinstance(key, jax.tree_util.DictKey) and key.key == layer for key in path) - - def _map_fn(path, value): - if not _should_keep(path, value): - return IGNORE - if isinstance(value, jax.ShapeDtypeStruct): - zeros_array = jnp.zeros(value.shape, value.dtype) - if value.sharding is not None: - try: - return jax.device_put(zeros_array, value.sharding) - except Exception as e: # pylint: disable=broad-except - max_logging.log(f"Error applying sharding for path {path}: {e}") - return zeros_array - return zeros_array - return value - - return jax.tree_util.tree_map_with_path(_map_fn, abstract_unboxed_params) - - -def main(argv: Sequence[str]) -> None: - jax.config.update("jax_default_prng_impl", "unsafe_rbg") - os.environ["TF_CPP_MIN_LOG_LEVEL"] = "0" - config = pyconfig.initialize(argv) - validate_config(config) - max_utils.print_system_information() - rng = jax.random.PRNGKey(1234) - quantization = LayerwiseQuantization(config, rng) - quantization.load_and_quantize() - - -def validate_config(config): - assert ( - config.load_full_state_path == "" - ), "Operation on full states not supported! Convert to parameter checkpoint first." - - -if __name__ == "__main__": - app.run(main) diff --git a/src/maxtext/utils/maxtext_utils.py b/src/maxtext/utils/maxtext_utils.py index fe34d9fac0..a8cc822758 100644 --- a/src/maxtext/utils/maxtext_utils.py +++ b/src/maxtext/utils/maxtext_utils.py @@ -1276,10 +1276,39 @@ def collect_intermediates_by_suffix(intermediate_outputs, *suffix_keys: str) -> return values -def get_intermediate_value(model, nested_key, default=None, clear=False): +def _find_and_remove_intermediates(state, suffix, clear=False): + """Recursively finds intermediate values matching suffix. + + If clear=True, removes them from the state. + Returns a list of (path, variable) tuples. """ - Retrieves an intermediate value from an NNX model. This functions has context about - where the intermediate value is located. + results = [] + + def _traverse(current_state, current_path): + keys_to_delete = [] + for k, v in list(current_state.items()): + new_path = current_path + (k,) + if isinstance(v, nnx.Intermediate): + if len(new_path) >= len(suffix) and new_path[-len(suffix) :] == suffix: + results.append((new_path, v)) + if clear: + keys_to_delete.append(k) + elif isinstance(v, (nnx.State, dict)): + _traverse(v, new_path) + if clear and not v: + keys_to_delete.append(k) + + for k in keys_to_delete: + del current_state[k] + + _traverse(state, ()) + return results + + +def get_intermediate_value(model, nested_key, default=None, clear=False): + """Retrieves an intermediate value from an NNX model. + + This functions has context about where the intermediate value is located. Args: model: The NNX model. @@ -1290,18 +1319,53 @@ def get_intermediate_value(model, nested_key, default=None, clear=False): Returns: The value associated with the nested key, or the default value if not found. """ - intermediate_value = default match nested_key: case "out_projection_activations": - if nested_key in model.decoder.layers["self_attention"]: - intermediate_value = model.decoder.layers["self_attention"][nested_key].get_value()[-1] - if clear: - del model.decoder.layers["self_attention"][nested_key] + suffixes = [ + ("self_attention", "out_projection_activations"), + ("GptOssAttention", "out_projection_activations"), + ] case _: - # Default case to handle any unknown nested keys raise ValueError(f"Incorrect nested_key: {nested_key}") - return intermediate_value + # Pop all intermediates to safely inspect and potentially clear them + intermediates = nnx.pop(model, nnx.Intermediate) + + found = [] + for suffix in suffixes: + found = _find_and_remove_intermediates(intermediates, suffix, clear=clear) + if found: + break + + # Put back the remaining intermediates + nnx.update(model, intermediates) + + if not found: + return default + + # Helper key function to sort paths numerically if indices are present + def path_sort_key(item): + path = item[0] + + def _to_int_if_possible(val): + if isinstance(val, int): + return val + if isinstance(val, str) and val.isdigit(): + return int(val) + return val + + return tuple(_to_int_if_possible(x) for x in path) + + found.sort(key=path_sort_key) + + values = [var.get_value()[-1] for path, var in found] + + if len(values) > 1: + # Multiple layers (sequential), stack them + return jnp.stack(values, axis=0) + else: + # Single layer (scanned or just 1 layer), return directly + return values[0] def update_state_param(state, target_path, value): @@ -1648,18 +1712,14 @@ def get_abstract_state_nnx(config, mesh, nnx_init_trainstate_fn, is_training=Tru # ourselves via nnx_construct_named_sharding, so auto-assignment is not needed here. abs_model = nnx.eval_shape(nnx_init_trainstate_fn) _, abs_var_state = nnx.split(abs_model) - named_sharding_state = sharding.nnx_construct_named_sharding( - abs_var_state, mesh - ) + named_sharding_state = sharding.nnx_construct_named_sharding(abs_var_state, mesh) abstract_state = jax.tree.map( lambda a, s: jax.ShapeDtypeStruct(a.shape, a.dtype, sharding=s), abs_var_state, named_sharding_state, ) - state_mesh_shardings = maxtext_utils_nnx.nnx_extract_named_sharding( - abstract_state - ) + state_mesh_shardings = maxtext_utils_nnx.nnx_extract_named_sharding(abstract_state) if is_training and config.shard_optimizer_over_data: # Add data to sharding for optimizer state diff --git a/tests/__init__.py b/tests/__init__.py index 46cd7ffa11..76d26765cb 100644 --- a/tests/__init__.py +++ b/tests/__init__.py @@ -16,6 +16,54 @@ Test initialization """ -import pathwaysutils +try: + import pathwaysutils -pathwaysutils.initialize() + pathwaysutils.initialize() +except ImportError: + import sys + from unittest.mock import MagicMock + + mock_pathwaysutils = MagicMock() + mock_pathwaysutils.__path__ = [] + mock_pathwaysutils.is_pathways_backend_used.return_value = False + sys.modules["pathwaysutils"] = mock_pathwaysutils + + mock_elastic = MagicMock() + mock_elastic.__path__ = [] + sys.modules["pathwaysutils.elastic"] = mock_elastic + + mock_manager = MagicMock() + sys.modules["pathwaysutils.elastic.manager"] = mock_manager + +try: + import tokamax + from tokamax._src.ops.experimental.tpu.splash_attention import splash_attention_kernel +except ImportError: + import sys + from unittest.mock import MagicMock + + mock_tokamax = MagicMock() + mock_tokamax.__path__ = [] + sys.modules["tokamax"] = mock_tokamax + sys.modules["tokamax._src"] = MagicMock() + sys.modules["tokamax._src.ops"] = MagicMock() + sys.modules["tokamax._src.ops.experimental"] = MagicMock() + sys.modules["tokamax._src.ops.experimental.tpu"] = MagicMock() + sys.modules["tokamax._src.ops.experimental.tpu.splash_attention"] = MagicMock() + +try: + import tensorflow +except ImportError: + import sys + from unittest.mock import MagicMock + + mock_tf = MagicMock() + mock_tf.__path__ = [] + sys.modules["tensorflow"] = mock_tf + sys.modules["tensorflow.io"] = MagicMock() + sys.modules["tensorflow.data"] = MagicMock() + sys.modules["tensorflow.compat"] = MagicMock() + sys.modules["tensorflow.compat.v1"] = MagicMock() + sys.modules["tensorflow.compat.v1.io"] = MagicMock() + sys.modules["tensorflow.compat.v1.io.gfile"] = MagicMock() diff --git a/tests/conftest.py b/tests/conftest.py index eec4afa225..b0eca11604 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -44,28 +44,14 @@ import jax import os import importlib.util +# Prevent TensorFlow (which might be imported by qwix or other libraries during test collection) +# from grabbing the GPU context and causing JAX NCCL communicator creation to fail. +try: + import tensorflow as tf -# Force early JAX initialization on GPU to prevent CUDA context conflicts with TensorFlow/PyTorch. -# If JAX initialization is deferred, TensorFlow/PyTorch (imported during test collection) -# might initialize CUDA first, causing JAX's subsequent NCCL communicator creation to fail -# with 'corrupted comm object detected'. -# Detect GPU environment using standard JAX env vars, GHA runner device types, -# and nvidia-docker visible device markers. -_jax_platforms = os.getenv("JAX_PLATFORMS", "").lower() -_device_type = os.getenv("INPUTS_DEVICE_TYPE", "").lower() -_has_gpu = ( - "cuda" in _jax_platforms - or "gpu" in _jax_platforms - or "cuda" in _device_type - or "gpu" in _device_type - or os.getenv("CUDA_VISIBLE_DEVICES") is not None - or os.getenv("NVIDIA_VISIBLE_DEVICES") is not None -) -if _has_gpu: - try: - _ = jax.devices() - except Exception: # pylint: disable=broad-exception-caught - pass + tf.config.set_visible_devices([], "GPU") +except Exception: # pylint: disable=broad-exception-caught + pass # --- Monkeypatch for absl.testing.parameterized --- # Context: Decorating a test method with @parameterized.named_parameters returns a custom diff --git a/tests/integration/decode_tests.py b/tests/integration/decode_tests.py index 0117dc1a6b..3af2f6d03a 100644 --- a/tests/integration/decode_tests.py +++ b/tests/integration/decode_tests.py @@ -62,7 +62,6 @@ class DecodeTests(unittest.TestCase): "max_target_length=128", "per_device_batch_size=1", "quantization=int8", - "quantize_kvcache=True", rf"tokenizer_path={os.path.join(MAXTEXT_ASSETS_ROOT, 'tokenizers', 'tokenizer.llama2')}", ], "pdb_lt_1": [ # tests decode with per_device_batch_size < 1 diff --git a/tests/integration/hlo_diff_test.py b/tests/integration/hlo_diff_test.py index cc18713749..8077a9e83f 100644 --- a/tests/integration/hlo_diff_test.py +++ b/tests/integration/hlo_diff_test.py @@ -59,7 +59,7 @@ def filter_line(line): not os.environ.get("GITHUB_ACTIONS"), reason="Skipping HLO diff test because it is not running in GitHub Actions", ) -@pytest.mark.tpu_backend +@pytest.mark.tpu_only class TestHloDiff: """Tests for HLO Graph Diff Verification.""" diff --git a/tests/unit/aqt_serve_roundtrip_nnx_test.py b/tests/unit/aqt_serve_roundtrip_nnx_test.py deleted file mode 100644 index 00d6825de7..0000000000 --- a/tests/unit/aqt_serve_roundtrip_nnx_test.py +++ /dev/null @@ -1,170 +0,0 @@ -# Copyright 2023-2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# https://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Round-trip test for the NNX serve-mode AQT checkpoint path. - -Builds a small NNX model in CONVERT mode with int8 quantization, runs a forward -to populate `qrhs.frozen`, saves the serve-mode-shape state to a local orbax -checkpoint, then reloads via `from_pretrained(quant_mode_str="serve")` and -checks that the loaded QTensor leaves match what was saved. - -This guards the chain of issues exercised by serve-mode reload (sharding helper -for QTensor, v[...] vs get_value() for composite values, Param-only filter -dropping aqt-typed leaves, Partitioned-unwrap for matching on-disk paths). -""" - -import os -import sys -import tempfile -import unittest - -import jax -import jax.numpy as jnp -import orbax.checkpoint as ocp -from flax import nnx -from flax.core.meta import Partitioned -from flax.linen import partitioning as nn_partitioning - -from maxtext.configs import pyconfig -from maxtext.utils import maxtext_utils, model_creation_utils, maxtext_utils_nnx -from maxtext.utils.globals import MAXTEXT_PKG_DIR -from maxtext.utils.layerwise_quantization import LayerwiseQuantization - - -def _wrap_value(node): - """Add `{"value": ...}` per-leaf wrap matching `_load_and_quantize_nnx` save format.""" - if isinstance(node, dict): - return {k: _wrap_value(v) for k, v in node.items()} - return {"value": node} - - -def _unbox(x): - return x.value if isinstance(x, Partitioned) else x - - -def _walk_qrhs(state): - """Yield (path_str, variable) pairs for every qrhs.frozen entry in an nnx.State.""" - for path, var in state.flat_state(): - keys = [str(getattr(k, "key", k)) for k in path] - if "qrhs" in keys and "frozen" in keys: - yield ".".join(keys), var - - -class ServeModeRoundTripTest(unittest.TestCase): - """End-to-end save+reload of a serve-mode NNX AQT checkpoint.""" - - def _init_cfg(self, ckpt_path, *, checkpoint_is_quantized): - """Build a pyconfig for save or reload.""" - # Use base.yml + gpt3-52k. The decoupled test config strips - # logical_axis_rules (e.g. "norm"), which the AQT serve-mode model - # construction needs. - base_yml = os.path.join(MAXTEXT_PKG_DIR, "configs", "base.yml") - args = [ - sys.argv[0], - base_yml, - "model_name=gpt3-52k", - "pure_nnx=true", - "enable_nnx=true", - "pure_nnx_decoder=true", - "max_target_length=64", - "max_prefill_predict_length=16", - "per_device_batch_size=1", - "scan_layers=true", - "quantization=int8", - "checkpoint_storage_use_ocdbt=false", - "checkpoint_storage_use_zarr3=false", - "skip_jax_distributed_system=true", - ] - if checkpoint_is_quantized: - args += [ - f"load_parameters_path={ckpt_path}", - "checkpoint_is_quantized=true", - "enable_checkpointing=true", # required by config validator when load_parameters_path is set - ] - else: - args += ["enable_checkpointing=false"] - return pyconfig.initialize(args) - - def test_save_then_reload_preserves_qrhs_frozen(self): - """Save a serve-mode-shape NNX checkpoint, then reload it and compare qvalue arrays.""" - with tempfile.TemporaryDirectory() as tmpdir: - ckpt_path = os.path.join(tmpdir, "quantized_ckpt") - - # Step 1: build CONVERT-mode model + run forward to populate qrhs.frozen. - cfg_save = self._init_cfg(ckpt_path, checkpoint_is_quantized=False) - mesh = maxtext_utils.get_mesh_from_config(cfg_save) - rngs = maxtext_utils_nnx.create_nnx_rngs(cfg_save) - with nn_partitioning.axis_rules(cfg_save.logical_axis_rules): - convert_model = model_creation_utils.from_config( - cfg_save, - mesh=mesh, - rngs=rngs, - model_mode="train", - quant_mode_str="convert", - ) - L = cfg_save.max_prefill_predict_length - tokens = jnp.zeros((1, L), dtype=jnp.int32) - pos = jnp.arange(L, dtype=jnp.int32)[None, :] - seg = jnp.ones((1, L), dtype=jnp.int32) - with nn_partitioning.axis_rules(cfg_save.logical_axis_rules): - _ = convert_model(tokens, pos, decoder_segment_ids=seg, enable_dropout=False, model_mode="train") - - # Step 2: capture the qrhs.frozen leaves we expect to round-trip, then save. - convert_state = nnx.state(convert_model).to_pure_dict() - serve_state = LayerwiseQuantization._strip_kernels_at_quantized_paths(convert_state) # pylint: disable=protected-access - saved_qrhs = {} - for path, var in _walk_qrhs(nnx.state(convert_model)): - qt = var.value if hasattr(var, "value") else var - saved_qrhs[path] = _unbox(qt.qvalue) - - # Replicate arrays across the mesh; orbax rejects SingleDeviceSharding - # once another test has initialized JAX-distributed state. - replicated = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec()) - serve_state = jax.tree.map( - lambda x: jax.device_put(x, replicated) if isinstance(x, jax.Array) else x, - serve_state, - ) - - orbax_checkpointer = ocp.PyTreeCheckpointer(use_ocdbt=False, use_zarr3=False) - orbax_checkpointer.save(ckpt_path, _wrap_value(serve_state), force=True) - self.assertGreater(len(saved_qrhs), 0, "Test config must produce at least one qrhs.frozen leaf") - - # Step 3: reload via from_pretrained in serve mode. - cfg_load = self._init_cfg(ckpt_path, checkpoint_is_quantized=True) - with nn_partitioning.axis_rules(cfg_load.logical_axis_rules): - loaded_model = model_creation_utils.from_pretrained( - cfg_load, - mesh=mesh, - model_mode="autoregressive", - quant_mode_str="serve", - ) - - # Step 4: assert every saved qrhs.frozen leaf matches what was persisted. - loaded_state = nnx.state(loaded_model) - loaded_qrhs = dict(_walk_qrhs(loaded_state)) - self.assertEqual(set(saved_qrhs.keys()), set(loaded_qrhs.keys())) - for path, saved_qv in saved_qrhs.items(): - var = loaded_qrhs[path] - qt = var.value if hasattr(var, "value") else var - loaded_qv = _unbox(qt.qvalue) - self.assertEqual(loaded_qv.shape, saved_qv.shape, f"shape mismatch at {path}") - self.assertEqual(loaded_qv.dtype, saved_qv.dtype, f"dtype mismatch at {path}") - self.assertTrue( - jnp.array_equal(loaded_qv.astype(jnp.int32), saved_qv.astype(jnp.int32)), - f"qvalue not preserved at {path}", - ) - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/unit/data_loader_test.py b/tests/unit/data_loader_test.py index c61217d598..52737c11c4 100644 --- a/tests/unit/data_loader_test.py +++ b/tests/unit/data_loader_test.py @@ -26,12 +26,15 @@ from maxtext.configs import pyconfig from maxtext.common.data_loader import DataLoader, RampUpDataLoader + from maxtext.utils import exceptions from maxtext.utils.maxtext_utils import create_device_mesh from maxtext.common.gcloud_stub import is_decoupled from maxtext.utils.rampup_batch import RampupBatchManager from tests.utils.test_helpers import get_test_config_path +pytestmark = pytest.mark.cpu_only + class DataLoaderTest(unittest.TestCase): diff --git a/tests/unit/grain_data_processing_test.py b/tests/unit/grain_data_processing_test.py index dc129b5ebc..1c22af1eac 100644 --- a/tests/unit/grain_data_processing_test.py +++ b/tests/unit/grain_data_processing_test.py @@ -29,10 +29,13 @@ from maxtext.configs import pyconfig from maxtext.input_pipeline import grain_data_processing from maxtext.input_pipeline import input_pipeline_interface + from maxtext.utils.globals import MAXTEXT_ASSETS_ROOT from maxtext.common.gcloud_stub import is_decoupled from tests.utils.test_helpers import get_test_base_output_directory, get_test_config_path, get_test_dataset_path +pytestmark = pytest.mark.cpu_only + class GrainBaseProcessingTest: """Base mixin with test_train_ds for all grain data processing tests. diff --git a/tests/unit/hf_data_processing_test.py b/tests/unit/hf_data_processing_test.py index 262c56ff9b..2a6cf93cd2 100644 --- a/tests/unit/hf_data_processing_test.py +++ b/tests/unit/hf_data_processing_test.py @@ -29,6 +29,10 @@ from maxtext.utils.globals import MAXTEXT_ASSETS_ROOT from tests.utils.test_helpers import get_test_config_path, get_test_base_output_directory +import pytest + +pytestmark = pytest.mark.cpu_only + class HfDataProcessingTest(unittest.TestCase): diff --git a/tests/unit/input_pipeline/olmo_data_grain_resume_test.py b/tests/unit/input_pipeline/olmo_data_grain_resume_test.py index bc1f547496..460ede691b 100644 --- a/tests/unit/input_pipeline/olmo_data_grain_resume_test.py +++ b/tests/unit/input_pipeline/olmo_data_grain_resume_test.py @@ -43,6 +43,10 @@ make_olmo_grain_data_loader, ) +import pytest + +pytestmark = pytest.mark.cpu_only + def _write_raw_uint32(tmpdir: str, name: str, values: np.ndarray) -> str: assert values.dtype == np.uint32 and values.ndim == 1 diff --git a/tests/unit/input_pipeline/olmo_data_grain_test.py b/tests/unit/input_pipeline/olmo_data_grain_test.py index 08a493d1c3..168ac57eaa 100644 --- a/tests/unit/input_pipeline/olmo_data_grain_test.py +++ b/tests/unit/input_pipeline/olmo_data_grain_test.py @@ -46,6 +46,10 @@ make_olmo_grain_data_loader, ) +import pytest + +pytestmark = pytest.mark.cpu_only + def _write_raw_uint32(tmpdir: str, name: str, values: np.ndarray) -> str: """Write a 1-D uint32 array as raw binary (no .npy header) — matches AI2.""" diff --git a/tests/unit/layerwise_quantization_nnx_test.py b/tests/unit/layerwise_quantization_nnx_test.py deleted file mode 100644 index bbd43f7964..0000000000 --- a/tests/unit/layerwise_quantization_nnx_test.py +++ /dev/null @@ -1,77 +0,0 @@ -# Copyright 2023–2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# https://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Unit tests for the NNX path of layerwise_quantization. - -Covers `_strip_kernels_at_quantized_paths` — the convert→serve shape converter -that drops the redundant full-precision kernel from quantized DenseGeneral -nodes while leaving non-quantized kernels (norms, embeddings) intact. -""" - -import unittest - -from maxtext.utils.layerwise_quantization import LayerwiseQuantization - - -class StripKernelsTest(unittest.TestCase): - - def test_drops_kernel_at_quantized_dense(self): - """A node with both `kernel` and `AqtDotGeneral_0` loses the kernel.""" - state = { - "decoder": { - "layers": { - "mlp": { - "wi": { - "kernel": "FULL_PRECISION_W", - "AqtDotGeneral_0": {"qrhs": {"frozen": "AQT_STATE"}}, - } - } - } - } - } - out = LayerwiseQuantization._strip_kernels_at_quantized_paths(state) # pylint: disable=protected-access - wi = out["decoder"]["layers"]["mlp"]["wi"] - self.assertNotIn("kernel", wi) - self.assertIn("AqtDotGeneral_0", wi) - self.assertEqual(wi["AqtDotGeneral_0"]["qrhs"]["frozen"], "AQT_STATE") - - def test_preserves_non_quantized_kernel(self): - """A non-quantized kernel (e.g. embedding, norm) survives.""" - state = { - "decoder": { - "decoder_norm": {"scale": "NORM_SCALE"}, - "logits_dense": {"kernel": "LOGITS_KERNEL"}, # no AqtDotGeneral_0 sibling - }, - "token_embedder": {"embedding": "EMB"}, - } - out = LayerwiseQuantization._strip_kernels_at_quantized_paths(state) # pylint: disable=protected-access - self.assertEqual(out["decoder"]["logits_dense"]["kernel"], "LOGITS_KERNEL") - self.assertEqual(out["decoder"]["decoder_norm"]["scale"], "NORM_SCALE") - self.assertEqual(out["token_embedder"]["embedding"], "EMB") - - def test_mixed_tree(self): - """Quantized + non-quantized at the same depth: only the quantized one strips.""" - state = { - "self_attention": { - "qkv_proj": {"kernel": "QKV", "AqtDotGeneral_0": "AQT"}, - "out": {"kernel": "OUT_FULL"}, # non-quantized output projection - } - } - out = LayerwiseQuantization._strip_kernels_at_quantized_paths(state) # pylint: disable=protected-access - self.assertNotIn("kernel", out["self_attention"]["qkv_proj"]) - self.assertEqual(out["self_attention"]["out"]["kernel"], "OUT_FULL") - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/unit/maxtext_utils_test.py b/tests/unit/maxtext_utils_test.py index 62e165c491..be2c3c4e47 100644 --- a/tests/unit/maxtext_utils_test.py +++ b/tests/unit/maxtext_utils_test.py @@ -19,7 +19,7 @@ from typing import Any, Sequence import unittest import pytest -from unittest.mock import MagicMock, Mock, patch +from unittest.mock import MagicMock, patch from dataclasses import dataclass, field import numpy as np import optax @@ -96,44 +96,139 @@ def test_fp8_stats_not_clipped_but_others_are(self): ) -class TestIntermediateValueRetrieval(unittest.TestCase): - """test class for IntermediateValueRetrieval""" +class MockAttention(nnx.Module): + """Mock attention module for testing intermediate extraction.""" - def setUp(self): - self.mock_model = MagicMock(name="Transformer") + def __init__(self, rngs): + pass - # 2. Create the Decoder Mock - self.mock_decoder = MagicMock(name="Decoder") - self.mock_model.decoder = self.mock_decoder - self.mock_layers = {} - self.mock_model.decoder.layers = self.mock_layers - self.self_attention = {} - self.mock_layers["self_attention"] = self.self_attention + def __call__(self, x, sow_val=None): + if sow_val is not None: + self.sow(nnx.Intermediate, "out_projection_activations", sow_val) + return x - def test_valid_intermediate_key(self): - expected_sowed_data = [0.1, 0.5, 0.9] - mock_sowed_variable = Mock(name="out_projection_activations") - mock_sowed_variable.get_value.return_value = (expected_sowed_data,) - self.mock_decoder.layers["self_attention"]["out_projection_activations"] = mock_sowed_variable +class MockDecoderLayer(nnx.Module): + """Mock decoder layer for testing intermediate extraction.""" - result = maxtext_utils.get_intermediate_value(self.mock_model, "out_projection_activations") + def __init__(self, rngs, attention_name="self_attention"): + if attention_name == "self_attention": + self.self_attention = MockAttention(rngs) + elif attention_name == "GptOssAttention": + self.GptOssAttention = MockAttention(rngs) + self.attention_name = attention_name - self.assertEqual(result, expected_sowed_data) + def __call__(self, x, sow_val=None): + attention = getattr(self, self.attention_name) + return attention(x, sow_val) - def test_returns_default_if_sow_did_not_happen(self): - """ - Simulate a scenario where the model ran, but this specific key - was NOT sowed (or the layer was skipped). - """ - result = maxtext_utils.get_intermediate_value(self.mock_model, "out_projection_activations", default="MyDefault") +class MockDecoderSequential(nnx.Module): + """Mock decoder sequential block for testing intermediate extraction.""" + def __init__(self, rngs, attention_name="self_attention"): + self.layers = nnx.List( + [ + MockDecoderLayer(rngs, attention_name), + MockDecoderLayer(rngs, attention_name), + ] + ) + + def __call__(self, x, sow_vals=None): + if sow_vals is None: + sow_vals = [None, None] + out = x + for layer, val in zip(self.layers, sow_vals): + out = layer(out, val) + return out + + +class MockDecoderScanned(nnx.Module): + """Mock decoder scanned block for testing intermediate extraction.""" + + def __init__(self, rngs, attention_name="self_attention"): + self.layers = MockDecoderLayer(rngs, attention_name) + self.attention_name = attention_name + + def __call__(self, x, sow_val=None): + if sow_val is not None: + attention = getattr(self.layers, self.attention_name) + attention.sow(nnx.Intermediate, "out_projection_activations", sow_val) + return x + + +class MockTransformer(nnx.Module): + """Mock transformer for testing intermediate extraction.""" + + def __init__(self, decoder_type, rngs, attention_name="self_attention"): + if decoder_type == "sequential": + self.decoder = MockDecoderSequential(rngs, attention_name) + elif decoder_type == "scanned": + self.decoder = MockDecoderScanned(rngs, attention_name) + + def __call__(self, x, sow_vals=None): + return self.decoder(x, sow_vals) + + +class TestIntermediateValueRetrieval(unittest.TestCase): + """test class for IntermediateValueRetrieval""" + + def test_valid_intermediate_key_sequential(self): + rngs = nnx.Rngs(0) + model = MockTransformer("sequential", rngs) + x = jnp.ones((1, 2)) + sow_vals = [jnp.array([0.1]), jnp.array([0.5])] + model(x, sow_vals) + + result = maxtext_utils.get_intermediate_value(model, "out_projection_activations") + expected = jnp.stack(sow_vals, axis=0) + self.assertTrue(jnp.allclose(result, expected)) + + def test_valid_intermediate_key_scanned(self): + rngs = nnx.Rngs(0) + model = MockTransformer("scanned", rngs) + x = jnp.ones((1, 2)) + sow_val = jnp.array([[0.1], [0.5]]) + model(x, sow_val) + + result = maxtext_utils.get_intermediate_value(model, "out_projection_activations") + self.assertTrue(jnp.allclose(result, sow_val)) + + def test_valid_intermediate_key_sequential_gpt_oss(self): + rngs = nnx.Rngs(0) + model = MockTransformer("sequential", rngs, attention_name="GptOssAttention") + x = jnp.ones((1, 2)) + sow_vals = [jnp.array([0.1]), jnp.array([0.5])] + model(x, sow_vals) + + result = maxtext_utils.get_intermediate_value(model, "out_projection_activations") + expected = jnp.stack(sow_vals, axis=0) + self.assertTrue(jnp.allclose(result, expected)) + + def test_valid_intermediate_key_scanned_gpt_oss(self): + rngs = nnx.Rngs(0) + model = MockTransformer("scanned", rngs, attention_name="GptOssAttention") + x = jnp.ones((1, 2)) + sow_val = jnp.array([[0.1], [0.5]]) + model(x, sow_val) + + result = maxtext_utils.get_intermediate_value(model, "out_projection_activations") + self.assertTrue(jnp.allclose(result, sow_val)) + + def test_returns_default_if_sow_did_not_happen(self): + rngs = nnx.Rngs(0) + model = MockTransformer("sequential", rngs) + x = jnp.ones((1, 2)) + model(x, None) + + result = maxtext_utils.get_intermediate_value(model, "out_projection_activations", default="MyDefault") self.assertEqual(result, "MyDefault") def test_unknown_key_raises_value_error(self): + rngs = nnx.Rngs(0) + model = MockTransformer("sequential", rngs) with self.assertRaises(ValueError) as cm: - maxtext_utils.get_intermediate_value(self.mock_model, "some_random_layer_name") + maxtext_utils.get_intermediate_value(model, "some_random_layer_name") self.assertEqual(str(cm.exception), "Incorrect nested_key: some_random_layer_name") @@ -1523,9 +1618,7 @@ class MockConfig: optimizer_memory_host_offload: bool = False parameter_memory_host_offload: bool = False param_scan_axis: int = 0 - logical_axis_rules: list = field( - default_factory=lambda: [["data", ["data"]], ["model", ["model"]]] - ) + logical_axis_rules: list = field(default_factory=lambda: [["data", ["data"]], ["model", ["model"]]]) class MockTrainState(nnx.Module): """Simulates a TrainState with params and optimizer state.""" diff --git a/tests/unit/quantizations_test.py b/tests/unit/quantizations_test.py index fff430e646..5aa252c9aa 100644 --- a/tests/unit/quantizations_test.py +++ b/tests/unit/quantizations_test.py @@ -14,314 +14,77 @@ """Tests for the quantizations""" -import functools -import os.path import sys -from typing import Any import unittest -from aqt.jax.v2 import aqt_tensor -from aqt.jax.v2.flax import aqt_flax -from flax import nnx -from flax.nnx import traversals import jax -from jax import lax from jax import numpy as jnp from jax.sharding import Mesh from maxtext.configs import pyconfig -from maxtext.utils.globals import MAXTEXT_CONFIGS_DIR from maxtext.common.common_types import DECODING_ACTIVE_SEQUENCE_INDICATOR -from maxtext.kernels.megablox import gmm -from maxtext.layers import nnx_wrappers, quantizations +from flax import nnx +from flax.nnx import traversals +from maxtext.layers import moe +from maxtext.layers import linears +from maxtext.layers import quantizations +from maxtext.kernels.megablox.ops import gmm +from maxtext.layers.initializers import nd_dense_init from maxtext.utils import maxtext_utils from maxtext.utils import model_creation_utils from tests.utils.test_helpers import get_test_config_path -import numpy as np import pytest -_QUERY_REGEX = ".*/query" -_VALUE_REGEX = ".*/value" - - -class QuantTestModule(nnx.Module): - """Test module for einsum.""" - - def __init__( - self, - quantization: quantizations.AqtQuantization, - data_type: Any, - rngs: nnx.Rngs, # pylint: disable=unused-argument - ): - self.quantization = quantization - self.identity = jnp.identity(2, dtype=data_type) - self.einsum = None - self.dot_general = None - - if self.quantization: - quant_dg, is_tiled, tiling_fn = None, False, None - if isinstance(self.quantization.quant_dg, dict): - quant_dg, is_tiled, tiling_fn = self._get_mixed_precision_cfg() - else: - quant_dg, is_tiled, tiling_fn = self.quantization.quant_dg, False, None - rhs_axis_metadata_wrapper = None - if self.quantization.quant_mode == aqt_flax.QuantMode.CONVERT: - rhs_axis_metadata_wrapper = None - else: - rhs_axis_metadata_wrapper = functools.partial( - quantizations._rhs_axis_metadata_wrapper, - mesh_axes=(), - is_tiled=is_tiled, - replicate_scale=self.quantization.replicate_scale, - ) - - aqt_dg_cls = aqt_flax.AqtDotGeneral( - quant_dg, - rhs_quant_mode=self.quantization.quant_mode, - lhs_freeze_mode=aqt_flax.FreezerMode.NONE, - rhs_freeze_mode=aqt_flax.FreezerMode.CALIBRATION_AND_VALUE, - rhs_axis_metadata_wrapper=rhs_axis_metadata_wrapper, - use_legacy_freezer=False, - tiling_fn=tiling_fn, - ) - aqt_dg_cls_nnx = nnx_wrappers.ToNNX(aqt_dg_cls, rngs=nnx.Rngs(params=0)) - aqt_einsum = aqt_flax.AqtEinsum( - cfg=quant_dg, - rhs_quant_mode=self.quantization.quant_mode, - lhs_freeze_mode=aqt_flax.FreezerMode.NONE, - rhs_freeze_mode=aqt_flax.FreezerMode.CALIBRATION_AND_VALUE, - rhs_axis_metadata_wrapper=rhs_axis_metadata_wrapper, - use_legacy_freezer=False, - tiling_fn=tiling_fn, - ) - aqt_einsum_nnx = nnx_wrappers.ToNNX(aqt_einsum, rngs=nnx.Rngs(params=0)) - - self.einsum = nnx.data(aqt_einsum_nnx) - self.dot_general = nnx.data(aqt_dg_cls_nnx) - else: - self.einsum = jnp.einsum - self.dot_general = lax.dot_general - - def __call__(self, inputs): - res_einsum = self.einsum("bc,ab->ac", inputs, self.identity) - res_dg = self.dot_general(inputs, inputs, (((), ()), ((), ())), precision=None) - return res_einsum, res_dg - -def _configure_quantization(quant_str="", quant_cfg_path="", mode_str="train", replicate_scale=False): +def _configure_quantization(quant_str="", mode_str="train"): config = pyconfig.initialize( [None, get_test_config_path()], enable_checkpointing=False, quantization=quant_str, - quant_cfg_path=quant_cfg_path, - replicate_quant_scale=replicate_scale, ) quant = quantizations.configure_quantization(config, mode_str) return quant -def _apply(quant_str=""): - rngs = nnx.Rngs(params=0) - inputs = jnp.ones((2, 2)) - data_type = inputs.dtype - quant = _configure_quantization(quant_str) - test_module = QuantTestModule(quant, data_type, rngs) - res_einsum, res_dg = test_module(inputs) - return inputs, res_einsum, res_dg - - class QuantizationTest(unittest.TestCase): """Tests for quantization.""" - def test_in_quant_mode(self): - quant = _configure_quantization(quant_str="int8", mode_str="convert") - self.assertTrue(quantizations.in_convert_mode(quant)) - self.assertFalse(quantizations.in_serve_mode(quant)) - def test_configure_quantization_is_null(self): for quant_mode in ["train", "serve", "convert"]: quant = _configure_quantization(quant_str="", mode_str=quant_mode) self.assertEqual(quant, None) - def test_configure_quantization_replicate_scale(self): - for quant_mode in ["train", "serve", "convert"]: - quant = _configure_quantization(quant_str="int8", mode_str=quant_mode) - self.assertEqual(quant.replicate_scale, False) - for quant_mode in ["train", "serve", "convert"]: - quant = _configure_quantization(quant_str="int8", mode_str=quant_mode, replicate_scale=True) - self.assertEqual(quant.replicate_scale, True) +class QuantizationConfigValidationTest(unittest.TestCase): + """Tests for quantization configuration validation.""" - @pytest.mark.cpu_only - def test_configure_quantization_is_int8(self): - for quant_mode in ["train", "serve", "convert"]: - quant = _configure_quantization(quant_str="int8", mode_str=quant_mode) - self.assertNotEqual(quant, None) - - @pytest.mark.tpu_only # b/421002974 - def test_aqt_quantization(self): - # Without quantization - inputs, res_einsum, res_dg = _apply() - self.assertTrue(jnp.array_equal(inputs, res_einsum)) - self.assertEqual(res_einsum.dtype, np.dtype(np.float32)) - self.assertTrue(jnp.array_equal(inputs, res_dg[0][0])) - self.assertEqual(res_dg.dtype, np.dtype(np.float32)) - - # With int8 quantization - inputs, res_einsum, res_dg = _apply(quant_str="int8") - self.assertTrue(jnp.greater(jnp.max(inputs), jnp.max(res_einsum))) - self.assertEqual(res_einsum.dtype, np.dtype(np.float32)) - self.assertTrue(jnp.greater(jnp.max(inputs), jnp.max(res_dg[0][0]))) - # self.assertEqual(res_dg.dtype, np.dtype(np.float32)) - - def test_mixed_precision_config_int8w(self): - quant = _configure_quantization( - quant_str="intmp", - quant_cfg_path=os.path.join(MAXTEXT_CONFIGS_DIR, "quantization", "int8_weight_only.json"), - ) - self.assertTrue(isinstance(quant.quant_dg, dict) and len(quant.quant_dg) == 1) - # pylint: disable=unsupported-membership-test - self.assertTrue(quantizations.DEFAULT in quant.quant_dg) - quant_cfg, _ = quant.quant_dg[quantizations.DEFAULT] - self.assertEqual(quant_cfg.fwd.dg_quantizer.lhs.numerics.dtype, None) - self.assertEqual(quant_cfg.fwd.dg_quantizer.rhs.numerics.bits, 8) - - def test_mixed_precision_config_scale(self): - quant = _configure_quantization( - quant_str="intmp", - quant_cfg_path=os.path.join( - MAXTEXT_CONFIGS_DIR, - "quantization", - "dense_llm_weight_only_scale.json", - ), - ) - self.assertTrue(isinstance(quant.quant_dg, dict) and len(quant.quant_dg) == 7) - # pylint: disable=unsupported-membership-test - self.assertTrue(quantizations.DEFAULT in quant.quant_dg) - quant_cfg, _ = quant.quant_dg[quantizations.DEFAULT] - self.assertEqual(quant_cfg.fwd.dg_quantizer.lhs.numerics.dtype, None) - self.assertEqual(quant_cfg.fwd.dg_quantizer.rhs.numerics.bits, 8) - quant_cfg, _ = quant.quant_dg[_QUERY_REGEX] - self.assertEqual(quant_cfg.fwd.dg_quantizer.lhs.numerics.dtype, None) - self.assertEqual(quant_cfg.fwd.dg_quantizer.rhs.numerics.bits, 4) - - def test_mixed_precision_config_subchannel(self): - quant = _configure_quantization( - quant_str="intmp", - quant_cfg_path=os.path.join( - MAXTEXT_CONFIGS_DIR, - "quantization", - "dense_llm_subchannel.json", - ), + def test_use_qwix_quantization_default(self): + # Verify that use_qwix_quantization defaults to True when initializing config + config = pyconfig.initialize( + [None, get_test_config_path()], + enable_checkpointing=False, + quantization="int8", ) - self.assertTrue(isinstance(quant.quant_dg, dict) and len(quant.quant_dg) == 7) - # pylint: disable=unsupported-membership-test - self.assertTrue(quantizations.DEFAULT in quant.quant_dg) - quant_cfg, tile_size = quant.quant_dg[quantizations.DEFAULT] - self.assertEqual(quant_cfg.fwd.dg_quantizer.lhs.numerics.bits, 8) - self.assertEqual(quant_cfg.fwd.dg_quantizer.rhs.numerics.bits, 8) - self.assertEqual(tile_size, -1) - quant_cfg, tile_size = quant.quant_dg[_QUERY_REGEX] - self.assertEqual(quant_cfg.fwd.dg_quantizer.lhs.numerics.bits, 8) - self.assertEqual(quant_cfg.fwd.dg_quantizer.rhs.numerics.bits, 4) - self.assertEqual(tile_size, 128) - - quant_cfg, tile_size = quant.quant_dg[_VALUE_REGEX] - self.assertEqual(quant_cfg.fwd.dg_quantizer.lhs.numerics.bits, 8) - self.assertEqual(quant_cfg.fwd.dg_quantizer.rhs.numerics.bits, 4) - self.assertEqual(tile_size, -1) - - def test_remove_quantized_params(self): - _params = { - "decoder": { - "decoder_norm": {"scale": 1.0}, - "layers": { - "mlp": { - "wi_0": {"kernel": 1.0}, - "wi_1": {"kernel": 1.0}, - "wo": {"kernel": 1.0}, - }, - "self_attention": { - "key": {"kernel": 1.0}, - }, - }, - "logits_dense": {"kernel": 1.0}, - }, - } - _aqt_vars = { - "decoder": { - "layers": { - "mlp": { - "wi_0": { - "AqtDotGeneral_0": { - "qrhs": { - "frozen": aqt_tensor.QTensor( - qvalue=[1.1, 1.0], - scale=[1.0], - scale_t=[1.0], - bias=1.0, - ) - } - } - }, - "wi_1": { - "AqtDotGeneral_0": { - "qrhs": { - "frozen": aqt_tensor.QTensor( - qvalue=[1.1, 1.0], - scale=[1.0], - scale_t=[1.0], - bias=1.0, - ) - } - } - }, - "wo": { - "AqtDotGeneral_0": { - "qrhs": { - "frozen": aqt_tensor.QTensor( - qvalue=[1.1, 1.0], - scale=[1.0], - scale_t=[1.0], - bias=1.0, - ) - } - } - }, - }, - "self_attention": { - "key": { - "AqtDotGeneral_0": { - "qrhs": { - "frozen": aqt_tensor.QTensor( - qvalue=[1.1, 1.0], - scale=[1.0], - scale_t=[1.0], - bias=1.0, - ) - } - } - } - }, - } - } - } - _expected = { - "decoder": { - "decoder_norm": {"scale": 1.0}, - "layers": { - "mlp": { - "wi_0": {"kernel": {}}, - "wi_1": {"kernel": {}}, - "wo": {"kernel": {}}, - }, - "self_attention": { - "key": {"kernel": {}}, - }, - }, - "logits_dense": {"kernel": 1.0}, - } - } - result = quantizations.remove_quantized_params(_params, _aqt_vars) - self.assertEqual(_expected, result) + self.assertTrue(config.use_qwix_quantization) + + def test_unsupported_quantization_without_qwix(self): + # Verify that setting use_qwix_quantization=False with non-native/non-TE quantization raises ValueError + with self.assertRaisesRegex(ValueError, "is unsupported because legacy AQT has been completely removed"): + pyconfig.initialize( + [None, get_test_config_path()], + enable_checkpointing=False, + quantization="int8", + use_qwix_quantization=False, + ) + + def test_supported_quantization_without_qwix(self): + # Verify that setting use_qwix_quantization=False with native FP8 or TE is allowed and does not raise + for quant_type in ["fp8", "nanoo_fp8", "fp8_gpu", "te_fp8_delayedscaling"]: + config = pyconfig.initialize( + [None, get_test_config_path()], + enable_checkpointing=False, + quantization=quant_type, + use_qwix_quantization=False, + ) + self.assertFalse(config.use_qwix_quantization) class QuantTest(unittest.TestCase): @@ -612,6 +375,7 @@ def test_fp8_te_nvfp4_quantization(self): ) @pytest.mark.tpu_only def test_gmm_kernel(group_sizes, k, n, tiling, dtype): + # pylint: disable=undefined-variable """Smoke-test + correctness check for the grouped matrix-multiply kernel. For each group i, gmm should compute @@ -652,5 +416,193 @@ def test_gmm_kernel(group_sizes, k, n, tiling, dtype): assert jnp.abs(quant_out - base_out).mean() / jnp.abs(base_out).mean() < 2e-1 +class QuantizationCoverageTest(unittest.TestCase): + """Explicit tests to ensure 100% test coverage of all quantization paths.""" + + def test_configure_quantization_paths(self): + # Test all configure_quantization paths on CPU (instantiation only) + config_fp8 = pyconfig.initialize( + [None, get_test_config_path()], + enable_checkpointing=False, + quantization="fp8", + use_qwix_quantization=False, + ) + quant_fp8 = quantizations.configure_quantization(config_fp8, "train") + self.assertIsNotNone(quant_fp8) + + config_nanoo = pyconfig.initialize( + [None, get_test_config_path()], + enable_checkpointing=False, + quantization="nanoo_fp8", + use_qwix_quantization=False, + ) + quant_nanoo = quantizations.configure_quantization(config_nanoo, "train") + self.assertIsNotNone(quant_nanoo) + + # Only run TE quantization config test if transformer_engine is installed + try: + import transformer_engine # pylint: disable=unused-import,import-outside-toplevel + + has_te = True + except (ImportError, OSError): + has_te = False + + if has_te: + config_te = pyconfig.initialize( + [None, get_test_config_path()], + enable_checkpointing=False, + quantization="te_fp8_delayedscaling", + use_qwix_quantization=False, + ) + quant_te = quantizations.configure_quantization(config_te, "train") + self.assertIsNotNone(quant_te) + + def test_configure_kv_quant(self): + config = pyconfig.initialize( + [None, get_test_config_path()], + enable_checkpointing=False, + quantize_kvcache=False, + ) + # Should not raise + quantizations.configure_kv_quant(config) + + config_fail = pyconfig.initialize( + [None, get_test_config_path()], + enable_checkpointing=False, + quantize_kvcache=True, + ) + with self.assertRaises(ValueError): + quantizations.configure_kv_quant(config_fail) + + @pytest.mark.tpu_only + def test_moe_quantization_coverage(self): + # Instantiates RoutedMoE on CPU to cover the AQT-free parameter initialization path in moe.py + config = pyconfig.initialize( + [None, get_test_config_path()], + enable_checkpointing=False, + quantization="int8", + use_qwix_quantization=True, + num_experts=2, + base_emb_dim=8, + base_mlp_dim=8, + base_moe_mlp_dim=8, # Required positive base value to derive positive moe_mlp_dim + parameter_memory_host_offload=True, # Cover the parameter offloading paths in linears.py + ) + + devices_array = maxtext_utils.create_device_mesh(config) + mesh = Mesh(devices_array, config.mesh_axes) + + quant = quantizations.configure_quantization(config, "train") + + with mesh: + moe_layer = moe.RoutedMoE( + config=config, + num_experts=config.num_experts, + num_experts_per_tok=1, + mesh=mesh, + kernel_init=nd_dense_init(1.0, "fan_in", "truncated_normal"), + kernel_axes=("expert", "embed_moe", "heads"), + rngs=nnx.Rngs(0), + quant=quant, + ) + + # In Flax NNX, parameters are fully initialized during instantiation. + self.assertIsNotNone(moe_layer.gate.kernel) + + # Execute a forward pass to cover DenseGeneral.__call__, RoutedMoE.__call__, + # sparse_matmul, and the custom quant_einsum wrapper in moe.py + # Batch size must be a multiple of the mesh axis size (e.g. devices_array.size) to be divisible under FSDP sharding + batch_size = max(4, devices_array.size) + inputs = jnp.ones((batch_size, 4, 8), dtype=jnp.float32) + outputs, _, _ = moe_layer(inputs) + self.assertEqual(outputs.shape, (batch_size, 4, 8)) + + def test_quantization_fallbacks(self): + # Cover the fallback return None path in _get_quant_config when an unsupported scheme is passed + config_invalid = pyconfig.initialize( + [None, get_test_config_path()], + quantization="fp8_gpu", + use_qwix_quantization=False, + ) + self.assertIsNone(quantizations.configure_quantization(config_invalid)) + + # Cover the implicit return None path in configure_kv_quant when quantize_kvcache is False + config_no_kv = pyconfig.initialize( + [None, get_test_config_path()], + quantize_kvcache=False, + ) + self.assertIsNone(quantizations.configure_kv_quant(config_no_kv)) + + def test_dense_general_parameter_offload_coverage(self): + # Covers parameter_memory_host_offload paths in linears.py + + dense_layer = linears.DenseGeneral( + in_features_shape=8, + out_features_shape=8, + parameter_memory_host_offload=True, + rngs=nnx.Rngs(0), + ) + inputs = jnp.ones((2, 8), dtype=jnp.float32) + outputs = dense_layer(inputs) + self.assertEqual(outputs.shape, (2, 8)) + + def test_configure_quantization_batch_split_schedule(self): + # Covers use_batch_split_schedule path in quantizations.py + config_bs = pyconfig.initialize( + [None, get_test_config_path()], + enable_checkpointing=False, + use_batch_split_schedule=True, + quantization="fp8_full", + use_qwix_quantization=True, + ) + quant = quantizations.configure_quantization(config_bs, "train") + self.assertIsInstance(quant, quantizations.QwixQuantization) + + config_bs_manual = pyconfig.initialize( + [None, get_test_config_path()], + enable_checkpointing=False, + use_batch_split_schedule=True, + quantization="fp8_full", + use_manual_quantization=True, + use_qwix_quantization=True, + ) + quant_manual = quantizations.configure_quantization(config_bs_manual, "train") + self.assertIsNone(quant_manual) + + @pytest.mark.tpu_only + def test_moe_gemma4_coverage(self): + # Covers GEMMA4 routing and expert scale fusion paths in moe.py + + config = pyconfig.initialize( + [None, get_test_config_path()], + enable_checkpointing=False, + decoder_block="gemma4", + model_call_mode="inference", + fuse_expert_scales=True, + num_experts=2, + base_emb_dim=8, + base_mlp_dim=8, + base_moe_mlp_dim=8, + ) + devices_array = maxtext_utils.create_device_mesh(config) + mesh = Mesh(devices_array, config.mesh_axes) + + with mesh: + moe_layer = moe.RoutedMoE( + config=config, + num_experts=config.num_experts, + num_experts_per_tok=1, + mesh=mesh, + kernel_init=nd_dense_init(1.0, "fan_in", "truncated_normal"), + kernel_axes=("expert", "embed_moe", "heads"), + rngs=nnx.Rngs(0), + ) + # Batch size must be a multiple of the mesh axis size (e.g. devices_array.size) to be divisible under FSDP sharding + batch_size = max(4, devices_array.size) + inputs = jnp.ones((batch_size, 4, 8), dtype=jnp.float32) + outputs, _, _ = moe_layer(inputs) + self.assertEqual(outputs.shape, (batch_size, 4, 8)) + + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/tfds_data_processing_test.py b/tests/unit/tfds_data_processing_test.py index fe1bb75f24..1952391868 100644 --- a/tests/unit/tfds_data_processing_test.py +++ b/tests/unit/tfds_data_processing_test.py @@ -32,6 +32,8 @@ from maxtext.input_pipeline import input_pipeline_interface from tests.utils.test_helpers import get_test_config_path, get_test_base_output_directory +pytestmark = pytest.mark.cpu_only + class TfdsDataProcessingTest(unittest.TestCase): diff --git a/tests/unit/tiling_test.py b/tests/unit/tiling_test.py index 899c1227e4..a302a21a65 100644 --- a/tests/unit/tiling_test.py +++ b/tests/unit/tiling_test.py @@ -268,6 +268,75 @@ def test_vocab_tiling_gradient_with_z_loss(self): "Gradients do not match for vocab tiling when z-loss is enabled.", ) + def test_vocab_tiling_gradient_with_z_loss_nnx(self): + """ + Tests loss and gradient correctness when z-loss is enabled, comparing + standard computation vs. vocabulary tiling computation with NNX enabled. + """ + cfg_non_tiling = pyconfig.initialize( + self.base_config, + run_name="grad_test_z_loss_no_tiling_nnx", + enable_checkpointing=False, + enable_dropout=False, + max_target_length=self.seq_len, + per_device_batch_size=self.batch_size, + logits_via_embedding=False, + base_num_decoder_layers=0, + dtype="float32", + matmul_precision="high", + num_vocab_tiling=1, + z_loss_multiplier=1e-4, # Enable z-loss + enable_nnx=True, + ) + quant_non_tiling = quantizations.configure_quantization(cfg_non_tiling) + devices_array_non_tiling = maxtext_utils.create_device_mesh(cfg_non_tiling) + mesh_non_tiling = Mesh(devices_array_non_tiling, cfg_non_tiling.mesh_axes) + model_non_tiling = models.transformer_as_linen( + cfg_non_tiling, mesh=mesh_non_tiling, quant=quant_non_tiling, model_mode=MODEL_MODE_TRAIN + ) + + rng_model, rng_targets = jax.random.split(self.rng) + + params = model_non_tiling.init( + {"params": rng_model, "dropout": rng_model}, + self.dummy_inputs, + self.dummy_inputs, + ) + + data = { + "targets": jax.random.randint(rng_targets, (self.batch_size, self.seq_len), 0, cfg_non_tiling.vocab_size), + "targets_segmentation": jnp.ones((self.batch_size, self.seq_len)), + } + + loss_non_tiling, grads_non_tiling = self.get_grads(cfg_non_tiling, params, data) + + cfg_tiling = pyconfig.initialize( + self.base_config, + run_name="grad_test_z_loss_with_tiling_nnx", + enable_checkpointing=False, + enable_dropout=False, + max_target_length=self.seq_len, + per_device_batch_size=self.batch_size, + logits_via_embedding=False, + base_num_decoder_layers=0, + dtype="float32", + matmul_precision="high", + num_vocab_tiling=4, + z_loss_multiplier=1e-4, # Enable z-loss + enable_nnx=True, + ) + loss_tiling, grads_tiling = self.get_grads(cfg_tiling, params, data) + + # Loss correctness test + assert jnp.allclose(loss_non_tiling, loss_tiling, rtol=self.rtol), "Losses do not match when z-loss is enabled (NNX)." + + # Gradient correctness test + self.assert_pytrees_all_close( + grads_non_tiling, + grads_tiling, + "Gradients do not match for vocab tiling when z-loss is enabled (NNX).", + ) + @pytest.mark.tpu_only def test_vocab_tiling_nnx_loss(self): """