Scaling Zero-Order Pretraining through Model Sharding

Francois Chaubard; Mykel J. Kochenderfer; Chris Ré

Figure 7: SOMA sweeps on logarithmic wall-clock axes. Left, Fixed-width experts scaling from

N\in\{1,2,8,32,256\}

, shown through a common cutoff of 183.6 estimated hours. Increasing number of experts (and ensemble size) improves training monotonically. Center,

n_{\rm pert}\in\{16,64,256,1024\}

. Right, effective batch

B\in\{16,64,256,1024\}

. The

B=256,1024

runs accumulate four and sixteen batches of 64, each with 64 independently sampled directions, giving 256 and 1,024 directions per update. SOMA

N=8

improves loss with more

n_{\rm pert}

and

B

. Settings are in Table 1 and Appendix A .

Quantity	Configuration
Training corpus / tokenizer	100B FineWeb-Edu bytes / byte vocabulary $V=256$
Context length	1,024 tokens
Expert validation set	The 1% split within each expert’s shard. The other 99% is used for training. Expert validation loss controls learning-rate decay.
Ensemble validation set	976 windows from 202 FineWeb-Edu documents, used for scaling curves. Scores are ensemble validation loss.
Test set	4,096 FineWeb-Edu documents screened for training overlap. Each contributes one 1,025-byte window. Scores are test loss. Construction and exclusions are in Appendix L .
Evaluation protocol	On the ensemble validation and test sets, routing uses the first 256 input bytes and loss scores the next 768 targets. Both losses are mean cross entropy in nats/byte.

Appendix

Table 1: Experimental settings and evaluation sets.

Figure 8: More experts do not lower loss at every training budget. These are the expert-count measurements on the ensemble validation set from Figure 7 , plotted against aggregate GPU-hours estimated from the recorded mean step times across all experts. Expert width is fixed and

n_{\rm pert}=B=64

. Greater

N

adds parameters and costs more per update. Section M compares the perturbation and batch allocations at fixed model size.

$N$	$d_{t}$	head per expert	ensemble parameters	expert body parameters
1	509	8.19k	8.33M	8.33M
2	361	8.19k	8.41M	4.20M
8	181	8.19k	8.55M	1.06M
256	32	8.19k (shared)	8.44M	32.9k

Appendix

Table 2: Parameter counts for the SPSA sharding comparison. Model size counts expert bodies and heads. All heads have width

d_{E}=32

. Wider bodies use input and output projections. The

N=1,2,8

heads are updated locally and counted separately. The

N=256

head is shared and frozen.

Figure 9: Semantic shards and warm initialization lower loss. Means and standard errors over three seeds for each condition at the shown training steps. The seed model’s training cost is excluded. Updating the tied head only through the decoder raises loss by 0.029 nats/byte, while freezing a random head performs much worse.

Figure 10: More shards and more perturbations reduce relative gradient variance. Lower is better. Centered estimator variance is divided by the squared exact gradient norm. Model size is roughly 8.44M in total (Section 3 ). Left,

N=256

has

118\times

lower variance than

N=1

. Right, doubling

n_{\rm pert}

halves variance and doubles update compute. Points use weights at 10,000 updates and the same two ensemble validation chunks on an RTX 5090. Dashed lines are predictions for the recorded perturbation distribution, without fitting. These tests measure variance at fixed weights, not learning speed. Appendix G.1 gives the calculation.

Figure 11: Keeping independent losses separate reduces gradient error at the same work. Lower is better. Both panels use SOMA

N=4

formed from SOMA

N=256

at 900k updates, with fixed weights and dense probes. The right panel uses one loss for four experts, two losses for two experts each, or four individual losses. Relative error is

\mathbb{E}\|\widehat{g}-g\|^{2}/\|g\|^{2}

, using the exact gradient

g

. Left fixes two ensemble validation chunks. At

n_{\rm pert}=64

, independent losses give

4.02\times

