├── .gitignore ├── .pre-commit-config.yaml ├── LICENSE ├── MANIFEST.in ├── README.md ├── abacus.py ├── arithmetic_eval_quicker.py ├── cramming ├── __init__.py ├── architectures │ ├── __init__.py │ ├── attention.py │ ├── components.py │ ├── construction.py │ ├── crammed_depthrecurrent.py │ ├── crammed_transformer.py │ ├── embeddings.py │ ├── huggingface_interface.py │ ├── losses.py │ └── sanity_check.py ├── backend │ ├── __init__.py │ ├── optimizers │ │ ├── __init__.py │ │ ├── optimizer_modifiers.py │ │ ├── progressive_batching.py │ │ └── schedulers.py │ ├── prepare_backend.py │ ├── torch_default.py │ └── utils.py ├── config │ ├── __init__.py │ ├── arch │ │ ├── __init__.py │ │ ├── albert.yaml │ │ ├── crammed-depthrecurrent.yaml │ │ ├── crammed-fakeRNN.yaml │ │ ├── crammed-janus.yaml │ │ ├── crammed-rnn.yaml │ │ ├── crammed-stack-janus.yaml │ │ ├── crammed-tiny.yaml │ │ ├── crammed-transformer.yaml │ │ ├── gpt2-base.yaml │ │ ├── hf-gpt2.yaml │ │ └── sanitycheck.yaml │ ├── cfg_eval.yaml │ ├── cfg_pretrain.yaml │ ├── data │ │ ├── __init__.py │ │ ├── arithmetic.yaml │ │ ├── c4-subset-processed.yaml │ │ ├── openweb.yaml │ │ ├── proofpile.yaml │ │ ├── sanity-check-1.yaml │ │ ├── sanity-check-2.yaml │ │ └── sources │ │ │ ├── ag_news.yaml │ │ │ ├── arithmetic.yaml │ │ │ ├── bookcorpus.yaml │ │ │ ├── c4.yaml │ │ │ ├── dash_books.yaml │ │ │ ├── fake.yaml │ │ │ ├── iwslt.yaml │ │ │ ├── local.yaml │ │ │ ├── no_code_stackexchange.yaml │ │ │ ├── openwebtext.yaml │ │ │ ├── oscar.yaml │ │ │ ├── proofpiledata.yaml │ │ │ ├── the_pile.yaml │ │ │ ├── the_pileCC.yaml │ │ │ ├── the_pile_dedup.yaml │ │ │ ├── the_pile_natural.yaml │ │ │ ├── the_pile_stream.yaml │ │ │ ├── uncorpus.yaml │ │ │ ├── uspto.yaml │ │ │ ├── wikibooks.yaml │ │ │ ├── wikinews.yaml │ │ │ ├── wikipedia.yaml │ │ │ ├── wikiquote.yaml │ │ │ ├── wikiversity.yaml │ │ │ └── wikivoyage.yaml │ ├── eval │ │ ├── __init__.py │ │ ├── pythia.yaml │ │ └── tasks │ │ │ ├── lambada_openai.yaml │ │ │ └── winogrande.yaml │ ├── hydra │ │ ├── __init__.py │ │ └── job_logging │ │ │ └── custom.yaml │ ├── impl │ │ ├── __init__.py │ │ ├── _default.yaml │ │ └── torch-default.yaml │ ├── train │ │ ├── __init__.py │ │ ├── common.yaml │ │ ├── cramming.yaml │ │ ├── janus-regime.yaml │ │ ├── optim │ │ │ ├── adafactor.yaml │ │ │ ├── adahessian.yaml │ │ │ ├── adam.yaml │ │ │ ├── adam8bit.yaml │ │ │ ├── adam_classic.yaml │ │ │ ├── adamscale.yaml │ │ │ ├── agd.yaml │ │ │ ├── lion.yaml │ │ │ ├── radam.yaml │ │ │ ├── sgd.yaml │ │ │ └── shampoo.yaml │ │ └── optim_mod │ │ │ ├── disabled.yaml │ │ │ ├── larc.yaml │ │ │ ├── lars.yaml │ │ │ ├── progressive.yaml │ │ │ └── sam.yaml │ └── wandb │ │ ├── default.yaml │ │ └── none.yaml ├── data │ ├── __init__.py │ ├── arithmetic_tokenizers.py │ ├── curriculum_sorting.py │ ├── deduplicate.py │ ├── pretraining_preparation.py │ ├── tokenizer_preparation.py │ └── utils.py └── utils.py ├── create_data_split.py ├── create_pos_or_variants.py ├── dataset_analysis.py ├── gen_eval_script.py ├── load_local_model.py ├── pretrain.py ├── pretty_plotter.py ├── pretty_plotter_big.py ├── pretty_plotter_sort.py ├── pyproject.toml ├── setup.cfg ├── shells ├── addition_ff.sh ├── addition_lt.sh ├── bitwise_or.sh ├── evaluation.sh ├── generate_and_tokenize_data.sh ├── multiplication.sh └── sorting.sh ├── sort_eval.py └── upload_processed_dataset.py /.gitignore: -------------------------------------------------------------------------------- 1 | outputs 2 | tables/*/*.csv 3 | tables/*/*.csv# 4 | tables/*.csv 5 | tables/*.csv# 6 | tables/*.ods 7 | *.png 8 | *.pdf 9 | 10 | # torchdynamo debug 11 | isolate 12 | repro.py 13 | 14 | checkpoints 15 | wandb-metadata.json 16 | 17 | torch_compile_debug/ 18 | 19 | dedup 20 | 21 | .vs/ 22 | 23 | *.pdf 24 | images 25 | 26 | *.temp.sh 27 | 28 | # Byte-compiled / optimized / DLL files 29 | __pycache__/ 30 | *.py[cod] 31 | *$py.class 32 | 33 | # C extensions 34 | *.so 35 | 36 | # Distribution / packaging 37 | .Python 38 | build/ 39 | develop-eggs/ 40 | dist/ 41 | downloads/ 42 | eggs/ 43 | .eggs/ 44 | lib/ 45 | lib64/ 46 | parts/ 47 | sdist/ 48 | var/ 49 | wheels/ 50 | pip-wheel-metadata/ 51 | share/python-wheels/ 52 | *.egg-info/ 53 | .installed.cfg 54 | *.egg 55 | MANIFEST 56 | 57 | # PyInstaller 58 | # Usually these files are written by a python script from a template 59 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 60 | *.manifest 61 | *.spec 62 | 63 | # Installer logs 64 | pip-log.txt 65 | pip-delete-this-directory.txt 66 | 67 | # Unit test / coverage reports 68 | htmlcov/ 69 | .tox/ 70 | .nox/ 71 | .coverage 72 | .coverage.* 73 | .cache 74 | nosetests.xml 75 | coverage.xml 76 | *.cover 77 | *.py,cover 78 | .hypothesis/ 79 | .pytest_cache/ 80 | 81 | # Translations 82 | *.mo 83 | *.pot 84 | 85 | # Django stuff: 86 | *.log 87 | local_settings.py 88 | db.sqlite3 89 | db.sqlite3-journal 90 | 91 | # Flask stuff: 92 | instance/ 93 | .webassets-cache 94 | 95 | # Scrapy stuff: 96 | .scrapy 97 | 98 | # Sphinx documentation 99 | docs/_build/ 100 | 101 | # PyBuilder 102 | target/ 103 | 104 | # Jupyter Notebook 105 | .ipynb_checkpoints 106 | 107 | # IPython 108 | profile_default/ 109 | ipython_config.py 110 | 111 | # pyenv 112 | .python-version 113 | 114 | # pipenv 115 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 116 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 117 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 118 | # install all needed dependencies. 119 | #Pipfile.lock 120 | 121 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow 122 | __pypackages__/ 123 | 124 | # Celery stuff 125 | celerybeat-schedule 126 | celerybeat.pid 127 | 128 | # SageMath parsed files 129 | *.sage.py 130 | 131 | # Environments 132 | .env 133 | .venv 134 | env/ 135 | venv/ 136 | ENV/ 137 | env.bak/ 138 | venv.bak/ 139 | 140 | # Spyder project settings 141 | .spyderproject 142 | .spyproject 143 | 144 | # Rope project settings 145 | .ropeproject 146 | 147 | # mkdocs documentation 148 | /site 149 | 150 | # mypy 151 | .mypy_cache/ 152 | .dmypy.json 153 | dmypy.json 154 | 155 | # Pyre type checker 156 | .pyre/ 157 | 158 | *.csv 159 | *.txt 160 | *.pth 161 | 162 | cramming-data/ 163 | sanity.sh 164 | log/ 165 | del.sh 166 | del.py 167 | sort_plots/ -------------------------------------------------------------------------------- /.pre-commit-config.yaml: -------------------------------------------------------------------------------- 1 | # precommit hooks from https://github.com/ashleve/lightning-hydra-template 2 | repos: 3 | - repo: https://github.com/pre-commit/pre-commit-hooks 4 | rev: v3.4.0 5 | hooks: 6 | # list of supported hooks: https://pre-commit.com/hooks.html 7 | - id: trailing-whitespace 8 | - id: end-of-file-fixer 9 | - id: check-yaml 10 | - id: check-added-large-files 11 | - id: debug-statements 12 | - id: detect-private-key 13 | 14 | # python code formatting 15 | - repo: https://github.com/psf/black 16 | rev: 22.3.0 17 | hooks: 18 | - id: black 19 | args: [--line-length, "140", "--fast"] # ;> 20 | 21 | # yaml formatting 22 | - repo: https://github.com/pre-commit/mirrors-prettier 23 | rev: v2.3.0 24 | hooks: 25 | - id: prettier 26 | types: [yaml] 27 | 28 | # python code analysis 29 | - repo: https://github.com/PyCQA/flake8 30 | rev: 4.0.1 31 | hooks: 32 | - id: flake8 33 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2024 Sean McLeish, Jonas Geiping 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /MANIFEST.in: -------------------------------------------------------------------------------- 1 | # added by check-manifest 2 | include *.py 3 | include *.yaml 4 | recursive-include cramming *.md 5 | recursive-include cramming *.yaml 6 | global-exclude *.pyc 7 | global-exclude __pycache__ 8 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Transformers Can Do Arithmetic with the Right Embeddings! [Link to arXiv paper](https://arxiv.org/abs/2405.17399) 2 | 3 | A joint project by: Sean McLeish, Arpit Bansal, Alex Stein, Neel Jain, John Kirchenbauer, Brian R. Bartoldson, Bhavya Kailkhura, Abhinav Bhatele, Jonas Geiping, Avi Schwarzschild and Tom Goldstein 4 | 5 | 6 | 7 | This repository contains code to replicate our research. It is a fork of the language model training framework [cramming](https://github.com/JonasGeiping/cramming) edited to for a next token prediction objective. 8 | 9 | We provide a standalone implementation of Abacus Embeddings in [abacus.py](abacus.py). 10 | 11 | ## Citing Our Work 12 | To cite our work, please use this bibtex. 13 | ``` 14 | @article{mcleish2024transformers, 15 | title={Transformers Can Do Arithmetic with the Right Embeddings}, 16 | author={Sean McLeish and Arpit Bansal and Alex Stein and Neel Jain and John Kirchenbauer and Brian R. Bartoldson and Bhavya Kailkhura and Abhinav Bhatele and Jonas Geiping and Avi Schwarzschild and Tom Goldstein}, 17 | journal={arXiv preprint arXiv:2405.17399}, 18 | year={2024} 19 | } 20 | ``` 21 | 22 | # Getting Started 23 | We developed in Python 3.10.4, to install run: 24 | ``` 25 | git clone git@github.com:mcleish7/arithmetic.git 26 | cd arithmetic 27 | pip install . 28 | ``` 29 | 30 | On some machines you will need to run: 31 | 1. `pip install multiprocess -U` 32 | 2. `pip install dill -U` 33 | 3. `pip install apache-beam -U` 34 | 35 | # Arithmetic 36 | ## Datasets 37 | We release our datasets on [Google Drive](https://drive.google.com/drive/folders/1DqjCrUM1cNV7069Zl25_qBw2Px2xAw9j?usp=sharing) both in zipped format. We recommend you work with the zipped version until it is correctly placed in your file system. 38 | 39 | Alternatively, you can make your own datasets using [create_data_split.py](create_data_split.py) using the commands from [shells/generate_and_tokenize_data.sh](shells/generate_and_tokenize_data.sh). 40 | 41 | ## File Structure 42 | We recommend creating another directory `cramming-data` inside of arithmetic. This is where the models, logs and data will be stored. 43 | 44 | You can either export you cramming base directory path to your `.bashrc` or you can replace `$cramming_base_dir` manually in the provided shells. 45 | ``` 46 | cd arithmetic 47 | mkdir cramming-data 48 | echo 'export cramming_base_dir=MY_BASE_DIR' >> ~/.bashrc 49 | source ~/.bashrc 50 | ``` 51 | For example, this may look like: `echo 'export cramming_base_dir=~/arithmetic/cramming-data' >> ~/.bashrc` 52 | 53 | For example our file system looks like: 54 | ``` 55 | cramming-generative 56 | └── cramming-data 57 | ├── addition-train-one 58 | │ ├── pretrain//