DOCSEARCHbyAlgolia
Submit Docs

arithmetic Research Repository

repository·main·Indexed 41 days ago

https://github.com/mcleish7/arithmetic

A research framework for enabling Transformers to perform arithmetic tasks using specialized embeddings, such as Abacus Embeddings. It includes tools for arithmetic-based next-token prediction, training frameworks, dataset generation, and evaluation visualization scripts.

Tokens
1K
Snippets
4
Records
9
Agent score
2%

What's inside arithmetic

  1. Configure the cramming_base_dir environment variable

    main

    The project uses a base directory for storing models, logs, and data. It is recommended to create a cramming-data directory inside the arithmetic folder and export it as an environment variable named cramming_base_dir. This allows the provided shell scripts to locate your data and outputs correctly.

    You can add this to your .bashrc to make it persistent.

    cd arithmetic
mkdir cramming-data
echo 'export cramming_base_dir=~/arithmetic/cramming-data' >> ~/.bashrc
source ~/.bashrc

  2. Analyze evaluation results with pretty_plotter scripts

    main

    The repository provides specialized plotting scripts to visualize evaluation grids:

    • General/Small Grids: Use pretty_plotter.py. Edit the main function in the script to include your model name.
    • Large 100x100 Grids: Use pretty_plotter_big.py.
    • Sorting Tasks: Use pretty_plotter_sort.py, which reads individual .txt files generated during testing and merges them into a single plot.

  3. Generate and tokenize arithmetic datasets

    main
    Datasets are available on Google Drive in zipped format. If you wish to create your own datasets, use the create_data_split.py script following the patterns defined in shells/generate_and_tokenize_data.sh.

  4. Run Multi-GPU training

    main

    To perform multi-GPU training, use torchrun instead of standard python. You must also set impl.fullgraph=false to ensure compatibility.

    torchrun --nproc_per_node=<NUM_GPUS> --standalone impl.fullgraph=false <other_args>

  5. Install the arithmetic repository

    main

    The project was developed using Python 3.10.4. To install the package, clone the repository and use pip to install the current directory.

    On some machines, you may need to manually install the following dependencies to ensure compatibility:

    • multiprocess
    • dill
    • apache-beam
    git clone git@github.com:mcleish7/arithmetic.git
cd arithmetic
pip install .

# Potential additional dependencies
pip install multiprocess -U
pip install dill -U
pip install apache-beam -U

  6. Configure single-GPU checkpointing

    main

    Checkpointing is currently implemented for single-GPU training. You can specify how often to save and what to name the files using impl.save_every_n_minutes and impl.save_intermediate_model_name.

    impl.save_every_n_minutes=60 impl.save_intermediate_model_name='last'

  7. Configure positional embeddings

    main

    You can control the type of positional embeddings used during training via the arch.embedding.pos_embedding flag.

    Absolute Embeddings

    • Learned: arch.embedding.pos_embedding=learned
    • Abacus: arch.embedding.pos_embedding=abacus. You can adjust the maximum length using arch.embedding.max_abacus_len=<VALUE> (default is 100).

    Relative Embeddings

    • NoPE: arch.embedding.pos_embedding=None
    • FIRE: Set arch.embedding.pos_embedding=None and specify arch.attention.rotary_embedding="fire". To introduce randomness, set arch.attention.max_length to a value larger than your sequence length.
    • RoPE: Set arch.attention.type="self-attention" and arch.attention.rotary_embedding=true.

  8. Configure Weights & Biases (WandB) logging

    main
    To log training runs to your Weights & Biases account, modify the wandb.entity and wandb.project parameters either directly in the command line or within the configuration file at cramming/config/wandb/default.yaml.

  9. Configure training hyperparameters and architecture

    main

    Use the following command-line flags to modify training behavior:

    FlagDescription
    arch.loss_reduction=noneGives samples and tokens equal importance in the loss calculation.
    arch.throttle=TrueDivides gradients in the recurrent block by the number of recurrences.
    arch.mask_before_equals=TrueMasks tokens before the equals sign.
    arch.forward_only_model_with_skip=TrueEnables skip connections inside the recurrent block.
    impl.fullgraph=falseRequired when using Multi-GPU training with torchrun.