diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py index f051e7a4be..92e52b771f 100644 --- a/src/maxtext/layers/moe.py +++ b/src/maxtext/layers/moe.py @@ -64,6 +64,9 @@ DISPATCH = "dispatch" COMBINE = "combine" +WI_0 = "wi_0" +WI_1 = "wi_1" +WO = "wo" @struct.dataclass @@ -497,6 +500,16 @@ def __init__( else: self._expert_parallelism_name = "expert" + if isinstance(self.quant, (quantizations.Fp8Quantization, quantizations.NANOOFp8Quantization)): + einsum_names = [WI_0, WI_1, WO] + if self.config.capacity_factor > 0: + einsum_names += [DISPATCH, COMBINE] + self.quant_einsums = nnx.Dict( + {name: quantizations.create_fp8_einsum(self.quant, self.dtype, self.rngs) for name in einsum_names} + ) + else: + self.quant_einsums = None + self.gate = GateLogit( in_features_shape=self.moe_expert_input_dim, out_features_shape=self.num_experts, @@ -2717,15 +2730,22 @@ def get_einsum( return jnp.einsum if self.quant: + op_id = einsum_name if einsum_name is not None else "einsum" - def aqt_einsum(*args, **kwargs): # pylint: disable=unused-argument + def quant_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 - - einsum_op = aqt_einsum + if self.quant_einsums is not None: + if op_id not in self.quant_einsums: + raise ValueError( + f"Einsum name '{op_id}' is not registered in quant_einsums. " + f"Available names: {list(self.quant_einsums.keys())}" + ) + return self.quant_einsums[op_id](*args, mutable=["_overwrite_with_gradient"]) + einsum = self.quant.einsum(mesh_axes=rhs_mesh_axes) # pytype: disable=attribute-error + return quantizations.apply_einsum_in_nnx(self, op_id, einsum, ["aqt"], *args) + + einsum_op = quant_einsum else: einsum_op = jnp.einsum return einsum_op @@ -2929,7 +2949,7 @@ def dense_matmul( with jax.named_scope("wi_0"): w0_kernel_axes = ("exp", None, "mlp") w0_kernel = self.maybe_all_gather_kernel_weight_in_expert_parallelism(w0_kernel, w0_kernel_axes) - layer_w0 = self.get_einsum(rhs_mesh_axes=w0_kernel_axes)( + layer_w0 = self.get_einsum(rhs_mesh_axes=w0_kernel_axes, einsum_name=WI_0)( mlp_up_einsum, dispatch, w0_kernel, precision=matmul_precision ) if self.config.mlp_bias: @@ -2943,7 +2963,7 @@ def dense_matmul( with jax.named_scope("wi_1"): w1_kernel_axes = ("exp", None, "mlp") w1_kernel = self.maybe_all_gather_kernel_weight_in_expert_parallelism(w1_kernel, w1_kernel_axes) - layer_w1 = self.get_einsum(rhs_mesh_axes=w1_kernel_axes)( + layer_w1 = self.get_einsum(rhs_mesh_axes=w1_kernel_axes, einsum_name=WI_1)( mlp_up_einsum, dispatch, w1_kernel, precision=matmul_precision ) if self.config.mlp_bias: @@ -2957,7 +2977,7 @@ def dense_matmul( with jax.named_scope("wo"): wo_kernel_axes = ("exp", "mlp", None) wo_kernel = self.maybe_all_gather_kernel_weight_in_expert_parallelism(wo_kernel, wo_kernel_axes) - intermediate_layer = self.get_einsum(rhs_mesh_axes=wo_kernel_axes)( + intermediate_layer = self.get_einsum(rhs_mesh_axes=wo_kernel_axes, einsum_name=WO)( mlp_down_einsum, layer_multiply, wo_kernel, @@ -3007,7 +3027,7 @@ def dense_matmul( ), ) with jax.named_scope("wi_0"): - layer_w0 = self.get_einsum(rhs_mesh_axes=self.wi_kernel_axes)( + layer_w0 = self.get_einsum(rhs_mesh_axes=self.wi_kernel_axes, einsum_name=WI_0)( "BSM,EMH -> BSEH", inputs, w0_kernel, precision=matmul_precision ) if self.config.mlp_bias: @@ -3016,7 +3036,7 @@ def dense_matmul( layer_w0 = layer_w0.astype(jnp.float32) layer_w0 = adc.checkpoint_name(adc.checkpoint_name(layer_w0, "mlpwi_0"), "moe_mlpwi_0") with jax.named_scope("wi_1"): - layer_w1 = self.get_einsum(rhs_mesh_axes=self.wi_kernel_axes)( + layer_w1 = self.get_einsum(rhs_mesh_axes=self.wi_kernel_axes, einsum_name=WI_1)( "BSM,EMH -> BSEH", inputs, w1_kernel, precision=matmul_precision ) if self.config.mlp_bias: @@ -3027,7 +3047,7 @@ def dense_matmul( layer_multiply = self.apply_ffn_activation(layer_w0, layer_w1) with jax.named_scope("wo"): - intermediate_layer = self.get_einsum(rhs_mesh_axes=self.wo_kernel_axes)( + intermediate_layer = self.get_einsum(rhs_mesh_axes=self.wo_kernel_axes, einsum_name=WO)( "BSEH,EHM -> BSEM", layer_multiply, wo_kernel, diff --git a/src/maxtext/layers/quantizations.py b/src/maxtext/layers/quantizations.py index a275a0afa8..29599cc54a 100644 --- a/src/maxtext/layers/quantizations.py +++ b/src/maxtext/layers/quantizations.py @@ -401,6 +401,10 @@ def dot_general_cls(self, mesh_axes: Tuple[str, ...] = ()): """Returns dot_general configured with aqt params.""" return nn.NANOOFp8DotGeneralOp + def einsum(self, dtype: DType = jnp.float32): + """Returns an einsum using the NANOO (fnuz) fp8 formats of AMD MI300/MI325.""" + return Fp8Einsum(dtype=dtype, e4m3_dtype=jnp.float8_e4m3fnuz, e5m2_dtype=jnp.float8_e5m2fnuz) + def _get_int8_quant_config(config): """Get int8 quantization configuration.""" @@ -768,6 +772,59 @@ def _apply_linen_module_in_nnx(linen_module_cls, op_id, *args, **kwargs): return linen_module_cls(name=op_id)(*args, **kwargs) +def create_fp8_einsum(quant: Quantization, dtype: DType, rngs: nnx.Rngs) -> nnx_wrappers.ToNNX: + """Creates an fp8 einsum that an `nnx.Module` can call. + + An fp8 einsum holds its scaling factors and amax histories in Linen variables, which can + only be created while the module is bound to a Linen scope. The bridge into NNX therefore + has to happen while the parent module is being built; creating the state on the first call + instead would grow the module graph inside the scanned layer loop, which NNX rejects. + + The state has a fixed shape, so a canonical pair of operands is enough to materialize it. + The returned einsum still accepts operands of any shape. + + Args: + quant: The fp8 quantization providing the einsum. + dtype: The computation dtype of the einsum. + rngs: The `nnx.Rngs` of the parent module. + """ + wrapper = nnx_wrappers.ToNNX(quant.einsum(dtype=dtype), rngs=rngs) # pytype: disable=attribute-error + dummy_operand = jnp.zeros((1, 1), dtype=dtype) + wrapper.lazy_init("ab,bc->ac", dummy_operand, dummy_operand) + return wrapper + + +def apply_einsum_in_nnx(parent: nnx.Module, op_id: str, einsum, mutable: Sequence[str], *args): + """Applies a quantized Linen einsum from within an NNX parent module. + + A Linen module cannot be called unbound, which is all an `nnx.Module` can offer it, so the + einsum is bridged into NNX on first use and the bridged instance is reused afterwards. + `op_id` must be unique per call site: two call sites sharing an id would also share + quantization state. Bridging here rather than while building the parent needs the operands, + so it only suits einsums whose state is shaped after them, such as AQT. + + Args: + parent: The NNX module hosting the einsum; it must carry an `rngs` attribute. + op_id: Stable identifier for the call site. + einsum: The einsum returned by a `Quantization`, either a Linen module or a callable. + mutable: Variable collections the einsum writes to. + *args: Arguments forwarded to the einsum. + """ + linen_einsum = einsum.func if isinstance(einsum, functools.partial) else einsum + if not isinstance(linen_einsum, nn.Module): + return einsum(*args) + + # The name must not start with an underscore: NNX treats such attributes as static, + # which cannot hold the bridged module's variables. + attr_name = f"quant_einsum_{op_id}" + wrapper = getattr(parent, attr_name, None) + if wrapper is None: + wrapper = nnx_wrappers.ToNNX(linen_einsum, rngs=parent.rngs) + wrapper.lazy_init(*args) + setattr(parent, attr_name, wrapper) + return wrapper(*args, mutable=list(mutable)) + + class NvidaFp8Provider(qwix.QtProvider): """Wraps nn.Fp8DirectDotGeneralOp with Qwix's provider interface.""" diff --git a/tests/integration/train_tests.py b/tests/integration/train_tests.py index 07f8fb446f..48adc26977 100644 --- a/tests/integration/train_tests.py +++ b/tests/integration/train_tests.py @@ -56,6 +56,17 @@ class TrainTests(unittest.TestCase): "sharding_tolerance=0.1", ] + # Routes the MoE layer through dense_matmul, which is what runs wherever the megablox and + # ragged kernels are unavailable. + _moe_model_overrides = [ + "decoder_block=mixtral", + "num_experts=4", + "num_experts_per_tok=2", + "base_moe_mlp_dim=32", + "sparse_matmul=False", + "megablox=False", + ] + CONFIGS = { "base": [ # short test for train.py with TFDS c4 None, @@ -135,6 +146,19 @@ class TrainTests(unittest.TestCase): rf"tokenizer_path={os.path.join(MAXTEXT_ASSETS_ROOT, 'tokenizers', 'tokenizer.llama2')}", ] + _small_model_overrides, + "moe": [ # tests a MoE model, to be combined with a quantization + None, + get_test_config_path(), + f"base_output_directory={_base_output_directory}", + "run_name=runner_test", + "dataset_type=synthetic", # use synthetic dataset_type to decrease training time + "steps=2", + "enable_checkpointing=False", + "enable_goodput_recording=False", + rf"tokenizer_path={os.path.join(MAXTEXT_ASSETS_ROOT, 'tokenizers', 'tokenizer.llama2')}", + ] + + _small_model_overrides + + _moe_model_overrides, "te_fp8_delayedscaling": [ # tests base config with te_fp8_delayedscaling None, get_test_config_path(), @@ -288,6 +312,26 @@ def test_gpu_fp8(self): def test_gpu_nanoo_fp8(self): train_main(TrainTests.CONFIGS["nanoo_fp8"] + ["attention=dot_product"]) + # The quantized MoE tests below carry no hardware marker on purpose. They cover how the + # quantized einsums are bound to the MoE layer, which breaks on every backend when it breaks, + # and both fp8 flavors are emulated in XLA rather than needing hardware support. + @pytest.mark.integration_test + def test_moe_int8(self): + train_main(TrainTests.CONFIGS["moe"] + ["quantization=int8"]) + + @pytest.mark.integration_test + def test_moe_fp8(self): + train_main(TrainTests.CONFIGS["moe"] + ["quantization=fp8"]) + + @pytest.mark.integration_test + def test_moe_nanoo_fp8(self): + train_main(TrainTests.CONFIGS["moe"] + ["quantization=nanoo_fp8"]) + + @pytest.mark.integration_test + def test_moe_fp8_token_dropping(self): + # capacity_factor > 0 adds the dispatch and combine einsums to the ones above. + train_main(TrainTests.CONFIGS["moe"] + ["quantization=fp8", "capacity_factor=1.25"]) + @pytest.mark.skip(reason="No runner with GPU arch >= 89 is available") @pytest.mark.integration_test @pytest.mark.gpu_only diff --git a/tests/unit/moe_test.py b/tests/unit/moe_test.py index c183e33d0c..a5813d62a1 100644 --- a/tests/unit/moe_test.py +++ b/tests/unit/moe_test.py @@ -33,7 +33,7 @@ from maxtext.layers import moe from maxtext.layers import nnx_wrappers from maxtext.layers.initializers import NdInitializer, nd_dense_init, variable_to_logically_partitioned -from maxtext.layers.quantizations import Fp8Quantization +from maxtext.layers.quantizations import configure_quantization, Fp8Quantization from maxtext.utils import max_logging, maxtext_utils from maxtext.utils.sharding import remove_expert_from_partition_spec from tests.utils.test_helpers import get_test_config_path @@ -1735,6 +1735,111 @@ def loss_fn(params, x): max_logging.log("\n" + diff_summary) +class GetEinsumTest(parameterized.TestCase): + """Tests for the quantized einsums RoutedMoE.get_einsum hands to dense_matmul.""" + + def _make_moe(self, quant): + """Builds a small RoutedMoE on the dense_matmul path with the given quantization.""" + cfg = pyconfig.initialize( + [None, get_test_config_path()], + run_name="get_einsum_test", + enable_checkpointing=False, + decoder_block="mixtral", + num_experts=4, + num_experts_per_tok=2, + base_emb_dim=64, + base_mlp_dim=32, + base_moe_mlp_dim=32, + dtype="float32", + weight_dtype="float32", + megablox=False, + sparse_matmul=False, + max_target_length=8, + per_device_batch_size=1, + ) + devices_array = maxtext_utils.create_device_mesh(cfg) + return moe.RoutedMoE( + config=cfg, + num_experts=cfg.num_experts, + num_experts_per_tok=cfg.num_experts_per_tok, + mesh=Mesh(devices_array, cfg.mesh_axes), + kernel_init=nd_dense_init(1.0, "fan_in", "truncated_normal"), + kernel_axes=("embed", "mlp"), + dtype=jnp.float32, + quant=quant, + rngs=nnx.Rngs(0), + ) + + def _quantization(self, quantization): + """Returns the quantization object the config string maps to.""" + return configure_quantization( + pyconfig.initialize( + [None, get_test_config_path()], + enable_checkpointing=False, + quantization=quantization, + ) + ) + + def test_fp8_einsum_is_bound(self): + model = self._make_moe(Fp8Quantization()) + einsum_fn = model.get_einsum(einsum_name=moe.WI_0) + result = einsum_fn("ab,bc->ac", jnp.ones((2, 3)), jnp.ones((3, 4))) + self.assertEqual(result.shape, (2, 4)) + + def test_aqt_einsum_is_bound(self): + model = self._make_moe(self._quantization("int8")) + self.assertIsNone(model.quant_einsums) + einsum_fn = model.get_einsum(einsum_name=moe.WI_0) + result = einsum_fn("ab,bc->ac", jnp.ones((2, 3)), jnp.ones((3, 4))) + self.assertEqual(result.shape, (2, 4)) + + def test_unregistered_quant_einsum_name_raises(self): + model = self._make_moe(Fp8Quantization()) + einsum_fn = model.get_einsum(einsum_name="not_registered") + with self.assertRaises(ValueError) as ctx: + einsum_fn("ab,bc->ac", jnp.ones((2, 3)), jnp.ones((3, 4))) + self.assertIn("not_registered", str(ctx.exception)) + self.assertIn("Available names", str(ctx.exception)) + + @parameterized.named_parameters( + ("fp8", "fp8", jnp.float8_e4m3fn), + ("nanoo_fp8", "nanoo_fp8", jnp.float8_e4m3fnuz), + ) + def test_fp8_einsum_quantizes_both_operands(self, quantization, e4m3_dtype): + """The bridged einsum is the plain one with both operands cast to the scheme's e4m3.""" + model = self._make_moe(self._quantization(quantization)) + lhs = jax.random.normal(jax.random.PRNGKey(0), (4, 16), dtype=jnp.float32) + rhs = jax.random.normal(jax.random.PRNGKey(1), (16, 8), dtype=jnp.float32) + + actual = model.get_einsum(einsum_name=moe.WI_0)("ab,bc->ac", lhs, rhs) + + # The scaling factors start at 1 and are only updated on the backward pass, so a forward + # call on a freshly built layer quantizes by a plain cast. + quantized = jnp.einsum( + "ab,bc->ac", lhs.astype(e4m3_dtype).astype(jnp.float32), rhs.astype(e4m3_dtype).astype(jnp.float32) + ) + np.testing.assert_array_equal(np.asarray(actual), np.asarray(quantized)) + self.assertFalse(np.array_equal(np.asarray(actual), np.asarray(jnp.einsum("ab,bc->ac", lhs, rhs)))) + + @parameterized.named_parameters(("fp8", "fp8"), ("nanoo_fp8", "nanoo_fp8")) + def test_quantized_dense_matmul_tracks_unquantized(self, quantization): + """A quantized MoE layer follows the same layer run unquantized, to within e4m3.""" + reference = self._make_moe(None) + model = self._make_moe(self._quantization(quantization)) + copy_weights(reference, model) + + inputs = jax.random.normal(jax.random.PRNGKey(42), (1, 8, reference.config.base_emb_dim), dtype=jnp.float32) + expected, _, _ = reference(inputs) + actual, _, _ = model(inputs) + + self.assertTrue(np.isfinite(actual).all()) + # e4m3 keeps three mantissa bits, and the layer is quantized at each of wi_0, wi_1 and wo, + # so the agreement is loose; it is the same threshold the qwix MoE test above uses. + relative_error = np.linalg.norm(actual - expected) / np.linalg.norm(expected) + self.assertLess(relative_error, 0.22) + self.assertFalse(np.allclose(actual, expected)) + + def make_moe(cfg, mesh): return moe.RoutedMoE( config=cfg, diff --git a/tests/unit/quantizations_test.py b/tests/unit/quantizations_test.py index 1c34d128be..d014a66ee0 100644 --- a/tests/unit/quantizations_test.py +++ b/tests/unit/quantizations_test.py @@ -717,6 +717,52 @@ def test_nnx_abstract_state_has_no_intermediates(self): self.assertNotIn("intermediates", state_dict) +class EinsumParent(nnx.Module): + """Minimal NNX parent for apply_einsum_in_nnx tests.""" + + def __init__(self, rngs: nnx.Rngs): + self.rngs = rngs + + +class MoEQuantizedEinsumTest(unittest.TestCase): + """Tests for MoE quantized einsum helpers.""" + + def test_nanoo_fp8_einsum_uses_fnuz_dtypes(self): + quant = quantizations.NANOOFp8Quantization() + einsum_mod = quant.einsum(dtype=jnp.float32) + self.assertEqual(einsum_mod.e4m3_dtype, jnp.float8_e4m3fnuz) + self.assertEqual(einsum_mod.e5m2_dtype, jnp.float8_e5m2fnuz) + + def test_create_fp8_einsum(self): + for quant_str in ("fp8", "nanoo_fp8"): + quant = _configure_quantization(quant_str=quant_str) + wrapper = quantizations.create_fp8_einsum(quant, jnp.float32, nnx.Rngs(0)) + lhs = jnp.ones((4, 8)) + rhs = jnp.ones((8, 16)) + result = wrapper("ab,bc->ac", lhs, rhs, mutable=["_overwrite_with_gradient"]) + self.assertEqual(result.shape, (4, 16)) + + def test_apply_einsum_in_nnx_plain_callable(self): + parent = EinsumParent(nnx.Rngs(0)) + lhs = jnp.ones((2, 3)) + rhs = jnp.ones((3, 4)) + result = quantizations.apply_einsum_in_nnx(parent, "plain", jnp.einsum, [], "ij,jk->ik", lhs, rhs) + expected = jnp.einsum("ij,jk->ik", lhs, rhs) + self.assertTrue(jnp.allclose(result, expected)) + + def test_apply_einsum_in_nnx_aqt_reuses_wrapper(self): + quant = _configure_quantization(quant_str="int8") + parent = EinsumParent(nnx.Rngs(0)) + lhs = jnp.ones((2, 2)) + rhs = jnp.ones((2, 2)) + einsum = quant.einsum(mesh_axes=()) + result1 = quantizations.apply_einsum_in_nnx(parent, "aqt_test", einsum, ["aqt"], "bc,ab->ac", lhs, rhs) + wrapper = getattr(parent, "quant_einsum_aqt_test") + result2 = quantizations.apply_einsum_in_nnx(parent, "aqt_test", einsum, ["aqt"], "bc,ab->ac", lhs, rhs) + self.assertIs(getattr(parent, "quant_einsum_aqt_test"), wrapper) + self.assertEqual(result1.shape, result2.shape) + + class StaticScaleTest(unittest.TestCase): """Tests for static scale extraction."""