[canary-tpu] grug MoE train crashes: megablox GMM output block (512,32) invalid for (524288,128)

[claude]

The TPU canary ferry (Grug MoE smoke test) crashes at compile time on the first train_step. On this run the crash was deterministic across all 4 task attempts, so max_task_failures was exceeded and the job failed.

The expert grouped-matmul fails to lower on TPU v6e: the megablox GMM pallas_call builds an output BlockSpec((512, 32)) for an output array of shape (524288, 128). Pallas requires the last block dim to be divisible by 128 or equal to the array dim; 32 is neither.

ValueError: The Pallas TPU lowering currently requires that the last two
dimensions of your block shape are divisible by 8 and 128 respectively, or be
equal to the respective dimensions of the overall array. Block spec for outputs
in pallas_call kernel at .../pallas/ops/tpu/megablox/gmm.py:393 has block shape
(Blocked(block_size=512), Blocked(block_size=32)), array shape (524288, 128)

Call path: levanter/grug/_moe/scatter.py:48 (expert MLP projection), then haliax.nn.ragged_dot._ragged_dot_megablox_impl (lib/haliax/src/haliax/nn/ragged_dot.py:93), then jax.experimental.pallas.ops.tpu.megablox.gmm. The output n=128 equals the canary model intermediate_dim.

Reproduce:

1. Run the TPU canary ferry python -m experiments.ferries.canary_ferry on v6e-4 (CANARY_MODEL: hidden_dim=32, intermediate_dim=128, num_experts=32, num_experts_per_token=2, num_layers=4).
2. The failure is intermittent across scheduled runs (history below); it reproduces when the megablox GMM lowers fresh instead of loading from the shared JAX compilation cache.

Expected: the megablox GMM lowers with a valid output block (last dim divisible by 128, or equal to n) and the canary trains to the regression gate.

Root-cause status (unresolved discrepancy)

The observed tn=32 does not match the checked-out tiling logic. _megablox_tile_size() at HEAD returns an n-tile of 1024, so _ragged_dot_megablox_impl passes tiling[2] = min(n=128, 1024) = 128, which would produce a valid (512, 128) output block. Static reading of the current code does not explain how tn=32 was emitted. This gap is the crux and needs a kernel owner.

Signals pointing away from a fresh code regression on the commit for this run:

• The prior scheduled canary (2026-08-27 06:30 UTC, f8972ac) succeeded. The three commits since (PR 8721 external eval/RL runtimes, PR 8717 echo dashboard, PR 8710 IAM grant) touch no JAX, haliax, levanter, or grug code.
• ragged_dot.py last changed 2026-08-25 (PR 8610), before two successful runs.
• Canary history alternates pass/fail (Aug 26 pass, Aug 27 06:30 pass, 03:42 and 17:21 fail).

The failing attempt logged a compilation-cache read failure immediately before the crash:

UserWarning: Error reading persistent compilation cache entry for
jit__create_lm_example: RuntimeError: Session is closed

Shared cache: gs://marin-eu-west4/tmp/ttl=30d/compilation-cache. Leading hypothesis: a latent megablox tiling incompatibility for these shapes, normally masked by compilation-cache hits and exposed when the cache read fails and the kernel lowers cold. Not confirmed.

Canary run context:

• Lane: tpu (GCP, TPU v6e-4, europe-west4)
• Iris job: /runner/iris-run-job-20260827-172216
• GHA run: https://github.com/marin-community/marin/actions/runs/33097960592
• W&B run: https://wandb.ai/marin-community/marin/runs/canary-tpu-33097960592-1
• Date: 2026-08-27

[claude]

🤖 Recurred on the scheduled TPU canary (2026-08-28). Same signature and call path: megablox GMM builds output BlockSpec((512, 32)) for array (524288, 128), and Pallas rejects it (last block dim neither divisible by 128 nor equal to 128). Crash is at compile time on the first train_step, deterministic across all 4 attempts, so max_task_failures was exceeded.

