Hugging Face's logo

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

Text Generation
PyTorch
Safetensors
English
elastic
elastic-models
foundation-models
hyperparameter-transfer
maximal-update-parameterization
pretraining
Model card Files
xet
 Community
ShellElastic-baseline-MatFormer-375M-906M / README.md
Can111's picture
Can111
Update tags
bb1501d verified
preview code
|
Raw History Blame Contribute Delete
13.8 kB
---
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
---
<div align="center">
<div align="center" class="dark:hidden"><img alt="ShellElastic" src="https://huggingface.co/spaces/Can111/ShellElastic/resolve/main/shellelastic-logo-light.svg" height="42"></div>
<div align="center" class="hidden dark:block"><img alt="ShellElastic" src="https://huggingface.co/spaces/Can111/ShellElastic/resolve/main/shellelastic-logo-dark.svg" height="42"></div>
<h3 align="center">Scalable Elastic Pretraining<br>with Simultaneous Maximal Updates and Hyperparameter Transfer</h3>
<p align="center"><small><b>Can&nbsp;Jin</b><sup>1*</sup> · <b>Yan&nbsp;Kang</b><sup>2</sup> · <b>Yifan&nbsp;Gong</b><sup>2</sup> · <b>Chongjian&nbsp;Ge</b><sup>2</sup> · <b>Mingze&nbsp;Xu</b><sup>2</sup> · <b>Yuanjun&nbsp;Xiong</b><sup>2</sup> · <b>Dimitris&nbsp;N.&nbsp;Metaxas</b><sup>1</sup> · <b>Hongwu&nbsp;Peng</b><sup>2</sup></small></p>
<p align="center"><small><sup>1</sup>Rutgers University &emsp; <sup>2</sup>Adobe Research</small></p>
<p align="center"><small><sup>*</sup>Corresponding author: <a href="mailto:can.jin@rutgers.edu">can.jin@rutgers.edu</a>. Work done during an internship at Adobe Research.</small></p>
<p align="center">
<a href="https://arxiv.org/pdf/2610.15124"><img alt="Paper" src="https://img.shields.io/badge/Paper-PDF-b31b1b?style=for-the-badge"></a>
<a href="https://jincan333.github.io/ShellElastic_page"><img alt="Project Page" src="https://img.shields.io/badge/Project-Page-2f6db5?style=for-the-badge"></a>
<a href="https://github.com/jincan333/ShellElastic"><img alt="Code" src="https://img.shields.io/badge/Code-GitHub-181717?style=for-the-badge&logo=github&logoColor=white"></a>
<a href="https://huggingface.co/spaces/Can111/ShellElastic"><img alt="Checkpoints" src="https://img.shields.io/badge/Checkpoints-Hugging%20Face-ffd21e?style=for-the-badge&logo=huggingface&logoColor=white"></a>
<a href="https://huggingface.co/spaces/Can111/ShellElastic"><img alt="Datasets" src="https://img.shields.io/badge/Datasets-Hugging%20Face-ffd21e?style=for-the-badge&logo=huggingface&logoColor=white"></a>
</p>
</div>
## 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](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}
}
```