Skip to content

release_notes.md

Latest commit

 

History

History
206 lines (147 loc) · 22.2 KB

release_notes.md

File metadata and controls

206 lines (147 loc) · 22.2 KB
Raw

MaxText release notes

PyPI Package

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

Unreleased

Releases

v0.2.4

Changes

  • 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).

Bug Fixes

  • 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).

Deprecations

  • 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.

v0.2.3

Changes

  • 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.

Deprecations

v0.2.2

Changes

  • 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.

Deprecations

  • 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.

v0.2.1

Changes

  • 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.

v0.2.0

Changes

Deprecations

  • 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

v0.1.0

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.