MaxText is available in PyPI and can be installed through pip. Please see our MaxText Installation Guide for setup instructions.

Last Updated: fb39de80e

• GraniteMoE Architecture & M3 Onboarding: Added GraniteMoE (granite-3.0-1b-a400m) support in M3, including pure Flax NNX decoder layers, checkpoint parameter mappings, conformance and real-weight TPU parity tests, and the model-bringup-m3 agent skill.
• M3 Model Conformance & Spec Discovery: Added CPU conformance specifications for Qwen3-0.6B and GLM-5.3-Flash, wired lazy dynamic spec discovery via pkgutil.iter_modules into conformance_test, and hardened pure-array NNX transform boundaries and scanned expert all-gather detection.
• MoE Router Precision & Qwen3 Shared Experts: Preserved float32 precision for router inputs in RoutedMoE to prevent premature downcasting and wired float32_gate_logits to the Qwen3 shared expert gating path for numerical stability.
• GLM-5.3-Flash Architecture & Training Integration: Consolidated the GLM-5.3-Flash model architecture submodules into a single modeling_glm5_3_flash.py module and relocated configurations to src/maxtext/m3/configs/glm5_3_flash_config.py per M3 single-file-per-model conventions; implemented the Layer-45 NextN Multi-Token Prediction (G8-mtp) block and training integration; added dense-mask DSA attention dispatch (dsa_attention_implementation="dense_mask"); established explicit return_aux contracts and heterogeneous super-block slot alignment; and implemented training loss modules with float32 masked cross-entropy, load-balancing loss, and AdamW optimizer parameter freezing in glm5_train.py alongside lazy FLOPs calculations.
• GLM-5.3-Flash Checkpoint Conversion & Parity Verification: Added single-owner checkpoint conversion specifications (iter_conversion_specs) with mixed-dtype float32 router parameters and host NumPy FP8 block-128 dequantization in glm5_ckpt.py; exported unit test suites to GitHub; and added end-to-end gradient and logits parity test suites across all three decoder block architectures (l0_kda_dense, l3_mla_moe, l4_kda_moe) and real model slices against HuggingFace references.
• Qwen3.8-Flash-Next Architecture & Tooling: Assembled the complete end-to-end Qwen3.8-Flash-Next model architecture with 4-stream Hyper-Connection residuals, heterogeneous 4-layer cycle blocks (3 GDN linear attention + 1 QSA sparse attention), Layer-2 Pre-Layer Embedding (PLE), and 512-expert Sparse MoE feed-forward blocks supporting both unrolled and scanned execution; added TPU kernel paths for GDN, QSA, and MoE; enabled FSDP batch sharding, scan remat, and per-chip MoE dispatch; implemented the Multi-Token Prediction (MTP) draft layer and integrated loss; and added training TFLOPs calculations and checkpoint conversion support.
• Hy4-Preview Architecture & FLOPs Accounting: Onboarded the Hy4-preview model architecture (hy4-preview.yml), combining Gated MLA, DeepSeek Sparse Attention (DSA) indexers, Identity Hyper-Connections (iHC), and routed MoE feed-forward blocks with HuggingFace checkpoint conversion and parameter mappings; added training TFLOPs calculations modeling Gated MLA projections, top-k sparse attention, and iHC stream expansion; and expanded FSDP weight prefetch guides (fsdp prefetch).
• GLM-5.3-Flash Architecture & Onboarding: Added full M3 modeling support for GLM-5.3-Flash (glm5.3-flash), including configuration and sharding foundations, Kimi Delta Attention (Glm5KdaAttention) and primitive modules (l2norm, Glm5RMSNormGated, Glm5ShortConv1d), full-attention sub-blocks (Glm5DsaMlaAttention) with k-pooled lightning indexing (Glm5Indexer), feed-forward and routed Mixture-of-Experts (MoE) stacks (glm5_moe.py), Manifold-Constrained Hyper-Connections (mHC), heterogeneous Transformer decoder blocks, 4-layer super-block scanning and rematerialization (glm5_scan.py), checkpoint conversion (glm5_ckpt.py), mini-slice creation utilities, and lazy factory model registration.
• Qwen3.8-Flash-Next Architecture: Added M3 modeling implementation and YAML configurations for Qwen3.8-Flash-Next (qwen3.8-flash-next.yml, qwen3.8-flash-next-mini.yml), including 4-stream Gated Residual (Hyper-Connection) NNX modules, Gated DeltaNet (GDN) linear attention with float32 chunked scan, Qwen Sparse Attention (QSA), Indexer with partial interleaved mRoPE, Sparse MoE blocks with 512 routed and 1 gated shared expert, and SplitMix64 n-gram hashing with Layer-2 Parameterized Linear Embedding (PLE) modules.
• Ming Architecture: Added Ming hybrid KDA/MLA text decoder with scanned execution, inhomogeneous layer cycling alternating between gated Multi-Head Latent Attention (MLA) and Kimi Delta Attention (KDA), DeepSeek-style routed MoE, and the ming-3.0-audio-tiny.yml model configuration.
• DeepSeek-V4-Pro (1600B) Support: Added DeepSeek-V4-Pro (1.6T parameters, 384 routed experts) model configuration (deepseek4-1600b.yml), HuggingFace parameter and shape mappings, checkpoint conversion, and end-to-end test pipelines (supported models and architectures).
• DeepSeek-V4 Sequence-Wise Router Balance Loss: Enabled the sequence-wise auxiliary load balancing loss for DeepSeek-V4 Sqrt(Softplus) routers (moe_use_megatron_seq_aux_loss: true), matching DeepSeek-V3 sequence-wise routing formulations and summing per-layer losses across decoder MoE layers.
• Qwen3.8 Architecture & Parameter Mappings: Added model configuration and HuggingFace parameter mapping for Qwen3.8-2.4T-A95B, parameterizing hf_prefix in Qwen3.5 mapping and shape builders to support Qwen3.8 variants.
• M3 Default Sharding: Added DefaultSharding for m3 (Modern MaxText Models) and aligned mesh axis names with legacy MaxText conventions.
• DeepSeek-V4 Layer Rematerialization: Honored remat_policy in DeepSeek-V4 train-mode layers by running unscanned prefix layers under jax.checkpoint within a length-1 scan barrier and applying layer-level rematerialization inside DeepSeek4ScannableBlock, avoiding block-level rematerialization to fit large configurations into device HBM.
• Lineage DeepSeek-V3 Modules & Mesh Sharding: Synced Lineage DeepSeek-V3 modules into maxtext/experimental/lineage, including the branched Splash Attention kernel and 1D expert-parallel sharding; looked up Lineage DeepSeek-V3 weight shardings dynamically by mesh topology; aligned the package with the code that produced the DeepSeek-V3 671B MLPerf Training v6.1 8k-chip result so it runs on released jax 0.11.1 / libtpu 0.0.48; and fed the decoder's final-norm output into the Lineage MTP block, mirroring mtp_norm_main_hidden_state (always on for the Lineage path).
• DeepSeek-V3 Lineage Architecture: Added Lineage DeepSeek-V3 modules under maxtext/experimental/lineage with the exported deepseek3-671b-lineage configuration and embed_router rule. Lineage is a very efficient implementation of DeepSeek-V3 targeting pdb=1 using a batch split strategy where the EP comms are fully overlapped (overlapped with splash). Also imported Lineage modules lazily in lineage_adapter when use_lineage=False, added optional model hooks for alternative DeepSeek-V3 implementations, and added lineage_expert_permutation for per-step expert load balancing in the Lineage MoE path.
• DeepSeek-V4 Multi-Token Prediction: Added Multi-Token Prediction (MTP) support for DeepSeek-V4 decoder layers along with HuggingFace parameter mapping and shape alignment utilities.
• WeaverOmni Architecture: Added the full 36-layer WeaverOmniTransformer backbone architecture and tests, supporting the weaver and weaver-mini-diffuser model configurations.
• Qwen3.5-9B Model Support: Added base model onboarding for Qwen3.5-9B, including model configurations, HuggingFace checkpoint conversion, decode workflows, and logit verification.
• Flax NNX Migration: Ported the remaining Linen sharding helpers into native MaxText implementations (sharding.logical_to_spec and sharding.with_logical_constraint), replacing Flax Linen dependencies across KV cache, attention ops, and decoder architectures. Additionally, converted FP8 ops (Fp8DirectDotGeneralOp, NANOOFp8DotGeneralOp, Fp8Einsum, QwixDotGeneral, QwixEinsum) and HuggingFace checkpoint parameter conversion (to_maxtext.py) to native Flax NNX modules, and deleted legacy ToLinen layer wrappers and factories.