lower error, against

4.00\times

predicted. Right selects 128 coordinates per expert and draws batches from a fixed pool of 32 chunks. Its reference is the pool mean gradient, with measured

\nu=8.32

. Perturbations are counted per expert. All loss groupings use the same forward evaluations. Right holds

n_{\rm pert}B=256

, so larger batches leave fewer perturbations. The total batch–perturbation budget is fixed at

n_{\rm pert}B=256

, so increasing batch size reduces the number of perturbation directions. Dashed lines are predictions. Markers show mean central-difference error. Bars are standard errors over 32 sets of random perturbations on the left and 16 batch-and-perturbation draws on the right.

Accounting method	GPU-hours	Calculation or recorded field
Recorded expert mean times	41.9k	$U\sum_{e=1}^{256}\bar{t}_{e}/3600$
W&B-reported physical total	49.6k	physical_gpu_hr
Fixed 0.168-second reference	49.7k	$256U(0.168)/3600$
Fixed 0.140-second normalization	41.4k	$256U(0.140)/3600$

Appendix

Table 3: Compute accounting for the same

N=256

checkpoint. The first row is used for the endpoint in Figures 1 and 15 . The other rows preserve the separate logged total and fixed-rate conventions. They do not identify different training endpoints.

Figure 12: DDPP synchronizes every global step. Nodes evaluate subsets of the batches and perturbations for one shared model. Dashed blue arrows carry local gradient estimates to the coordinator, and solid orange arrows return the combined estimate. FSO instead updates each expert using its own loss and exchanges no per-step updates.

Configuration	$n_{\rm pert}$	Updates	GPU-hours	Test loss
SOMA $N=2$	64	65.0k	150	1.76
Monolithic SPSA	64	30.0k	152	2.00
Monolithic SPSA	256	1.32k	150	2.11
Monolithic SPSA	1024	839	148	2.00
Monolithic EGGROLL	—	7.40k	141	2.21

Appendix

Table 4: Measured test losses near 150 aggregate GPU-hours for models of approximately 8.44M parameters. Updates are checkpoint counters, per expert for SOMA. Test loss is in nats/byte. Device costs include communication and waiting, so this compares the measured training systems rather than arithmetic work alone. Values are rounded to three significant figures. EGGROLL uses its population of

2^{20}

and

B=256

.

$N$	SOMA	BPTT	SOMA	BPTT	SOMA	BPTT
	50 GPU-hours		75 GPU-hours		100 GPU-hours
$1$	$2.1237$	$1.0198$	$2.0590$	$1.0206$	$2.0343$	1.0091 ∗
$2$	$1.9363$	$1.0711$	$1.8734$	$1.0653$	$1.8257$	1.0535
$8$	$2.0880$	$1.1826$	$2.0456$	$1.1797$	$2.0179$	1.1776

Appendix

Table 5: BPTT reaches lower ensemble validation loss at these budgets. Values are in nats/byte on the ensemble validation set, using roughly 8.44M models (Section 3 ). Values interpolate measured losses, except the marked BPTT

N=1

entry, which uses its last checkpoint at 97.3 GPU-hours. Hardware and training recipes differ.

Learning rate	Initial $\varepsilon$	$N=1$	$N=2$
		Ensemble validation loss
0.00125	0.00250	$2.2768\pm 0.0040$	$2.2370\pm 0.0034$
0.00250	0.00250	$2.3247\pm 0.0149$	$2.2708\pm 0.0025$
0.00500	0.00250	$2.4179\pm 0.0267$	$2.3379\pm 0.0073$
0.00125	0.00125	$2.2764\pm 0.0040$	$2.2370\pm 0.0034$
0.00125	0.00500	$2.2811\pm 0.0048$	$2.2404\pm 0.0033$

Appendix

Table 6: Learning-rate and

\varepsilon

sensitivity at an equal aggregate GPU compute budget per configuration. Entries are mean

