All pull requests

Pull requests

• fp8: per-layer delayed scaling for the backward cotangent (bwd_delayed_scaling)

  #5467·ultrons opened Sep 30, 2026Collaborator·

  ·
  5

• fix(dependencies): avoid recursive copy of gs://maxtext-test-assets subdirectories in Dockerfiles

  #5466·KevinYenky opened Sep 30, 2026Collaborator·

  ·
  1

  4

• Replicate MoE router weight under fsdp-as-dp-for-attn

  #5465·NuojCheng opened Sep 30, 2026Collaborator·

  ·
  4

• MTP: reuse main token embeddings and merge the shared output head (flags)

  #5462·RissyRan opened Sep 30, 2026Collaborator·

  ·
  4

• Add per-phase qk_diag_skip switches to the Tokamax splash kernel and a sa_qk_diag_skip MaxText flag.

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

  ·
  1

• Push all MaxText Docker image tags from BuildKit instead of gcloud.

  #5460·copybara-service[bot] opened Sep 29, 2026Bot·

  ·
  1

  

• fix(ci): explicitly add push=true to buildx outputs

  #5458·SurbhiJainUSC opened Sep 29, 2026Collaborator·

  ·
  2

• MoE: save routing index maps (moe_routing) and bf16 output from SC gather-reduce

  #5457·RissyRan opened Sep 29, 2026Collaborator·

  ·

• [NNX] Delete unused ToLinen wrappers and move fp8 and HF conversion to NNX

  #5456·ecnal-cienet opened Sep 29, 2026Collaborator·

  ·
  3

• ci: push automated changelog updates to a branch

  #5453·SurbhiJainUSC opened Sep 29, 2026Collaborator·

  ·
  1

• Reject scanned layer counts that are not whole layer cycles

  #5449·ecnal-cienet opened Sep 29, 2026Collaborator·

  ·
  1

• Fix slow Qwix fp8/int8 training under FSDP

  #5448·ecnal-cienet opened Sep 29, 2026Collaborator·

  ·
  3

• Use a unique run_name per process in MaxText train integration tests.

  #5447·copybara-service[bot] opened Sep 29, 2026Bot·

  ·
  1

  

• [Pallas:SC] Migrate from plsc.bitcast to pltpu.bitcast(dim=-1)

  #5446·copybara-service[bot] opened Sep 29, 2026Bot·

  ·

• convaiinnovations/laya

  #5445·hengtaoguo opened Sep 29, 2026Collaborator·

  ·
  9

• Enables restoring checkpoints from an external directory (directory: str | None = None) without mutating the configured checkpoint save directory:

  #5441·copybara-service[bot] opened Sep 29, 2026Bot·

  ·
  1

  

• Plumb multi-host RL profile configuration from MaxText into Tunix core.

  #5438·copybara-service[bot] opened Sep 29, 2026Bot·

  ·
  1

• [Draft] feat(weight_converter): enable BF16 trainer to MoE-only FP8 sampler weight sync for Qwen 3.5

  #5437·wenxindongwork opened Sep 29, 2026Collaborator·

  ·
  13

• fix(checkpointing): support disabling Orbax fingerprinting and bypass for host-offloaded optimizer state

  #5436·niting opened Sep 29, 2026Collaborator·

  ·
  2

• [Draft] Enable MoE-only FP8 checkpoint loading and layerwise stream quantization

  #5432·wenxindongwork opened Sep 29, 2026Collaborator·

  ·
  7

• Shard the MoE router and shared experts on expert under fsdp-as-dp-for-attn

  #5431·NuojCheng opened Sep 29, 2026Collaborator·

  ·
  7

• serve_fp8_weight: static FP8 activation scales for dense layers and fused MoE experts

  #5430·Shuwen-Fang opened Sep 29, 2026Collaborator·

  ·
  3

• Add Pallas cost_estimate to forward fused_conv1d_gdn kernel

  #5427·suhithr opened Sep 28, 2026Collaborator·

  ·
  4

• [trainerperf] Merge forward and backward scans in vocab_tiling_nnx_loss

  #5424·suhithr opened Sep 28, 2026Collaborator·

  ·
  2

• test(qwen3-next): run the packing equivalence tests on a GatedDeltaNet stack

  #5422·A9isha opened Sep 28, 2026Collaborator·

  ·