Can111/ShellElastic-baseline-MatFormer-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.

MatFormer baseline, 375M–906M scale

Final weights of the MatFormer 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). MatFormer nests only the FFN hidden width (Section 4.1), and the paper trains it with µ-Elastic (Section 3.2 and Appendix D). model.safetensors holds the supernet with its five sub-networks (Low, Medium, High, Extra, Max); modeling_elastic.py runs any of them and extracts it as a standalone dense model.

Sub-networks

Sub-network	Model width	Attention heads	FFN width	Parameters
Low	1,024	8	1,024	371,852,289 (371.9M)
Medium	1,024	8	2,176	485,098,497 (485.1M)
High	1,024	8	3,456	610,927,617 (610.9M)
Extra	1,024	8	4,864	749,339,649 (749.3M)
Max	1,024	8	6,528	912,917,505 (912.9M)

Parameter counts include the learned readout scalar α; rounded to 0.1M, they equal the entries of Table 10 of the paper. model.num_parameters(tier) counts the extracted sub-network, whose readout weights already contain α, and reports 1 fewer. model.safetensors stores the supernet: 1,722,418,177 parameters (6.89 GB in float32).

Results

Sub-network	Val. loss, this checkpoint	Val. loss, paper	5-shot acc. (%), this checkpoint	5-shot acc. (%), paper
Low	3.0720	3.072	39.2	39.7
Medium	3.0044	3.005	40.0	40.3
High	2.9742	2.975	40.3	40.4
Extra	2.9619	2.963	41.1	41.2
Max	2.9621	2.963	41.0	41.3
Avg.	2.9949	2.996	40.3	40.6

This checkpoint: the final validation loss of each sub-network of these weights (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 MatFormer at the 375M–906M scale (Table 2). Avg.: the mean over the five sub-networks. Appendix G of the paper gives the per-task accuracies.

Architecture

MatFormer (Devvrit et al., 2024) keeps the model width fixed and nests only the FFN hidden width (Section 3.2). This baseline has the settings of Appendices D and E.1 (Tables 6, 7 and 8):

Nesting	model width D_model = 1,024 in every sub-network; sub-network s uses the first s·h0 FFN hidden units, with base FFN width h0 = 128 and S = {8, 17, 27, 38, 51} (FFN widths 1,024, 2,176, 3,456, 4,864, 6,528)
Shared at full width	attention (8 heads of size d_head = 128), normalization, the embedding and a single readout U, scaled by one learned scalar α (Appendix D)
FFN	SwiGLU; the down projection multiplies its output by 1/e, where e = h0/D_model = 1/8 takes the place of the FFN expansion (Table 6)
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)

model.safetensors stores every hidden matrix as the two tensors of Equation 23, the frozen initialization init (W^init) and the trained delta delta (Δ); the token embedding, the norm gains and the readout are single tensors. With the multipliers of Table 6, sub-network s uses W^init + Δ for the attention matrices and for the FFN up and gate projections (their first s·h0 rows), s^−1/2·W^init[:, :s·h0] + s^−1·Δ[:, :s·h0] for the FFN down projection, and α·U as its readout. With i = 0, …, 31 numbering the layers:

Tensor	Shape
embed_tokens.weight	[50432, 1024]
layers.<i>.{attn_norm,mlp_norm}.weight, layers.<i>.attn.{q,k}_norm.weight, norm.weight	[1024]
layers.<i>.attn.{q,k,v,gate,o}_proj.{init,delta}	[1024, 1024]
layers.<i>.mlp.{gate,up}_proj.{init,delta}	[6528, 1024]
layers.<i>.mlp.down_proj.{init,delta}	[1024, 6528]
lm_head.weight	[50432, 1024]: U
lm_head.logit_scale	[1]: α

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 tokenized training and validation data of these runs are in the dataset repository Can111/ShellElastic-data.

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

Hyperparameters (Section 4.1, Appendices D and E.1, and Table 9):

Setting	Value
Parameterization	µ-Elastic (Section 3.2, Table 6) with d0 = 128 and m = D_model/d0 = 8
Optimizer	AdamW with decoupled weight decay
Base learning rate η0	2×10⁻³
Base initialization σ0	3.3×10⁻³, truncated normal at ±3 standard deviations
Base weight decay λ0	0.01; none on norm gains and α
β1, β2	0.95, 0.95
AdamW ε	10⁻¹⁴
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-MatFormer-375M-906M", local_dir="ShellElastic-baseline-MatFormer-375M-906M")) and run, from that directory:

import torch
from modeling_elastic import ElasticForCausalLM

model = ElasticForCausalLM.from_pretrained(".")       # float32 on the CPU
# model = ElasticForCausalLM.from_pretrained(".", 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, tier="medium")                    # [1, 5, 50432], float32; tiers: low, medium, high, extra, max
logits, loss = model(ids[:, :-1], tier="low", labels=ids[:, 1:])   # mean next-token cross-entropy
out = model.generate(ids, max_new_tokens=20, tier="max")           # greedy decoding (no key/value cache)
weights = model.extract("low")                        # the Low sub-network as a dense model (state dict)
model.save_subnetwork("low", "low_model")             # its config.json and model.safetensors

extract(tier) returns a sub-network in the format of the released Standalone models: an ordinary dense decoder with the multipliers of the sub-network applied, the FFN divisor e folded into down_proj and the learned readout scalar into lm_head. save_subnetwork writes it with its config.json; the folder loads with modeling_dense.py from the Standalone repositories (for example Can111/ShellElastic-baseline-Standalone-375M-906M). python example.py --tier medium runs a sub-network on the CPU and extracts it (--save DIR writes the extracted model).

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 sub-network 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_elastic.py: the reference implementation (PyTorch only)
• example.py: loads the supernet on the CPU and runs one sub-network and extracts it
• config.json: the architecture of the supernet
• model.safetensors: the weights of the supernet (float32)

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}
}