Integrating TE Collective GEMM

Author: dinodeep

src/maxtext/configs/base.yml

@@ -159,6 +159,8 @@ weight_sparsity_m: null
 weight_sparsity_update_step: 10
 # The first number of steps before updating the sparsity masks.
 weight_sparsity_start_step: 50
+# Whether to use Transformer Engine's collective GEMM overlap algorithm with quantization in MLP layers. Note: only valid with tensor-sequence parallelism.
+use_te_comm_gemm_overlap: false
 
 decoder_block: "llama2" # which style of decoderblock to use.
 # global parameter scale needs to be a power of 2. if you want finer grained control of the model sizes

src/maxtext/configs/types.py

@@ -470,6 +470,9 @@ class Quantization(BaseModel):
       50,
       description=("The first number of steps before updating the sparsity masks."),
   )
+  use_te_comm_gemm_overlap: bool = Field(
+      False, description="If True, uses Transformer Engine's collective GEMM overlap algorithm."
+  )
 
 
 class ModelArchitecture(BaseModel):
@@ -2555,6 +2558,26 @@ def validate_ragged_buffer_factor(self):
           "  2. Ragged sort with ring of experts (use_ring_of_experts=True AND use_ragged_sort=True)"
       )
 
+  def _validate_use_te_comm_gemm_overlap(self):
+    """Validates that use_te_comm_gemm_overlap is used with supported settings to enable TE Collective GEMM ops."""
+    te_has_distributed_env = jax.local_device_count() == 1 and jax.distributed.is_initialized()
+
+    if self.hardware != "gpu_multiprocess" or not te_has_distributed_env:
+      raise ValueError(
+          "TE Collective GEMM operations are only supported for hardware=gpu_multiprocess with rank per GPU."
+      )
+
+    tsp_size = self.ici_tensor_sequence_parallelism * self.dcn_tensor_sequence_parallelism
+    tp_size = self.ici_tensor_parallelism * self.dcn_tensor_parallelism
+
+    if tsp_size == 1 or tp_size > 1:
+      raise ValueError("TE Collective GEMM operations are only supported for TSP size > 1 and TP size == 1.")
+
+    if not self.quantization.startswith("te_"):
+      raise ValueError(
+          "TE Collective GEMM operations are only supported for TE quantization recipes (i.e. starting with 'te_')."
+      )
+
   @model_validator(mode="after")
   def set_derived_and_validate_values(self) -> "MaxTextConfig":
     """
@@ -3453,6 +3476,9 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
 
     self._validate_check_vma_is_supported()
 
+    if self.use_te_comm_gemm_overlap:
+      self._validate_use_te_comm_gemm_overlap()
+
     # Final string-to-enum conversions if they haven't been coerced by pydantic yet.
     if isinstance(self.decoder_block, str):
       self.decoder_block = DecoderBlockType(self.decoder_block.lower())

src/maxtext/layers/quantizations.py

@@ -951,6 +951,8 @@ def __init__(self, config):
 
     self._recipe = TransformerEngineQuantization._get_recipe(config.quantization)
 
+    self._perform_collective_gemm = config.use_te_comm_gemm_overlap
+
   def __hash__(self):
     return hash((self.quant_mode, self._recipe))
 
@@ -1001,11 +1003,13 @@ def _wrap(self, f, name=None):
     2. Wraps the given function in a Flax linen module. This module does not store any Flax
     parameters but can store Flax variables for quantizers if required by the recipe.
 
-    3. When the wrapper is called, it provides an additional argument to the given function `f`,
-    'generate_quantizer_set' as the first argument. 'generate_quantizer_set' is a function that
-    can be called to generate a TransformerEngine/JAX quantizer set object used in
-    TransformerEngine/JAX APIs. 'generate_quantizer_set' will generate quantizers based on the
-    recipe of this TransformerEngineQuantizer object.
+    3. When the wrapper is called, it provides two additional arguments to the given function `f`,
+    'generate_quantizer_set' as the first argument and 'generate_collective_op_set' as the second argument.
+    'generate_quantizer_set' is a function that can be called to generate a TransformerEngine/JAX quantizer
+    set object used in TransformerEngine/JAX APIs based on the recipe of this TransformerEngineQuantizer
+    object. Similarly, 'generate_collective_op_set' is a function that can be called to generate a
+    TransformerEngine/JAX collective operation set object used in TransformerEngine/JAX APIs based
+    the kernel's mesh axes.
 
     Args:
       f: The function to wrap. The first argument must be 'generate_quantizer_set'.
@@ -1016,6 +1020,7 @@ def _wrap(self, f, name=None):
     """
 
     import transformer_engine.jax  # pylint: disable=import-outside-toplevel # pytype: disable=import-error
