Add Qwen3.5 weight mapping for tpu-inference's vLLM (torchax) path

[wenxindongwork]

Summary

Weight mapping from MaxText's Qwen3.5 (text) model to tpu-inference's native Qwen3.5, which runs on the vLLM (torchax) path, for RL weight sync through Tunix's VllmSampler.update_params.

• src/maxtext/integration/tunix/weight_mapping/qwen3_5.py: QWEN3_5_VLLM_MAPPING. The target is vLLM's canonical TP=1 parameter layout (module paths such as language_model.model.layers.N.self_attn.qkv_proj.weight); tpu-inference converts canonical arrays into its tp/backend-dependent internal layout (VllmModelWrapper.load_canonical_weights, companion PR), so the mapping knows nothing about TP size, KV-head replication or the MoE backend.
• Tunix's key mapping is one-to-one, so every fusion/reorder lives in preprocess_src_state: GDN per-key-head interleaved in_proj_qkvz/in_proj_ba → [Q|K|V|Z]/[B|A], query/key/value → qkv_proj (q keeps the attention output gate), wi_0/wi_1 → w13_weight, shared-expert gate_up_proj, conv1d [K,1,C] → [C,1,K], transposes. to_hf_mapping is then a regex rename (targets match both language_model.model.layers.N… and model.layers.N…, and experts(.routed_experts)?.w13_weight). Unscanned (layers_{i}) and inhomogeneous scanned (layers.layer_{b}) parameter trees are accepted; the GDN geometry comes from the HF config or is inferred from shapes. Each layer is materialized under jax.jit with even out_shardings on the source mesh (eager execution on fsdp-sharded kernels came back replicated and did not fit next to the trainer and sampler weights).
• preprocess_src_state is plumbed through VllmWeightMapping and TunixMaxTextAdapter (Tunix's MappingConfig.build/from_model pick it up).
• Registry fix: StandaloneVllmWeightMapping now checks the qwen3.5 prefix before qwen3, which used to swallow it and silently return the dense Qwen3 mapping.

Testing

• tests/post_training/unit/qwen3_5_vllm_weight_mapping_test.py (CPU): re-layouts cross-checked against MaxText's own HF export hooks (QWEN3_5_MAXTEXT_TO_HF_PARAM_HOOK_FN(saving_to_hf=True)), geometry inference, mapping coverage and target-regex resolution against real vLLM names, scanned vs unscanned equivalence.
• tests/post_training/integration/qwen3_5_vllm_weight_sync_check.py (manual, TPU): dump loads Qwen/Qwen3.5-35B-A3B from HF on the torchax path and records the runner state + greedy generations; sync loads gs://maxtext-model-checkpoints/qwen3.5-35b-a3b/unscanned/0/items (bf16), starts Tunix's VllmSampler, captures in-process reference generations with the HF weights, overwrites all weights with random values (or --random-init for vLLM's dummy loader), runs update_params, and compares every tensor and the generations.
• Result on v7x-8 (TP=8, bf16): all 613 weight tensors bit-identical to the HF-loaded model; greedy generations 4/4 identical to the in-process HF-weight reference; update_params ~150 s (first call, JIT-dominated).
• Same result (4/4 identical generations) with the production Qwen3.5 serving configuration on the sampler side (TP=8, --enable-expert-parallel, enable_dp_attention + attn_dp_size=4, prefix caching, async scheduling, chunked prefill, production env/LIBTPU flags): the mapping is unchanged, tpu-inference produces the EP/attn-DP layouts.

Companion PRs: vllm-project/tpu-inference#3477 (load_canonical_weights), google/tunix#2014 (VllmSampler.update_params torchax branch).

[gemini-code-assist]

Code Review

This pull request introduces weight mapping support from MaxText's Qwen3.5 model to tpu-inference's vLLM (torchax) Qwen3.5 implementation, including integration and unit tests. The review feedback highlights several improvement opportunities: refining the state flattening logic to prevent nested dictionaries from breaking the process, adding a fallback sharding strategy in _even_sharding to avoid out-of-memory errors on larger clusters, safely handling cases where decoder.logits_dense.kernel is missing, and using context managers (with statements) for file operations in the integration tests to prevent resource leaks.

[gemini-code-assist]

The check isinstance(state, Mapping) and all(isinstance(k, str) for k in state) is too broad because it matches nested dictionaries (where the top-level keys are strings). When a nested dictionary is passed, it will enter this branch, call state.items(), and then silently skip all nested dictionaries because they do not have a shape attribute. This completely breaks flattening for nested parameter trees.

To fix this, ensure that the mapping is flat (i.e., does not contain nested mappings) before using state.items(), or simply let tree_flatten_with_path handle it.

Suggested change

	elif isinstance(state, Mapping) and all(isinstance(k, str) for k in state):
	items = state.items()
	elif isinstance(state, Mapping) and all(isinstance(k, str) and not isinstance(v, Mapping) for k, v in state.items()):
	items = state.items()

[gemini-code-assist]

Currently, _even_sharding only shards a dimension if it is divisible by the total mesh size n. On larger clusters (where DP > 1) or with non-power-of-two mesh sizes, the total mesh size n can be quite large (e.g., 48, 96, 128). In such cases, none of the dimensions of a tensor (like w13 with shape [64, 2816, 5120]) might be divisible by n, causing the tensor to be fully replicated across all devices. This can easily lead to Out-Of-Memory (OOM) errors.

As a robust fallback, if no dimension is divisible by the total mesh size n, we can try to shard over the last axis of the mesh (typically the "model" axis, representing the TP size). This is much more likely to succeed and still provides significant memory savings.

def _even_sharding(mesh: jax.sharding.Mesh, shape) -> jax.sharding.NamedSharding:
  """Shards the first dimension divisible by the device count over the whole mesh (else replicates)."""
  n = mesh.size
  spec = [None] * len(shape)
  for i, dim in enumerate(shape):
    if dim % n == 0 and dim >= n:
      spec[i] = tuple(mesh.axis_names)
      break
  else:
    # Fallback: try to shard over the model axis (last axis) if total mesh size doesn't divide any dimension
    model_axis = mesh.axis_names[-1]
    model_size = mesh.shape[model_axis]
    for i, dim in enumerate(shape):
      if dim % model_size == 0 and dim >= model_size:
        spec[i] = model_axis
        break

[gemini-code-assist]

decoder.logits_dense.kernel is unconditionally accessed from globals_. If a model configuration ties embeddings (or for some other reason does not have separate LM head weights), this will raise a KeyError.

It is safer to check if "decoder.logits_dense.kernel" in globals_ before attempting to materialize and update it.

Suggested change

	out.update(
	_materialize(
	lambda g: {"base.decoder.logits_dense.kernel": jnp.transpose(g["decoder.logits_dense.kernel"])},
	{"decoder.logits_dense.kernel": globals_["decoder.logits_dense.kernel"]},
	)
	)
	if "decoder.logits_dense.kernel" in globals_:
	out.update(
	_materialize(
	lambda g: {"base.decoder.logits_dense.kernel": jnp.transpose(g["decoder.logits_dense.kernel"])},
	{"decoder.logits_dense.kernel": globals_["decoder.logits_dense.kernel"]},
	)
	)

[gemini-code-assist]

Opening files directly with open(...) without a with statement can lead to resource leaks because the file descriptor is not guaranteed to be closed immediately.

Use with open(...) as f: context managers to ensure files are properly closed.

Suggested change

	json.dump(meta, open(os.path.join(out, "vllm_state_meta.json"), "w"), indent=1)
	json.dump(_generate(llm), open(os.path.join(out, "vllm_ref_generations.json"), "w"), indent=1)
	with open(os.path.join(out, "vllm_state_meta.json"), "w") as f:
	json.dump(meta, f, indent=1)
	with open(os.path.join(out, "vllm_ref_generations.json"), "w") as f:
	json.dump(_generate(llm), f, indent=1)

[gemini-code-assist]

Opening meta_path directly with open(...) without a with statement can lead to resource leaks.

Use a with open(...) as f: context manager to ensure the file is properly closed.

Suggested change

	for m in json.load(open(meta_path)):
	with open(meta_path) as f:
	meta_data = json.load(f)
	for m in meta_data:

[gemini-code-assist]

Opening ref_path directly with open(...) without a with statement can lead to resource leaks.

Use a with open(...) as f: context manager to ensure the file is properly closed.

Suggested change

	ref = json.load(open(ref_path))
	with open(ref_path) as f:
	ref = json.load(f)