diff --git a/src/dependencies/extra_deps/post_train_github_deps.txt b/src/dependencies/extra_deps/post_train_github_deps.txt index cdc56f4053..437d934bbe 100644 --- a/src/dependencies/extra_deps/post_train_github_deps.txt +++ b/src/dependencies/extra_deps/post_train_github_deps.txt @@ -1,3 +1,3 @@ google-tunix @ https://github.com/google/tunix/archive/c4ec573d29e4c3a3955b348256d464b119c8a6d1.zip -tpu-inference @ https://github.com/vllm-project/tpu-inference/archive/7ecc401e6faefe3f793f3e4a8e6e2dbc49eca868.zip -vllm @ git+https://github.com/vllm-project/vllm@0ba2aa35a81dcc3246b26291368b53fa2389c7d7 +tpu-inference @ https://github.com/vllm-project/tpu-inference/archive/b67ae5f8f234fd559cf4b376840e7bb9ae6d3275.zip +vllm @ git+https://github.com/vllm-project/vllm@d626108b1841888ec90aced33367149a6bbc7e4b diff --git a/src/maxtext/integration/vllm/maxtext_vllm_adapter/adapter.py b/src/maxtext/integration/vllm/maxtext_vllm_adapter/adapter.py index cd842eccac..ce062c11f8 100644 --- a/src/maxtext/integration/vllm/maxtext_vllm_adapter/adapter.py +++ b/src/maxtext/integration/vllm/maxtext_vllm_adapter/adapter.py @@ -128,6 +128,23 @@ def generate_maxtext_config(vllm_config: VllmConfig) -> pyconfig.HyperParameters f"num_lanes={num_lanes}, tp={tp}, attn_dp={attn_dp}, ep={ep}, moe_mlp_tp_size={moe_mlp_tp_size}" ) + # The native tpu-inference model paths derive use_ep from + # parallel_config.enable_expert_parallel, but the mesh built by + # ShardingConfigManager.from_vllm_config takes expert parallelism only from + # additional_config's sharding_strategy. MaxText derives use_ep from the mesh, so + # --enable-expert-parallel alone leaves the experts unsharded here while the native + # implementation of the same model runs expert-parallel. Warn rather than fail, since + # the run is still correct, only sharded differently than the flag suggests. + expert_shard_degree = ep * sharding_config.attn_dp_expert_size + if vllm_config.parallel_config.enable_expert_parallel and expert_shard_degree == 1: + max_logging.warning( + "--enable-expert-parallel was requested but the mesh has no expert shards " + f"(expert={ep}, attn_dp_expert={sharding_config.attn_dp_expert_size}), so MaxText will shard the MoE over " + "the MLP dimension instead of the expert dimension. The vLLM flag does not reach the JAX mesh; set " + 'expert_parallelism in the vLLM additional_config sharding_strategy, e.g. \'{"sharding": ' + '{"sharding_strategy": {"expert_parallelism": }}}\', to actually shard experts.' + ) + # Replicate the number of KV heads if its less than the total degree of model parallelism if kv_tp_size % num_kv_heads == 0 and num_kv_heads < kv_tp_size: max_logging.log( @@ -152,8 +169,15 @@ def generate_maxtext_config(vllm_config: VllmConfig) -> pyconfig.HyperParameters while (padded_hidden_size // moe_mlp_tp_size) < (2 * num_lanes): padded_hidden_size = next_power_of_two(padded_hidden_size + 1) - max_logging.log( - f"Padding moe_intermediate_size from {hidden_size} to {padded_hidden_size} to match MLP MoE requirements." + # This inflates every expert weight, so it is a real memory/FLOP cost rather than a + # cosmetic reshape: at moe_mlp_tp_size=4 a 512-wide MoE is padded to 1024 (2x the MoE + # weights), and at moe_mlp_tp_size=8 to 2048 (4x). Log it at WARNING so it is visible + # under vLLM's logging configuration, which does not surface absl INFO records. + max_logging.warning( + f"Padding moe_intermediate_size from {hidden_size} to {padded_hidden_size} to match MLP MoE requirements " + f"(moe_mlp_tp_size={moe_mlp_tp_size}, 2*num_lanes={2 * num_lanes}). This multiplies the MoE weights and " + f"MoE FLOPs by {padded_hidden_size / hidden_size:g}x. Consider sharding the MoE over the expert axis " + f"instead, by setting expert_parallelism in the vLLM additional_config sharding_strategy." ) overrides["padded_base_moe_mlp_dim"] = padded_hidden_size diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py index 82d4b01e47..de70870da2 100644 --- a/src/maxtext/layers/moe.py +++ b/src/maxtext/layers/moe.py @@ -3290,6 +3290,7 @@ def fused_moe_matmul( # pylint: disable=import-outside-toplevel # pytype: disable=import-error from tpu_inference.layers.common.fused_moe_gmm import fused_moe_func + from tpu_inference import envs as tpu_inference_envs except ImportError as e: raise ImportError("fused_moe_matmul requires the tpu-inference package.") from e @@ -3333,6 +3334,18 @@ def fused_moe_matmul( use_ep=use_ep, activation=activation, scoring_fn=scoring_fn, + # Forward the same environment-backed kernel knobs that tpu-inference passes on its + # own serving path (tpu_inference/layers/common/moe.py). Without these, env vars such + # as ONEHOT_MOE_PERMUTE_THRESHOLD and VLLM_MOE_CHUNK_SIZE are silently ignored when a + # MaxText model is served through vLLM, so the two paths run the same kernel with + # different configurations. In particular, ONEHOT_MOE_PERMUTE_THRESHOLD selects the + # one-hot permute path instead of the SparseCore ragged_gather_reduce kernel, which + # is what makes expert parallelism usable on TPU generations whose SparseCore has + # fewer SIMD lanes than the kernel requires (e.g. v5p). + enable_rs_kernel=tpu_inference_envs.ENABLE_RS_KERNEL, + use_gmm_fused_rs_kernel=tpu_inference_envs.USE_GMM_FUSED_RS_KERNEL, + onehot_moe_permute_threshold=tpu_inference_envs.ONEHOT_MOE_PERMUTE_THRESHOLD, + moe_chunk_size=tpu_inference_envs.VLLM_MOE_CHUNK_SIZE, ) # Reshape output 2D [T, D] -> 3D [B, S, D]