\pm

sample standard deviation across three optimization seeds, in nats/byte. Selection uses each seed’s unrounded validation losses.

Optimization seed	Monolithic SPSA	SOMA ( $N=2$ )	Difference
1	2.2641	2.2273	0.0368
2	2.2666	2.2204	0.0462
3	2.2576	2.2218	0.0358
Mean	2.2628	2.2232	0.0396
Sample standard deviation	0.0047	0.0036	0.0057

Appendix

Table 7: Test loss after selecting each seed’s settings using only the ensemble validation set. All losses are nats/byte, and the difference is monolithic SPSA minus SOMA. Standard deviations measure optimization-seed variation with fixed starting checkpoints, not test-document uncertainty.

$N$	Batch size	Cosine distance	Batch size	Cosine distance
	Initialization		10,000 updates
1	8000	0.00002	8000	0.01045
2	8000	0.00002	8000	0.00503
8	8000	0.00002	4000	0.00492
256	8000	0.00013	1600	0.00759

Appendix

Table 8: BPTT reference batch sizes and cosine distance to an independent batch of the same size. Trained

N=1

is just above

0.01

. Using the independent references changes every plotted value by less than

4\%

.

Checkpoint	Training compute	Validation	Test
	(GPU-hours)	(nats/byte)	(nats/byte)
8.44M parameters
SOMA $N=2$ , $n_{\rm pert}=64$	150	2.047	2.066
Monolithic SPSA, $n_{\rm pert}=64$	152	2.240	2.254
Monolithic SPSA, $n_{\rm pert}=256$	150	2.348	2.363
Monolithic SPSA, $n_{\rm pert}=1024$	148	2.271	2.290

Appendix

Table 9: Frozen next-byte evaluation on WikiText-103 for approximately 8.44M-parameter models near 150 aggregate training GPU-hours. Training compute uses each checkpoint’s recorded RTX 5090 accounting. Models and checkpoints are selected before observing WikiText scores.

Compute (operations)	Estimated ensemble validation loss
$2\times 10^{18}$	1.8046–1.8057
$4\times 10^{18}$	1.7947–1.8000
$8\times 10^{18}$	1.7848–1.7887

Appendix

Table 10: The saved 0.272M curves favor

(N,n_{\rm pert},B)=(8,64,1024)

at these budgets. Losses use the ensemble validation set from Figure 13 . Ranges span linear- and log-compute interpolation.

Figure 13: Larger batches reach 1.90 with fewer GPU-hours in these runs. All curves use

N=8

and a 0.272M model (Section 3 ). For

n_{\rm pert}=64,B=1024

, the first plotted point at or below 1.90 costs 28.7 GPU-hours, compared with 44.5 for

n_{\rm pert}=1024,B=64

. The larger batch uses

1.55\times

fewer GPU-hours. Costs sum recorded mean step times across the eight experts. The dotted line marks 1.90. Batches of 256 and 1024 accumulate four and sixteen independently sampled batches of 64 with independently sampled directions. These runs are separate from the 8.44M comparison in Figure 1 . The controlled continuation in Table 11 finds nearly the same loss for three allocations.

$n_{\rm pert}$	$B$	Test loss (nats/byte)
64	1024	2.376587
256	256	2.376575
1024	64	2.376625

Appendix

Table 11: Nearly equal test loss after 512 further updates at equal arithmetic operations. All settings continue the same

N=8

, 0.272M model. Values average two paired seeds on all 4,096 test documents.

Figure 14: More shards can lower loss under a tight per-expert budget. Color gives the expert count of the lowest-loss saved checkpoint that fits both budgets. We compare 44 checkpoints from

N=1,2,8,256

on the ensemble validation set of Figure 7 , including both

N=256

recipes through 900k and 4.16M updates.

N=1

is never selected. Average work per expert is total GPU-hours divided by

N

