diff --git a/src/maxtext/integration/vllm/maxtext_vllm_rollout.py b/src/maxtext/integration/vllm/maxtext_vllm_rollout.py index 9dc530a82b..1cd7853f96 100644 --- a/src/maxtext/integration/vllm/maxtext_vllm_rollout.py +++ b/src/maxtext/integration/vllm/maxtext_vllm_rollout.py @@ -38,7 +38,9 @@ from flax.traverse_util import flatten_dict, unflatten_dict from tunix.generate import mappings +from tunix.generate import utils as tunix_gen_utils from tunix.generate.vllm_sampler import VllmConfig, VllmSampler +from tunix.rl import reshard as tunix_reshard from tunix.rl.rollout import base_rollout, vllm_rollout from maxtext.integration.vllm.convert_utils import _sharding_summary @@ -46,11 +48,11 @@ WeightConverter, MODEL_TO_CONVERSION_RULES, ) +from maxtext.integration.vllm.torchax_converter.base import BaseMaxTextToVLLMConverter from maxtext.integration.vllm.torchax_converter.gemma4_moe import Gemma4MaxTextToVLLMConverter from maxtext.integration.vllm.torchax_converter.qwen35_moe import Qwen35MaxTextToVLLMConverter from maxtext.integration.vllm.torchax_converter.qwen3_moe import Qwen3MaxTextToVLLMConverter - # Sentinel distinguishing "this model has no entry" from "this model has an # entry whose value is None", which means direct-sync-only. _NO_RULE_TABLE = object() @@ -73,10 +75,27 @@ def _create_model_converter( mesh: jax.sharding.Mesh, use_hf_mapping: bool = False, use_weight_converter: bool = False, + use_standalone_converter: bool = False, + sharding_hints: Optional[dict] = None, debug: bool = False, ): """Instantiate the converter for a MaxText model name.""" tp = config.rollout_tensor_parallelism + if use_standalone_converter: + # Standalone torchax converters emit the tpu-inference runner's internal + # layout keyed by its state names; MaxTextVllmSampler syncs them with + # `_sync_standalone_converted`. Requires vLLM to run its *native* model + # (no MaxTextForCausalLM overrides). + if model_name.startswith("qwen3.5"): + return Qwen35MaxTextToVLLMConverter( + config=config, + mesh=mesh, + vllm_attn_dp=sharding_hints.get("attn_dp_size", 1) if sharding_hints else 1, + vllm_use_ep=sharding_hints.get("enable_expert_parallel", False) if sharding_hints else False, + ) + if model_name.startswith("gemma4"): + return Gemma4MaxTextToVLLMConverter(config=config, mesh=mesh) + raise NotImplementedError(f"use_standalone_converter: no standalone torchax converter for {model_name}") if not use_hf_mapping and not use_weight_converter: # Default MaxText-to-MaxText sync uses direct transfer_state_directly with unroll return None @@ -463,6 +482,17 @@ def update_params( filter_types: Optional[Tuple[Any, ...]] = None, ): """Update the vLLM runner weights from a MaxText state tree.""" + if isinstance(self._converter, BaseMaxTextToVLLMConverter): + try: + return self._sync_standalone_converted(updated_weights) + except BaseException: + logging.error("MaxTextVllmSampler standalone sync failed:\n%s", traceback.format_exc()) + for handler in logging.getLogger().handlers: + try: + handler.flush() + except Exception: # pylint: disable=broad-except + pass + raise if self._converter is None: if self._direct_maxtext_sync: updated_weights = unroll_qwen_scanned_weights( @@ -486,6 +516,92 @@ def update_params( pass raise + def _sync_standalone_converted(self, updated_weights): + """Standalone torchax-converter sync path. + + The converter emits tensors in the tpu-inference runner's *internal* layout, + keyed by its state names, so this bypasses Tunix's mapped/direct transfers: + tear down the KV cache, convert, reshard each tensor onto its existing + sharding (chunked, Pathways-aware) and assign into the runner's flat state + dict in place. + """ + runner = self._model_runner + state = runner.state + if not isinstance(state, dict): + raise TypeError( + "Standalone torchax converters target the vLLM (torchax) model " + "implementation, whose runner state is a flat dict; got " + f"{type(state).__name__}. Remove MaxTextForCausalLM overrides so " + "vLLM runs its native model." + ) + + if self.llm is not None: + self.llm.reset_prefix_cache() + self.llm.collective_rpc("delete_kv_cache") + elif self._driver is not None: + self._driver.llm_engine.reset_prefix_cache() + self._driver.llm_engine.collective_rpc("delete_kv_cache") + jax.effects_barrier() + + start = time.time() + pure = updated_weights.to_pure_dict() if hasattr(updated_weights, "to_pure_dict") else updated_weights + converted = self._converter.convert(pure) + + src = {k: v for k, v in converted.items() if k in state} + version_aliases = sorted(set(converted) - set(src)) + if version_aliases: + logging.info( + "Standalone sync: %d converted tensors have no runner target (vLLM version aliases), e.g. %s", + len(version_aliases), + version_aliases[:3], + ) + uncovered = [k for k in state if k not in src and not k.rsplit(".", 1)[-1].startswith("_") and "rotary_emb" not in k] + if uncovered: + logging.warning( + "Standalone sync: %d runner tensors NOT covered by the converter (stale weights!), e.g. %s", + len(uncovered), + uncovered[:5], + ) + + spec = {k: state[k] for k in src} + expected = {k: (tuple(v.shape), v.dtype) for k, v in spec.items()} + chunk = getattr(self.config, "reshard_chunk_size", None) + delete_dst = getattr(self.config, "delete_dst_buffers", True) + reshard_in_chunks = getattr(tunix_gen_utils, "_reshard_in_chunks", None) + if chunk and reshard_in_chunks is None: + logging.warning("Standalone sync: this Tunix has no _reshard_in_chunks; falling back to one reshard call.") + chunk = None + if chunk: + resharded = reshard_in_chunks( + src_flat=dict(src), + spec_flat=spec, + reshard_fn=tunix_reshard.reshard_pytree, + chunk_size=chunk, + delete_spec_buffers=delete_dst, + ) + else: + shardings = {k: v.sharding for k, v in spec.items()} + if delete_dst: + tunix_gen_utils._delete_target_buffers(spec, src) # pylint: disable=protected-access + resharded = tunix_reshard.reshard_pytree(src, shardings) + + for k in src: + new = resharded[k] + shape, dtype = expected[k] + if tuple(new.shape) != shape or new.dtype != dtype: + raise ValueError( + f"{k}: converter produced {tuple(new.shape)}/{new.dtype}, the runner expects {shape}/{dtype}; " + "the converter's layout is out of date with tpu-inference." + ) + state[k] = new + runner.state_leaves = state + logging.info("Standalone sync: updated %d/%d runner tensors in %.1fs", len(src), len(state), time.time() - start) + + if self.llm is not None: + self.llm.collective_rpc("reinitialize_kv_cache") + elif self._driver is not None: + self._driver.llm_engine.collective_rpc("reinitialize_kv_cache") + class MaxTextVllmRollout(vllm_rollout.VllmRollout): """VllmRollout that uses MaxTextVllmSampler for weight synchronization. @@ -542,6 +658,21 @@ def __init__( getattr(maxtext_config, "use_weight_converter", False) or vllm_additional_config.get("use_weight_converter", False) ) + use_standalone_converter = bool( + getattr(maxtext_config, "use_standalone_converter", False) + or vllm_additional_config.get("use_standalone_converter", False) + ) + # Sampler sharding the standalone converter must mirror: attention DP from + # the sharding_strategy blob, expert parallelism from the vLLM engine kwargs. + strategy = {} + sharding_blob = vllm_additional_config.get("sharding") if isinstance(vllm_additional_config, dict) else None + if isinstance(sharding_blob, dict): + strategy = sharding_blob.get("sharding_strategy") or {} + rollout_vllm_kwargs = getattr(rollout_config, "rollout_vllm_kwargs", None) or {} + sharding_hints = { + "attn_dp_size": (int(strategy.get("attn_dp_size") or 1) if strategy.get("enable_dp_attention", False) else 1), + "enable_expert_parallel": bool(rollout_vllm_kwargs.get("enable_expert_parallel", False)), + } # Accepted from either spelling, matching use_weight_converter above, so a # debug run can be triggered by editing the same JSON blob. self._weight_sync_debug = bool( @@ -553,6 +684,8 @@ def __init__( mesh=mesh, use_hf_mapping=use_hf, use_weight_converter=use_weight_converter, + use_standalone_converter=use_standalone_converter, + sharding_hints=sharding_hints, debug=self._weight_sync_debug, ) diff --git a/src/maxtext/integration/vllm/torchax_converter/qwen35_moe.py b/src/maxtext/integration/vllm/torchax_converter/qwen35_moe.py index 808fd9e809..984676e10a 100644 --- a/src/maxtext/integration/vllm/torchax_converter/qwen35_moe.py +++ b/src/maxtext/integration/vllm/torchax_converter/qwen35_moe.py @@ -25,6 +25,17 @@ class Qwen35MaxTextToVLLMConverter(BaseMaxTextToVLLMConverter): """Converts MaxText Qwen3.5 (Scanned Block) layout to vLLM execution layout.""" + def __init__(self, config, mesh, vllm_attn_dp: int = 1, vllm_use_ep: bool = False): + super().__init__(config, mesh) + self.vllm_attn_dp = max(1, int(vllm_attn_dp or 1)) + self.vllm_use_ep = bool(vllm_use_ep) + assert ( + self.vllm_tp % self.vllm_attn_dp == 0 + ), f"rollout_tensor_parallelism={self.vllm_tp} must be divisible by attn_dp_size={self.vllm_attn_dp}" + # Attention (and GDN / shared-expert column) projections are sharded over + # the per-attention-group tensor axis. + self.attn_shards = self.vllm_tp // self.vllm_attn_dp + def convert(self, model_state: dict, **kwargs): """Converts model_state parameters to vLLM format.""" logging.info("\n%sStarting Qwen 3.5 Conversion (Hybrid MoE)...%s", GREEN, RESET) @@ -67,6 +78,16 @@ def _convert_global(self, params): params["base"]["decoder"]["logits_dense"]["kernel"], (1, 0) ) + def _replicate_kv_heads(self, kv): + """[D, n_kv, dh] -> [D, n_kv * replicas, dh] with each head repeated + consecutively, as vLLM's QKVParallelLinear lays KV heads out when + tp > num_kv_heads (rank r reads head r // replicas).""" + n_kv = kv.shape[1] + if self.attn_shards <= n_kv: + return kv + assert self.attn_shards % n_kv == 0, f"attention shards={self.attn_shards} must be a multiple of num_kv_heads={n_kv}" + return jnp.repeat(kv, self.attn_shards // n_kv, axis=1) + def _convert_attn(self, params): """Converts attention weights.""" decoder = params["base"]["decoder"] @@ -102,13 +123,17 @@ def _unstack_rep(x): self.vllm_state[f"{prefix}.input_layernorm.weight"] = pre_ln[rep] self.vllm_state[f"{prefix}.post_attention_layernorm.weight"] = post_ln[rep] - q, k, v = q_layers[rep], k_layers[rep], v_layers[rep] + # q carries the attention output gate ([q | gate] per head); k/v are + # replicated up to one head per shard when tp > num_kv_heads. + q = q_layers[rep] + k = self._replicate_kv_heads(k_layers[rep]) + v = self._replicate_kv_heads(v_layers[rep]) q_T = jnp.transpose(q, (1, 2, 0)) k_T = jnp.transpose(k, (1, 2, 0)) v_T = jnp.transpose(v, (1, 2, 0)) - tp_size = self.vllm_tp + tp_size = self.attn_shards q_tp_shards = jnp.split(q_T.reshape(-1, q.shape[0]), tp_size, axis=0) k_tp_shards = jnp.split(k_T.reshape(-1, k.shape[0]), tp_size, axis=0) v_tp_shards = jnp.split(v_T.reshape(-1, v.shape[0]), tp_size, axis=0) @@ -163,7 +188,7 @@ def _unstack_rep(x): v = t_r[:, 2 * D_k : 2 * D_k + V_per_K * D_v, :].reshape(H_v * D_v, -1) z = t_r[:, 2 * D_k + V_per_K * D_v :, :].reshape(H_v * D_v, -1) - tp_size = self.vllm_tp + tp_size = self.attn_shards q_shards = jnp.split(q, tp_size, axis=0) k_shards = jnp.split(k, tp_size, axis=0) v_shards = jnp.split(v, tp_size, axis=0) @@ -233,7 +258,9 @@ def _convert_moe(self, params): wi_1 = jnp.transpose(routed["wi_1"], (1, 0, 2, 3)) num_reps, num_experts, d_model, d_inner = wi_0.shape - tp_size = self.vllm_tp + # GMM_EP (expert parallelism) shards experts and keeps [gate | up] whole; + # GMM_TP interleaves per-TP gate/up chunks. + tp_size = 1 if self.vllm_use_ep else self.vllm_tp # vLLM's TPU Grouped GEMM kernel requires 128-alignment per expert chunk chunk_size = d_inner // tp_size @@ -264,7 +291,9 @@ def _convert_moe(self, params): shared = mlp_block["shared_expert"] sh_gate_layers = jnp.unstack(jnp.transpose(shared["wi_0"]["kernel"], (1, 2, 0)), axis=0) sh_up_layers = jnp.unstack(jnp.transpose(shared["wi_1"]["kernel"], (1, 2, 0)), axis=0) - sh_down_layers = jnp.unstack(jnp.transpose(shared["wo"]["kernel"], (1, 2, 0)), axis=0) + # wo.kernel is [F, layer, D]; per-layer [F, D] is already the runner's + # [in, out] layout for down_proj. + sh_down_layers = jnp.unstack(shared["wo"]["kernel"], axis=1) if "shared_expert_gate" in mlp_block: sh_gate_router_layers = jnp.unstack(jnp.transpose(mlp_block["shared_expert_gate"]["kernel"], (1, 2, 0)), axis=0) @@ -274,30 +303,29 @@ def _convert_moe(self, params): p = f"vllm_model.language_model.model.layers.{i}" self.vllm_state[f"{p}.mlp.gate.weight"] = router_weights[rep] - # vLLM nests the routed-expert weights one level deeper than the - # router (`mlp.experts.routed_experts.*`, not `mlp.experts.*`) -- - # missing that segment means these never match spec_flat during - # sync, so the routed experts (the bulk of a MoE model's params) - # silently stay at their random dummy-init value. + # Current vLLM nests the expert tensors under a `routed_experts` + # submodule; older versions keep them on the FusedMoE layer directly. + # Emit both names (same array, no copy) and let the structural sync + # pick whichever the target has. self.vllm_state[f"{p}.mlp.experts.routed_experts.w13_weight"] = w13_layers[rep] self.vllm_state[f"{p}.mlp.experts.routed_experts.w2_weight"] = down_layers[rep] + self.vllm_state[f"{p}.mlp.experts.w13_weight"] = w13_layers[rep] + self.vllm_state[f"{p}.mlp.experts.w2_weight"] = down_layers[rep] if has_shared: sh_g, sh_u = sh_gate_layers[rep], sh_up_layers[rep] - sh_per_tp = sh_g.shape[0] // self.vllm_tp + sh_per_tp = sh_g.shape[0] // self.attn_shards shared_gate_up = jnp.concatenate( [ - sh_g.reshape(self.vllm_tp, sh_per_tp, sh_g.shape[1]), - sh_u.reshape(self.vllm_tp, sh_per_tp, sh_u.shape[1]), + sh_g.reshape(self.attn_shards, sh_per_tp, sh_g.shape[1]), + sh_u.reshape(self.attn_shards, sh_per_tp, sh_u.shape[1]), ], axis=1, ).reshape(-1, sh_g.shape[1]) - # Same (in_features, out_features) convention as the GDN in_proj/out_proj - # weights above -- neither of these was transposed into it yet. self.vllm_state[f"{p}.mlp.shared_expert.gate_up_proj.weight"] = jnp.transpose(shared_gate_up, (1, 0)) - self.vllm_state[f"{p}.mlp.shared_expert.down_proj.weight"] = jnp.transpose(sh_down_layers[rep], (1, 0)) + self.vllm_state[f"{p}.mlp.shared_expert.down_proj.weight"] = sh_down_layers[rep] if "shared_expert_gate" in mlp_block: self.vllm_state[f"{p}.mlp.shared_expert_gate.weight"] = sh_gate_router_layers[rep]