fp4 support (#14625)

Signed-off-by: qiyuw <qiyuw@nvidia.com>
Co-authored-by: qiyuw <qiyuw@nvidia.com>
Co-authored-by: gautham-kollu <gkollu@nvidia.com>

Author: 3 people

nemo/lightning/pytorch/plugins/mixed_precision.py

@@ -69,6 +69,9 @@ class DtypeConfig:
     # fp8 related
     fp8: str = None
     fp8_recipe: str = "delayed"
+    # fp4 related
+    fp4: str = None
+    fp4_recipe: str = "nvfp4"
     first_last_layers_bf16: bool = False
     fp8_margin: int = 0
     fp8_amax_history_len: int = 1
@@ -116,6 +119,9 @@ def __init__(
         fp8_multi_head_attention: bool = False,
         fp8_params: bool = None,
         fp8_param_gather: bool = None,
+        # fp4 related
+        fp4: str = None,
+        fp4_recipe: str = "nvfp4",
         fp16_loss_scale: float = None,
         fp16_initial_loss_scale: float = 4294967296,
         fp16_min_loss_scale: float = 1.0,
@@ -161,6 +167,8 @@ def __init__(
             fp8_multi_head_attention=fp8_multi_head_attention,
             fp8_param=fp8_param_gather,
             fp8_param_gather=fp8_param_gather,
+            fp4=fp4,
+            fp4_recipe=fp4_recipe,
             num_layers_at_start_in_bf16=num_layers_at_start_in_bf16,
             num_layers_at_end_in_bf16=num_layers_at_end_in_bf16,
             reuse_grad_buf_for_mxfp8_param_ag=reuse_grad_buf_for_mxfp8_param_ag,