Step0 reforward test scripts - #4925
Conversation
…rift results - Implemented 25-intermediate activation tensor capture harness for Qwen3.5 MoE decoder layer (tests/unit/qwen3_5_layer_dump_test.py). - Added CLI comparison and NPZ analysis tool (tests/analyze_qwen3_5_layer_dump.py). - Added SPS benchmark runner on Cloud TPU v5p (tests/run_sps_qwen3_5_dump.py). - Documented 25-tensor numerical drift benchmark results (docs/qwen3_5_kernel_drift_results.md). - Added 'dcp' axis support to vllm.yml and types.py for inference mesh compatibility. TAG=agy CONV=a6a5e1d4-a4a7-4cab-be35-bbf37e64f5e2
- Added automated comparative evaluation on Cloud TPU v5p for Splash block sizes (512 vs 128) and exact softmax transcendental math - Fixed sm_scale calculation in forward_serve_vllm to use self.query_scale or 1/sqrt(head_dim) - Updated documentation with empirical findings and MXU hardware behavior TAG=agy CONV=a6a5e1d4-a4a7-4cab-be35-bbf37e64f5e2
… with authentic RPA and Pallas MoE
…ic summary and numpy metric optimizations
…ayers/attentions.py
…tedMoE - Compute SwiGLU activation functions and intermediate elementwise products in Float32 to reduce truncation drift. - Ensure gate logits use Float32 for Qwen3 decoder blocks when float32_gate_logits is configured.
- Support Tokamax Splash Attention vs RPA and Batched RPA comparison on TPU. - Add sweep options for base-2 vs base-e natural exponential and reciprocal fusion.
- Implement isolated TPU benchmarks comparing Tokamax GMM v2 vs Pallas Fused MoE. - Add diagnostic scripts to measure spectral error amplification through MLP projections.
…kernel configs - Enable Tokamax Splash Option A (sa_use_base2_exp=False, sa_fuse_reciprocal=True). - Enable Tokamax GMM v2 (wi_tile_fwd_batch_seq=256, sparse_matmul=True) in BF16 and FP32. - Configure FP32 routing logits and weighted combination precision flags.
…and next plan - Document Attention & MoE kernel parity findings and 1-ULP quantization limits. - Add comprehensive parity improvement story with progressive diff tables. - Add multi-layer error mitigation plan and roadmap in docs/next_plan.md. - Record updated 25-intermediate tensor drift metrics in docs/qwen3_5_kernel_drift_results.md.
Drop the sps_ prefix now that these scripts run directly against locally-attached TPU chips rather than the Shared Pathways Service proxy. Co-Authored-By: Claude Sonnet 5 <[email protected]>
…xecution Drop the pathwaysutils/isc_pathways SPS connection boilerplate in favor of running directly against locally-attached Cloud TPU v5p chips. Also fix inference-mesh construction: tensor-parallel degree is capped at num_kv_heads (not data-parallel across all devices), so the mesh is now built by hand from a device slice and inputs are explicitly re-placed (device_put) onto it, instead of maxtext_utils.create_device_mesh, which requires the ICI product to equal the full visible device count and was producing shard_map device-set mismatches in the RPA kernel. Co-Authored-By: Claude Sonnet 5 <[email protected]>
Builds the full Qwen3.5 model end-to-end on identical seeded weights and compares final logits between the training path (Splash Attention + Tokamax GMM v2 MoE) and the inference path (vLLM RPA attention + fused MoE kernel), reporting the top-1/top-5 argmax agreement and KL divergence that actually determine greedy-decoding parity, rather than raw intermediate-tensor distance metrics alone. Co-Authored-By: Claude Sonnet 5 <[email protected]>
Refresh attention and MoE kernel repro results and the 1-layer intermediate tensor drift breakdown with numbers from the local-TPU-VM runs (post SPS proxy removal and mesh fixes), and add a standalone MoE kernel repro results doc. Co-Authored-By: Claude Sonnet 5 <[email protected]>
Fold learnings.md, docs/next_plan.md, and docs/parity_improvement_story.md into docs/train_infer_logit_parity.md alongside the full-model logit parity results, removing the fragmented/stale duplicates. Also drop tests/analyze_qwen3_5_layer_dump.py, superseded by the drift results doc. Co-Authored-By: Claude Sonnet 5 <[email protected]>
|
Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA). View this failed invocation of the CLA check for more information. For the most up to date status, view the checks section at the bottom of the pull request. |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces comprehensive standalone reproduction, diagnostic, and parity verification scripts for Qwen3.5 attention and MoE kernels, along with corresponding documentation. It also updates logical axis rules for vLLM context parallelism and adjusts activation casting in linear layers and MoE routing. The review feedback highlights several critical issues in the newly added tests and linear layers: unconditionally casting activations to float32 in linear layers bypasses configuration settings and increases memory usage; attempting to subscript nnx.Param objects directly in weight synchronization will raise errors; and multiple JAX device mismatch bugs exist in the test suites due to sharing inputs and parameter states directly between different device meshes without explicit device placement.
| if cfg.activations_in_float32: | ||
| x = x.astype(jnp.float32) | ||
| x = _convert_to_activation_function(act_fn)(x) | ||
| x = _convert_to_activation_function(act_fn)(x.astype(jnp.float32)) |
There was a problem hiding this comment.
Unconditionally casting x to jnp.float32 before applying the activation function defeats the purpose of the cfg.activations_in_float32 configuration. This will force float32 activations for all models using this linear layer, leading to increased memory usage (HBM/VMEM) and potentially slower execution when running in lower precision (e.g., bfloat16). It should remain conditional on cfg.activations_in_float32.
| x = _convert_to_activation_function(act_fn)(x.astype(jnp.float32)) | |
| if cfg.activations_in_float32: | |
| x = x.astype(jnp.float32) | |
| x = _convert_to_activation_function(act_fn)(x) |
| and hasattr(src_moe, "wi_0") | ||
| and hasattr(src_moe, "wi_1") | ||
| ): | ||
| wi_fused = jnp.concatenate([src_moe.wi_0[...], src_moe.wi_1[...]], axis=-1) |
There was a problem hiding this comment.
src_moe.wi_0 and src_moe.wi_1 are nnx.Param (or nnx.Variable) objects in Flax NNX, which do not support direct slicing/indexing with [...]. Attempting to do src_moe.wi_0[...] will raise a TypeError: 'Param' object is not subscriptable (or similar). You should access their underlying JAX arrays using the .value attribute instead.
| wi_fused = jnp.concatenate([src_moe.wi_0[...], src_moe.wi_1[...]], axis=-1) | |
| wi_fused = jnp.concatenate([src_moe.wi_0.value, src_moe.wi_1.value], axis=-1) |
| # Force weight synchronization | ||
| nnx.update(infer_model, nnx.state(train_model, nnx.Param)) |
There was a problem hiding this comment.
Directly updating infer_model with the parameter state of train_model via nnx.update will copy/alias the parameters that are placed on train_mesh's devices. Since infer_model runs on infer_mesh (which has a different device topology/count), executing infer_model will trigger a JAX device mismatch error (ValueError: Received incompatible devices for jitted computation).
To fix this, you must explicitly place the parameter state onto infer_mesh's devices using jax.device_put with a replicated sharding before calling nnx.update, similar to the working pattern implemented in tests/run_qwen3_5_logit_parity.py.
| # Force weight synchronization | |
| nnx.update(infer_model, nnx.state(train_model, nnx.Param)) | |
| # Force weight synchronization and place them on the inference mesh devices | |
| infer_replicated_sharding = NamedSharding(infer_mesh, P()) | |
| train_param_state = nnx.state(train_model, nnx.Param) | |
| train_param_state_for_infer = jax.tree_util.tree_map( | |
| lambda x: jax.device_put(x, infer_replicated_sharding), train_param_state | |
| ) | |
| nnx.update(infer_model, train_param_state_for_infer) |
| print(" Executing Training Model...") | ||
| logits_train = train_model(token_ids, positions, model_mode=MODEL_MODE_TRAIN) | ||
|
|
||
| print(" Executing Inference Model...") |
There was a problem hiding this comment.
The inputs token_ids and positions are placed on train_mesh. Passing them directly to infer_model (which is compiled and executed on infer_mesh) will cause a JAX device mismatch error. You must explicitly place separate copies of the inputs onto infer_mesh using jax.device_put before calling infer_model.
infer_token_ids = jax.device_put(token_ids, infer_replicated_sharding)
infer_positions = jax.device_put(positions, infer_replicated_sharding)
infer_out = infer_model(infer_token_ids, infer_positions, model_mode=MODEL_MODE_PREFILL)| # Synchronize weights | ||
| infer_moe.gate.kernel = train_moe.gate.kernel | ||
| if hasattr(train_moe.gate, "bias") and train_moe.gate.bias is not None: | ||
| infer_moe.gate.bias = train_moe.gate.bias | ||
| infer_moe.wi_0 = train_moe.wi_0 | ||
| infer_moe.wi_1 = train_moe.wi_1 | ||
| infer_moe.wo = train_moe.wo |
There was a problem hiding this comment.
Directly aliasing the parameter variables from train_moe to infer_moe causes infer_moe to use parameters that are placed on train_mesh's devices. Since infer_moe is instantiated on infer_mesh, executing it will trigger a JAX device mismatch error.
To resolve this, you should split the infer_moe state, place the synchronized weights onto infer_mesh's devices using jax.device_put, and then merge them back, breaking the aliasing and ensuring correct device placement.
| # Synchronize weights | |
| infer_moe.gate.kernel = train_moe.gate.kernel | |
| if hasattr(train_moe.gate, "bias") and train_moe.gate.bias is not None: | |
| infer_moe.gate.bias = train_moe.gate.bias | |
| infer_moe.wi_0 = train_moe.wi_0 | |
| infer_moe.wi_1 = train_moe.wi_1 | |
| infer_moe.wo = train_moe.wo | |
| # Synchronize weights and place them on the inference mesh devices | |
| infer_moe.gate.kernel = train_moe.gate.kernel | |
| if hasattr(train_moe.gate, "bias") and train_moe.gate.bias is not None: | |
| infer_moe.gate.bias = train_moe.gate.bias | |
| infer_moe.wi_0 = train_moe.wi_0 | |
| infer_moe.wi_1 = train_moe.wi_1 | |
| infer_moe.wo = train_moe.wo | |
| infer_replicated_sharding = NamedSharding(infer_mesh, P()) | |
| _infer_graphdef, _infer_state = nnx.split(infer_moe) | |
| _infer_state = jax.tree_util.tree_map( | |
| lambda x: jax.device_put(x, infer_replicated_sharding), _infer_state | |
| ) | |
| infer_moe = nnx.merge(_infer_graphdef, _infer_state) |
| out_train_2d = out_train.reshape(num_tokens, emb_dim) | ||
|
|
||
| print(" [2/3] Executing Inference MoE (Fused MoE via tpu-inference)...") | ||
| out_infer_fused, _, _ = infer_moe(inputs_3d) |
There was a problem hiding this comment.
inputs_3d is placed on train_mesh. Passing it directly to infer_moe (which runs on infer_mesh) will cause a JAX device mismatch error. You must explicitly place a copy of the inputs onto infer_mesh's devices using jax.device_put before calling infer_moe.
| out_infer_fused, _, _ = infer_moe(inputs_3d) | |
| infer_inputs_3d = jax.device_put(inputs_3d, infer_replicated_sharding) | |
| out_infer_fused, _, _ = infer_moe(infer_inputs_3d) |
2 participants
Description
Start with a short description of what the PR does and how this is a change from
the past.
The rest of the description includes relevant details and context, examples:
If the change fixes a bug or a Github issue, please include a link, e.g.,:
FIXES: b/123456
FIXES: #123456
You can also provide a comma-separated list. If you don't want to close a bug but
simply to reference it, use BUGS, e.g.:
BUGS: b/123456
Notice 1: Once all tests pass, the "pull ready" label will automatically be assigned.
This label is used for administrative purposes. Please do not add it manually.
Notice 2: For external contributions, our settings currently require an approval from a MaxText maintainer to trigger CI tests.
Tests
Please describe how you tested this change, and include any instructions and/or
commands to reproduce.
Checklist
Before submitting this PR, please make sure (put X in square brackets):
gemini-reviewlabel.