Quantize activations before EP all gather

run_qwen3_next_80b_xpk.sh

@@ -0,0 +1,152 @@
+#!/bin/bash
+set -e
+
+# Activate Python virtual environment
+source /home/shuwenf_google_com/venv-maxtext/bin/activate
+
+# --- Environment Variables ---
+export PROJECT_ID="tpu-prod-env-one-vm"
+export CLUSTER_NAME="bodaborg-v6e-256-lcscld-c"
+export ZONE="southamerica-west1-a"
+
+# --- Configuration & Automated Image Build ---
+TIMESTAMP=$(date +%m%d%H%M%S)
+export WORKLOAD_IMAGE="gcr.io/tpu-prod-env-one-vm/param3_21jul:shuwenf_${TIMESTAMP}"
+export WORKLOAD_NAME="shuwenf-qn80b-${TIMESTAMP}"
+export DEVICE_TYPE="v6e-256"
+export NUM_SLICES=1
+export PRIORITY="very-high"
+export NUM_STEPS=15
+export MAX_RESTARTS=${MAX_RESTARTS:-0}
+export MODEL_NAME="qwen3-next-80b-a3b"
+export BASE_OUTPUT_DIR="gs://shuwenf-hlo-dumps/qwen3-next-80b-profiles/run-${TIMESTAMP}"
+
+echo "========================================================================"
+echo "Building and uploading Docker runner image for full 15-step execution..."
+echo "Target Image: ${WORKLOAD_IMAGE}"
+echo "========================================================================"
+
+(
+  cd /home/shuwenf_google_com/maxtext &&   sudo CLOUD_IMAGE_NAME="${WORKLOAD_IMAGE}"   BASE_IMAGE="gcr.io/tpu-prod-env-one-vm/param3_21jul:latest"   bash src/dependencies/scripts/docker_upload_runner.sh
+)
+
+echo "Docker image upload complete: ${WORKLOAD_IMAGE}"
+
+# --- XLA Flags ---
+XLA_FLAGS_ARRAY=(
+  "--xla_tpu_scheduler_percent_shared_memory_limit=35"
+  "--xla_msa_enable_sync_slice_replacement=false"
+  "--xla_tpu_enable_sparse_core_collective_offload_2d_all_gather=true"
+  "--xla_msa_enable_sync_copy_replacement=false"
+  "--xla_tpu_scoped_vmem_limit_kib=81000"
+  "--xla_tpu_enable_sparse_core_collective_offload_all_gather=true"
+  "--xla_tpu_enable_sparse_core_collective_offload_all_reduce=true"
+  "--xla_tpu_offload_gather_to_sparsecore=true"
+  "--xla_tpu_dvfs_p_state=7"
+  "--xla_tpu_disable_sparse_core_collective_offload_remover=true"
+  "--xla_tpu_enable_async_collective_fusion=true"
+  "--xla_tpu_overlap_compute_collective_tc=true"
+  "--xla_tpu_enable_async_collective_fusion_multiple_steps=true"
+  "--xla_tpu_enable_latency_hiding_scheduler=true"
+  "--xla_latency_hiding_scheduler_rerun=10"
+  "--xla_tpu_all_gather_collective_matmul_mode=post_spmd_conservative"
+  "--xla_tpu_reduce_scatter_collective_matmul_mode=post_spmd_conservative"
+  "--xla_latency_hiding_scheduler_enable_selective_resources=true"
+  "--xla_tpu_enable_ilp_latency_hiding_scheduler=true"
+  "--xla_tpu_enable_all_experimental_scheduler_features=true"
+  "--xla_tpu_enable_scheduler_memory_pressure_tracking=true"
+  "--xla_tpu_host_transfer_overlap_limit=24"
+  "--xla_tpu_aggressive_opt_barrier_removal=ENABLED"
+  "--xla_lhs_prioritize_async_depth_over_stall=DISABLED"
+  "--xla_tpu_enable_ag_backward_pipelining=true"
+  "--xla_should_allow_loop_variant_parameter_in_chain=ENABLED"
+  "--xla_should_add_loop_invariant_op_in_chain=ENABLED"
+  "--xla_max_concurrent_host_send_recv=100"
+)
+export XLA_FLAGS="${XLA_FLAGS_ARRAY[*]}"
+
+# --- MaxText Workload Overrides ---
+MAXTEXT_ARGS_ARRAY=(
+  "model_name=${MODEL_NAME}"
+  "base_output_directory=${BASE_OUTPUT_DIR}"
+  "run_name=param-3"
+  "dataset_type=synthetic"
+  "dataset_name=synthetic"
+  "dtype=bfloat16"
+  "allow_split_physical_axes=True"
+  "ici_expert_parallelism=4"
+  "use_ring_of_experts=True"
+  "custom_mesh=hybrid_ring_64x4"
+  "use_ragged_sort=True"
+  "use_random_routing=True"
+  "per_device_batch_size=4"
+  "opt_type=muon"
+  "muon_consistent_rms=0.2"
+  "muon_weight_decay=0.1"
+  "learning_rate=1e-5"
+  "max_target_length=2048"
+  "ragged_buffer_factor=1.5"
+  "remat_policy=full"
+  "reuse_example_batch=1"
+  "decoder_layer_input=offload"
+  "context=device"
+  "ici_fsdp_parallelism=-1"
+  "steps=15"
+  "shard_exp_on_fsdp=True"
+  "sharding_tolerance=0.5"
+  "sa_q_layout=SEQ_MINOR"
+  "sa_k_layout=HEAD_DIM_MINOR"
+  "sa_v_layout=HEAD_DIM_MINOR"
+  "sa_block_q=2048"
+  "sa_block_kv=2048"
+  "sa_block_kv_compute=1024"
+  "sa_block_q_dkv=2048"
+  "sa_block_kv_dkv=2048"
+  "sa_block_kv_dkv_compute=1024"
+  "hardware=tpu"
+  "skip_jax_distributed_system=False"
+  "attention=flash"
+  "use_tokamax_splash=True"
+  "sa_use_fused_bwd_kernel=True"
+  "use_tokamax_gmm=True"
+  "use_gmm_v2=True"
+  "sparse_matmul=True"
+  "megablox=True"
+  "optimizer_memory_host_offload=True"
+  "parameter_memory_host_offload=False"
+  "enable_checkpointing=False"
+  "async_checkpointing=False"
+  "tokenizer_type=tiktoken"
+  "tokenizer_path=tokenizer_74B/"
+  "override_model_config=true"
+  "mhc_expansion_rate=4"
+  "profiler=xplane"
+  "profiler_steps=5"
+  "skip_first_n_steps_for_profiler=2"
+  "enable_tpu_profiling_options=True"
+  "upload_all_profiler_results=true"
+)
+MAXTEXT_ARGS="${MAXTEXT_ARGS_ARRAY[*]}"
+
+# Clean container temporary log setup safely
+RUN_COMMAND="set -e && rm -rf /tmp/tpu_logs/* 2>/dev/null || true; mkdir -p /tmp/tpu_logs && export LIBTPU_INIT_ARGS=\"${XLA_FLAGS}\" && export JAX_PLATFORMS='tpu,cpu' && export ENABLE_PJRT_COMPATIBILITY='true' && export JAX_DISTRIBUTED_INITIALIZE_TIMEOUT=1800 && export PYTHONPATH=/deps:/deps/src:/deps/src/maxtext/src && python3 src/maxtext/trainers/pre_train/train.py src/maxtext/configs/base.yml ${MAXTEXT_ARGS}"
+
+# --- XPK Workload Creation ---
+echo "Creating XPK workload: ${WORKLOAD_NAME} on cluster: ${CLUSTER_NAME}"
+
+python3 -m xpk.main workload create   --cluster="${CLUSTER_NAME}"   --project="${PROJECT_ID}"   --zone="${ZONE}"   --priority="${PRIORITY}"   --max-restarts="${MAX_RESTARTS}"   --device-type="${DEVICE_TYPE}"   --num-slices="${NUM_SLICES}"   --docker-image="${WORKLOAD_IMAGE}"   --enable-debug-logs   --workload="${WORKLOAD_NAME}"   --command="${RUN_COMMAND}"
+
+LOGS_URL="https://console.cloud.google.com/logs/query;query=resource.type%3D%22k8s_container%22%0Aresource.labels.project_id%3D%22${PROJECT_ID}%22%0Aresource.labels.location%3D%22southamerica-west1%22%0Aresource.labels.cluster_name%3D%22${CLUSTER_NAME}%22%0Aresource.labels.namespace_name%3D%22default%22%0Aresource.labels.pod_name%3A%22${WORKLOAD_NAME}-slice-job-0-0-%22%0Aseverity%3E%3DDEFAULT;storageScope=project;duration=P1D?project=${PROJECT_ID}"
+GKE_URL="https://console.cloud.google.com/kubernetes/service/southamerica-west1/${CLUSTER_NAME}/default/${WORKLOAD_NAME}/details?project=${PROJECT_ID}"
+TB_URL="https://tensorboard.corp.google.com/?logdir=${BASE_OUTPUT_DIR}/param-3/tensorboard"
+
+echo "========================================================================"
+echo "📋 Pantheon Cloud Logging (Worker 0 Logs):"
+echo "${LOGS_URL}"
+echo ""
+echo "☸️ GKE Workload Details:"
+echo "${GKE_URL}"
+echo ""
+echo "📊 GCS TensorBoard Link:"
+echo "${TB_URL}"
+echo "========================================================================"

