Improve training setup and continued pretraining docs

README.md

@@ -28,6 +28,7 @@ In this repo, you'll find code for:
   * [Loguru](#loguru)
   * [Weights & Biases](#weights--biases)
 - [Experimental: Retrieval-Augmented Models](#experimental-retrieval-augmented-models)
+- [Dataset Licensing](#dataset-licensing)
 - [See Also](#see-also)
 - [License](#license)
 - [Citing OpenChatKit](#citing-openchatkit)
@@ -70,6 +71,20 @@ conda install mamba -n base -c conda-forge
 mamba env create -f environment.yml 
 ```
 
+The default environment targets Linux hosts with NVIDIA GPUs. If the solver fails on package conflicts, make sure conda's strict channel priority is disabled before creating the environment:
+
+```shell
+conda config --set channel_priority false
+```
+
+The Llama 2 32K training modules require FlashAttention. After activating the environment, install a FlashAttention build that matches your CUDA and PyTorch versions, for example:
+
+```shell
+pip install flash-attn --no-build-isolation
+```
+
+If you do not have an NVIDIA GPU or CUDA-capable Linux environment, use the inference-only paths or a platform-specific PyTorch environment instead of the root training environment.
+
 6. Activate the new conda environment.
 
 ```shell
@@ -142,6 +157,18 @@ As the training loop runs, checkpoints are saved to the `model_ckpts` directory
 
 Please see [the training README](training/README.md) for more details about customizing the training run.
 
+## Continued pre-training
+
+For continued pre-training from the prepared Llama-2-7B-32K-beta checkpoint, use `training/continued_pretrain_llama-2-7b-32k.sh`. The script uses the same distributed training entrypoint as the fine-tuning examples and can be pointed at any local or remote JSONL dataset with a `text` field:
+
+```shell
+PRETRAIN_DATASETS=/path/to/corpus.jsonl:1 \
+PRETRAIN_TOTAL_STEPS=1000 \
+bash training/continued_pretrain_llama-2-7b-32k.sh
+```
+
+Multiple datasets can be supplied as a comma-separated list with sampling weights, for example `dataset_a.jsonl:0.8,dataset_b.jsonl:0.2`. The default sequence length is 32K, matching the long-context model; lower `PRETRAIN_SEQ_LENGTH`, `PRETRAIN_BATCH_SIZE`, or `PRETRAIN_GRADIENT_ACCUMULATE_STEP` if your hardware cannot fit the default configuration.
+
 ## Converting trained weights to Hugging Face format
 
 Before you can use this model to perform inference, it must be converted to the Hugging Face format. Run this command from the root of the repo to do so.
@@ -311,6 +338,14 @@ Zurich is located in Switzerland.
 >>>
 ```
 
+# Dataset Licensing
+
+This repository's code is Apache 2.0 licensed. Dataset licensing is separate from code licensing and should be checked at the dataset source before training or redistribution.
+
+The OIG dataset used by the chat-model examples is published on Hugging Face as `laion/OIG` with an Apache 2.0 dataset license. The LAION OIG dataset card also notes that OIG combines components from multiple sources, including permissively licensed data, CC-BY-SA data such as Wikipedia-derived content, and web-crawled data. Users should inspect the specific files they use and comply with the upstream license or usage terms for each component.
+
+The Llama 2 32K examples use files from `togethercomputer/Long-Data-Collections`. Those files are downloaded from Hugging Face at training time, so users should review that dataset card and any referenced source licenses before using the data commercially or redistributing derivatives.
+
 # See Also
 * [docs/GPT-NeoXT-Chat-Base-20B.md](docs/GPT-NeoXT-Chat-Base-20B.md). OpenChatKit also provides a larger, 20B parameter chat model that was trained on GPT-NeoXT-Chat-Base-20B from Eleuther AI.
 

environment.yml

@@ -8,21 +8,21 @@ dependencies:
   - cudatoolkit=11.8.0
   - cupy=12.1.0
   - faiss-gpu=1.7.2
-  - fastparquet=0.5.0
+  - fastparquet=2023.7.0
   - nccl=2.18.3.1
   - pip=23.2
   - pyarrow=12.0.1
   - python=3.10.9
   - python-snappy=0.6.1
   - pytorch=2.0.1
   - pytorch-cuda=11.8
-  - snappy=1.1.9
+  - snappy=1.1.10
   - torchaudio=2.0.2
   - torchvision=0.15.2
   - pip:
       - accelerate==0.21.0
       - boto3
-      - datasets==2.13.1
+      - datasets==2.15.0
       - loguru==0.6.0
       - netifaces==0.11.0
       - pandas==2.0.3

training/README.md

@@ -4,6 +4,23 @@ This directory contains code for training a chat model using OpenChatKit. The ma
 
 To customize training, make a copy of the script and modify the arguments.
 
+## Continued Pre-training
+
+`continued_pretrain_llama-2-7b-32k.sh` is a template for continuing training from the prepared `Llama-2-7B-32K-beta` checkpoint before any downstream fine-tuning. It uses the same `dist_clm_train.py` entrypoint as the Llama fine-tuning scripts, but exposes the dataset, step count, learning rate, sequence length, and batch settings through environment variables.
+
+Example:
+
+```bash
+PRETRAIN_DATASETS=/path/to/corpus.jsonl:1 \
+PRETRAIN_TOTAL_STEPS=1000 \
+PRETRAIN_CHECKPOINT_STEPS=100 \
+bash training/continued_pretrain_llama-2-7b-32k.sh
+```
+
+The dataset must be JSONL readable by Hugging Face `datasets` with a `text` field. Multiple datasets can be passed as a comma-separated list with sampling weights, such as `corpus_a.jsonl:0.75,corpus_b.jsonl:0.25`.
+
+The template defaults to 32K sequence length and an eight-stage pipeline. Override `PRETRAIN_SEQ_LENGTH`, `PRETRAIN_BATCH_SIZE`, `PRETRAIN_MICRO_BATCH_SIZE`, and `PRETRAIN_GRADIENT_ACCUMULATE_STEP` when adapting it to smaller hardware.
+
 ## Arguments
 
 Environment vars that should be set:

training/continued_pretrain_llama-2-7b-32k.sh

@@ -0,0 +1,94 @@
+#!/usr/bin/env bash
+
+DIR=$(cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd)
+
+netif=${NET_IFACE:-lo}
+export GLOO_SOCKET_IFNAME=${netif}
+export NCCL_SOCKET_IFNAME=${netif}
+export MODEL_NAME=${MODEL_NAME:-llama-2-7b-32k-continued-pretrain}
+
+export SHOW_DATA=${SHOW_DATA:-1}
+
+BASE_MODEL=${BASE_MODEL:-"${DIR}/../pretrained/Llama-2-7B-32K-beta/togethercomputer_Llama-2-7B-32K-beta"}
+
+TOTAL_STEPS=${PRETRAIN_TOTAL_STEPS:-1000}
+CHECKPOINT_STEPS=${PRETRAIN_CHECKPOINT_STEPS:-100}
+CHECKPOINT_PATH=${PRETRAIN_CHECKPOINT_PATH:-"${DIR}/../model_ckpts/${MODEL_NAME}"}
+SEQ_LENGTH=${PRETRAIN_SEQ_LENGTH:-32768}
+BATCH_SIZE=${PRETRAIN_BATCH_SIZE:-4}
+MICRO_BATCH_SIZE=${PRETRAIN_MICRO_BATCH_SIZE:-1}
+GRADIENT_ACCUMULATE_STEP=${PRETRAIN_GRADIENT_ACCUMULATE_STEP:-1}
+LR=${PRETRAIN_LR:-1e-5}
+
+DATASETS=${PRETRAIN_DATASETS:-"${DIR}/../data/OIG/files/unified_oscar_en_sample_dialog.jsonl:1"}
+
+validate_cli_arg() {
+    local name=$1
+    local value=$2
+
+    if [[ "${value}" == *[$'\n\r']* || "${value}" =~ [\;\&\|\<\>\`\$\\] ]]; then
+        printf 'Invalid %s: shell metacharacters and control characters are not allowed.\n' "${name}" >&2
+        exit 2
+    fi
+}
+
+validate_cli_arg PRETRAIN_DATASETS "${DATASETS}"
+
+ARGS=(
+    --model-name "${BASE_MODEL}"
+    --tokenizer-name "${BASE_MODEL}"
+    --project-name together
+    --model-type llama
+    --optimizer adam
+    --seed 42
+    --load-pretrained-model true
+    --task-name "${DATASETS}"
+    --checkpoint-path "${CHECKPOINT_PATH}"
+    --total-steps "${TOTAL_STEPS}"
+    --warmup-steps 0
+    --train-warmup-steps 0
+    --checkpoint-steps "${CHECKPOINT_STEPS}"
+    --lr "${LR}"
+    --seq-length "${SEQ_LENGTH}"
+    --batch-size "${BATCH_SIZE}"
+    --micro-batch-size "${MICRO_BATCH_SIZE}"
+    --gradient-accumulate-step "${GRADIENT_ACCUMULATE_STEP}"
+    --dist-url tcp://127.0.0.1:7033
+    --num-layers 4
+    --embedding-dim 4096
+    --world-size 8
+    --pipeline-group-size 8
+    --data-group-size 1
+    --job-id 0
+    --net-interface "${netif}"
+    --fp16
+    --dp-backend nccl
+    --dp-mode allreduce
+    --pp-mode gpipe
+    --profiling no-profiling
+)
+
+run_rank() {
+    local cuda_id=$1
+    local rank=$2
+    python "${DIR}/dist_clm_train.py" "${ARGS[@]}" --cuda-id "${cuda_id}" --rank "${rank}"
+}
+
+(trap 'kill 0' SIGINT; \
+run_rank 0 0 \
+    & \
+run_rank 1 1 \
+    & \
+run_rank 2 2 \
+    & \
+run_rank 3 3 \
+    & \
+run_rank 4 4 \
+    & \
+run_rank 5 5 \
+    & \
+run_rank 6 6 \
+    & \
+run_rank 7 7 \
+    & \
+wait)