Add FP4 [E2M1] support to Maxtext [Model convergence not tested yet].

PiperOrigin-RevId: 952967390

Author: Google-ML-Automation

src/maxtext/configs/types.py

@@ -84,6 +84,7 @@ class QuantizationType(str, Enum):
   """Supported quantization schemes."""
 
   NONE = ""
+  FP4 = "fp4"
   INT4 = "int4"
   INT8 = "int8"
   INTMP = "intmp"

src/maxtext/layers/quantizations.py

@@ -824,6 +824,9 @@ def make_qt_rule(dtype) -> list[qwix.QtRule]:
     case "int8":
       return make_qt_rule(jnp.int8)
 
+    case "fp4" | "fp4_e2m1":
+      return make_qt_rule(jnp.float4_e2m1fn)
+
     case "fp8_e5m2":
       return make_qt_rule(jnp.float8_e5m2)
 
@@ -841,7 +844,7 @@ def make_qt_rule(dtype) -> list[qwix.QtRule]:
 def get_qt_provider(config):
   """Get quantization rules based on the config."""
   match config.quantization:
-    case "int4" | "int8" | "fp8" | "fp8_e5m2" | "fp8_e4m3" | "fp8_full":
+    case "int4" | "int8" | "fp4" | "fp4_e2m1" | "fp8" | "fp8_e5m2" | "fp8_e4m3" | "fp8_full":
       return qwix.QtProvider(get_quantization_rule(config))
     case "fp8_gpu":
       return NvidaFp8Provider(get_quantization_rule(config))