• Halo Sliding-Window Context Parallelism & Per-Layer CP: Added Halo Sliding-Window Context Parallelism (AttentionOp.tpu_halo_flash_attention, local_context_parallel_strategy="halo") and per-layer CP strategy resolution across hybrid sliding-window models (Gemma 2/3/4, GPT-OSS, OLMo 3), sequence-sharding Q, K, and V across the context axis and exchanging only the block-aligned W-token K/V halo tail via lax.ppermute across contiguous and load-balanced dual_chunk_swap layouts.
• GPU Tokamax Attention Enhancements: Supported non-causal attention, sliding window masks, and attention logit soft-capping in the GPU Tokamax ragged attention path (AttentionOp).
• Threaded DiLoCo Data-Parallel Submesh Replicas: Added train loop and input pipeline hooks for training data-parallel replicas on submeshes in threaded Streaming DiLoCo, including setup_train_loop(..., mesh=None), centralized shard index calculations (input_pipeline_utils.get_dataloading_shard), and per-process replica configuration fields.
• JAX Ring Attention: Implemented JAX Ring Attention across devices using SPMD ppermute communication, leveraging JAX Flash Attention forward and backward passes for distributed Context Parallelism.
• Threaded Streaming DiLoCo Transfer & Optimization: Added FragmentTransfer and Nesterov outer-step utilities for threaded streaming DiLoCo distributed training, supporting shard-local 1-D parameter fragment extraction, asynchronous mesh transfer, and in-place layer parameter application without cross-device resharding overhead.
• Splash Attention Ops: Promoted Splash Attention kernels out of experimental ops into standard MaxText attention operations.
• M3 Profiling Training Loop & Trace Scopes: Added a lightweight M3 training loop for workload profiling (maxtext/m3/train/loop.py) featuring JIT-compiled synthetic batch generation, eliminated per-step NNX graph traversal overhead by caching model and optimizer state with nnx.cached_partial, and restored hierarchical jax.named_scope tracking for NNX modules in XProf traces.
• Splash Attention FP8 Scaling: Added the experimental sa_q_prescaled_log2e (folds log2(e) into the float32 Q softmax scale before the FP8 cast, so Q is rounded to FP8 once and in-kernel Q scaling is skipped) and sa_dkv_reuse_scaled_q (the fused dKV kernel reuses the log2(e)-scaled Q for dK) Splash Attention options, both off by default.
• Splash Attention Stream Chunking Optimization: Added sa_ignore_segment_ids to skip uniform segment mask loading and comparisons in Splash Attention when running continuous stream-chunked c4_mlperf datasets, eliminating redundant VPU overhead while preserving position IDs and loss masking.
• YaRN Cis Optimization: Built YaRN rotary position embedding freqs_cis values directly from complex(cos(theta), sin(theta)) instead of jnp.exp(1j * theta), avoiding XLA complex exponential expansion and redundant exponential evaluations under explicit sharding.
• Multi-Token Prediction Hidden State Normalization: Added mtp_norm_main_hidden_state configuration flag to apply decoder normalization before the output head and feed the normalized hidden state directly into the Multi-Token Prediction (MTP) block, aligning with Megatron-LM DeepSeek-V3 behavior.
• JAX Flash & Splash Attention: Added the JAX Flash and Splash Attention backward pass supporting Multi-Head Attention (MHA), Grouped-Query Attention (GQA), and Multi-Query Attention (MQA) with block sparsity, causal masking, soft capping, and packed-sequence segment IDs, and exposed sa_qk_diag_skip, sa_sv_diag_skip, and sa_qk_diag_grid configuration flags in Tokamax Splash Attention.
• YaRN RoPE Optimizations: Applied YaRN RoPE rescaling to DeepSeek-V4 compressed attention (CSA/HCA) layers matching the HuggingFace configuration. Also added rope_direct_position_freqs to compute YaRN RoPE angles directly from position frequencies in 4D float32 without lookup tables, and added rope_dot_pairwise using a 4D block-skew contraction matrix to preserve rank-4 activations during pairwise rotation and eliminate layout copy overhead.
• DiLoCo Distributed Training: Added split_mesh_along_axis and remove_mesh_axis_from_rules to derive per-replica views by splitting meshes along the diloco axis and stripping the axis from logical rules, and introduced diloco_bucketize_non_scanned to optionally bucketize non-scanned parameter matrices across layer fragments along their longest unsharded axis to balance data transfer.

• MaxText Training Engine Kernel & Gradient Architecture: Separated forward-backward execution from running gradient accumulation into dedicated kernels donating sums in grad_dtype, divided gradients once by total denominators in float32, added cast_grads_after_all_reduce, refactored parameter freezing via trainable_parameters_mask with public optimizers.get_path_mask_fn, and introduced step guards (skip_step_on_nan) and spike protection (skip_step_on_spikes).
• MaxText Training Engine Checkpointing & GRPO AOT Compilation: Hardened checkpoint save/restore by migrating MaxTextTrainingEngine to Orbax v1 training.Checkpointer with explicit checkpoint_dir support, async save timeouts, and pinned-host optimizer leaf restoration into device memory; and added GRPO ahead-of-time (AOT) compilation with per-device HBM utilization reporting.
• Tunix & vLLM Weight Synchronization: Ported weight conversion and weight synchronization functionalities to align with Tunix APIs, supporting unrolled tensor rank expert detection, bidirectional plan validation, and target-free FP8 MoE rollout quantization.
• Tunix Chunked Log-Probabilities & RL Configuration: Wired up chunked log-probability computation (skip_lm_head and compute_final_logits) in TunixMaxTextAdapter, and declared RLConfig.train_steps as an explicit validated Pydantic field across RL workflows.
• DeepSeek-V4 LoRA Support: Added DeepSeek-V4 module paths to lora_module_path.yml covering attention projections, shared experts, and HCA/CSA compressor projections, and preserved lazily created LoRA parameters across per-layer rematerialization scans.
• Tunix Distributed RL Trainer Alignment: Updated AbstractTrainingEngine and MaxTextTrainingEngine to declare explicit argument signatures without **kwargs and migrated trainer imports to tunix.dist_rl.
• DeepSeek-V4 SFT Chat Encoding: Added native DeepSeek-V4 chat encoding and message formatting to SFT preprocessing in maxtext/input_pipeline/encoding_dsv4.py, supporting thinking-mode turn splitting and multi-turn tool calling across Grain and HuggingFace SFT pipelines.
• vLLM Rollout Weight Sync & Conversion: Added Decode Context Parallelism (DCP) weight synchronization for Qwen3.5 in MaxTextVllmSampler and WeightConverter, preserving user-configured tensor/expert parallelism overrides during rollout initialization, and added rollout_weight_dtype for target-free weight conversion.
• RL Profiling & Managed ML Diagnostics: Integrated the Tunix RL trainer with Tunix's per-role, invocation-counted profiler (RLProfileConfig) through the new rl.profiler_start_invocation, rl.profiler_num_invocations, rl.profiler_timeout_secs, and rl.profiler_max_num_hosts keys (the last also traces disaggregated sampler hosts on Pathways), producing one <role>_invocation_<n> trace per role. With managed_mldiagnostics=true, the trainer creates the Managed ML Diagnostics run before the RLCluster, uploads the trainer config (plus differing sampler.* keys on multi-slice topologies), and installs MLDiagScalarBackend as a metrics backend. Added managed_mldiagnostics_metric_only_run to create metric-only runs.
• Decode Context Parallelism for vLLM: Added Decode Context Parallelism (DCP) support for Qwen3.5 in MaxTextForCausalLM via the vLLM adapter.
• Block-Diffusion Fine-Tuning: Extended the block-diffusion objective to Continued Fine-Tuning (CFT) and completion-only Supervised Fine-Tuning (SFT), integrating with Tunix diffusion training and corrupting only configured supervision scopes while preventing future-prompt leakage.
• Continuous Profiling: Added the enable_continuous_profiling configuration flag across profilers and post-training pipelines (SFT, DPO, RL, and Distillation) to capture profiles for long-running workloads.

• Wan 2.1 Diffusion Pipelines & Inference: Added full diffusion inference pipelines (wan_pipeline.py), checkpoint conversion utilities (to_maxtext_diffusers.py), and a unified sample.py inference entrypoint for Wan 2.1 models, while promoting diffusers to shared base requirements across all environments.
• Wan 2.1 Diffusion Transformer: Implemented the Wan 2.1 Diffusion Transformer (DiT) model family under src/maxtext/m3/models/wan/, including the 3D causal VAE (autoencoder_kl_wan.py), UMT5 text encoder (text_encoder_wan.py), attention kernels, and ahead-of-time (AOT) cache support for wan2.1-t2v-14b.
• Ming Whisper Audio Integration & Checkpoint Conversion: Added audio input processing and inference preprocessing for Ming Whisper under src/maxtext/multimodal/processor_ming.py, and added HuggingFace audio checkpoint parameter mappings and golden logit verification utilities for Ming audio models.
• Diffusion Schedulers & Multimodal VAE Processors: Added M3 diffusion model configurations, continuous flow matching (scheduling_flow_match.py) and UniPC multistep schedulers (scheduling_unipc_multistep.py), multimodal VAE processors (processor_vae.py), and added diffusers to pre-training requirements for Wan 2.1 diffusion models.
• Qwen3.5-9B Multimodal Support: Enabled use_multimodal=true for Qwen3.5-9B by adding its Qwen3.5 vision-encoder (vision_encoder_block: qwen3_5) and MRoPE settings to qwen3.5-9b.yml and registering the model with the multimodal processor and decoder alongside the other Qwen3.5 variants.
• Qwen3-VL Integration: Added M3-style qwen3-vl-2b modeling implementation and layer correctness comparison tests.
• Wan 2.1 Video Diffusion Scaffolding: Added M3 architecture skeleton and module stubs for diffusion models and Wan 2.1 text-to-video (wan2.1-t2v-14b), including Diffusers schedulers and pipeline utilities.

