https://github.com/kempnerinstitute/tmrc
Transformer Model Research Codebase (TMRC)
https://github.com/kempnerinstitute/tmrc
Last synced: 8 months ago
JSON representation
Transformer Model Research Codebase (TMRC)
- Host: GitHub
- URL: https://github.com/kempnerinstitute/tmrc
- Owner: KempnerInstitute
- License: mit
- Created: 2024-08-01T19:35:21.000Z (almost 2 years ago)
- Default Branch: develop
- Last Pushed: 2025-01-27T02:28:46.000Z (over 1 year ago)
- Last Synced: 2025-01-27T03:23:16.793Z (over 1 year ago)
- Language: Python
- Homepage: https://kempnerinstitute.github.io/tmrc/
- Size: 5.31 MB
- Stars: 2
- Watchers: 3
- Forks: 0
- Open Issues: 14
-
Metadata Files:
- Readme: README.md
- License: LICENSE
Awesome Lists containing this project
README
# TMRC
_Transformer model research codebase_
TMRC (Transformer Model Research Codebase) is a simple, explainable codebase to train transformer-based models. It was developed with simplicity and ease of modification in mind, particularly for researchers. The codebase will eventually be used to train foundation models and experiment with architectural and training modifications.
## Documentation
[TMRC Documentation](https://kempnerinstitute.github.io/tmrc/)
## Installation
- Step 1: Load required modules
If you are using the Kempner AI cluster, load required modules:
```bash
module load python/3.12.5-fasrc01
module load cuda/12.4.1-fasrc01
module load cudnn/9.1.1.17_cuda12-fasrc01
```
If you are not using the Kempner cluster, install torch and cuda dependencies following instructions on the [PyTorch website](https://pytorch.org). TMRC has been tested with torch `2.5.0+cu124` and Python `3.12`.
- Step 2: Create a Conda environment
```bash
conda create -n tmrc_env python=3.12
conda activate tmrc_env
```
- Step 3: Clone the repository
```bash
git clone git@github.com:KempnerInstitute/tmrc.git
```
- Step 4: Install the package
```bash
cd tmrc
pip install poetry
poetry install
```
## Running Experiments
- Step 1: Login to Weights & Biases to enable experiment tracking
```bash
wandb login
```
### Single-GPU Training
- Step 2: Request compute resources. For example, on the Kempner AI cluster, to request an H100 80GB GPU run
```bash
salloc --partition=kempner_h100 --account= --nodes=1 --ntasks=1 --cpus-per-task=24 --mem=375G --gres=gpu:1 --time=00-07:00:00
```
If you are not using the Kempner AI cluster, you can run experiments on your local machine (if you have a GPU) or on cloud services like AWS, GCP, or Azure. TMRC should automatically find the available GPU. If there are no GPUs available, it will run on CPU (though this is not recommended, since training will be prohibitively slow for any reasonable model size).
- Step 3: Activate the Conda environment
```bash
conda activate tmrc_env
```
- Step 4: Launch training
```bash
python src/tmrc/core/training/train.py
```
### Multi-node multiple-GPU Training
- Step 2: Request compute resources. For example, on the Kempner AI cluster, to request eight H100 80GB GPUs on two nodes run
```bash
salloc --partition=kempner_h100 --account= --nodes=2 --ntasks-per-node=4 --ntasks=8 --cpus-per-task=24 --mem=375G --gres=gpu:4 --time=00-07:00:00
```
- Step 3: Activate the Conda environment
```bash
conda activate tmrc_env
```
- Step 4: Launch training
```bash
srun python src/tmrc/core/training/train.py
```
> [!NOTE]
> For distributed training, TMRC uses `Distributed Data Parallelism (DDP)` by default. For larger models, to use `Fully Sharded Data Parallelism (FSDP)`, set `distributed_strategy` to `fsdp` in the `training` part of the config file or see the next section on how to have a custom config file.
### Configuration
By default, the training script uses the configuration defined in `configs/training/default_train_config.yaml`.
To use a custom configuration file
python src/tmrc/core/training/train.py --config-name YOUR_CONFIG
> [!NOTE]
> The `--config-name` parameter should be specified without the `.yaml` extension.
> [!TIP]
> Configuration files should be placed in the `configs/training/` directory. For example, if your config is named `my_experiment.yaml`, use `--config-name my_experiment`
Make sure to change the `path` under `datasets` block in the config file.
## Build the documentation locally
- Step 1: Install the required packages
```bash
poetry install --with dev
```
- Step 2: Build the documentation
```bash
cd docs
make html
```
- Step 3: Open the documentation in your browser
```bash
open _build/html/index.html
```