DOCs/Docs/TRAINING.md

# Andromeda Model Training Standard Operating Procedure

This document provides instructions on how to train the Andromeda model end-to-end using the provided code. The training procedure consists of three main scripts: `build_dataset.py`, `model.py`, and `train_distributed.py`. Follow the steps below to train the Andromeda model.

## Prerequisites

Before starting the training process, ensure that you have the following requirements:

- Python 3.7 or higher
- PyTorch 1.9 or higher
- Transformers library
- Datasets library
- Accelerate library
- Wandb library (optional, for logging)

## Step 1: Building the Dataset

The first step is to build the dataset required for training. The `build_dataset.py` script processes the training data and prepares it for training. Follow the instructions below to build the dataset:

1. Open the `build_dataset.py` script.
2. Set the configuration parameters in the `CFG` class according to your requirements:
   - `HF_ACCOUNT_REPO`: Replace with your Hugging Face API key.
   - `TOKENIZER`: Choose the tokenizer model to use (e.g., "EleutherAI/gpt-neox-20b").
   - `DATASET_NAME`: Choose the dataset to process (e.g., "tiiuae/falcon-refinedweb").
   - `SEQ_LEN`: Set the desired sequence length.
3. Save the changes to the script.
4. Open a terminal or command prompt and navigate to the directory containing the `build_dataset.py` script.
5. Run the following command to execute the script:
   ```
   python build_dataset.py
   ```
6. The script will process the dataset and push it to your Hugging Face account repository specified by `HF_ACCOUNT_REPO`.

## Step 2: Defining the Andromeda Model

The second step is to define the Andromeda model architecture. The `model.py` script contains the model definition and configuration. Follow the instructions below to configure the Andromeda model:

1. Open the `model.py` script.
2. Set the configuration parameters in the `AndromedaTokenizer` and `Andromeda` classes according to your requirements:
   - `tokenizer`: Configure the tokenizer with the desired parameters.
   - `Andromeda`: Configure the Andromeda model with the desired architecture.
3. Save the changes to the script.

## Step 3: Training the Andromeda Model

The final step is to train the Andromeda model using the `train_distributed.py` script. Follow the instructions below to start the training process:

1. Open the `train_distributed.py` script.
2. Set the configuration parameters in the `TrainAndromeda.CFG` class according to your requirements:
   - `BATCH_SIZE`: Set the batch size for training.
   - `GRADIENT_ACCUMULATE_EVERY`: Set the number of gradient accumulation steps.
   - `LEARNING_RATE`: Set the learning rate for the optimizer.
   - `WEIGHT_DECAY`: Set the weight decay for the optimizer.
   - `SEQ_LEN`: Set the desired sequence length.
   - `USE_DEEPSPEED`: Set to `True` if using DeepSpeed for optimization.
   - `USE_FSDP`: Set to `True` if using Fully Sharded Data Parallelism.
   - `USE_PRETOKENIZED`: Set to `True` if using a pre-tokenized dataset.
   - `USE_ACTIVATION_CHECKPOINTING`: Set to `True` if using activation checkpointing.
   - `RESUME_FROM_CHECKPOINT`: Set to the path of a checkpoint to resume training from.
   - `CHECKPOINTING_STEPS`: Set the number of steps between checkpoints.
   - `OUTPUT_DIR`: Set the output directory for saving the model checkpoints and logs.
   - `ENTITY_NAME`: Set the Wandb entity name for logging (optional).
3. Save the changes to the script.
4. Open a terminal or command prompt and navigate to the directory containing the `train_distributed.py` script.
5. Run the following command to start the training:
   ```
   python train_distributed.py
   ```
6. The script will train the Andromeda model using the specified configuration and dataset.
7. During training, the progress will be displayed in the terminal, and logs will be saved to the specified output directory.

# Other Training methods

First:

`Accelerate Config`

Enable Deepspeed 3: 

`Accelerate launch train_distributed_accelerate.py`