• Qwix Quantization & M3 Integration: Added Qwix multi-pass quantization (MultipassQtProvider) to MaxText, introduced M3 precision and quantization transforms (matmul_precision), and enabled multi-pass Qwix quantization with default multi-device FSDP state sharding for M3 models.
• DeepSeek-V3 671B FP8 TPU7x Recipes: Added pre-training recipes for DeepSeek-V3 671B FP8 (FSDP + EP) on TPU7x targeting 2,048 chips (2k), 4,096 chips (4k), and 8,192 chips (8k) under benchmarks/deepseek3_671b_tpu7x/, featuring FP8 Splash Q/K, tuned GMM v2 tiles, and SparseCore all-reduce offload.
• FSDP Weight Prefetch Optimizations: Pinned FSDP prefetch all-gather and reduce-scatter collectives to SparseCore 1 on TPU7x via jax.experimental.compute_on, automatically registered max_async_overlap_transform, optimized FSDP prefetch under expert parallelism, supported host-offloaded layer inputs (decoder_layer_input=offload), and updated documentation (fsdp prefetch).
• Muon Optimizer Support for Non-Leading Scan Axes: Extended the Muon optimizer to support scanned parameters where param_scan_axis != 1, fixed flag handling, and added comprehensive unit test coverage.
• Tensor-Axis Expert Parallelism (tp-as-ep): Added the tp-as-ep.yml custom mesh and sharding rule preset, sharding routed experts over (expert, tensor) to eliminate tensor-axis collectives from MoE blocks while retaining standard layouts for other layers (custom mesh and rule).
• Transformer Engine MoE Support & Fixed Routing Benchmarks: Integrated Transformer Engine (TE) MoE block support for DeepSeek-V4 changes, and added fixed random-logit routing (moe_benchmark_routing: fixed_random_logits) to balance expert utilization during GPU Transformer Engine benchmarking without full router scoring overhead (benchmark and performance).
• Pipeline Parallelism under Explicit Sharding: Enabled pipeline parallelism under shard_mode=explicit, deferring per-microbatch weight-gradient reductions in JAX using unreduced tags, applying PP-as-DP construction rules, and integrating explicit pipeline execution tests (shard mode performance).
• Hy4 Indexer Activation Rematerialization: Saved topk_indices across activation rematerialization boundaries via checkpoint_name in Hy4Indexer, allowing JAX dead-code elimination to prune redundant projections and top-k operations during backward passes.
• Llama4-Maverick v5p-256 Expert Parallelism Recipe: Switched the Llama4-Maverick v5p-256 benchmark recipe from FSDP32 x TP4 to FSDP8 x EP16 with dropless fallback, Pallas ragged-sort permute kernels, SparseCore collective offloading, and tuned Tokamax Splash and Megablox tiles.
• M3 Named Scope Profiling: Added jax.named_scope tracking across M3 modules and training loops to ensure NNX module and layer names appear distinctly in XProf and TPU execution profiles.
• Native Microscaling (MX) FP8 and FP4 Training: Added native MXFP8 and MXFP4 training support via jax.nn.scaled_dot_general in src/maxtext/layers/quantizations.py.
• Transformer Engine GPU Kernels & Grouped GEMM: Integrated Transformer Engine (TE) mHC Triton kernels for GPU and added an opt-in TE grouped GEMM backend (use_te_grouped_gemm) for sparse MoE execution with per-expert bias fusion.
• FSDP Prefetch Inhomogeneous Block Overlap: Extended prefetch_fsdp_defer_grad_reduce_scatter to inhomogeneous decoder blocks (fsdp_prefetch_inhomogeneous.py, e.g. Qwen3.5), overlapping each sub-layer's weight-gradient reduce-scatters with preceding sub-layer backward compute.
• Training Memory Lifetime Reductions: Released LM head logits cotangents together with input and weight cotangents via tie_lm_head_cotangents to prevent retaining [batch, seq, vocab] tensors across backward passes, and eliminated saving logits in the forward residual of cross-entropy loss computations.
• Circular Pipeline Parallelism & BSW Schedule Overlap: Overlapped circular pipeline FSDP collectives using a shifted Bidirectional Shift-Wheel (BSW) schedule that reorders prologue bubble iterations and single/mixed-slot phases to hide forward and backward all-gather and reduce-scatter communications behind independent compute; added carried w_curr optimization with @jax.custom_vjp to eliminate redundant all-gathers; introduced BSW phase-split, pipeline_bsw_no_carry, and compute-dtype options; and preserved cp-as-ep context sharding on MoE expert weights in BSW.
• FSDP Weight Prefetching & Sub-Layer Double-Buffering: Added sub-layer FSDP weight prefetching for inhomogeneous blocks (fsdp_prefetch_inhomogeneous) by overlapping sub-layer all-gathers and reduce-scatters across block boundaries; double-buffered gathered weights across sub-layer scan trips to eliminate loop-carry copies in Qwen3.5; supported FSDP weight prefetching under expert parallelism with co-sharded axes; reused Qwen3-Next nested block scanning for Qwen3.5 (fsdp prefetch); and eliminated MoE weight relayout transposes during FSDP prefetching via 2D flat routed-expert weights (moe_flat_fsdp_weights).
• GPU Collective Overlap & cuDNN GMM Fusion: Introduced a custom GPU scheduling pass (gpu_collective_overlap_pass.py) to enhance XLA latency-hiding scheduling using feedback-directed optimization (FDO) profiles, and added Transformer Engine MoE cuDNN GMM fusion overrides.
• FSDP Weight Prefetching & Gathered Weight Caching: Added prefetch_fsdp_weights for scanned NNX decoder layers using a custom_vjp double-buffering pipeline that overlaps layer k+1 weight all-gathers with layer k computation in forward and layer k-1 all-gathers during backward recomputation (fsdp prefetch); supported symmetric jax.lax.scan(..., unroll=2); extended prefetching to honor layer-level rematerialization policies; and introduced save_fsdp_gathered_weights to retain gathered FSDP kernels in compute dtype across forward and backward passes.
• Kernel-Order Weight Gradients & LM Head Optimization: Introduced weight_grad_in_kernel_order in DenseGeneral, lm_head_weight_grad_in_kernel_order, and dense_weight_grad_in_kernel_order to route matmuls through jax.sharding.auto_axes under explicit sharding, allowing XLA to fold activation transposes into gradient matmuls and recover all-reduce to reduce-scatter fusion; contracted tied LM head logits directly over embedding feature axes (einsum("...e,ve->...v")) to eliminate redundant type casts and transposes; and added opt-in lm_head_fsdp_as_vocab_tp to shard untied LM heads along vocabulary dimensions using FSDP axes for small-batch efficiency.
• Explicit Sharding Expansion & Performance Guide: Onboarded GPT-3, GPT-OSS, Llama4, and Gemma 4 decoders to explicit sharding with loss parity against auto sharding (bit-for-bit at step 0 and under tensor parallelism, within bf16 reassociation tolerance thereafter), and published a comprehensive performance guide detailing HLO differences, step-time comparisons, and explicit sharding tuning flags (optimization, shard mode performance).
• TransformerEngine MoE Sharding & Quantization: Added TransformerEngine MoE checkpoint policies, exposed optional quantization before the TE MoE FSDP weight all-gather (te_moe_quantize_before_fsdp_all_gather), passed explicit mesh resources, and enabled configurable dispatch rematerialization (te_moe_dispatch_remat_location).
• MoE Logical Sharding & 1D Expert Parallelism: Expressed joint expert and FSDP placement with existing logical axes via the expert-on-fsdp-ep custom mesh and rules, and updated sharding configurations to support 1D expert parallelism.
• Lineage DeepSeek-V3 Quantization & MLPerf Performance: Introduced fp8_full_splash_qk_e3_combine_per_hop Lineage quantization combining FP8 Splash Attention Q/K with per-hop FP8 combine reduce-scatter; fused per-hop token absmax into SparseCore FP8 combine reduce-scatter; exposed FP8 Q/K Splash and Splash-skip flags; enabled async TensorCore ragged all-to-all in Lineage DeepSeek-V3; supported gradient accumulation (gradient_accumulation_steps > 1) with explicit ragged_buffer_factor overrides; enabled warm_eval_input_reshard_before_run_start and eval_cache_prefill_in_background by default in the deepseek3-671b-lineage recipe and set use_stream_chunking: true so Lineage Splash Attention skips the no-op segment mask on stream-chunked c4_mlperf training rows; supported fractional eval sharding (0 < eval_per_device_batch_size < 1) using expert axes as context parallelism; capped mllog eval_samples at the 1024-sequence validation set size and reported fp8 for lowest_numerical_precision_in_attn when Splash Q/K FP8 is enabled; and removed lineage_expert_permutation: "lpt" from the default Lineage DeepSeek-V3 recipe.
• TensorCore MoE Ragged Sort & Backward Quantization: Enhanced TensorCore ragged sort with [start, start+cap) windowed routing, uint32 packed sorting, and unique_indices weight gathering, and fused MoE combine backward quantization prior to TC unsort-gather while skipping 3D DLHS zero-initialization.
• Quantization & Qwix Optimizations: Disabled Qwix weight-gradient matmul tiling by default via qwix_tile_weight_grad_matmuls, supported use_sliced_mla_proj with Qwix interception quantization by preserving dot-general interception across sliced projection kernels, and added a configuration flag to keep backward passes in BF16 for Q projections during quantized training.
• Pre-Training Initialization Warmup: Added warmup_programs_in_init to execute precompiled train and eval programs on synthetic batches during setup to warm JIT dispatch caches before timed execution, with options to skip dropless training warmup, and added log_step_diagnostics for step-level timing breakdown.
• TensorCore Ragged Sort & 3D Layout Pipeline for Ring-of-Experts MoE: Added a TensorCore ragged sort implementation for Ring-of-Experts MoE, applied MoE routing weights to activations alongside the ragged sort, maintained the sorted MoE token buffer in 3D layout through GMM v2 kernels, and enabled token all-gather and combine in the TensorCore kernels' native 3D layout.
• In-Place MoE Gradient Accumulation: Added moe_accumulate_wi_dlhs to accumulate wi_0 and wi_1 backward DLHS gradients in-place via gmm_v2 partial_sum with fused DLHS scale, and introduced moe_accumulate_chunk_wgrad to accumulate weight gradients across token chunks in-place via tgmm_v2 partial_sum.
• GMM v2 Native Transpose RHS: Added native transpose_rhs support to Tokamax GMM v2 block specs and tiled matmuls to consume expert weights in their existing layout and eliminate RHS layout copies.
• FP8 Quantization: Added logits_proj_bwd_quant_calibration_method allowing custom backward quantization calibration overrides for logits projections (quantization), and added the lineage_quantization (none or fp8_full) option to support quantization for Lineage DeepSeek-V3 routed experts.
• MLPerf DeepSeek-V3 Optimization: Added the fsdp-as-dp-for-attn-rf custom mesh and sharding rule to shard MoE router weights on FSDP and overlap gradient communication with GMM backward execution (custom mesh and rule), and added warm_eval_input_reshard_before_run_start and eval_cache_prefill_in_background flags to overlap eval input caching with initial training steps.

