Can111/ShellElastic-baseline-SampledStandalone-375M-906M

Scalable Elastic Pretraining
with Simultaneous Maximal Updates and Hyperparameter Transfer

Can Jin^1* · Yan Kang² · Yifan Gong² · Chongjian Ge² · Mingze Xu² · Yuanjun Xiong² · Dimitris N. Metaxas¹ · Hongwu Peng²

¹Rutgers University   ²Adobe Research

^*Corresponding author: can.jin@rutgers.edu. Work done during an internship at Adobe Research.

About ShellElastic

Foundation models need to support diverse inference budgets. Elastic models train several nested sub-networks inside one weight-shared supernet. ShellElastic is an elastic transformer that organizes hidden dimensions into a dense core and successive shells, so that every hidden matrix is block lower triangular. One full-width forward pass contains the exact hidden states of every sub-network, so joint training keeps a fixed computation graph and costs about as much per step as training the widest sub-network alone.

µ-Elastic extends maximal update parameterization (µP) to ShellElastic and other elastic architectures. Its scaling rules keep all sub-networks of a supernet maximally updated simultaneously, and hyperparameters tuned on a small non-elastic proxy stay near-optimal for elastic families and larger supernets. A single ShellElastic run produces sub-networks that outperform MatFormer and fully elastic baselines at matched parameter counts, and approach separately trained models that together take several times the steps.

Sampled standalone baseline, 375M–906M scale

Final weights of the Sampled standalone baseline at the 375M–906M scale of the paper (Section 4.1). These weights are a baseline trained for this paper, not an official release of MatFormer (Devvrit et al., 2024). Sampled standalone, the standalone baseline of MatFormer, trains the five Standalone models within the steps of one elastic run (Section 4.1). Each folder of this repository holds one of the five models (Low, Medium, High, Extra, Max).

Models

Model	Folder	Model width	Attention heads	FFN width	Parameters
Low	low/	832	13	2,288	377,529,152 (377.5M)
Medium	medium/	960	15	2,640	487,711,680 (487.7M)
High	high/	1,088	17	2,992	611,787,840 (611.8M)
Extra	extra/	1,216	19	3,344	749,757,632 (749.8M)
Max	max/	1,344	21	3,696	901,621,056 (901.6M)

Parameter counts include every weight of the model; rounded to 0.1M, they equal the entries of Table 10 of the paper. Each folder holds one model: config.json and model.safetensors (float32).

Results

Model	Val. loss, this checkpoint	Val. loss, paper	5-shot acc. (%), this checkpoint	5-shot acc. (%), paper
Low	3.2015	3.208	37.8	37.7
Medium	3.1748	3.175	37.5	37.8
High	3.1457	3.141	37.6	38.0
Extra	3.1197	3.115	38.1	38.7
Max	3.0908	3.093	39.5	39.0
Avg.	3.1465	3.147	38.1	38.2

This checkpoint: the final validation loss of each model (per-token cross-entropy in nats on held-out C4 text, at step 50,000) and its mean 5-shot accuracy over the 13 downstream tasks of Appendix F. Paper: the main results of the paper for Sampled standalone at the 375M–906M scale (Table 2). Avg.: the mean over the five models. Appendix G of the paper gives the per-task accuracies.

Architecture

Each folder holds one dense decoder-only transformer, one of the five Standalone architectures of the 375M–906M scale (Appendix E.1, Tables 7 and 8): width 832 to 1,344, FFN expansion 2.75 and d_head = 64.

Attention	heads of size d_head = 64 (13 to 21 heads, Models table); separate query, key and value projections
FFN	SwiGLU with expansion e = 2.75
Normalization	RMSNorm over the model width
Shared settings	32 pre-norm layers with residual-branch multiplier 1; RMSNorm with learnable gains and ε_N = 10⁻⁵; query–key RMSNorm per head; rotary embeddings with base 10,000 (the code uses the rotate-half form); attention logit scale d_head^−1/2 and an elementwise sigmoid gate on the attention output, sigmoid(x W_G^⊤) (Equation 4); SwiGLU FFN; no biases; untied embeddings over the 50,432-token GPT-NeoX vocabulary; sequence length 2,048 (Table 8 and Section 4.1)

The stored weights already include the constant forward multipliers of the model (the FFN output multiplier 1/e in down_proj), so every folder runs as an ordinary dense transformer of this architecture. Tensors of a model of width d and FFN width f, with i = 0, …, 31 numbering the layers:

Tensor	Shape
embed_tokens.weight	[50432, d]
layers.<i>.{attn_norm,mlp_norm}.weight, layers.<i>.attn.{q,k}_norm.weight, norm.weight	[d]
layers.<i>.attn.{q,k,v,gate,o}_proj.weight	[d, d]
layers.<i>.mlp.{gate,up}_proj.weight	[f, d]
layers.<i>.mlp.down_proj.weight	[d, f]
lm_head.weight	[50432, d]

Training

