Can111/ShellElastic-baseline-FullyElastic-119M-285M

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.

Fully elastic baseline, 119M–285M scale

Final weights of the Fully elastic baseline at the 119M–285M scale of the paper (Section 4.1). These weights are a baseline trained for this paper, not an official release of slimmable networks (Yu et al., 2018), of which Fully elastic is a transformer adaptation. Fully elastic nests the residual, attention and FFN widths by slicing shared dense matrices (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	384	6	1,536	118,997,381 (119.0M)
Medium	448	7	1,792	154,427,845 (154.4M)
High	512	8	2,048	194,314,757 (194.3M)
Extra	576	9	2,304	238,658,117 (238.7M)
Max	640	10	2,560	287,457,925 (287.5M)

Parameter counts include the five learned readout scalars α_s; 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 α_s, and reports 5 fewer. model.safetensors stores the supernet: 542,556,805 parameters (2.17 GB in float32).

Results

Sub-network	Val. loss, this checkpoint	Val. loss, paper	5-shot acc. (%), this checkpoint	5-shot acc. (%), paper
Low	3.4298	3.429	34.6	34.9
Medium	3.3549	3.354	36.0	35.5
High	3.3162	3.315	36.1	35.7
Extra	3.2970	3.295	36.1	35.9
Max	3.2914	3.290	36.4	36.0
Avg.	3.3379	3.336	35.9	35.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 25,000) and its mean 5-shot accuracy over the 13 downstream tasks of Appendix F. Paper: the main results of the paper for Fully elastic at the 119M–285M scale (Table 2). Avg.: the mean over the five sub-networks. Appendix G of the paper gives the per-task accuracies.

Architecture

Fully elastic models, a transformer adaptation of slimmable networks (Yu et al., 2018), slice shared dense matrices along the residual, attention and FFN widths (Section 3.2). This baseline has the settings of Appendices D and E.1 (Tables 5, 7 and 8):

Nesting	base width g = 64 and S = {6, 7, 8, 9, 10}: sub-network s has model width s·g (384 to 640) and FFN width e·s·g, and reads the top-left slice of every hidden matrix (Appendix D)
Attention	heads of size d_head = 64; sub-network s uses s·g/d_head of them (6 to 10)
Normalization	RMSNorm over the active width (Table 8)
FFN	SwiGLU with expansion e = 4; the down projection multiplies its output by 1/e (Table 5)
Embedding and readout	the embedding is sliced (sub-network s uses its first s·g columns); one shared readout, of which sub-network s reads the first s·g columns, with a learned scalar α_s per sub-network (Appendix D)
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 and the shared readout as the two tensors of Equation 23, the frozen initialization init (W^init) and the trained delta delta (Δ); the token embedding and the norm gains are single tensors. Sub-network s uses the top-left slice of every hidden matrix as W_s = s^−1/2·W^init[slice] + s^−1·Δ[slice] (Equation 23; query, key and value are separate tensors), and as its readout α_s·(s^−1/2·W^init[:, :s·g] + s^−1·Δ[:, :s·g]) of the readout tensors (Table 5). With i = 0, …, 31 numbering the layers:

Tensor	Shape
embed_tokens.weight	[50432, 640]
layers.<i>.{attn_norm,mlp_norm}.weight, layers.<i>.attn.{q,k}_norm.weight, norm.weight	[640]
layers.<i>.attn.{q,k,v,gate,o}_proj.{init,delta}	[640, 640]
layers.<i>.mlp.{gate,up}_proj.{init,delta}	[2560, 640]
layers.<i>.mlp.down_proj.{init,delta}	[640, 2560]
lm_head.init, lm_head.delta	[50432, 640]
lm_head.logit_scale	[5]: α_s of Low to Max

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 25,000 steps of 128 sequences, 6.55B tokens (Section 4.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 5) with d0 = 128 and m = g/d0 = 1/2
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 α_s
β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 20,000 to step 25,000
Batch size	128 sequences (262,144 tokens)
Sequence length	2,048
Training steps	25,000 (6.55B 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-FullyElastic-119M-285M", local_dir="ShellElastic-baseline-FullyElastic-119M-285M")) 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-119M-285M). 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}
}