• Dynamic Safetensors Loading & Multi-Host Sharding: Fixed and optimized source_checkpoint_layout=safetensors_dynamic loading across gpt-oss, Qwen3.5, Qwen3-Next, and Gemma; shared the single-axis stacking rule (tensor_handling.single_stack_axis) between loaders and converters; sharded reads along the leading axis on multi-host runs for contiguous reads; and enabled dynamic safetensors loading across standalone entry points outside training restore.
• Elastic Training Robustness & Batch Size Modes: Added elastic_global_batch_size_mode ("scaled" or "constant_padded") to keep global batch sizes invariant across elastic slice resizes via zero-segmentation row padding; recovered elastic training after losing the slice holding the default device by setting process-wide jax_default_device and re-raising slice-down errors as retryable JaxRuntimeErrors; and switched elastic backend detection to generic feature-based attribute checks.
• Colocated ArrayRecord Shard Caching: Sharded ArrayRecord record counting across colocated hosts on the iterator's CPU mesh, caching pattern listings and shard counts in colocated sidecars to eliminate redundant GCS reads during startup and elastic resizes.
• ML Goodput Monitoring & Logging Tools: Used get_goodput_job_name in setup_checkpoint_logger to honor custom Goodput job name overrides when initializing Orbax CloudLogger, and added the maxtext-goodput agent skill for automated Goodput monitoring and dashboard analysis.
• Pure-Dict NNX Checkpoints & M3 Sharding: Added support for pure-dict NNX checkpoints without intermediate 'value' wrapping dictionaries, and automated 1D NamedSharding across fsdp/fsdp_transpose for unpartitioned M3 variables in model_creation_utils.
• Elastic Training Sharding & In-Memory Snapshots: Enabled in-memory elastic snapshotting alongside colocated Python data input, sharded colocated elastic dataset iterators across hosts of the current mesh, and restricted incomplete elastic checkpoint cleanup strictly to GCS destinations when checkpointing is enabled to avoid removing valid checkpoints on local and persistent filesystems.
• Orbax V1 Goodput Logging Integration: Required orbax-checkpoint>=0.12.8 to leverage public logger support in orbax.checkpoint.v1.training.Checkpointer and automated RestoreStepStatistics emission during load_checkpointables(), replacing private Orbax internal logger attachments.
• In-Memory Snapshots for Elastic Training: Wired up elastic_backup_kind: snapshot (the default) so elastic runs keep a copy of the train state and data-iterator position in pinned host memory every elastic_snapshot_interval steps (new, default 10) and restore from it when it is at least as recent as the latest checkpoint. Snapshot mode requires dataset_type=grain with grain_use_elastic_iterator=True (or synthetic data) and is rejected at config time with DiLoCo, enable_rampup_batch_size, or colocated_python_data_input; such runs must set elastic_backup_kind: checkpoint.
• NNX & Flat MoE Checkpoint Resharding: Updated reshard_checkpoint.py to build the abstract train-mode NNX model for the target mesh (restoring every weight training saves, including MoEBiasVar), support flat MoE layouts (moe_flat_fsdp_weights=True), read each stored chunk once instead of once per device to cut host memory on large meshes, and preserve each weight's stored dtype. convert_deepseek_family_ckpt.py gained --weight_dtype keep to retain source dtypes (e.g., DeepSeek-V3's float32 e_score_correction_bias).
• Goodput Monitoring for RL & Distillation: Instrumented the RL and Distillation trainers with ML Goodput recording and monitoring hooks (job start/end, TPU init, training preparation, and step start events via GoodputTrainingHooks), honoring the existing enable_goodput_recording, monitor_goodput, goodput_job_name, and enable_pathways_goodput flags, which are now also listed in post_train/rl.yml.
• Multi-Tier Checkpointing Restore Documentation: Documented procedures for restoring Multi-Tier Checkpointing jobs from specific backup checkpoints on Google Cloud Storage (multi tier checkpointing).
• Goodput Monitoring & Logging: Integrated the Orbax v1 logger into Goodput monitoring, and added sliding window tracking for goodput metrics to accurately evaluate sustained throughput and performance trends over time.
• Pathways Checkpointing: Added opt-in colocated-Python checkpoint reads for Pathways single-controller runs across all parameter restoration paths (including metadata read) and introduced the train_lib_tpu build target.
• Elastic Training: Unified all elastic training modes to use elastic_retry as the single entry point.

• Treescope Guide & Inspection Notebook: Added documentation and an interactive tutorial notebook for inspecting attention masks on packed sequences with and without segment IDs using Treescope (monitoring and debugging, debug with treescope).
• M3 Scan Support & Dependency Updates: Added M3 scan support and bumped the required Flax version (update dependencies).
• Cluster Toolkit Benchmark Migration: Migrated the MaxText benchmark runner and configurations from XPK to Cluster Toolkit (gcluster), adding ClusterConfig in benchmarks/ctk_configs.py, gcluster job submit orchestration with kubectl completion polling in benchmarks/maxtext_ctk_runner.py, and retaining backward-compatible XPK re-export shims; benchmark_runner.py gained ctk/gcluster subcommands with --location, --compute_type, --topology, and --mount flags.
• Dependency Upgrades: Upgraded tokamax to 0.0.16 for TPU environments (the GPU extra stays pinned at 0.0.14) (install maxtext), raised the orbax-checkpoint floor to >=0.12.8, and raised google-cloud-mldiagnostics to >=1.0.9.
• Decision Models & Scoring Harness: Added vllm_score for JEV-style decision-making LLMs, enabling next-token probability distribution extraction, candidate continuation scoring, and batch evaluation (tutorials, decision models, inference).
• Configuration Support: Updated pyconfig to accept configuration files with .yaml extensions alongside .yml.
• Documentation & Guides: Updated ReadTheDocs links across documentation and tutorial notebooks to point to default versioned builds via redirects.
• Dependencies & Compatibility: Updated pre-training dependency requirements to require flax>=0.12.10 to ensure compatibility with JAX 0.11.2 (update dependencies).