Data. DCLM-Baseline (Li et al., 2024), a filtered Common Crawl corpus, with the GPT-NeoX-20B tokenizer (Black et al., 2022) and a sequence length of 2,048 (Section 4.1). A run at this scale takes 50,000 steps of 256 sequences, 26.2B tokens (Appendix E.1). The five models were trained in one run of this length (Appendix E.1). The tokenized training and validation data of these runs are in the dataset repository Can111/ShellElastic-data.

Recipe. Trained as described in the paper (Section 4.1, Appendices E.1 and E.2, and Table 9).

Hyperparameters (Section 4.1, Appendix E.1 and Table 9; the paper rescales the optimizer HPs of this baseline with the token-budget rule of Mlodozeniec et al. (2026), and the table gives the rescaled values of this run):

Setting	Value
Parameterization	µP with d0 = 128 and m = width/128 (Table 8)
Optimizer	AdamW with decoupled weight decay
Base learning rate η0	4.47×10⁻³ (rescaled)
Base initialization σ0	3.3×10⁻³, truncated normal at ±3 standard deviations
Base weight decay λ0	0.0224 (rescaled); none on norm gains
β1, β2	0.75, 0.75 (rescaled)
AdamW ε	4.47×10⁻¹⁵ (rescaled)
Learning-rate schedule	warmup-stable-decay
Learning-rate warmup	linear, first 1,000 steps
Learning-rate decay	linear to 0, from step 40,000 to step 50,000
Batch size	256 sequences (524,288 tokens)
Sequence length	2,048
Training steps	50,000 (26.2B tokens)
Gradient clipping	global norm 4
Step skipping	after step 5,000, a step is skipped if its gradient norm exceeds 20× the clip value
z-loss coefficient λz	10⁻⁴ (Equation 8); all reported losses are the cross-entropy alone (Appendix E.2)
Precision	bf16 mixed precision with fp32 master weights and fp32 gradient reduction

Usage

The repository is self-contained: the reference implementation needs only PyTorch, plus safetensors to read the weights (pip install -r requirements.txt), and runs on CPU and GPU. Download the repository (for example with huggingface_hub.snapshot_download("Can111/ShellElastic-baseline-SampledStandalone-375M-906M", local_dir="ShellElastic-baseline-SampledStandalone-375M-906M")) and run, from that directory:

import torch
from modeling_dense import DenseDecoder

model = DenseDecoder.from_pretrained("low")           # float32 on the CPU; folders: low, medium, high, extra, max
# model = DenseDecoder.from_pretrained("max", device="cuda", dtype=torch.bfloat16)
ids = torch.tensor([[510, 5347, 273, 6181, 310]])     # GPT-NeoX-20B ids of "The capital of France is"
logits = model(ids)                                   # [1, 5, 50432], float32
logits, loss = model(ids[:, :-1], labels=ids[:, 1:])  # mean next-token cross-entropy
out = model.generate(ids, max_new_tokens=20)          # greedy decoding (no key/value cache)

python example.py --tier max runs one model on the CPU (--dtype bfloat16 runs it in bfloat16).

Inputs are token ids of the GPT-NeoX-20B tokenizer, for example from tokenizers.Tokenizer.from_pretrained("EleutherAI/gpt-neox-20b"); the tokenizer is not bundled here. The vocabulary has 50,432 entries (Table 7), of which the tokenizer uses the first 50,277.

Validation

Loaded in bfloat16 (the paper uses bf16 mixed precision, Table 9), the weights and code in this repository reproduce the final validation loss in the this checkpoint column of Results for the Max model to within 2×10⁻⁴ nats.

Files

• README.md: this card
• LICENSE: the Creative Commons Attribution 4.0 International license (CC BY 4.0)
• requirements.txt: torch and safetensors
• modeling_dense.py: the reference implementation (PyTorch only)
• example.py: loads one model on the CPU and runs it
• low/, medium/, high/, extra/, max/: config.json and model.safetensors (float32) of each model

Limitations

These are research checkpoints of language models pretrained on DCLM-Baseline, a filtered Common Crawl corpus, for a study of elastic pretraining. They have not been instruction-tuned, aligned or evaluated for safety, are not intended for deployment, and can produce incorrect, biased or offensive text. Use them for research only.

License

Everything in this repository (the weights, the configurations and the code) is released under CC-BY-4.0 (LICENSE).

Third-party terms: the training data, DCLM-Baseline, are distributed under CC-BY-4.0 and derived from Common Crawl; the validation text is from C4, distributed under ODC-BY together with the Common Crawl terms of use. Neither is included in this repository; the tokenized data of these runs are in the dataset repository Can111/ShellElastic-data. The GPT-NeoX-20B tokenizer (Apache-2.0) is not bundled here.

Citation

@article{jin2026shellelastic,
  title   = {{ShellElastic}: Scalable Elastic Pretraining with Simultaneous Maximal Updates and Hyperparameter Transfer},
  author  = {Jin, Can and Kang, Yan and Gong, Yifan and Ge, Chongjian and Xu, Mingze and Xiong, Yuanjun and Metaxas, Dimitris N. and Peng, Hongwu},
  journal = {arXiv preprint arXiv:2610.15124},
  year    = {2026},
  url     = {https://arxiv.org/abs/2610.15124}
}