+    import transformer_engine.jax.cpp_extensions as tex  # pylint: disable=import-outside-toplevel # pytype: disable=import-error
     from transformer_engine.common import recipe  # pylint: disable=import-outside-toplevel # pytype: disable=import-error
 
     default_recipe = self._recipe
@@ -1041,9 +1046,20 @@ def generate_quantizer_set(
             n_groups=n_groups,
         )
 
+      def generate_collective_op_set(self, mesh_axes: Tuple[str, ...] = ()):
+        """Inspect the kernel's mesh axes to determine the type of collective operation to use for collective GEMM."""
+
+        if len(mesh_axes) >= 1:
+          if mesh_axes[0] == "embed" and mesh_axes[-1] == "mlp":
+            return tex.CollectiveOpSet.create(tex.CollectiveOp.ALL_GATHER)
+          elif mesh_axes[0] == "mlp" and mesh_axes[-1] == "embed":
+            return tex.CollectiveOpSet.create(tex.CollectiveOp.REDUCE_SCATTER)
+
+        return tex.noop_collective_op_set
+
       @nn.compact
       def __call__(self, *args, **kwargs):
-        return f(self.generate_quantizer_set, *args, **kwargs)
+        return f(self.generate_quantizer_set, self.generate_collective_op_set, *args, **kwargs)
 
     TEWrapper.__name__ = f"TEWrapper_{name if name else f.__name__}"
 
@@ -1052,17 +1068,22 @@ def __call__(self, *args, **kwargs):
   def dot_general_cls(self, mesh_axes: Tuple[str, ...] = ()):
     """Placeholder for dot_general implementation in subclasses."""
     import transformer_engine.jax  # pylint: disable=import-outside-toplevel # pytype: disable=import-error
+    import transformer_engine.jax.cpp_extensions as tex  # pylint: disable=import-outside-toplevel # pytype: disable=import-error
 
-    def te_dot_general(generate_quantizer_set, x, kernel, dims, **kwargs):
+    def te_dot_general(generate_quantizer_set, generate_collective_op_set, x, kernel, dims, **kwargs):
       contracting_dims, batch_dims = dims
       assert batch_dims == ((), ()), "Batch dimensions must be empty for TransformerEngine dot."
 
       quantizer_set = generate_quantizer_set()
+      collective_op_set = (
+          generate_collective_op_set(mesh_axes) if self._perform_collective_gemm else tex.noop_collective_op_set
+      )
       return transformer_engine.jax.dense.dense(
           x,
           kernel,
           contracting_dims=contracting_dims,
           quantizer_set=quantizer_set,
+          collective_op_set=collective_op_set,
       )
 
     return self._wrap(te_dot_general, "dot_general")

src/maxtext/trainers/pre_train/train.py

@@ -78,9 +78,7 @@
 def get_first_step(model, state):
   if isinstance(model, nn.Module):
     return int(state.step)
-  if hasattr(
-      state, "inner_state"
-  ):  # DiLoCoTrainState (NNX DiLoCo): step is the optimizer step var
+  if hasattr(state, "inner_state"):  # DiLoCoTrainState (NNX DiLoCo): step is the optimizer step var
     return int(state.step.get_value())
   return int(state.optimizer.step.get_value())
 
@@ -784,11 +782,7 @@ def train_loop(config, recorder, state=None):
     # the Zero-1 opt overlay doesn't apply through the diloco wrapper.
     params_shardings = state_mesh_shardings.params
   else:
