Record: Parallel Muon + Parameter Banking — 81.87ms/step, val_bpb 1.1247 (3-seed mean)

[abaybektursun]

Novel Contribution: Parameter Banking + Parallel Muon

This submission introduces Parameter Banking, a weight layout restructuring that enables batched optimizer operations, combined with an adapted Parallel Muon communication strategy. Together, these provide a 3.4% training throughput improvement that is architecture-agnostic and composes with any Muon-based training stack. The approach has since been adopted by subsequent competition submissions (e.g., PR #549).

Pure systems optimization — model architecture and hyperparameters are unchanged.

3-Seed Results (8×H100 80GB SXM, PyTorch 2.9.1+cu128, 600s)

Seed	step_avg	steps	int6 sliding val_bpb	artifact
1337	81.86 ms	7,331	1.1241	15,830,960 bytes
42	81.88 ms	7,328	1.1253	15,819,728 bytes
2025	81.86 ms	7,330	1.1247	15,796,052 bytes
Mean	81.87 ms	7,330	1.1247 (std 0.0006)	~15.8 MB

Technical Approach

1. Parameter Banking (novel)

We restructure 66 separate nn.Linear weight matrices into 4 contiguous 3D nn.Parameter tensors, grouped by shape:

• qo_bank: (22, 512, 512) — Q + Out projections
• kv_bank: (22, 256, 512) — K + V projections
• mlp_up_bank: (11, 1536, 512) — MLP up
• mlp_down_bank: (11, 512, 1536) — MLP down

Forward pass uses F.linear(x, bank[layer_idx]) — compiles identically to nn.Linear under torch.compile. Verified: banked forward+backward = 72.33ms vs baseline 72.59ms.

The key benefit: Newton-Schulz orthogonalization (used by Muon) becomes a single torch.bmm over the batch dimension, replacing 66 sequential small GEMMs. This reduces optimizer time from 19.7ms to 1.3ms (15× faster).

2. Parallel Muon (adapted from arXiv:2511.07464)

Standard DDP is incompatible with parameter banking: bank gradients aggregate across all 11 layers and are only available at end of backward, destroying compute-communication overlap (+4ms regression).

Our solution removes DDP for banked parameters and schedules communication explicitly:

1. Launch async reduce_scatter for all banks (biggest first)
2. all_reduce + Adam step on small replicated params (while bank RS is in-flight)
3. Wait for RS, local batched NS on each GPU's shard, async all_gather

This follows the DDP-free communication pattern from modded-nanogpt, adapted to work with our banking structure.

Engineering notes

Approach	Result	Lesson
Non-surgery batching (keep 66 params, batch in optimizer)	85.73ms	Gather/scatter kernel overhead offsets speedup
DDP with banks	88.8ms (+4ms)	Bank grads only available at end of backward
Polar Express (arXiv:2505.16932)	82ms, 16.2MB	PE weights compress ~190KB worse than NS
Parameter Banking + Parallel Muon	81.87ms, 15.8MB	Architecture-agnostic, composable

Compatibility analysis

Base PR	Speed	Score	Finding
#315 (EMA only)	-3.4%	-0.0006 BPB	Extra steps improve EMA monotonically
#374 (Tight SWA)	-3.5%	+0.001	SWA averages warmdown weights; extra steps don't enter the window
#401 (EMA+SWA)	-2.8%	+0.0005	Same SWA dilution
#398 (TTT)	-2.3%	+0.004	More-converged model has less room for TTT adaptation

Key finding: The throughput advantage translates to quality gains exclusively for EMA-based models, where every additional step monotonically refines the exponential moving average.

Credits

• Architecture: PR #315 by @jfprincz (11L Partial RoPE + LN Scale + EMA + XSA4)
• Parallel Muon scheduling: Adapted from arXiv:2511.07464 and modded-nanogpt

🤖 Generated with Claude Code