This run corroborates the compilation-cache hypothesis in the issue body. Every attempt logged the cache read failure immediately before lowering cold:

jax/_src/compiler.py:815: UserWarning: Error reading persistent compilation cache
entry for 'jit__create_lm_example': RuntimeError: Session is closed

Shared cache gs://marin-eu-west4/tmp/ttl=30d/compilation-cache (attempts 0/1/2/3 all hit it). The pattern holds: the megablox tiling incompatibility is masked by cache hits and exposed when the read fails and the kernel lowers fresh. The tn=32 vs. expected-128 discrepancy from the checked-out tiling still stands and still needs a kernel owner.

Canary run context:

• Lane: tpu (GCP, TPU v6e-4, europe-west4)
• Iris job: /runner/iris-run-job-20260828-181501
• GHA run: https://github.com/marin-community/marin/actions/runs/33198349960
• W&B run: https://wandb.ai/marin-community/marin/runs/canary-tpu-33198349960-1
• Date: 2026-08-28

[rjpower]

@weaverbot fix this please

[loom-oa-dev]

On it — https://loom.oa.dev/s/yjmq2r7w
marin-style-nightly-updates · partitioned-data-loader-design · partitioned-data-loader-research · partitioned-data-loader-spec

• 🟢 Aug 28 19:48Z Investigating GitHub issue [canary-tpu] grug MoE train crashes: megablox GMM output block (512,32) invalid for (524288,128) #8731 and reproducing the Grug MoE megablox tile mismatch.
• 🟠 Aug 28 19:53Z attention — PR Fix Megablox tiling across ragged-dot gradients #8719 contains the tested fix for [canary-tpu] grug MoE train crashes: megablox GMM output block (512,32) invalid for (524288,128) #8731 and is green; it needs review and merge: Fix Megablox tiling across ragged-dot gradients #8719

[loom-oa-dev]

🤖 Research

The invalid Megablox grouped-matmul output block is produced in the left-hand-side gradient path. Haliax passes the concrete forward tile (512, 128, 32) into JAX Megablox's custom VJP; grad_lhs swaps the contraction and output dimensions, then reuses that tuple and produces output block (512, 32) for an array with 128 columns. The compilation-cache read failures coincided with cold lowering in the reported runs, but the reused forward tile causes the compiler error.

Proposed Fix

PR #8719 already changes Haliax to pass a callable tiling policy, so each forward or backward grouped matmul computes a tile for its own shape; the left-hand-side gradient selects (512, 32, 128). No duplicate implementation was opened. The PR has green GitHub checks, exact-shape v6e-4 parity with XLA for loss, forward output, and both gradients (maximum absolute difference 0.0), and a canary run that displayed 238/238 while its validator recorded 237 completed steps and final finite loss 5.7976. It is open and needs review and merge. The incident record is at Echo.

[claude]

🤖 Recurred on the scheduled TPU canary (2026-08-29). Same signature, same call path: the megablox GMM builds output BlockSpec((512, 32)) for array (524288, 128), and Pallas rejects it (last block dim neither divisible by 128 nor equal to 128). Crash is at compile time on the first train_step, deterministic across all 4 attempts (0/1/2/3), so max_task_failures was exceeded and the job failed.

The cold-lowering pattern held again: attempt 1 logged the compilation-cache read failure immediately before lowering fresh.

jax/_src/compiler.py:815: UserWarning: Error reading persistent compilation cache
entry for 'jit__create_lm_example': RuntimeError: Session is closed

Shared cache gs://marin-eu-west4/tmp/ttl=30d/compilation-cache. This matches the confirmed root cause: the (512, 32) block comes from the left-hand-side gradient VJP reusing the forward tile (512, 128, 32) after the contraction/output dims swap, not from the forward tiling.