• Aliased pltpu.reciprocal to pl.reciprocal for vLLM TPU inference on JAX 0.11.0 so non-float32 query kernels do not fail with AttributeError.
• Fixed JetStream chat-template tokenization and encoding compatibility for HuggingFace transformers>=5.0.0.
• Fixed colocated Python ArrayRecord shard count exchange on Pathways by reading host count directly from reported shard count rows instead of inferring from mesh device process indices.
• Limited GCS glob listings in gcs_glob_pattern to the pattern's literal prefix, preventing multi-hour whole-bucket scans during GCS glob evaluations and elastic checkpoint recovery.
• Fixed PyPI and OS package vulnerabilities in MaxText Docker container images.
• Hardened reinforcement learning exact-format regex to eliminate catastrophic backtracking on adversarial or malformed generation sequences.
• Fixed vLLM KV-sharing cache writebacks where untouched sharing aliases overwrote donor caches on models with layer KV sharing (e.g., Gemma 4).
• Kept FSDP prefetch reduce-scatter minor-dimension padding unreduced to prevent ShardingTypeError under newer JAX.
• Fixed multi-host reshard_train_state with scanned NNX PRNG keys (scan_layers=True) under ZeRO-1 by moving typed keys as raw data to avoid multi-controller ValueError exceptions.
• Fixed Qwen3-Next shared expert gate under explicit sharding with tensor parallelism to resolve ambiguous sharding conflicts and ShardingTypeError.
• Masked left-padding tokens before conv1d in Qwen3NextGatedDeltaNet to prevent causal convolution bleed into real tokens, and guarded unsharded recurrent state casts against NoneType crashes.
• Fixed vLLM memory leak in from_pretrained by promptly releasing the restored NNX checkpoint tree during weight alignment, preventing redundant copies of MoE weights from persisting in device HBM.
• Fixed vLLM fused MoE weights under expert parallelism by packing gate and up projections into a [gate | up] layout to prevent corrupted decodes during EP rollout steps.
• Honored remat_policy=none across NNX decoder loops and scannable blocks (DeepSeek4, Gemma4, Qwen3-Next, Qwen3.5), resolving full-rematerialization fallbacks when remat_policy is "none" or None (fsdp prefetch).
• Fixed two FP8 MoE training failures under JAX HiJAX unreduced lowering in quantizations.py and megablox/ops.py.
• Added fallback to pkgutil.get_data in MaxTextConfig._load_mesh_config_from_yaml to resolve custom mesh and rule YAML files when running MaxText from zipped archives.
• Read logical axis names directly in from_pretrained checkpoint alignment via variable out_sharding metadata to prevent ambient logical rules from causing incorrect weight repetition or alignment failures during vLLM and Tunix loading.
• Aligned Engram n-gram hashing and causal short convolution (ShortConv) with DUAL_CHUNK_SWAP Context Parallel load balancing, skipping halo exchanges for width-1 kernels and preserving hash tables as NumPy arrays under JIT.
• Fixed driver host OOM during Zero-1 train state resharding under Pathways single controller (enable_single_controller=true) by replacing host-staged jax.device_put with an on-device jitted identity reshard (sharding.reshard_train_state).
• Fixed out-of-memory errors during RL training on models with nested-scanned linear-attention layers by adding local_layers to logical_axis_rules in post_train/rl.yml, and added scanned-layer fallback support to direct-sync coverage validation for Tunix and vLLM rollouts.
• Fixed DeepSeek-V4 router projection precision by enabling float32_gate_logits: true in deepseek4-284b.yml and deepseek4-tiny.yml to match official fp32 gate projections.
• Fixed silent token dropping on the ragged all-to-all MoE path by wiring up in-layer dropless fallback (moe_dropless_fallback=layer) to honor worst-case receive buffers.
• Fixed GPU while-loop double-buffering when gradient accumulation is disabled (gradient_accumulation_steps=1) by applying the model loop unroll context to single-step training.
• Fixed lost training progress during elastic scale-ups by taking a snapshot immediately before the attempt is interrupted.
• Fixed colocated data iterator checkpoint persistence under Pathways by storing iterator state directly within the Orbax-managed checkpoint directory.
• Fixed Multi-Token Prediction (MTP) layer sharding under pipeline parallelism by applying logical_axis_rules_pp_act_as_dp, summed per-expert routing counts across microbatches for loss-free MoE bias updates, and honored user-configured prng_impl.
• Fixed scanned Qwen3-Next and Qwen3.5 decoder blocks raising NotImplementedError during vLLM decoding by properly routing per-layer attention metadata dictionaries to sub-layers during static unrolling.
• Fixed DeepSeek SFT loss masking zeroing out completion EOS tokens when pad_token_id == eos_token_id by remapping padding to <｜▁pad▁｜> during data preprocessing.
• Fixed gmm_v2 reading uninitialized VMEM in RHS scale rows when tile_k > size_k on partial K tiles, preventing NaN accumulator poisoning in quantized and unquantized MoE execution.
• Fixed MoE execution on CPU backends by automatically bypassing TPU-specific Megablox Pallas kernels and falling back to jax.lax.ragged_dot.
• Fixed DeepSeek-V4 decode drifting from train-mode forward passes after padded prefill by deriving prefill true length from segment IDs in compressed attention, properly sizing compressed AR cache, and correcting sliding-window decode masks measured from token positions across DeepSeek-V4, Gemma2, and GPT-OSS.
• Fixed DeepSeek router numerics by applying SwiGLU input clamping (mlp_activations_limit) to shared experts, retaining float32 precision for routed gate bias (bias_dtype) and source non-bf16 tensors in deepseek_dequantize.py, and zero-initializing dropped token rows in combine ragged all-to-all when ragged_buffer_factor > 0.
• Fixed XLA arithmetic fusion breakage and numerical drift in Gemma4 under ShardMode.AUTO by strictly restricting physical out_sharding and explicit sharding constraints to ShardMode.EXPLICIT.
• Fixed native_fp8_dot_general quantization by using per-token activation scaling in per_tensor mode and properly skipping quantization for modules specified in unquantized_modules.
• Fixed safetensors_dynamic checkpoint loading keeping single-key weights (embedding, final norm, LM head) in the source checkpoint dtype instead of casting them to the target leaf's weight_dtype, and fixed -Instruct model names being stripped before the HF_IDS lookup (which silently loaded the base model's weights); exact model names are now required.
• Fixed TPU SparseCore collective offload stalls during gradient accumulation by applying the _xla_loop_unroll_strategy="double-buffer" attribute only on GPU hardware, preserving TensorCore compute_on collective placement on TPU.
• Fixed duplicated evaluation sequences when multiple eval hosts share a single C4 evaluation file in tfds_data_processing_c4_mlperf.py.
• Fixed Profiler.deactivate() raising RuntimeError: No profile started on early StopTraining teardown when profiling windows have already closed.
• Fixed distillation training by aligning _train_step signatures, slicing steps properly, and propagating distill_beta to the teacher model.
• Fixed to_maxtext conversion of scanned DeepSeek checkpoints (MoE gate-bias stacking and MoEBiasVar/Tid2EidVar key detection), added DeepSeek MTP-layer parameter mappings, and made Grain iterator worker cleanup robust in multihost_dataloading.py; migrated the DeepSeek end-to-end tests to to_maxtext.
• Fixed scheduled execution failures in the dpo_qwen3_demo_no_eval.ipynb example notebook.
• Fixed loss calculation zeroing the language model loss when indexer loss is disabled or unscaled (indexer_loss_scaling_factor == 0).
• Fixed DeepSeek-V4 logit checker forward pass execution under mesh JIT, token space decode assertion, and Linen side collection restoration (Tid2EidVar).
• Fixed post-training hooks swallowing non-StopIteration data loading exceptions, ensuring unexpected loader errors surface immediately (#5523).
• Fixed scanned decoders with inhomogeneous layer cycles silently dropping remainder layers or crashing when num_decoder_layers is not an exact multiple of inhomogeneous_layer_cycle_interval, raising an explicit ValueError.
• Fixed gmm_v2 partial K tail handling when tile_k does not divide size_k by masking stale VMEM in the LHS operand, preventing NaN/Inf poisoning of the accumulator and block absmax quantization.
• Fixed MLPerf parallelism, precision, and config filename MLLog disclosure logging for Lineage DeepSeek-V3 runs.
• Fixed SFT evaluation invoking the loss function with is_train=True, ensuring dropout, MTP loss, and sparsity batch_stats updates are disabled during eval.

• Added automatic fallback resolution to src/maxtext/configs/models/deprecated/ with a DeprecationWarning when loading deprecated model configuration files.
• Deprecated XPK launcher paths and scripts (run_distill_xpk.sh, xpk_olmo3_7b_stage1.sh, and maxtext_xpk_runner.py) in favor of Cluster Toolkit wrappers (run maxtext via xpk).
• Removed the pallas_tpu_ops dependency from MaxText.
• Migrated Flash Attention callers from deprecated Pallas Triton kernels in Tokamax to the pure Triton kernel implementation.
• Migrated SparseCore Pallas ragged gather-reduce kernels from plsc.bitcast to lax.bitcast_convert_type ahead of plsc.bitcast deprecation.
• Removed obsolete force_on_policy_ratio and log_sampler_trainer_agreement configuration flags from RLSettings following agreement fusion unification in Tunix.
• Removed experimental Lineage hooks from decoder, MTP, and MoE layers, and removed the lineage_integration configuration flag.
• Deprecated calling from_config() without explicit rngs and deprecated models.transformer_as_linen in favor of native Flax NNX module initialization.

• DeepSeek-V3 Lineage Integration: Integrated Lineage DeepSeek-V3 (dsv3.py) into MaxText, storing routed expert output weights in Lineage's layout and sowing router bias updates along the use_lineage path.
• DeepSeek-V4 & Multi-Token Prediction: Supported DeepSeek-V4 checkpoint conversion and parameter mapping with 2D transposed shape alignment, added DeepSeek-V4 indexer loss, and implemented separate final normalization for Multi-Token Prediction (MTP).
• Weaver Architecture: Added Weaver Mixture-of-Transformers block and standardized Weaver model naming across configs.
• Qwen3.5 & Qwen3-Next: Added qwen3.5-35b-a3b-fp8 and qwen3.5-397b-a17b-fp8 model configs with HuggingFace parameter mapping and weight conversion, and optimized Qwen3-Next shared expert construction to only instantiate when shared_experts > 0.
• Cosmos 3 Bring-Up: Added model definitions and bring-up support for Cosmos 3 Nano and Cosmos 3 Super Reasoner.
• Normalization Enhancements: Added RMSNorm normalization layer support for Apple Envy decoder blocks and OLMo 3 MLP blocks (DecoderBlockType.OLMO3).
• GPT-OSS Attention Sink: Added attention-sink support for GPT-OSS models.
• M3 Core Architecture: Added m3 package scaffolding with architecture rules and layout verification tests, introducing m3 core RoPE implementation, mesh creation, and sharding rules.

• JAX Flash Attention Optimization: Enhanced JAX Flash Attention kernel with a specialized high-performance execution path for long sequences (splitting at a 4,096 sequence length threshold).
• DeepSeek-V4 CSA & Splash Attention: Added fused Pallas TPU StreamIndex score kernel, QK attention head chunking for Compressed Self-Attention (CSA) memory reduction, and HCA static compilation for Splash Attention on DeepSeek-V4.
• Context Parallelism (CP): Sharded GatedDeltaNet sequences under ici_context_parallelism, enabled native block-sparse execution for indexer masks under CP, removed all-to-all communication overhead from TSP/CP segmentation masks, and split embedding logical names in attention modules.
• Multimodal RoPE (MRoPE): Eliminated redundant JAX recompilations in get_mrope_input_positions using inline NumPy operations, and normalized MRoPE position shapes for vLLM across hybrid cache utilities and embeddings.
• DiLoCo Distributed Training: Added comprehensive documentation and tutorials for DiLoCo pre-training (core concepts, diloco, tutorials), and migrated DiLoCo launch scripts and tutorials from XPK to Cluster Toolkit (diloco pretraining).
• Quantized Multi-Token Prediction: Enabled quantized Multi-Token Prediction (MTP) for accelerated pre-training.

• Reinforcement Learning (GRPO) Infrastructure: Plumbed GRPO flags to AgenticGRPOLearner, added RL sequence packing and vLLM block-size configuration knobs, added free_kv_cache_during_weight_sync, reward_num_workers, and gc_collect_after_weight_sync settings, made extract_answer linear in response length, and added a Gemma 4 26B GRPO tutorial (post training index, rl gemma4 26b, rl gemma4 e4b).
• LoRA & Distillation Workflows: Refactored LoRA resharding to modernize NNX variable access and graph traversal, migrated full fine-tuning scripts from TFDS to Grain (full finetuning), added launch scripts and configs for Distillation Training, and integrated MaxText's trainer backend into a distributed GSM8K example.
• vLLM Rollout Enhancements: Added Mamba prefix caching for Qwen3 GDN layers in vLLM, indexed KV caches via layer_name_to_kvcache_index in the vLLM adapter, and implemented MoE 128-lane chunking and padding for GMM layout.

• Multimodal SFT Pipeline: Added Grain-based multimodal Supervised Fine-Tuning (SFT) data processing pipeline (multimodal), end-to-end video SFT supporting variable frame sizes and durations, and ChartNet SFT configuration with string response handling in multimodal input pipelines.
• Omni Pipeline & Qwen3-VL Inference: Added custom multimodal processor and end-to-end pipeline for Omni multimodal models, and enabled vllm_decode inference support for Qwen3-VL (inference).

• Explicit Sharding & ZeRO-1 Expansion: Onboarded explicit sharding and ZeRO-1 gradient accumulation across Qwen3.5, Qwen3-Next, Qwen3, Qwen2, Kimi-K2, Mistral, Mixtral, and Gemma model families; supported partitioned optimizer state sharding propagation and MaskedNode filtering in ZeRO-1; deferred data-parallel gradient all-reduce to update() under GA with backward cotangents scaled to GA=1.
• Sharded Muon Optimizer: Integrated a sharding-aware, distributed Muon optimizer into MaxText with generalized weight dimension extraction, support for Qwen3 and GPT-OSS MoE blocks, and a configuration flag to govern applying Muon to MoE routers.
• FP8 Quantization & Dequantization: Implemented an FP8 weight-only dynamic dequantization engine with direct FP8 and scale tensor ingestion in to_maxtext, native FP8 inference (serve_fp8_weight) for DenseGeneral and GMM v2, weight-gradient-arm cotangent calibration (drhs_grad_quantization_calibration_method), quantized logits projection (quantization), pre-quantized weights for TPU inference fused MoE kernels, and options to keep MoE router projections unquantized for stability (moe configuration, quantization).
• Quantized Ring-of-Experts Pipeline: Enabled token-quantized all-gather pipelines with Ring of Experts and added quantized all-gather for the Ring of Experts combine backward (moe configuration, quantization).
• MoE Dropless Fallback & Buffer Resilience: Implemented configurable MoE dropless fallback (moe_dropless_fallback: None|step|layer) with memory-neutral step fallback, first-phase ragged buffers, host-driven step replay retry mechanisms for buffer overflows, evaluation-specific buffer scaling (eval_ragged_buffer_factor), and evaluation tile sizes and QKV layouts (moe configuration).
• MoE Routing & Kernel Optimizations: Added dense custom VJPs for router top-k selection (eliminating backward scatters), optimized top-2 expert group score computation, aligned sigmoid-router auxiliary losses and expert bias updates with Megatron-LM, added collective psum reductions for expert bias load balancing, and supported extracting and replaying router expert decisions between inference and training (moe configuration).
• MoE Kernels & Sharding: Integrated TransformerEngine (TE) MoEBlock, supported Tokamax GMM v2 heuristic tiling functions, forwarded TPU inference MoE kernel knobs, saved routed expert inputs across remat boundaries via moe_x_sorted, added nested scanning over hybrid attention with per-layer remat in Qwen3-Next, and added explicit tracking for shard_embed_moe_on_fsdp (moe configuration).
• SparseCore & Pallas Accelerations: Added moe_pin_sparse_core_all_gathers to offload FSDP and Expert Parallel all-gathers to SparseCores, lowered plsc.bitcast to tpu.bitcast when layout passes are enabled, and refactored the MaxText mHC kernel to utilize emit_pipeline.

• Orbax v1 Checkpointing: Refactored MaxText's checkpointing architecture to natively utilize the Orbax v1 API (retiring legacy v0 implementations), added dequantize-on-load parameter restoration, and added support for restoring checkpoint weights stored outside the standard "params" collection.
• Goodput & Pathways Persistence: Added goodput_job_name configuration option for Goodput monitoring, registered Pathways persistence array handlers with bounded host staging, and introduced standalone checkpointer benchmarking loop features and configuration flags.
• Checkpointing Documentation: Migrated multi-tier and emergency checkpointing guides from XPK to Cluster Toolkit (emergency checkpointing, multi tier checkpointing).

• MLPerf Training Compliance & Optimizations: Added full MLPerf training logging compliance, accelerated step-0 initialization time via per-file c4_mlperf train sharding and data-free setup, added host-local eval batch caching, implemented continuous stream chunking for C4 MLPerf datasets, and added the dsv3-mlperf-4k custom mesh and sharding rule.
• Data Pipeline Enhancements: Added Megatron MMap data format support for Grain data loaders with an integrated mmap index builder (data input pipeline, data input grain, data input megatron mmap), allowed Grain to ingest TFDS configurations, optimized data input pipelines for evaluation, and set default dataset_type to synthetic.
• Documentation & Cluster Toolkit Migration: Migrated core navigation, getting started, and execution guides from XPK to Cluster Toolkit (getting started, install maxtext, architecture overview), updated container image registry references from GCR to Artifact Registry (run maxtext elastic training, run maxtext via pathways, lora on multi host), and published TPU pre-training and post-training Docker images for MaxText in the build tutorial (build maxtext).
• Dependencies & Developer Tooling: Upgraded JAX dependencies to 0.11.1 (update dependencies), Tokamax to 0.0.14, and DrJAX to 0.2.1; added custom GCS path support for ML Diagnostics telemetry; added Gemini CLI development skills; streamlined multi-worker XLA dumping (--dump_xla_all_hosts); and defined default rule constants in types.py with base configuration parity testing.

• Fixed FP8 and quantized MoE numerical paths across dense and sparse matmul implementations.
• Fixed MoE weight tracking, vLLM HBM memory allocation, and restored wo in Tunix's zero-pad key set during RL rollouts.
• Fixed RLConfig mesh initialization AttributeError and resolved defects on the Tunix adapter and engine packing path.
• Fixed vLLM sampler initialization failures and corrected cache_heads and paged_kv_heads logical axis order in vllm.yml.
• Unblocked Qwen3.5 hybrid decode through the vLLM adapter and fixed the standalone torchax converter for TPU inference rollouts.
• Fixed standalone checkpointer restore loop to use Orbax v1 and ensured fatal checkpointing errors raise RuntimeError.
• Resolved Pyrefly typing errors related to JAX scalar types across trainers and Megablox kernels.
• Derived cp_size dynamically from active logical axis rules and computed local block padding for dynamic Splash indexer masks under Context Parallelism.
• Fixed GDN Mamba block table indexing under prefix caching and resolved argument mismatches in compact Mamba block override calls.
• Fixed dropped routing-weight gradients in ring_ragged_unsort and aligned dropless-replay remat overrides with the NNX decoder.
• Configured layout passes on SparseCore Pallas ragged gather and ragged gather-reduce kernels.
• Eliminated XLA backward adjoint interior padding in reorder_sequence.
• Fixed duplicate XLA compilation of jit(train_step).
• Fixed vocab tiling hidden-state cotangent scaling.
• Fixed FP8 custom-gradient state overwrite during NNX training.
• Fixed NNX throughput regression when running with scan_layers=False.
• Fixed XLA fusion breakage and numerical drift in Qwen3 under AUTO sharding mode.
• Fixed edge cases in AttentionOp.generate_attention_mask for DeepSeek-V4 compressed attention across chunked prefill, packed sequences, and standalone environments.
• Fixed Qwen3-VL tokenizer type mismatch and resolved Grain IterDataset crashes on multimodal data.
• Ensured weight_dtype is propagated correctly in Qwen3-Next sparse MoE shared expert gate.
• Fixed EOS token detection and model routing for Qwen vision models in multimodal evaluation.
• Fixed GCS path resolution by skipping redundant exists() checks on Google Cloud Storage paths.
• Allowed warm starts from load_parameters_path without requiring checkpointing enabled.
• Fixed MoE Data Parallel reduce-scatter and logical axis sharding in the vLLM serving path.
• Fixed Ahead-Of-Time (AOT) compilation for ragged kernels and supported internal compilation and orchestration workflows.
• Fixed C4 MLPerf evaluation input pipeline handling for fractional evaluation batch sizes.
• Resolved Grain dataloader starvation and process leak on slice recovery.
• Fixed chat template formatting and SFT prompt/completion masking for Gemma 4 reasoning.
• Fixed ImportError on optional tensorflow_text in post-training workflows.
• Explicitly set enable_checkpointing when loading state in DiLoCo tests.
• Fixed AttributeError in MHC-lite during shape evaluation and sharding construction.
• Fixed inference query head rule mapping and aligned QKV batch dimensions for cuDNN flash attention.

• Flax Linen Removal: Removed legacy Flax Linen module code, dead code, and shared-primitive imports, deprecating and removing pure_nnx, enable_nnx, and pure_nnx_decoder configuration flags, Linen decoder/attention layers, and *_as_linen model wrappers across trainers, inference dispatch, quantization, and sharding (distillation, knowledge distillation, lora).
• Replaced deprecated google-cloud-sdk package with google-cloud-cli in Dockerfiles (run maxtext via xpk, rl on multi host, sft on multi host).
• Made TensorFlow an optional dependency in requirements (update dependencies, data input tfds, install maxtext).

• Flax NNX Migration: Enabled pure_nnx, enable_nnx, and pure_nnx_decoder configurations by default (PR #3526), migrating MaxText primarily on Flax NNX (PR #2885).

• Dependency Upgrades: Upgraded JAX to version 0.10.2 for pre-training and 0.11.0 for post-training.

• Model Support & Architecture:

  • DeepSeek-V4: Full model integration, decoders, and configuration stack (PR #4153), added HyperHead, aligned Sinkhorn implementation (PR #4337), and added checkpoint conversion support (PR #4336). See the user guide for more details.
  • Qwen3-VL: Added support for Qwen3-VL models (PR #4293, PR #4517) and Qwen3-VL-4B (PR #4263).
  • Apple Envy MoE: Added model configurations and support for Apple Envy Switch architectures.
  • Chunked MoE: Added chunked MoE support via num_moe_token_chunks to reduce memory footprint (PR #4499).
  • Block Diffusion: Added block-diffusion pre-training support (PR #4776), model-independent block corruption utilities (PR #4737), and causal-block attention across Dense, Splash, and Tokamax kernels (PR #4743).

• LoRA & QLoRA: Added native LoRA and QLoRA support for Gemma4, Gemma3, Qwen3, and Llama3, along with interactive tutorials (PR #3969, PR #4265, PR #4068, PR #3968, PR #3970, PR #4417).

• Context Parallelism (CP), Ring Attention:

  • Added Ulysses and USP CP strategy and packing (PR #4687, PR #4825, PR #4836), Tokamax load-balanced Ring Attention (PR #4266, PR #4537, PR #4622), and sequence packing for USP and All-Gather CP (PR #4230, PR #4887).
  • DeepSeek MoE & MLA: Added Ring Attention with DSA Sparse Indexer PR #4767, auxiliary loss-free and sequence-wise load balancing PR #4753, MLA QK head chunking PR #4564, optimized generate_mask PR #4437, and Approximate Top-K PR #4243.
  • Positional Embeddings: Added YaRN RoPE config PR #4238, standardized MRoPE to BS3 convention for multimodal training PR #4709, and fixed Qwen3.5 partial rotary factor handling.
  • Kernels & Megacore: Added configurable attention_for_vit kernels PR #4232 and enabled Megacore for Splash Attention dkv backward PR #4755.

• Quantization & Performance: Added FP4 [E2M1] (PR #4495) and experimental attention quantization (PR #4487); enabled TE Collective GEMMs (PR #4470) and overlap (PR #4307), MoE comms with collective matmul (PR #4295), Tokamax GMM v2 (MoE configuration guide), and double-buffered inner scans during gradient accumulation (PR #4316).

• Checkpointing: Added support for Multi-tier checkpointing in Pathways.

• Goodput & Elasticity:

  • Added Goodput support for Pathways Elasticity & Slice Efficiency, including record_slice_state() to query live slice counts (PR #4840).
  • Implemented checkpoint-based elasticity using set-based slice tracking (PR #4245).

• Post Training:

  • Added reward_functions_path and reward_functions CLI knobs for custom rewards (PR #4149) to RL training.
  • Updated tutorials with AgenticGRPOLearner for async RL training (PR #4181) and added GRPO Gemma4-e4b tutorial (PR #4427).
  • Added RL support for Qwen3 30B and GPT-OSS 20B. See the Qwen3 30B RL tutorial and GPT-OSS 20B RL tutorial for recipes.
  • Added support for DPO along with tutorials (PR #4362).

• Usability & Infrastructure:

  • Added wandb logging support (PR #3053).
  • Added Hugging Face Grain streaming integration and onboarding guide (PR #4486).
  • Added Simple-evals runner support for gpt-oss model family (PR #4644).
  • Added scripts to run vanilla DiLoCo on MaxText (PR #4095).
  • Added option to enable on-demand profiling server in ML Diagnostics (PR #4131).

• Post-Training:

  • Resolved Gemma 3/4 RL rollout gibberish issue by unrolling scanned weights for vLLM adapter (PR #4536, PR #4519, PR #4404).
  • Fixed RL LR schedule defaults (PR #4225), added drop_remainder=True to prevent shape mismatches on tail batches during GRPO training (PR #4252) and resolved Qwen3.5 MRoPE/Kv-cache rollout issues (PR #4177).

• Compilation:

  • Fixed double-compilation in train_step by matching input sharding (PR #4174).
  • Truncated out_sharding on extra pspec dimensions (PR #4769) and restricted GMM quantization to fp8_full (PR #4842).

• Model-Specific Fixes:

  • Qwen3.5: Applied partial MRoPE for Qwen3.5 (PR #4764).
  • Mixtral: Fixed EP throughput via configurable expert-axis batch sharding (PR #4179).

• NNX, MoE & MTP:

  • Resolved silent zero-loss (PR #4525) and targets_segmentation bugs (PR #4756) in Multi-Token Prediction (MTP).
  • Preserved scanned layer intermediates for MoE load-balancing loss in NNX (PR #4829).
  • Relanded Qwix quantization on NNX (PR #4198) and fixed Qwix LoRA mesh sharding (PR #4866).

• Tensor Transpose Parallelism Removed: Completely removed the tensor_transpose physical mesh axis and deleted ici_tensor_transpose_parallelism and dcn_tensor_transpose_parallelism configuration options.
• Flax Linen Deprecation Warning: Flax Linen is now deprecated in favor of Flax NNX; running with pure_nnx=False or enable_nnx=False will issue a deprecation warning.

• Upgraded JAX to version 0.10.0 for pre-training and 0.10.1 for post-training.
• New vLLM-Powered Evaluation Framework: Introduced an eval framework for running lm-eval, evalchemy, and custom benchmarking against MaxText checkpoints. See the evaluation guide for details.
• Added support for pre-training new models:
  • Qwen3.5: Qwen3.5 35B & 397B is now supported.
  • Qwen3-Omni: Support for multimodal SFT (PR #3863).
• Direct Preference Optimization (DPO/ORPO) Support: Full support for DPO and ORPO alignment pipelines. See the DPO tutorial for details.
• Reinforcement Learning (RL) Recipe: Added a pre-configured RL recipe for Qwen3-30b-a3b.
• Iterative Quality Monitoring (RL): Added intermediate evaluation hooks to automatically run quality benchmarks during RL training (every eval_interval steps), optimized with a new eval_batch_size configuration knob.
• Developer Extensibility: Added dataset_processor_path CLI knob for custom dataset integration, and refactored shared post-training hooks to simplify custom SFT, DPO, and RL workflow development.
• Generalized Learn-to-Init (LTI) for Distillation: Enhanced post-training distillation capabilities with generalized LTI support.
• Added support for recording elastic goodput events during training to track efficiency (PR #3901).
• Installation Updates: Updated the [tpu-post-train] installation command to require UV_TORCH_BACKEND=cpu(see Installation Guide).
• Zero1 AOT Compilation: Added zero1 support to Ahead-Of-Time (AOT) compilation in train compile, improving compilation capabilities for zero1 config.
• MoE Performance Optimization: Integrated ragged gather reduce into Mixture of Experts (MoE) layers to optimize memory and performance by replacing ragged scatter and supporting backward pass.
• Added E2E scripts to run checkpoint conversion, pre-training and post-training (SFT, RL) with Gemma3-4B model.
• Bug Fixes and Usability Enhancements:
  • Attention Masking Fix in RL: Fixed an issue in TunixMaxTextAdapter where queries at non-pad positions could attend to pad-position keys during training, which was corrupting log-probabilities and affecting GRPO training reward trajectories (PR #4016).
  • JAX/NNX Gradient Mutation Fix: Refactored post-training loops (train_distill, train_sft, train_rl) to use jax.value_and_grad with explicit NNX state split/merge instead of nesting nnx.value_and_grad inside nnx.jit (PR #3652).
  • Qwen3-MoE Checkpoint Conversion: Fixed checkpoint conversion issues for Qwen3-MoE models (PR #3868).
  • Duplicate Configuration Failures Fix: Allowed identical config overrides and handled configuration exceptions cleanly (PR #3933).
• Documentation Improvements: Updated Getting started guide, including new guides for the evaluation framework and the DPO tutorial.

• Deleted legacy DPO implementation in favor of the integrated DPO trainer.
• Removed stack trace collection feature.

• Upgraded JAX to version 0.9.2, improving support for both pre-training and post-training.
• Introduced simplified APIs for accessing MaxText models.
• Included maxtext_with_gepa.ipynb, a new notebook demonstrating AIME prompt optimization using the GEPA framework within MaxText.
• Added support for Kimi-K2 models and the MuonClip optimizer. Users can explore this with the kimi-k2-1t config (see user guide for details).
• Kimi-K2-Thinking, Kimi-K2.5 (text), and Kimi-K2.6 (text) are now supported. See Run_Kimi.md for details.
• DeepSeek-V3.2 is now supported, including DeepSeek Sparse Attention for handling long contexts. Use the deepseek3.2-671b config to try it out (refer to the user guide for more information).
• Support has been added for Gemma 4 multi-modal models (26B MoE and 31B dense). These can be used with the gemma4-26b and gemma4-31b configs. See Run_Gemma4.md for further details.
• Support has been added for Gemma 4 inference using MaxText on vLLM plugin.
• Enhanced RL capabilities with support for the open-r1/OpenR1-Math-220k dataset and nvidia/OpenMathReasoning.
• Added more evaluation modes for RL like majority voting and pass@1 estimation.
• Sync weights to vllm prior to pre RL evaluation.
• More robust usage of math-verify in RL.
• MaxText's Supervised Fine-Tuning (SFT) now supports non-instruct models.
• Added support for tensor parallelism using the Fused MoE kernel for MaxText on vLLM inference.
• Added support for MaxText to vllm converters for Qwen3 and Gemma4 family of models.
• validate_converter.py now runs on multislice environment to test larger models with utilities to compare maxtext and vllm weights.

• Legacy MaxText.* shims have been removed. Please refer to src/MaxText/README.md for details on the new command locations and how to migrate.
• Sequence parallelism has been deprecated, please use context parallelism instead.
• The flag expert_shard_attention_option is deprecated, use custom_mesh_and_rule=ep-as-cp for the same functionality.

• Use the new maxtext[runner] installation option to build Docker images without cloning the repository. This can be used for scheduling jobs through Cluster Toolkit on GKE. See the MaxText installation instructions for more info.
• Config can now be inferred for most MaxText commands. If you choose not to provide a config, MaxText will now select an appropriate one.
• Configs in MaxText PyPI will now be picked up without storing them locally.
• New features from DeepSeek-AI are now supported: Conditional Memory via Scalable Lookup (Engram) and Manifold-Constrained Hyper-Connections (mHC). Try them out with our deepseek-custom starter config.
• MaxText now supports customizing your own mesh and logical rules. Two examples guiding how to use your own mesh and rules for sharding are provided in the custom_mesh_and_rule directory.

• New tpu-post-train target in PyPI. Please also use this installation option for running vllm_decode. See the MaxText installation instructions for more info.
• Qwen3-Next is now supported.
• New MaxText structure! MaxText has been restructured according to RESTRUCTURE.md. Please feel free to share your thoughts and feedback.
• Muon optimizer is now supported.
• DeepSeek V3.1 is now supported. Use existing configs for DeepSeek V3 671B and load in V3.1 checkpoint to use model.
• New RL and SFT Notebook tutorials are available.
• The ReadTheDocs documentation site has been reorganized.
• Multi-host support for GSPO and GRPO is now available via new RL tutorials.
• A new guide, What is Post Training in MaxText?, is now available.
• Ironwood TPU co-designed AI stack announced. Read the blog post on its co-design with MaxText.
• Optimized models tiering documentation has been refreshed.
• Added Versioning. Check out our first set of release notes!
• Post-Training (SFT, RL) via Tunix is now available.
• Vocabulary tiling (PR) is now supported in MaxText! Adjust config num_vocab_tiling to unlock more efficient memory usage.
• The GPT-OSS family of models (20B, 120B) is now supported.

• Many MaxText modules have changed locations. Core commands like train, decode, sft, etc. will still work as expected temporarily. Please update your commands to the latest file locations
• install_maxtext_github_deps installation script replaced with install_maxtext_tpu_github_deps
• tools/setup/setup_post_training_requirements.sh for post training dependency installation is deprecated in favor of pip installation

Our first MaxText PyPI package is here! MaxText is a high performance, highly scalable, open-source LLM library and reference implementation written in pure Python/JAX and targeting Google Cloud TPUs and GPUs for training. We are excited to make it easier than ever to get started.

Users can now install MaxText through pip, both for local development and through stable PyPI builds. Please see our MaxText Installation Guide for more setup details.

Going forward, this page will document notable changes as we release new versions of MaxText.