All pull requests

Pull requests

• Fix remaining fixable critical and high docker image vulnerabilities

  #5512·KevinYenky opened Oct 2, 2026Collaborator·4 tasks done·

  ·
  3

• Require flax>=0.12.10 for pre-training so MaxText runs with JAX 0.11.2.

  #5509·copybara-service[bot] opened Oct 2, 2026Bot·

  ·
  1

  

• Feat(Checkpoint): Add colocated_python checkpointing

  #5508·YixuanWang-99 opened Oct 2, 2026Collaborator·

  ·
  3

• [DRAFT] All-gather / combine MoE tokens in the TC kernels' 3D layout

  #5507·gobbleturk opened Oct 2, 2026Collaborator·4 tasks·

  ·
  2

• Train loop and input pipeline hooks for training a data-parallel replica on a submesh.

  #5506·copybara-service[bot] opened Oct 2, 2026Bot·

  ·
  1

  

• Threaded (non-SPMD) streaming DiLoCo trainer.

  #5505·copybara-service[bot] opened Oct 2, 2026Bot·

  ·
  1

  

• DiLoCo: threaded streaming DiLoCo recipe and docs.

  #5504·copybara-service[bot] opened Oct 2, 2026Bot·

  ·
  1

  

• Profiling and logging support for threaded DiLoCo.

  #5503·copybara-service[bot] opened Oct 2, 2026Bot·

  ·
  1

  

• Quantize the GMM bwd gradient once under fixed calibration

  #5502·RissyRan opened Oct 2, 2026Collaborator·2 of 4 tasks·

  ·
  4

• MoE: accumulate wi_0 and wi_1 backward DLHS in-place

  #5501·RissyRan opened Oct 2, 2026Collaborator·4 tasks done·

  ·
  3

• MoE: accumulate weight gradients across token chunks in place

  #5500·RissyRan opened Oct 2, 2026Collaborator·4 tasks done·

  ·
  5

• FP8: add logits_proj_bwd_quant_calibration_method

  pull ready
  #5499·RissyRan opened Oct 2, 2026Collaborator·4 tasks done·

  ·
  8

• [DRAFT] Keep the TC-sorted MoE token buffer in 3D layout through the gmm_v2 k…

  #5498·gobbleturk opened Oct 2, 2026Collaborator·4 tasks·

  ·
  4

• [DRAFT] Apply MoE routing weights to the activation with the TC ragged sort

  #5497·gobbleturk opened Oct 2, 2026Collaborator·4 tasks·

  ·
  2

• [DRAFT] Add TensorCore ragged sort for ring-of-experts MoE

  #5496·gobbleturk opened Oct 2, 2026Collaborator·4 tasks·

  ·
  3

• Forward-pass velocity numerical checker for WeaverOmniTransformer

  #5495·parsley9877 opened Oct 1, 2026Collaborator·3/3·4 tasks done·

  ·
  5

• fix(tests): follow-ups to #5411 for decoupled and checkpoint output isolation

  pull ready
  #5491·gulsumgudukbay opened Oct 1, 2026Collaborator·4 tasks done·

  ·
  15

• add new oob test script

  #5490·pv97 opened Oct 1, 2026·4 tasks·

  ·
  6

• Add checkpoint conversion and parameter mapping for WeaverOmniTransformer

  #5488·parsley9877 opened Oct 1, 2026Collaborator·2/3·4 tasks·

  ·
  5

• [Draft] [GPU] DSv4 Support for TE MoE Block

  #5487·dinodeep opened Oct 1, 2026Collaborator·4 tasks done·

  ·
  7

• Migrate MaxText benchmark runner and configs from XPK to Cluster Toolkit

  #5483·KevinYenky opened Oct 1, 2026Collaborator·4 tasks done·

  ·
  31

• Restore Orbax checkpoint Cloud Logging for ML Goodput after the Orbax V1 migration

  #5478·lydhr opened Sep 30, 2026Collaborator·4 tasks done·

  ·
  9

• Add WeaverOmniTransformer full 36-layer backbone architecture and tests

  #5476·parsley9877 opened Sep 30, 2026Collaborator·1/3·4 tasks done·

  ·
  2

• Fix MaxText GPU Docker image build for the Ubuntu 24.04 / CUDA 13.4 base image.

  #5475·copybara-service[bot] opened Sep 30, 2026Bot·

  ·
  1

  

• Gemma4 26B MoE on TPU v6e: step-time optimizations (4.80 -> 4.00 s/step, one commit per optimization)

  #5474·csgoogle opened Sep 30, 2026Contributor·

  ·
  3