diff --git a/run_qwen3_next_80b_xpk.sh b/run_qwen3_next_80b_xpk.sh new file mode 100755 index 0000000000..22a2711862 --- /dev/null +++ b/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 "========================================================================" diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index 08d22bae24..6bb0765530 100644 --- a/src/maxtext/configs/base.yml +++ b/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). diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 68684cb38a..a9bb27a089 100644 --- a/src/maxtext/configs/types.py +++ b/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}") diff --git a/src/maxtext/kernels/megablox/ops.py b/src/maxtext/kernels/megablox/ops.py index bf533e353c..e890c90db2 100644 --- a/src/maxtext/kernels/megablox/ops.py +++ b/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, diff --git a/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py b/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py index 1826ef576b..cdd89179c2 100644 --- a/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py +++ b/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: diff --git a/src/maxtext/kernels/ragged/ragged_sort.py b/src/maxtext/kernels/ragged/ragged_sort.py index f3def646b9..b0ced6f607 100644 --- a/src/maxtext/kernels/ragged/ragged_sort.py +++ b/src/maxtext/kernels/ragged/ragged_sort.py @@ -102,18 +102,21 @@ def _ring_ragged_sort_fwd(hidden_states_local, topk_indices_local): shard_output_start = group_offsets[experts_start] shard_output_end = group_offsets[experts_end] - if buffer_size is None or buffer_size >= num_tokens_local * topk: - local_buffer_size = num_tokens_local * topk - x = ragged_gather( + def _gather(indices, start, end): + return ragged_gather( hidden_states_local, - token_indices_sorted, - shard_output_start[None], - shard_output_end[None], + indices, + start, + end, enforce_fallback=enforce_gather_fallback, flops_override=gather_flops_override, bytes_accessed_override=gather_bytes_accessed_override, use_single_sparsecore=use_single_sparsecore, ) + + if buffer_size is None or buffer_size >= num_tokens_local * topk: + local_buffer_size = num_tokens_local * topk + x = _gather(token_indices_sorted, shard_output_start[None], shard_output_end[None]) else: local_buffer_size = buffer_size # We only gather up to the available buffer size or the actual number of @@ -128,16 +131,7 @@ def _ring_ragged_sort_fwd(hidden_states_local, topk_indices_local): local_buffer_size, axis=0, ) - x = ragged_gather( - hidden_states_local, - sliced_indices, - jnp.int32(0)[None], - gather_end[None], - enforce_fallback=enforce_gather_fallback, - flops_override=gather_flops_override, - bytes_accessed_override=gather_bytes_accessed_override, - use_single_sparsecore=use_single_sparsecore, - ) + x = _gather(sliced_indices, jnp.int32(0)[None], gather_end[None]) out = (x, group_sizes_local, topk_argsort_revert_indices) @@ -174,17 +168,11 @@ def _ring_ragged_sort_bwd(res, g_out): # rather than materializing a (mostly-zero) dense buffer ourselves. n = topk_argsort_revert_indices.shape[0] - if local_buffer_size >= n: - valid_rows_mask = (topk_argsort_revert_indices >= shard_output_start) & ( - topk_argsort_revert_indices < shard_output_end - ) - # The forward scatter-add over `token_indices_sorted` is equivalent to a - # gather-reduce: each input token has exactly `topk` contributions located - # at sorted positions `topk_argsort_revert_indices[t*topk:(t+1)*topk]`. - # `topk_weights` is set to ones because this op has no per-row weighting. - grad_hidden_states = ragged_gather_reduce( + def _gather_reduce(indices, valid_rows_mask): + """`topk_weights` is set to ones because this op has no per-row weighting.""" + return ragged_gather_reduce( g_x, - topk_argsort_revert_indices, + indices, topk_weights=jnp.ones((n,), dtype=jnp.float32), valid_rows_mask=valid_rows_mask, reduce_group_size=topk, @@ -193,6 +181,12 @@ def _ring_ragged_sort_bwd(res, g_out): bytes_accessed_override=gather_reduce_bytes_accessed_override, use_single_sparsecore=use_single_sparsecore, ) + + if local_buffer_size >= n: + valid_rows_mask = (topk_argsort_revert_indices >= shard_output_start) & ( + topk_argsort_revert_indices < shard_output_end + ) + grad_hidden_states = _gather_reduce(topk_argsort_revert_indices, valid_rows_mask) else: # Buffering: g_x has size `local_buffer_size` (packed). # The revert indices are global [0, n), but they must map to the local @@ -207,18 +201,7 @@ def _ring_ragged_sort_bwd(res, g_out): # Clamp invalid indices to 0 to prevent compile-time/run-time out-of-bounds # in JAX. These clamped values will be ignored due to `valid_rows_mask`. safe_indices = jnp.where(valid_rows_mask, shifted_indices, 0) - - grad_hidden_states = ragged_gather_reduce( - g_x, - safe_indices, - topk_weights=jnp.ones((n,), dtype=jnp.float32), - valid_rows_mask=valid_rows_mask, - reduce_group_size=topk, - enforce_fallback=enforce_gather_reduce_fallback, - flops_override=gather_reduce_flops_override, - bytes_accessed_override=gather_reduce_bytes_accessed_override, - use_single_sparsecore=use_single_sparsecore, - ) + grad_hidden_states = _gather_reduce(safe_indices, valid_rows_mask) return grad_hidden_states, None _ring_ragged_sort.defvjp(_ring_ragged_sort_fwd, _ring_ragged_sort_bwd) diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py index f6e89f204d..cf451d57ec 100644 --- a/src/maxtext/layers/moe.py +++ b/src/maxtext/layers/moe.py @@ -112,6 +112,17 @@ def _truncate_matrix(all_shards_group_sizes: jax.Array, buffer_size: int) -> jax return jnp.diff(clamped_cumsum_extended, axis=0) +def _map_qarray(fn, value: jax.Array | qpl.QArray) -> jax.Array | qpl.QArray: + """Applies `fn` to `value`. + + If `value` is a QArray, applies `fn` to `qvalue` and `scale` independently + and rewraps the result, instead of requiring the caller to unpack/repack it. + """ + if isinstance(value, qpl.QArray): + return qpl.QArray(qvalue=fn(value.qvalue), scale=fn(value.scale)) + return fn(value) + + def _sort_activations( inputs: jax.Array, sort_indices: jax.Array, @@ -868,7 +879,10 @@ def permute( # reshape inputs (batch, sequence, emb) to (batch * sequence, emb) inputs_shape = inputs.shape bsz_times_seq_len = inputs_shape[0] * inputs_shape[1] - inputs_2d = jnp.reshape(inputs, (bsz_times_seq_len, inputs_shape[2])) + if isinstance(inputs, qpl.QArray): + inputs_2d = inputs.reshape(bsz_times_seq_len, inputs.shape[-1]) + else: + inputs_2d = jnp.reshape(inputs, (bsz_times_seq_len, inputs_shape[2])) weights, selected_experts = self.get_topk(gate_logits, pre_bias_logits, rngs, input_ids) lb_loss = None if self.config.load_balance_loss_weight > 0.0 and not self.is_hash_routing: @@ -891,6 +905,7 @@ def permute( inputs_2d = inputs_2d * router_scores.reshape(bsz_times_seq_len, -1) num_expert_parallelism = self.get_expert_parallelism_size() + # The ragged-kernel path inside permute()/unpermute() is only correct for # the ring-of-experts strategy: each shard's output is masked to its own # [start, end) range within a globally-sorted layout. When ring of experts @@ -915,22 +930,33 @@ def permute( else: buffer_size = None - sorted_inputs, group_size, sorted_selected_experts = ring_ragged_sort( - inputs_2d, - topk_indices_2d, - self.config.num_experts, - self.num_experts_per_tok, - self._expert_parallelism_name, - num_expert_parallelism, - buffer_size=buffer_size, - enforce_gather_fallback=self.config.ragged_gather_fallback, - enforce_gather_reduce_fallback=self.config.ragged_gather_reduce_fallback, - gather_flops_override=self.config.ragged_gather_cost_estimate_flops, - gather_reduce_flops_override=self.config.ragged_gather_reduce_cost_estimate_flops, - gather_bytes_accessed_override=self.config.ragged_gather_cost_estimate_bytes_accessed, - gather_reduce_bytes_accessed_override=self.config.ragged_gather_reduce_cost_estimate_bytes_accessed, - use_single_sparsecore=self.config.ragged_sort_use_single_sparsecore, - ) + def _permute_ring_ragged_sort(tensor): + return ring_ragged_sort( + tensor, + topk_indices_2d, + self.config.num_experts, + self.num_experts_per_tok, + self._expert_parallelism_name, + num_expert_parallelism, + buffer_size=buffer_size, + enforce_gather_fallback=self.config.ragged_gather_fallback, + enforce_gather_reduce_fallback=self.config.ragged_gather_reduce_fallback, + gather_flops_override=self.config.ragged_gather_cost_estimate_flops, + gather_reduce_flops_override=self.config.ragged_gather_reduce_cost_estimate_flops, + gather_bytes_accessed_override=self.config.ragged_gather_cost_estimate_bytes_accessed, + gather_reduce_bytes_accessed_override=self.config.ragged_gather_reduce_cost_estimate_bytes_accessed, + use_single_sparsecore=self.config.ragged_sort_use_single_sparsecore, + ) + + if isinstance(inputs_2d, qpl.QArray): + sorted_qvalue, group_size, sorted_selected_experts = _permute_ring_ragged_sort(inputs_2d.qvalue) + if inputs_2d.scale.shape[0] == inputs_2d.qvalue.shape[0]: + sorted_scale, _, _ = _permute_ring_ragged_sort(inputs_2d.scale) + else: + sorted_scale = inputs_2d.scale + sorted_inputs = qpl.QArray(qvalue=sorted_qvalue, scale=sorted_scale) + else: + sorted_inputs, group_size, sorted_selected_experts = _permute_ring_ragged_sort(inputs_2d) else: flatten_selected_experts = jnp.ravel(selected_experts) @@ -938,10 +964,12 @@ def permute( flatten_selected_experts = (flatten_selected_experts - roll_to_expert_id) % self.num_experts sorted_selected_experts = jnp.argsort(flatten_selected_experts) # sort inputs for number of selected experts - replicated_inputs_2d = jnp.repeat(inputs_2d, self.num_experts_per_tok, axis=0) - sorted_inputs = _sort_activations(replicated_inputs_2d, sorted_selected_experts, use_custom_sort_vjp).astype( - self.dtype + replicated_inputs_2d = _map_qarray(lambda t: jnp.repeat(t, self.num_experts_per_tok, axis=0), inputs_2d) + sorted_inputs = _map_qarray( + lambda t: _sort_activations(t, sorted_selected_experts, use_custom_sort_vjp), replicated_inputs_2d ) + if not isinstance(sorted_inputs, qpl.QArray): + sorted_inputs = sorted_inputs.astype(self.dtype) group_size = jnp.bincount(flatten_selected_experts, length=self.num_experts) num_tokens = bsz_times_seq_len * self.num_experts_per_tok @@ -1477,6 +1505,8 @@ def extract_vma(tensor): # Parses the varying mesh axes from JAX's type string for a tensor inside shard_map. # jax.typeof(t) renders as e.g. 'f32[128,256]{V:(expert, fsdp)}'; this extracts # ('expert', 'fsdp'). Returns () if the tensor has no varying axes. + if isinstance(tensor, qpl.QArray): + tensor = tensor.qvalue type_str = str(jax.typeof(tensor)) if "{V:" in type_str: start = type_str.index("{V:") + 3 @@ -1487,7 +1517,7 @@ def extract_vma(tensor): lhs_vma_axes = extract_vma(inputs) rhs_vma_axes = extract_vma(kernel) - if inputs.shape[0] != expert_assignments.shape[0]: + if (inputs.qvalue if isinstance(inputs, qpl.QArray) else inputs).shape[0] != expert_assignments.shape[0]: raise ValueError("The number of input tokens must match the number of expert assignments!") tokamax_group_sizes = get_tokamax_group_sizes(group_sizes, inputs, kernel) @@ -1663,9 +1693,24 @@ def roe_ag_and_route(x, logits, pre_bias_logits, num_ep, expert_shard_id, rngs, # The ring-of-experts strategy first duplicates the inputs to all # expert shards, and then routes within each shard. - # Duplicate inputs to all expert shards. - x, logits, pre_bias_logits = tuple( - jax.lax.all_gather(z, axis_name=self._expert_parallelism_name, tiled=True) for z in (x, logits, pre_bias_logits) + # Duplicate inputs to all expert shards + rule = qpl.get_current_rule("gmm") + + if self.config.quantize_before_ep_all_gather and rule and rule.act_qtype and not isinstance(x, qpl.QArray): + x_q = qpl.quantize( + x, + rule.act_qtype, + channelwise_axes=[] if rule.disable_channelwise_axes else [0], + calibration_method=rule.act_calibration_method, + ) + x_qvalue = jax.lax.all_gather(x_q.qvalue, axis_name=self._expert_parallelism_name, tiled=True) + x_scale = jax.lax.all_gather(x_q.scale, axis_name=self._expert_parallelism_name, tiled=True) + x = qpl.QArray(qvalue=x_qvalue, scale=x_scale) + else: + x = jax.lax.all_gather(x, axis_name=self._expert_parallelism_name, tiled=True) + + logits, pre_bias_logits = tuple( + jax.lax.all_gather(z, axis_name=self._expert_parallelism_name, tiled=True) for z in (logits, pre_bias_logits) ) # "Route" tokens within each shard. diff --git a/tests/unit/moe_test.py b/tests/unit/moe_test.py index 7781f4b797..49b99851d5 100644 --- a/tests/unit/moe_test.py +++ b/tests/unit/moe_test.py @@ -15,6 +15,7 @@ import unittest from absl.testing import parameterized +import chex import pytest from flax import nnx @@ -506,6 +507,46 @@ def get_moe_output(self, variables, hidden_states, cfg, mesh): output = jax.jit(model.apply)(moe_variables, hidden_states) # pylint: disable=not-callable return output + def get_quantized_moe_model(self, cfg, mesh): + """Builds a RoutedMoE wrapped with the fp8_full qwix quantization rule.""" + model = moe.get_routed_moe( + name="MoeBlock", + config=cfg, + num_experts=cfg.num_experts, + num_experts_per_tok=cfg.num_experts_per_tok, + mesh=mesh, + kernel_init=nd_dense_init(1.0, "fan_in", "truncated_normal"), + kernel_axes=("embed", "mlp"), + intermediate_dim=cfg.mlp_dim, + dtype=cfg.dtype, + ) + # Similar to `quantizations.get_fp8_full_qwix_rule_w_sparsity`. + quantization_rule = [ + qwix.QtRule( + module_path=".*", + weight_qtype=jnp.float8_e4m3fn, + act_qtype=jnp.float8_e4m3fn, + bwd_qtype=jnp.float8_e5m2, + weight_calibration_method=cfg.weight_quantization_calibration_method, + act_calibration_method=cfg.act_quantization_calibration_method, + bwd_calibration_method=cfg.bwd_quantization_calibration_method, + op_names=("gmm", "ragged_dot"), + ), + ] + return qwix.quantize_model(model, qwix.QtProvider(quantization_rule)) + + def get_moe_loss_and_grad(self, model, variables, hidden_states): + """Computes (loss, output) and grads w.r.t. params and the input hidden states.""" + + def loss_fn(params, x): + out, lb_loss, _ = model.apply({"params": params}, x) + loss = jnp.mean(out.astype(jnp.float32) ** 2) + if lb_loss is not None: + loss = loss + lb_loss.astype(jnp.float32) + return loss, out + + return jax.jit(jax.value_and_grad(loss_fn, argnums=(0, 1), has_aux=True))(variables["params"], hidden_states) + @pytest.mark.tpu_only def test_megablox(self): cfg = pyconfig.initialize( @@ -1623,6 +1664,87 @@ def loss_fn(params, x): diff_summary = compare_tree(tree_ref, tree_tgt, relative_norm_diff_threshold) max_logging.log("\n" + diff_summary) + def _build_ep_all_gather_test_cfg(self, quantize_before_ep_all_gather): + return pyconfig.initialize( + [None, get_test_config_path()], + run_name="quantize_before_ep_all_gather_equivalence_test", + enable_checkpointing=False, + model_name="mixtral-8x7b", + weight_dtype="float32", + dtype="bfloat16", + per_device_batch_size=2, + max_target_length=256, + float32_gate_logits=True, + ici_expert_parallelism=4, + sparse_matmul=True, + megablox=True, + use_tokamax_gmm=True, + use_gmm_v2=True, + use_ring_of_experts=True, + use_ragged_sort=True, + quantization="fp8_full", + use_qwix_quantization=True, + weight_quantization_calibration_method="absmax", + act_quantization_calibration_method="absmax", + bwd_quantization_calibration_method="absmax", + quantize_before_ep_all_gather=quantize_before_ep_all_gather, + wi_tile_fwd_batch_seq=128, + wi_tile_dlhs_batch_seq=128, + wi_tile_dlhs_embed_dim=256, + wi_tile_drhs_batch_seq=128, + wo_tile_fwd_batch_seq=128, + wo_tile_fwd_embed_dim=256, + wo_tile_dlhs_batch_seq=128, + wo_tile_dlhs_mlp_dim=256, + wo_tile_drhs_batch_seq=128, + ) + + def _run_ep_all_gather_test(self, quantize_before_ep_all_gather, rng_model, hidden_states): + cfg = self._build_ep_all_gather_test_cfg(quantize_before_ep_all_gather) + mesh = Mesh(maxtext_utils.create_device_mesh(cfg), cfg.mesh_axes) + model = self.get_quantized_moe_model(cfg, mesh) + with jax.set_mesh(mesh), nn_partitioning.axis_rules(cfg.logical_axis_rules): + variables = model.init({"params": rng_model, "dropout": rng_model}, hidden_states) + return self.get_moe_loss_and_grad(model, variables, hidden_states) + + @pytest.mark.skip_on_tpu7x + @pytest.mark.tpu_only + def test_quantize_before_ep_all_gather_equivalence(self): + """Quantizing activations before the ring-of-experts EP all-gather should be + numerically equivalent to quantizing them later inside the gmm call: both + sides quantize the same per-token values (channelwise/absmax scale), just + at a different point in the pipeline. Unlike test_gmm_grad_equivalence + (which compares quantized vs. unquantized and expects real quantization + noise), any meaningful gap here is a bug in the ahead-of-time relocation, + not expected noise -- so this uses a much tighter tolerance. + """ + rng = jax.random.PRNGKey(4567) + rng_model, rng_hidden_states = jax.random.split(rng) + cfg = self._build_ep_all_gather_test_cfg(quantize_before_ep_all_gather=False) + # Normal distribution for realistic variance/negative values, so the + # quantization scale != 1.0 and scale-dropping bugs are actually caught. + hidden_states = jax.random.normal( + rng_hidden_states, + (int(cfg.per_device_batch_size) * jax.device_count(), cfg.max_target_length, cfg.base_emb_dim), + dtype=cfg.dtype, + ) + + # Reference: quantize dynamically inside the gmm call (activations stay + # bf16 through the EP all-gather and ragged sort). Target: quantize ahead + # of the EP all-gather (this branch's default). Same RNG for both, so the + # initial weights match. + (_, output_ref), (grads_ref, x_grad_ref) = self._run_ep_all_gather_test(False, rng_model, hidden_states) + (_, output_tgt), (grads_tgt, x_grad_tgt) = self._run_ep_all_gather_test(True, rng_model, hidden_states) + + tree_ref = {"output": output_ref, "state_grad": x_grad_ref, "var_grad": grads_ref} + tree_tgt = {"output": output_tgt, "state_grad": x_grad_tgt, "var_grad": grads_tgt} + # Use an absolute+relative tolerance (not compare_tree's pure relative-norm + # metric): the hidden-state gradient's reference norm is near the float32 + # noise floor for this loss, so a relative-only metric blows up on + # ordinary floating-point reordering noise even when both sides agree to + # ~1e-9 in absolute terms. + chex.assert_trees_all_close(tree_tgt, tree_ref, atol=1e-6, rtol=1e-3) + def make_moe(cfg, mesh): return moe.RoutedMoE(