The fix (PR #8719) is still open and unmerged (REVIEW_REQUIRED as of this run). The canary will keep failing on cache-miss runs until it lands.

Canary run context:

• Lane: tpu (GCP, TPU v6e-4, europe-west4)
• Iris job: /runner/iris-run-job-20260829-122144
• GHA run: https://github.com/marin-community/marin/actions/runs/33252261455
• W&B run: https://wandb.ai/marin-community/marin/runs/canary-tpu-33252261455-1
• Date: 2026-08-29

[claude]

🤖 Recurred on the scheduled TPU canary (2026-08-30), the third straight daily recurrence logged here (Aug 28 / 29 / 30). Same signature and call path: the megablox GMM builds output BlockSpec((512, 32)) for array (524288, 128), and Pallas rejects the 32-wide trailing block (neither divisible by 128 nor equal to 128). The crash is at compile time on the first train_step, deterministic across all 4 attempts (0/1/2/3), so max_task_failures was exceeded and the job failed.

Matches the confirmed root cause: the (512, 32) block is the left-hand-side gradient, where megablox's VJP (_gmm_bwd) reuses the forward tile (512, 128, 32) for grad_lhs = gmm(..., transpose_rhs=True). The transpose makes that output n = rhs.shape[1] = intermediate_dim = 128, so the reused tn=32 is applied to a 128-wide output. The forward down-projection itself lowers fine because its output width is hidden_dim=32 and tn=32 equals the array dim there.

No Session is closed cache-read warning this run; the kernel simply lowered cold with no cache hit for the entry, which is enough to trigger the deterministic failure.

The fix (PR #8719) is still open and unmerged. As of this run it is mergeable with green checks but REVIEW_REQUIRED. The canary will keep failing on every cache-miss run until it lands.

Canary run context:

• Lane: tpu (GCP, TPU v6e-4, europe-west4)
• Iris job: /runner/iris-run-job-20260830-112054
• GHA run: https://github.com/marin-community/marin/actions/runs/33308652303
• W&B run: https://wandb.ai/marin-community/marin/runs/canary-tpu-33308652303-1
• Date: 2026-08-30

[claude]

🤖 Recurred on the scheduled TPU canary (2026-08-31) — fourth straight daily recurrence (Aug 28 / 29 / 30 / 31). Same signature and call path: the megablox GMM builds output BlockSpec((512, 32)) for array (524288, 128), and Pallas rejects the 32-wide trailing block (neither divisible by 128 nor equal to 128). The crash is at compile time on the first train_step, deterministic across attempts 1/2/3, so max_task_failures was exceeded and the job failed.

Matches the confirmed root cause: the (512, 32) block is the left-hand-side gradient, where megablox's VJP reuses the forward tile (512, 128, 32) for grad_lhs = gmm(..., transpose_rhs=True). The transpose makes that output n = rhs.shape[1] = intermediate_dim = 128, so the reused tn=32 is applied to a 128-wide output.

No Session is closed cache-read warning this run; the kernel simply lowered cold on a cache miss, which is enough to trigger the deterministic failure.

This run is the first on the re-vendored runtimes from PR #8804 (external runtime advance, merged 2026-08-31 05:47Z). The new lib/haliax/src/haliax/nn/ragged_dot.py still passes a static tuple tiling (min(m,512), min(k,…), min(n,1024)), not the per-shape callable policy from the fix, so the bug is unchanged and #8719 still applies cleanly to the new tree (MERGEABLE).

The fix (PR #8719) is still open and unmerged — BLOCKED / REVIEW_REQUIRED as of this run. The canary will keep failing on every cache-miss run until it lands.

Canary run context:

• Lane: tpu (GCP, TPU v6e-4, europe-west4)
• Iris job: /runner/iris-run-job-20260831-124909
• GHA run: https://github.com/marin-community/marin/actions/runs/33393578343
• W&B run: https://wandb.ai/marin-community/marin/runs/canary-tpu-33393578343-1
• Date: 2026-08-31