--- license: cc-by-4.0 language: - en datasets: - mlfoundations/dclm-baseline-1.0 - Can111/ShellElastic-data library_name: pytorch pipeline_tag: text-generation tags: - elastic-models - foundation-models - hyperparameter-transfer - maximal-update-parameterization - pretraining ---
ShellElastic

Scalable Elastic Pretraining
with Simultaneous Maximal Updates and Hyperparameter Transfer

Can Jin1* · Yan Kang2 · Yifan Gong2 · Chongjian Ge2 · Mingze Xu2 · Yuanjun Xiong2 · Dimitris N. Metaxas1 · Hongwu Peng2

1Rutgers University   2Adobe Research

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

Paper Project Page Code Checkpoints Datasets

## 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..{attn_norm,mlp_norm}.weight`, `layers..attn.{q,k}_norm.weight`, `norm.weight` | [1024] | | `layers..attn.{q,k,v,gate,o}_proj.{init,delta}` | [1024, 1024] | | `layers..mlp.{gate,up}_proj.{init,delta}` | [6528, 1024] | | `layers..mlp.down_proj.{init,delta}` | [1024, 6528] | | `lm_head.weight` | [50432, 1024]: U | | `lm_head.logit_scale` | [1]: α | ## Training **Data.** [DCLM-Baseline](https://huggingface.co/datasets/mlfoundations/dclm-baseline-1.0) (Li et al., 2024), a filtered Common Crawl corpus, with the [GPT-NeoX-20B tokenizer](https://huggingface.co/EleutherAI/gpt-neox-20b) (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](https://huggingface.co/datasets/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: ```python 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](https://huggingface.co/EleutherAI/gpt-neox-20b), 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](https://creativecommons.org/licenses/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](https://huggingface.co/datasets/Can111/ShellElastic-data). The GPT-NeoX-20B tokenizer (Apache-2.0) is not bundled here. ## Citation ```bibtex @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} } ```