diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index 3c5d20fef7..f0465752c2 100644 --- a/src/maxtext/configs/base.yml +++ b/src/maxtext/configs/base.yml @@ -134,6 +134,7 @@ dtype: "bfloat16" # 'nanoo_fp8' for 8-bit floating-point gemms on amd mi300/mi325 gpus. # 'fp8_full' for fp8 quantization with static scaling. # 'serve_fp8_weight' to load and use FP8 checkpoint directly at inference time. +# 'jax_mxfp8' / 'jax_mxfp4' for OCP MX training through jax.nn.scaled_dot_general. quantization: "" # used to configure constant_bound_config in aqt lib for static scaling, e.g. constant_bound_config='0.5, 0.5, 0.5, 0.5, 0.5, 0.5' constant_bound_config: "" diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 1eb27808af..04b9a7398c 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -103,6 +103,8 @@ class QuantizationType(str, Enum): TE_FP8_DS = "te_fp8_delayedscaling" TE_FP8_CS = "te_fp8_currentscaling" TE_MXFP8 = "te_mxfp8" + JAX_MXFP8 = "jax_mxfp8" + JAX_MXFP4 = "jax_mxfp4" TE_NVFP4 = "te_nvfp4" TE_NVFP4_NO_RHT = "te_nvfp4_no_rht" SERVE_FP8_WEIGHT = "serve_fp8_weight" @@ -4448,6 +4450,16 @@ def validate_shard_embed_moe_on_fsdp(self) -> "MaxTextConfig": ) return self + def _validate_jax_mx_quantization(self) -> None: + """Rejects jax_mxfp8 and jax_mxfp4 settings that would silently leave GEMMs unquantized or fail obscurely in JAX.""" + name = self.quantization.value + if self.use_qwix_quantization: + raise ValueError(f"quantization='{name}' does not support use_qwix_quantization=True.") + if self.shard_mode == ShardMode.EXPLICIT: + raise ValueError(f"quantization='{name}' does not support shard_mode='explicit'.") + if self.num_experts > 1: + raise ValueError(f"quantization='{name}' does not support MoE models (num_experts > 1).") + @model_validator(mode="after") def set_derived_and_validate_values(self) -> "MaxTextConfig": """ @@ -4744,6 +4756,8 @@ def get_num_target_devices(): if self.quantization not in ( "fp8", "nanoo_fp8", + "jax_mxfp8", + "jax_mxfp4", ) and not self.quantization.startswith("te_"): logger.warning( "WARNING: AQT quantization is deprecated and will be removed in a future release. " @@ -6068,6 +6082,9 @@ def _stream_chunked(tokenize_data: bool) -> bool: "counterparts) to 1, or turn off `shard_optimizer_over_data`." ) + if self.quantization in ("jax_mxfp8", "jax_mxfp4"): + self._validate_jax_mx_quantization() + # Diloco params # Resolve dcn_diloco_parallelism=-1 if left unspecified, using the same convention as dcn_data_parallelism. # num_diloco_replicas must be computed after this resolution, so we resolve it here rather than diff --git a/src/maxtext/layers/quantizations.py b/src/maxtext/layers/quantizations.py index 301a85f2b3..3b94b24a6e 100644 --- a/src/maxtext/layers/quantizations.py +++ b/src/maxtext/layers/quantizations.py @@ -162,6 +162,45 @@ def einsum(self, dtype: DType = jnp.float32): """Placeholder for einsum implementation in subclasses.""" +class JaxBlockScaledDotGeneral(nnx.Module): + """Stateless OCP MX dot_general: `jax.nn.scaled_dot_general` with one config for both operands and the gradient.""" + + def __init__(self, mode: str): + self.mode = mode + + def __call__( + self, + lhs: jax.Array, + rhs: jax.Array, + dimension_numbers: jax.lax.DotDimensionNumbers, + precision: jax.lax.PrecisionLike = None, + preferred_element_type: jax.typing.DTypeLike | None = None, + ) -> jax.Array: + # jax.nn.scaled_dot_general takes no precision argument. + del precision + output_dtype = jnp.result_type(lhs, rhs) if preferred_element_type is None else preferred_element_type + config = jax.nn.get_scaled_dot_general_config(self.mode) + return jax.nn.scaled_dot_general( + lhs, rhs, dimension_numbers, preferred_element_type=output_dtype, configs=[config, config, config] + ) + + +@dataclass +class JaxBlockScaledQuantization(Quantization): + """Configures OCP MX quantization through `jax.nn.scaled_dot_general`.""" + + quant_mode = "train" + mode: str + + def dot_general_cls(self, mesh_axes: Tuple[str, ...] = ()): + del mesh_axes + return functools.partial(JaxBlockScaledDotGeneral, mode=self.mode) + + def einsum(self, *args, **kwargs): + del args, kwargs + raise ValueError(f"quantization='jax_{self.mode}' does not support einsum.") + + def infer_scale_granularity(scale_shape: Tuple[int, ...]) -> Tuple[str, int | None]: """Infers quantization scheme ('per_tensor', 'per_channel', or 'block_wise') from scale shape.""" nontrivial_axes = [i for i, s in enumerate(scale_shape) if s > 1] @@ -899,6 +938,8 @@ def _get_quant_config(config): """Set quantization params based on user configuration.""" if not config.quantization or config.quantization == "": return None + if config.quantization in ("jax_mxfp8", "jax_mxfp4"): + return config.quantization if config.quantization == "int8": return _get_int8_quant_config(config) if config.quantization == "intmp": @@ -959,7 +1000,9 @@ def configure_quantization(config: Config, quant_mode_str: str = "train"): return None quant_cfg = _get_quant_config(config) if quant_cfg: - if quant_cfg == "fp8": + if quant_cfg in ("jax_mxfp8", "jax_mxfp4"): + return JaxBlockScaledQuantization(mode=quant_cfg.removeprefix("jax_")) + elif quant_cfg == "fp8": return Fp8Quantization() elif quant_cfg == "nanoo_fp8": return NANOOFp8Quantization() diff --git a/tests/unit/jax_mx_quantization_test.py b/tests/unit/jax_mx_quantization_test.py new file mode 100644 index 0000000000..22efa5d6b1 --- /dev/null +++ b/tests/unit/jax_mx_quantization_test.py @@ -0,0 +1,130 @@ +# 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. + +"""Tests for the native JAX MX recipes, quantization=jax_mxfp8 and jax_mxfp4.""" + +import sys +import unittest + +from flax import nnx +import jax +import jax.numpy as jnp +import numpy as np + +from maxtext.configs import pyconfig +from maxtext.layers import linears +from maxtext.layers import quantizations +from tests.utils.test_helpers import get_test_config_path + + +def _supports(mode: str) -> bool: + try: + jax.nn.get_scaled_dot_general_config(mode) + except ValueError: + return False + return True + + +def _modes() -> list[str]: + return [mode for mode in ("mxfp8", "mxfp4") if _supports(mode)] + + +def _config(**kwargs): + """A tiny dense model whose contracting extents are all multiples of 32, with jax_mxfp8 by default.""" + init_kwargs = { + "run_name": "test", + "dataset_type": "synthetic", + "enable_checkpointing": False, + "enable_goodput_recording": False, + "steps": 1, + "per_device_batch_size": 1, + "use_qwix_quantization": False, + "skip_jax_distributed_system": True, + "base_emb_dim": 64, + "base_num_query_heads": 2, + "base_num_kv_heads": 2, + "head_dim": 32, + "base_mlp_dim": 128, + "base_num_decoder_layers": 1, + "max_target_length": 32, + "quantization": "jax_mxfp8", + } | kwargs + return pyconfig.initialize([sys.argv[0], get_test_config_path()], **init_kwargs) + + +class JaxMxQuantizationTest(unittest.TestCase): + """The quantization object and the layer it builds.""" + + def test_configure_quantization(self): + for mode in _modes(): + with self.subTest(mode=mode): + quant = quantizations.configure_quantization(_config(quantization=f"jax_{mode}")) + self.assertIsInstance(quant, quantizations.JaxBlockScaledQuantization) + self.assertEqual(quant.mode, mode) + + def test_layer_matches_scaled_dot_general(self): + x = jax.random.normal(jax.random.key(1), (2, 32, 64), jnp.float32).astype(jnp.bfloat16) + for mode in _modes(): + with self.subTest(mode=mode): + self._check_layer_matches_scaled_dot_general(mode, x) + + def _check_layer_matches_scaled_dot_general(self, mode, x): + """DenseGeneral with the MX backend computes exactly what jax.nn.scaled_dot_general computes.""" + layer = linears.DenseGeneral( + in_features_shape=64, + out_features_shape=128, + weight_dtype=jnp.float32, + dtype=jnp.bfloat16, + quant=quantizations.JaxBlockScaledQuantization(mode=mode), + rngs=nnx.Rngs(0), + ) + kernel = layer.kernel[...].astype(jnp.bfloat16) + config = jax.nn.get_scaled_dot_general_config(mode) + + def reference(lhs): + return jax.nn.scaled_dot_general( + lhs, kernel, (((2,), (0,)), ((), ())), preferred_element_type=jnp.bfloat16, configs=[config, config, config] + ) + + def loss(fn, lhs): + return fn(lhs).astype(jnp.float32).sum() + + y = layer(x) + self.assertEqual(y.dtype, jnp.bfloat16) + np.testing.assert_array_equal(np.asarray(y, np.float32), np.asarray(reference(x), np.float32)) + dx = jax.grad(loss, argnums=1)(layer, x) + dx_ref = jax.grad(loss, argnums=1)(reference, x) + np.testing.assert_array_equal(np.asarray(dx, np.float32), np.asarray(dx_ref, np.float32)) + + +class JaxMxConfigValidationTest(unittest.TestCase): + """Settings that would silently leave GEMMs unquantized or fail obscurely in JAX are rejected at config time.""" + + def assert_rejected(self, regex, **kwargs): + # pydantic's ValidationError, which pyconfig raises, is a ValueError. + with self.assertRaisesRegex(ValueError, regex): + _config(**kwargs) + + def test_qwix_flag(self): + self.assert_rejected("does not support use_qwix_quantization=True", use_qwix_quantization=True) + + def test_explicit_sharding(self): + self.assert_rejected("does not support shard_mode='explicit'", shard_mode="explicit") + + def test_moe(self): + self.assert_rejected("does not support MoE models", num_experts=8, num_experts_per_tok=2, base_moe_mlp_dim=128) + + +if __name__ == "__main__": + unittest.main()