All pull requests

Pull requests

• [AGENT DRAFT] Reduce MLPerf step-0 time: per-file c4_mlperf train sharding and data…

  #5402·gobbleturk opened Sep 28, 2026Collaborator·

  ·
  3

• Align sigmoid-router MoE aux loss, full-batch expert bias update, and comm precision mllog with Megatron-LM

  #5401·zxhe-sean opened Sep 28, 2026Collaborator·

  ·
  2

• Use 3-pass bf16 (Precision.HIGH / bf16x3) for GDN matmuls & bump VMEM limit to 0.9

  #5399·Rohan-Bierneni opened Sep 27, 2026Collaborator·

  ·
  3

• Use sliding window to track goodput metrics over time

  #5398·lydhr opened Sep 27, 2026Collaborator·

  ·
  4

• Exclude cuda-bindings and CUTLASS DSL from decoupled requirements

  #5396·ksapru opened Sep 27, 2026·

  ·
  1

• DiLoCo: fragment extraction/application and outer step for threaded streaming DiLoCo.

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

  ·
  1

  

• DiLoCo sharding: split a mesh along the diloco axis and strip that axis from logical rules.

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

  ·
  1

  

• DiLoCo fragmenter: optionally bucketize large non-scanned params across layer fragments.

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

  ·
  1

  

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