-    params_shardings, state_mesh_shardings = (
-        sharding.maybe_update_params_sharding_with_opt(
-            config, state_mesh_shardings
-        )
-    )
+    params_shardings, state_mesh_shardings = sharding.maybe_update_params_sharding_with_opt(config, state_mesh_shardings)
 
   p_train_step, p_eval_step = train_utils.jit_train_and_eval_step(
       config,
@@ -828,9 +822,7 @@ def train_loop(config, recorder, state=None):
   if isinstance(model, nn.Module):
     setup_params = state.params
   elif config.enable_diloco:
-    setup_params = (
-        state.params
-    )  # DiLoCoTrainState.params: the outer (global) params
+    setup_params = state.params  # DiLoCoTrainState.params: the outer (global) params
   else:
     _, setup_params, _ = nnx.split(state.model, nnx.Param, ...)
   metric_logger_instance.write_setup_info_to_tensorboard(setup_params)
@@ -935,6 +927,9 @@ def initialize(argv: Sequence[str]) -> tuple[pyconfig.HyperParameters, Any]:
   if config.use_vertex_tensorboard or os.environ.get("UPLOAD_DATA_TO_TENSORBOARD"):
     vertex_tensorboard_manager.configure_vertex_tensorboard(config)
 
+  if config.use_te_comm_gemm_overlap:
+    max_utils.bootstrap_transformer_engine_cgemm(config)
+
   # Create the Goodput recorder
   recorder = create_goodput_recorder(config)
 

src/maxtext/utils/max_utils.py

@@ -1185,6 +1185,22 @@ def print_non_trivial_mesh_axis(mesh):
       print(f"{mesh_axis}: {axis_size}", flush=True)
 
 
+def bootstrap_transformer_engine_cgemm(config):
+  """Potentially initialize NCCL communicators for Collective GEMM operations if
+  the environment is distributed and has the appropriate config."""
+  import transformer_engine.jax.cpp_extensions as tex  # pylint: disable=import-outside-toplevel # pytype: disable=import-error
+
+  tsp_size = config.ici_tensor_sequence_parallelism * config.dcn_tensor_sequence_parallelism
+
+  # Setup NCCL buffers for GPU Collective GEMM operations
+  tex.collective_gemm_bootstrap(
+      jax.device_count(),
+      jax.local_device_count(),
+      jax.process_index(),
+      tsp_size,
+  )
+
+
 @contextmanager
 def maybe_get_transformer_engine_context(config):
   """Runs a transformer engine context engine manager for GPUs only."""
@@ -1212,7 +1228,7 @@ def transformer_engine_context():
     mesh_resource = MeshResource(  # pytype: disable=wrong-arg-types
         dp_resource="data",
         tp_resource="tensor",
-        # tpsp_resource = "tensor_sequence", #TODO(Phuong): add this back when upstreaming CGEMM
+        tpsp_resource="tensor_sequence",
         fsdp_resource="fsdp",
         pp_resource=None,
         cp_resource="context",

tests/integration/train_tests.py

@@ -293,6 +293,22 @@ def test_gpu_nanoo_fp8(self):
   def test_gpu_te_fp8_delayedscaling(self):
     train_main(TrainTests.CONFIGS["te_fp8_delayedscaling"] + ["attention=dot_product"])
 
+  @pytest.mark.skip(reason="No runner with GPU arch >= 89 is available")
+  @pytest.mark.integration_test
+  @pytest.mark.gpu_only
+  def test_gpu_te_fp8_delayedscaling_tsp_cgemm(self):
+    if jax.process_count() <= 1:
+      pytest.skip("Requires rank-per-GPU launch (JAX_PROCESS_COUNT > 1)")
+    if jax.local_device_count() != 1:
+      pytest.skip(f"Requires rank-per-GPU launch (local_device_count==1), " f"got {jax.local_device_count()}")
+    if not jax.distributed.is_initialized():
+      pytest.skip("Requires jax.distributed.initialize() (hardware=gpu_multiprocess)")
+
+    train_main(
+        TrainTests.CONFIGS["te_fp8_delayedscaling"]
+        + ["attention=dot_product", "ici_tensor_sequence_parallelism=2", "use_te_comm_gemm_overlap=true"]
+    )
+
   @pytest.mark.skip(reason="No runner with GPU arch >= 89 is available")
   @pytest.mark.integration_test
   @pytest.mark.gpu_only