. At one GPU per expert, it approximates average elapsed training time. Labels give the selected checkpoint loss at the marked budgets. Work is estimated from recorded update rates and excludes initialization. The

N=256

ensemble uses 0.168 seconds per expert update. Region boundaries come from the saved checkpoints. Training recipes follow the fixed-size controls in Figure 1 , detailed in Appendix I . Model size follows Section 3 .

Figure 15: More shards reduce inference cost, but training progress changes their ranking. Left, the four settings from Figure 1 at 100 estimated training GPU-hours, with losses interpolated between checkpoints. Right, recorded

N=1,2,256

checkpoints at 152, 150, and 41.9k estimated training GPU-hours, using the recorded-update accounting in Table 3 . Their training budgets differ. Both panels use the same test set. Bars show 95% document-bootstrap intervals and exclude interpolation uncertainty on the left. Arithmetic operations count the selected experts and exclude routing. Parameter counts and arithmetic operations follow Sections 3 and 7 .

Model	Batch	Throughput (M tokens/s)	Batch latency (ms)	Peak GPU memory (GiB)
Monolithic SPSA	768	0.555	1,060	16.5
SOMA $N=2$	768	0.468	1,260	12.3
SOMA $N=8$	1,024	0.257	3,060	15.0
SOMA $N=256$	4,096	2.36	1,330	8.30

Appendix

Table 12: Complete warm inference at the fastest validated settings, including routing when needed. Each token is one byte, with 768 evaluated tokens per 1,024-byte request. Values are medians of seven repetitions on the same GPU. Peak allocated memory includes intermediates. Compilation and warmup are excluded.

Model	Parameters	Updates	Estimated training wall-clock hours	Aggregate training GPU-hours	Test loss
SOMA $N=8$	8.55M	53.5k–70.0k	104	700	1.71
SOMA $N=256$	8.44M	2.18M	100	22.0k	1.75
SOMA $N=256$	8.44M	4.16M	191	41.9k	1.68

Appendix

Table 13: Test loss and training budgets for the approximately equal-total-size inference comparison. All rows use top-

k

routing with

k=4

. The SOMA

N=8

and later SOMA

N=256

checkpoints are timed in Table 12 . The middle row supplies the nearest measured SOMA

N=256

checkpoint to the training duration of SOMA

N=8

. Wall-clock and aggregate time retain the accounting used in Figure 1 , excluding initialization, evaluation and checkpointing.

Figure 16: Top-

k

routing with

k=4

captures most of the gain at the same training stage. Each curve varies

k

in a saved 4M-update ensemble. All experts have the same width and use

n_{\rm pert}=B=64

. SOMA

N=8

and SOMA

N=256

reach these checkpoints after approximately 176 and 192 elapsed hours, respectively. All scores use the same ensemble validation set and average the selected experts’ probabilities. The dashed line marks

k=4

.

Figure 17: Averaging helps more when experts make different errors. The horizontal axis measures error similarity after removing shared token difficulty. The vertical axis measures how much lower the averaged prediction’s ensemble validation loss is than the mean loss of the two experts alone. For example, 0.10 means a reduction of 0.10 nats/byte. Each point averages about 1,300 expert pairs with similar error correlations, using 749,568 ensemble validation tokens. Pairs share experts. This tests averaging pairs of experts, not top-4 routing.

Scaling Zero-Order Pretraining through Model Sharding

Abstract

Figures & tables

Appendix

Zeroth-Order Optimization

Probe-Space Preconditioning for Fast and Stable Zero-Order Training

GRZO: Group-Relative Zeroth-Order Optimization for Large Language Model Fine-Tuning

Sample-Adaptive Zeroth-Order Post-Training

Abstract

Zeroth-Order Optimization

Explore similar work

Probe-Space Preconditioning for Fast and Stable Zero-Order Training

GRZO: Group-Relative Zeroth-Order Optimization for Large Language Model Fine-Tuning

Sample-Adaptive Zeroth-Order Post-Training