src/maxtext/configs/base.yml

@@ -226,6 +226,8 @@ load_balance_loss_weight: 0.0 # weight for the load balance loss
 use_random_routing: false # whether to use random routing for debug/test purpose
 use_custom_sort_vjp: true # whether to use a custom VJP sort for efficient backward pass processing in sparse matmul
 use_ring_of_experts: false # whether to use ring of experts for sparse matmul expert parallelism
+quantize_before_ep_all_gather: true # whether to quantize activations before the ring-of-experts EP all-gather
+                                    # (fp8 collective + ragged sort) vs. quantizing later inside the gmm call
 num_moe_emb_chunks: 0 # number of chunks for overlapping token all-gather and GMM computation along embedding dimension
 # If true, peel the 'expert' mesh axis off the MoE dispatch/MLP batch dim so the expert GEMM
 # stays expert-parallel (AllToAll); false keeps 'expert' on the batch dim (activation_batch_moe).

src/maxtext/configs/types.py

@@ -867,6 +867,13 @@ class MoEGeneral(BaseModel):
       False,
       description="Whether to use Ring of Experts for sparse matmul expert parallelism.",
   )
+  quantize_before_ep_all_gather: bool = Field(
+      True,
+      description=(
+          "Whether to quantize activations before the Ring of Experts EP all-gather (so the "
+          "collective and ragged sort move fp8, not bf16), vs. quantizing later inside the gmm call."
+      ),
+  )
   moe_dispatch_no_expert_sharding: bool = Field(
       False,
       description=(
@@ -2841,6 +2848,16 @@ 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_quantize_before_ep_all_gather(self):
+    """Validates quantize_before_ep_all_gather is used with supported settings."""
+    if self.quantize_before_ep_all_gather and not (
+        self.use_ring_of_experts and self.use_qwix_quantization and self.use_gmm_v2
+    ):
+      raise ValueError(
+          "quantize_before_ep_all_gather=True is only supported with use_ring_of_experts=True and "
+          "qwix quantization, and gmm v2 kernel"
+      )
+
   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()
@@ -3948,6 +3965,8 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
       if self.use_batch_split_schedule:
         raise ValueError("GMM v2 is not supported with a batch split schedule.")
 
+    self._validate_quantize_before_ep_all_gather()
+
     for val in self.compress_ratios:
       if val != 0 and val < 4:
         raise ValueError(f"compress_ratio must be 0 (disabled) or >= 4, got {val}")

src/maxtext/kernels/megablox/ops.py

@@ -27,6 +27,8 @@
 from maxtext.layers import quantizations
 import qwix
 import qwix.pallas as qpl
+from qwix._src.core import numerics as qwix_numerics
+from qwix._src.core.qarray import call_with_generic_broadcast
 import tokamax
 
 
@@ -102,13 +104,18 @@ def gmm(
           act_calibration_method="absmax",
       )
 
+  lhs_scale = None
+  if isinstance(lhs, qpl.QArray):
+    lhs_scale = lhs.scale
+    lhs = lhs.qvalue
+
   gmm_fwd_bwd = lambda *args: _gmm_fwd(*args)[0]  # pylint: disable=C3001
   gmm_fwd_bwd = jax.custom_vjp(
       gmm_fwd_bwd,
       nondiff_argnums=(3, 4, 7, 8, 9, 10, 11, 12, 13, 14, 15),
   )
   gmm_fwd_bwd.defvjp(_gmm_fwd, functools.partial(_gmm_bwd, lhs.dtype, rhs.dtype))
-  return gmm_fwd_bwd(
+  out = gmm_fwd_bwd(
       lhs,
       rhs,
       group_sizes,
@@ -127,6 +134,9 @@ def gmm(
       use_gmm_v2,
       partial_sum,
   )
+  if lhs_scale is not None:
+    out = call_with_generic_broadcast(jnp.multiply, out, lhs_scale.astype(out.dtype))
+  return out
 
 
 # ==============================================================================
@@ -202,15 +212,7 @@ def _gmm_fwd(
     out = _fwd_run_tokamax_v1(lhs, rhs, group_sizes, preferred_element_type, transpose_rhs, use_manual_quantization)
   elif use_tokamax_backend and use_gmm_v2:
     out = _fwd_run_tokamax_v2(
-        lhs,
-        rhs,
-        group_sizes,
-        preferred_element_type,
-        tiling,
-        group_offset,
-        partial_sum,
-        transpose_rhs,
-        quantization_rule,
+        lhs, rhs, group_sizes, preferred_element_type, tiling, group_offset, partial_sum, transpose_rhs, quantization_rule
     )
   else:
     out = _fwd_run_megablox(
@@ -238,7 +240,7 @@ def _fwd_quantize_activation_and_weight(
     transpose_rhs: bool,
 ) -> tuple[jnp.ndarray | qpl.QArray, jnp.ndarray | qpl.QArray]:
   """Handles act and weight quantization for GMM forward inputs."""
-  if quantization_rule.act_qtype and not isinstance(lhs, qpl.QArray) and not use_gmm_v2:
+  if quantization_rule.act_qtype and not isinstance(lhs, qpl.QArray) and qwix_numerics.should_quantize(lhs.dtype):
     lhs = qpl.quantize(  # pyrefly: ignore[bad-assignment]
         lhs,
         quantization_rule.act_qtype,
@@ -382,24 +384,37 @@ def _fwd_run_tokamax_v2(
     rhs_operand = rhs_operand.qvalue
     rhs_scale = _fwd_prepare_rhs_scale(rhs, transpose_rhs=transpose_rhs)
 
+  lhs_operand = lhs.qvalue if isinstance(lhs, qpl.QArray) else lhs
+  maybe_quantize_lhs = not isinstance(lhs, qpl.QArray) and qwix_numerics.should_quantize(lhs_operand.dtype)
+
+  lhs_scale = _fwd_prepare_lhs_scale(quantization_rule) if maybe_quantize_lhs else None
+
   custom_fwd_tiling = gmm_v2.TileSizes(
       tile_m=tiling[0],
       tile_k=tiling[1],
       tile_n=tiling[2],
   )
 
-  return gmm_v2.gmm_v2(
-      lhs=lhs,  # pyrefly: ignore[bad-argument-type]
+  eff_pref_dtype = preferred_element_type if qwix_numerics.should_quantize(lhs_operand.dtype) else jnp.bfloat16
+
+  out = gmm_v2.gmm_v2(
+      lhs=lhs_operand,  # pyrefly: ignore[bad-argument-type]
       rhs=rhs_operand,  # pyrefly: ignore[bad-argument-type]
       group_sizes=group_sizes,
       rhs_scale=rhs_scale,
       tile_info=custom_fwd_tiling,
-      preferred_element_type=preferred_element_type,
+      preferred_element_type=eff_pref_dtype,
       partial_sum=partial_sum,
       group_offset=group_offset,
-      lhs_scale=_fwd_prepare_lhs_scale(quantization_rule),
+      maybe_quantize_lhs=maybe_quantize_lhs,
+      lhs_scale=lhs_scale,
   )
 
+  if isinstance(lhs, qpl.QArray):
+    out *= lhs.scale.astype(out.dtype)
+
+  return out
+
 
 def _fwd_run_megablox(
     lhs: jnp.ndarray,
@@ -560,7 +575,12 @@ def _bwd_prepare_inputs(
 
   # GMM2 FWD performs lhs quantization inside kernel, lhs is stored as unquantized dtype
   # in the residual tuple. In BWD, we explicitly quantize lhs.
-  if quantization_rule and quantization_rule.act_qtype and not isinstance(lhs, qpl.QArray):
+  if (
+      quantization_rule
+      and quantization_rule.act_qtype
+      and not isinstance(lhs, qpl.QArray)
+      and qwix_numerics.should_quantize(lhs.dtype)
+  ):
     lhs = qpl.quantize(  # pyrefly: ignore[bad-assignment]
         lhs,
         quantization_rule.act_qtype,

src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py

@@ -525,7 +525,12 @@ def _matmul(is_first_k_step: bool, is_last_k_step: bool):
           # Perform lhs quantization. Note that for every block_lhs,
           # same computation will be performed tiles_n//mxu_size times.
           # But we can let compiler perform CSE and avoid recomputation.
-          if should_use_external_scale:
+          if tiled_lhs.dtype == lhs_q_dtype:
+            # lhs block already arrives quantized, the real dequant
+            # scale is applied externally, so just pass an identity scale here. Comparing against lhs_q_dtype
+            block_lhs_q = block_lhs
+            block_scale = jnp.array(1.0, dtype=acc_ref.dtype)
+          elif should_use_external_scale:
             assert lhs_scale is not None
             assert lhs_scale_inv is not None
             block_lhs_q = jnp.clip(block_lhs * lhs_scale_inv, -dtype_max, dtype_max).astype(lhs_q_dtype)
@@ -1206,7 +1211,11 @@ def make_gmm_configs(
   )
 
   lhs_q_dtype = None
-  if maybe_quantize_lhs and rhs_cfgs.should_dequantize_after_matmul:
+  if jnp.issubdtype(lhs.dtype, jnp.integer) or jnp.issubdtype(lhs.dtype, jnp.float8_e4m3fn):
+    # lhs arrives already quantized (e.g. pre-quantized ahead-of-time by the
+    # caller): use its dtype as-is, no in-kernel quantization/scale needed.
+    lhs_q_dtype = lhs.dtype
+  elif maybe_quantize_lhs and rhs_cfgs.should_dequantize_after_matmul:
     # Choose lhs quantization dtype based on TPU hardware support.
     is_rhs_float = jnp.issubdtype(rhs_quant_dtype, jnp.floating)  # pyrefly: ignore[bad-argument-type]
     tpu_info = pltpu.get_tpu_info()
@@ -1241,8 +1250,12 @@ def make_gmm_configs(
       has_scale=has_lhs_scale,
   )
 
-  if out_dtype is None:
-    out_dtype = lhs.dtype
+  if out_dtype is None or jnp.issubdtype(out_dtype, jnp.float8_e4m3fn):
+    # The raw quantized-domain matmul output isn't yet rescaled -- writing it
+    # directly as fp8 would lose precision before the scale multiply happens
+    # (either inside this kernel via lhs_scale/block_scale, or externally by
+    # the caller for a pre-quantized lhs). Floor to bf16 as a safe intermediate.
+    out_dtype = jnp.bfloat16.dtype
 
   if acc_dtype is None:
     if lhs_cfgs.quant_dtype is None: