All pull requests

Pull requests

• Fix MaxTextTrainingEngine optimizer host offload, logical_axis_rules ownership, and fp32 grad accumulation

  #5390·suhithr opened Sep 26, 2026Collaborator·

  ·
  3

• [DO NOT MERGE] Wire up optimizer_memory_host_offload in MaxTextTrainingEngine and fix AOT trace mesh context

  #5387·suhithr opened Sep 26, 2026Collaborator·

  ·

• Trim small pre-GMM routing and ragged sort operations in MoE RoE.

  #5385·copybara-service[bot] opened Sep 26, 2026Bot·

  ·
  1

• Weaver MoT Block

  #5384·parsley9877 opened Sep 25, 2026Collaborator·

  ·
  3

• c4_mlperf: read and cache each host's own eval batch once

  gemini-review
  #5382·ultrons opened Sep 25, 2026Collaborator·

  ·
  4

• MoE: dense custom VJP for the router top-k selection (no scatter in the backward)

  #5381·ultrons opened Sep 25, 2026Collaborator·

  ·
  2

• MoE: first-phase ragged buffer, dropless fallback and precompiled replay/eval programs

  #5380·ultrons opened Sep 25, 2026Collaborator·

  ·
  3

• [MaxText] Add spatial minor [k, n, g] TGMM kernel

  #5379·Shuwen-Fang opened Sep 25, 2026Collaborator·

  ·
  10

• [MLPerf] Cherry-pick #5293, #5207, #5355, #5360 from main

  #5378·igorts-git opened Sep 25, 2026Collaborator·

  ·
  2

• Restore custom NNX collections in model_creation_utils.from_pretrained

  gemini-review
  #5376·snehalv2002 opened Sep 25, 2026Collaborator·

  ·
  6

• Add TE Groupe GEMM as GPU GEMM for MoE

  #5374·shurale-nkn opened Sep 25, 2026Contributor·

  ·
  3

• Add unit test for NNXDecoder dropless remat_policy_override

  #5373·Shuwen-Fang opened Sep 25, 2026Collaborator·

  ·
  3

• Add fsdp-as-dp-for-attn-cp-as-ep-for-moe mesh rule and its eval companion

  #5372·NuojCheng opened Sep 25, 2026Collaborator·

  ·
  8

• [MaxTextEngine] Fix Qwen3.5-397B MLPerf 6.1 standalone trainer harness, MoE aux loss parity, and optimizer host offload (b/565501755)

  #5371·uwlei opened Sep 24, 2026Collaborator·

  ·
  19

• Adding DualPipe Style Schedule to Improve MoE Performance

  #5368·abhinavgoel95 opened Sep 24, 2026Contributor·

  ·
  3

• Add Qwen3.8-2.4T-A95B model configuration and parameter mappings

  #5366·JamesDeng42 opened Sep 24, 2026Collaborator·

  ·
  4

• fix(data): Fail the job on JAX runtime errors from the data iterator

  #5365·ecnal-cienet opened Sep 24, 2026Collaborator·

  ·
  1

• bespokelabs/Bespoke-Nimble-9B Lora PoC

  #5356·hengtaoguo opened Sep 23, 2026Collaborator·

  ·
  1

• Qwen3.5-9B base model

  #5354·hengtaoguo opened Sep 23, 2026Collaborator·

  ·
  11

• feat(training-engine): optional per-leaf gradient norms

  #5352·A9isha opened Sep 23, 2026Collaborator·

  ·
  2

• fix(vllm): enable Qwen3.5 hybrid GDN paged attention and KV cache spec for vLLM serving

  #5349·Shuwen-Fang opened Sep 23, 2026Collaborator·

  ·
  6

• MoE: moe_x_sorted remat location (save the routed expert input; skip the backward dispatch all-gather + ragged sort)

  #5348·ultrons opened Sep 23, 2026Collaborator·

  ·
  3

• fix(dependencies): fix Docker image vulnerabilities and migrate base images to Airlock

  #5347·chiajunglien opened Sep 23, 2026Collaborator·

  ·
  11

• Add elastic reshard and recovery to the MaxText pre-train loop.

  #5343·copybara-service[bot] opened Sep 22, 2026Bot·

  ·
  

• Add in-memory elastic snapshot utilities for Pathways multi-slice training.

  #5342·copybara-service[bot] opened Sep 22, 2026Bot·

  ·