Decouple recipe and scaling mode

Signed-off-by: Jeremy Berchtold <[email protected]>

Author: jberchtold-nvidia

tests/jax/test_custom_call_compute.py

@@ -36,6 +36,7 @@
     ScalingMode,
     QuantizerFactory,
     QuantizeLayout,
+    QuantizerParams,
 )
 from transformer_engine.jax.quantize import helper
 from transformer_engine.jax.activation import activation
@@ -188,9 +189,11 @@ def test_act_grad_with_tensor_scaling_fp8(
         )
 
         quantizer = QuantizerFactory.create(
-            scaling_mode=scaling_mode,
-            q_dtype=output_type,
-            q_layout=QuantizeLayout.ROWWISE,
+            QuantizerParams(
+                scaling_mode=scaling_mode,
+                q_dtype=output_type,
+                q_layout=QuantizeLayout.ROWWISE,
+            )
         )
 
         prim_out, (prim_grad,) = value_n_grad_primitive_func(x, activation_type, quantizer)
@@ -219,9 +222,11 @@ def test_act_forward_with_tensor_scaling_fp8(
 
         te_quantizer, jax_quantizer = QuantizerFactory.create(
             n_quantizers=2,
-            scaling_mode=scaling_mode,
-            q_dtype=output_type,
-            q_layout=q_layout,
+            q_params=QuantizerParams(
+                scaling_mode=scaling_mode,
+                q_dtype=output_type,
+                q_layout=q_layout,
+            ),
         )
 
         te_output = tex.act_lu(x, activation_type, te_quantizer)
@@ -244,7 +249,9 @@ def test_act_forward_with_block_scaling_fp8(
         self.activation_type = activation_type
 
         quantizer = QuantizerFactory.create(
-            scaling_mode=ScalingMode.MXFP8_1D_SCALING, q_dtype=output_type, q_layout=q_layout
+            QuantizerParams(
+                scaling_mode=ScalingMode.MXFP8_1D_SCALING, q_dtype=output_type, q_layout=q_layout
+            )
         )
 
         output = tex.act_lu(x, activation_type, quantizer)
@@ -378,7 +385,7 @@ def test_norm_grad_with_tensor_scaling_fp8(
             pytest.skip("RMSNorm and zero_centered_gamma is not supported!")
 
         quantizer = QuantizerFactory.create(
-            scaling_mode=scaling_mode, q_dtype=out_dtype, q_layout=q_layout
+            QuantizerParams(scaling_mode=scaling_mode, q_dtype=out_dtype, q_layout=q_layout)
         )
         self._test_norm_grad(
             n, hidden, norm_type, zero_centered_gamma, epsilon, inp_dtype, quantizer
@@ -406,7 +413,12 @@ def _test_norm_forward(
         gamma = jnp.asarray(gamma, inp_dtype)
 
         quantizer, ref_quantizer = QuantizerFactory.create(
-            n_quantizers=2, scaling_mode=scaling_mode, q_dtype=out_dtype, q_layout=q_layout
+            n_quantizers=2,
+            q_params=QuantizerParams(
+                scaling_mode=scaling_mode,
+                q_dtype=out_dtype,
+                q_layout=q_layout,
+            ),
         )
         if norm_type == "layernorm":
             beta = jax.random.uniform(subkeys[2], (hidden,), jnp.float32, -1, 1)
@@ -562,9 +574,11 @@ def test_qdq(self, in_dtype, input_shape, q_dtype, scaling_mode, q_layout, flatt
 
         # Quantizer is created once as some quantization approaches use state from previous iterations (e.g. delayed scaling)
         quantizer = QuantizerFactory.create(
-            scaling_mode=scaling_mode,
-            q_dtype=q_dtype,
-            q_layout=q_layout,
+            QuantizerParams(
+                scaling_mode=scaling_mode,
+                q_dtype=q_dtype,
+                q_layout=q_layout,
+            )
         )
         # Adding dimension to test if padding is done correctly when flatten 3D to 2D
         if flatten_axis == -2:
@@ -587,7 +601,8 @@ def test_quantize_bitwise(
         input = jax.random.uniform(key, input_shape, in_dtype)
 
         te_quantizer, jax_quantizer = QuantizerFactory.create(
-            n_quantizers=2, q_dtype=q_dtype, scaling_mode=scaling_mode, q_layout=q_layout
+            n_quantizers=2,
+            q_params=QuantizerParams(q_dtype=q_dtype, scaling_mode=scaling_mode, q_layout=q_layout),
         )
 
         jax_output = _jax_quantize(input, quantizer=jax_quantizer, flatten_axis=flatten_axis)
@@ -619,7 +634,10 @@ def test_quantize_dbias(
         input = jax.random.uniform(key, input_shape, in_dtype)
 
         jax_quantizer, te_quantizer = QuantizerFactory.create(
-            n_quantizers=2, q_dtype=out_dtype, scaling_mode=scaling_mode, q_layout=q_layout
+            n_quantizers=2,
+            q_params=QuantizerParams(
+                q_dtype=out_dtype, scaling_mode=scaling_mode, q_layout=q_layout
+            ),
         )
 
         te_output, te_dbias = jit(
@@ -649,7 +667,10 @@ def _test_quantize_dact_dbias(
         dz = jax.random.uniform(subkeys[1], input_shape, in_dtype, -1, 1)
 
         jax_quantizer, te_quantizer = QuantizerFactory.create(
-            n_quantizers=2, q_dtype=out_dtype, scaling_mode=scaling_mode, q_layout=q_layout
+            n_quantizers=2,
+            q_params=QuantizerParams(
+                q_dtype=out_dtype, scaling_mode=scaling_mode, q_layout=q_layout
+            ),
         )
         is_casted_output = te_quantizer is not None
 

tests/jax/test_layer.py

@@ -27,6 +27,8 @@
     ScalingMode,
     is_fp8_available,
     update_collections,
+    UsageContext,
+    UsageType,
 )
 
 
@@ -354,7 +356,13 @@ def test_backward(
                     test_others,
                     test_layer,
                 )
-                if QuantizeConfig.SCALING_MODE == ScalingMode.DELAYED_TENSOR_SCALING:
+                if (
+                    QuantizeConfig.RECIPE_MANAGER is not None
+                    and QuantizeConfig.RECIPE_MANAGER.get_quantizer_params(
+                        UsageContext(UsageType.X)
+                    ).scaling_mode
+                    == ScalingMode.DELAYED_TENSOR_SCALING
+                ):
                     _, updated_quantize_meta = flax.core.pop(
                         updated_state[0], QuantizeConfig.COLLECTION_NAME
                     )

transformer_engine/jax/flax/module.py

@@ -31,7 +31,15 @@
     jax_scaled_masked_softmax,
     jax_scaled_upper_triang_masked_softmax,
 )
-from ..quantize import QuantizerFactory, QuantizeConfig, QuantizeMeta, QuantizeMetaSet, ScalingMode
+from ..quantize import (
+    QuantizerFactory,
+    QuantizeConfig,
+    QuantizeMeta,
+    QuantizeMetaSet,
+    ScalingMode,
+    UsageContext,
+    UsageType,
+)
 from ..sharding import get_non_contracting_logical_axes
 
 PRNGKey = Any
@@ -356,7 +364,13 @@ def generate_quantize_meta(quantizer_name: str):
             ).value
             return QuantizeMeta(scale=scale, amax_history=amax_history)
 
-        if QuantizeConfig.SCALING_MODE == ScalingMode.DELAYED_TENSOR_SCALING:
+        if (
+            QuantizeConfig.RECIPE_MANAGER is not None
+            and QuantizeConfig.RECIPE_MANAGER.get_quantizer_params(
+                UsageContext(UsageType.X)
+            ).scaling_mode
+            == ScalingMode.DELAYED_TENSOR_SCALING
+        ):
             x_meta = generate_quantize_meta("x")
             kernel_meta = generate_quantize_meta("kernel")
             grad_meta = generate_quantize_meta("grad")