diff --git a/.gitignore b/.gitignore index 50b9fa06..4ad6f92f 100644 --- a/.gitignore +++ b/.gitignore @@ -5,3 +5,6 @@ build/ *.log *.report.rank* *.records.log.rank* + +__pycache__/ +/data/ diff --git a/README.md b/README.md index ed6f0dd9..abd8070b 100644 --- a/README.md +++ b/README.md @@ -50,7 +50,7 @@ Build Options: > Both options are optional and can be disabled for CPU-only builds. -## ✨ InfiniTrain Overview +## ✨ InfiniTrain Overview ### ✔ Support Matrix @@ -96,54 +96,160 @@ For example, the `llama3` example produces a binary named `llama3`. To view available runtime options: -```bash -./llama3 --help -``` - -### Getting Started - -The following examples demonstrate **LLaMA 3 supervised fine-tuning (SFT)** using InfiniTrain. - -#### Single-node Training Example - -```bash -./llama3 \ - --device cuda \ - --input_bin [training_data_path] \ - --llmc_filepath [model_path] \ - --num_iteration 10 - +```bash +./build/llama3 --help ``` -#### Multi-nodes Training Example (3D parallel) - -```bash -./infini_run \ - --nnodes=2 \ - --nproc_per_node=1 \ +### Getting Started + +#### Prepare Datasets and Weights + +Run the asset preparation script from the repository root. Prepared files are +written to `data/` by default. + +```bash +# MNIST dataset +./scripts/assets/prepare-infinitrain-assets.sh mnist + +# GPT-2 124M weights, tokenizer, and tokenized TinyShakespeare data +./scripts/assets/prepare-infinitrain-assets.sh gpt2 + +# LLaMA 3.2 1B weights and tokenized TinyShakespeare data +HF_TOKEN=hf_xxx ./scripts/assets/prepare-infinitrain-assets.sh llama3 +``` + +Preparing LLaMA requires access to the gated +`meta-llama/Llama-3.2-1B` repository. Accept its license on Hugging Face and +provide `HF_TOKEN`, or authenticate with `hf auth login`, before running the +command. The complete LLaMA preparation requires approximately 8.5 GB of free +disk space, including the downloaded checkpoint and converted FP32 weights. + +Use `DATA_DIR` to write the assets elsewhere, or prepare all supported assets +in one invocation: + +```bash +DATA_DIR=/path/to/data \ +HF_TOKEN=hf_xxx \ +./scripts/assets/prepare-infinitrain-assets.sh all +``` + +#### Model Examples + +The generated files can be passed directly to the corresponding executables: + +##### MNIST + +```bash +./build/mnist \ + --device cpu \ + --dataset data/mnist +``` + +##### GPT-2 124M + +```bash +./build/gpt2 \ + --device cuda \ + --input_bin data/gpt2/tiny_shakespeare_train.bin \ + --input_val_bin data/gpt2/tiny_shakespeare_val.bin \ + --tokenizer_bin data/gpt2/gpt2_tokenizer.bin \ + --llmc_filepath data/gpt2/gpt2_124M.bin \ + --num_iteration 10 +``` + +##### LLaMA 3.2 1B + +```bash +./build/llama3 \ + --device cuda \ + --input_bin data/llama3/tiny_shakespeare_train.bin \ + --input_val_bin data/llama3/tiny_shakespeare_val.bin \ + --llmc_filepath data/llama3/llama3.2_1B_fp32.bin \ + --num_iteration 10 +``` + +### Launch Modes + +GPT-2 and LLaMA training support both thread-based and process-based launches. +The examples below use LLaMA, but the same launch modes also apply to GPT-2. + +#### Direct Launch + +Running a model executable directly uses one process and one device by default. +Set `--nthread_per_process` to use multiple execution threads and devices in the +same process: + +```bash +./build/llama3 \ + --device cuda \ + --input_bin data/llama3/tiny_shakespeare_train.bin \ + --llmc_filepath data/llama3/llama3.2_1B_fp32.bin \ + --nthread_per_process 8 \ + --num_iteration 10 +``` + +#### Single-node Multi-process Launch + +Use `infini_run` to start multiple training processes on one node. Each process +uses one execution thread by default: + +```bash +./build/infini_run \ + --nnodes=1 \ + --nproc_per_node=8 \ + ./build/llama3 \ + --device cuda \ + --input_bin data/llama3/tiny_shakespeare_train.bin \ + --llmc_filepath data/llama3/llama3.2_1B_fp32.bin \ + --num_iteration 10 +``` + +#### Multi-node Multi-process Launch + +Run the following command on every node with the same rendezvous settings and +a distinct `node_rank`: + +```bash +./build/infini_run \ + --nnodes=2 \ + --nproc_per_node=4 \ --node_rank=[rank_id] \ --rdzv_endpoint=[master_addr]:29500 \ --rdzv_id=[job_id] \ - ./llama3 \ - --device cuda \ - --input_bin [training_data_path] \ - --llmc_filepath [model_path] \ - --num_iteration 10 \ - --nthread_per_process 8 \ - --batch_size 40 \ - --total_batch_size 10240 \ - --tensor_parallel 2 \ - --pipeline_parallel 2 \ - --sequence_parallel -``` + ./build/llama3 \ + --device cuda \ + --input_bin data/llama3/tiny_shakespeare_train.bin \ + --llmc_filepath data/llama3/llama3.2_1B_fp32.bin \ + --num_iteration 10 \ + --tensor_parallel 2 \ + --pipeline_parallel 2 \ + --sequence_parallel +``` + +`--nproc_per_node` and `--nthread_per_process` can be combined. The total +training world size is: + +```text +world_size = nnodes × nproc_per_node × nthread_per_process +``` ### Parallelism Strategies -#### Distributed Data Parallelism (DDP) - -```bash ---nthread_per_process 8 # ddp_size = nthread_per_process / (tensor_parallel × pipeline_parallel) -``` +#### Distributed Data Parallelism (DDP) + +For a direct launch with TP and PP disabled, the following starts eight +data-parallel workers in one process: + +```bash +--nthread_per_process 8 # 8-way DDP when TP=1 and PP=1 +``` + +For all launch modes, the data-parallel size is derived from the total world +size after accounting for tensor and pipeline parallelism: + +```text +data_parallel_size = world_size / (tensor_parallel × pipeline_parallel) +``` #### Tensor Parallelism (TP) @@ -210,4 +316,4 @@ Multiple parallelism strategies (DDP, TP, SP, PP) can be freely combined to scal optimizations. Integrated a CTest + GTest based testing infrastructure to strengthen the - framework's automated test workflow. \ No newline at end of file + framework's automated test workflow. diff --git a/scripts/assets/prepare-infinitrain-assets.sh b/scripts/assets/prepare-infinitrain-assets.sh new file mode 100755 index 00000000..4df1abdd --- /dev/null +++ b/scripts/assets/prepare-infinitrain-assets.sh @@ -0,0 +1,232 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Prepare datasets / weights that can be read directly by InfiniTrain examples. +# +# Usage: +# ./scripts/assets/prepare-infinitrain-assets.sh gpt2 +# ./scripts/assets/prepare-infinitrain-assets.sh llama3 +# ./scripts/assets/prepare-infinitrain-assets.sh mnist +# ./scripts/assets/prepare-infinitrain-assets.sh all +# +# Optional environment variables: +# DATA_DIR=/path/to/data +# PYTHON=python3 +# HF_TOKEN=hf_xxx +# FORCE=1 +# SKIP_LLAMA3_WEIGHTS=1 +# +# Output layout: +# data/ +# ├── gpt2/ +# │ ├── gpt2_124M.bin +# │ ├── gpt2_tokenizer.bin +# │ ├── tiny_shakespeare_train.bin +# │ └── tiny_shakespeare_val.bin +# ├── llama3/ +# │ ├── llama3.2_1B_fp32.bin +# │ ├── tiny_shakespeare_train.bin +# │ └── tiny_shakespeare_val.bin +# └── mnist/ +# ├── train-images-idx3-ubyte +# ├── train-labels-idx1-ubyte +# ├── t10k-images-idx3-ubyte +# └── t10k-labels-idx1-ubyte + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd -- "${SCRIPT_DIR}/../.." && pwd)" +DATA_DIR="${DATA_DIR:-${REPO_ROOT}/data}" +CACHE_DIR="${DATA_DIR}/.cache" +PYTHON="${PYTHON:-python3}" +FORCE="${FORCE:-0}" +SKIP_LLAMA3_WEIGHTS="${SKIP_LLAMA3_WEIGHTS:-0}" + +GPT2_DIR="${DATA_DIR}/gpt2" +LLAMA3_DIR="${DATA_DIR}/llama3" +MNIST_DIR="${DATA_DIR}/mnist" + +TARGET="${1:-all}" + +case "${TARGET}" in + gpt2|llama3|mnist|all) ;; + *) + echo "Usage: $0 {gpt2|llama3|mnist|all}" + exit 2 + ;; +esac + +mkdir -p "${CACHE_DIR}" "${GPT2_DIR}" "${LLAMA3_DIR}" "${MNIST_DIR}" + +log() { + printf '\n[%s] %s\n' "$(date '+%H:%M:%S')" "$*" +} + +die() { + echo "ERROR: $*" >&2 + exit 1 +} + +need_cmd() { + command -v "$1" >/dev/null 2>&1 || die "required command not found: $1" +} + +download_file() { + local url="$1" + local dst="$2" + + if [[ -s "${dst}" && "${FORCE}" != "1" ]]; then + echo "skip existing: ${dst}" + return 0 + fi + + mkdir -p "$(dirname "${dst}")" + local tmp="${dst}.part" + rm -f "${tmp}" + + echo "download: ${url}" + curl -fL --retry 4 --retry-delay 2 --connect-timeout 20 \ + -o "${tmp}" "${url}" + mv "${tmp}" "${dst}" +} + +prepare_gpt2() { + log "Preparing GPT-2 dataset / weights" + + # InfiniTrain's GPT-2 LLMC loader currently accepts the FP32 v3 file. + # These artifacts are the same llm.c starter-pack files used by TinyInfiniTrain. + local base="https://huggingface.co/datasets/karpathy/llmc-starter-pack/resolve/main" + + local files=( + "gpt2_124M.bin" + "gpt2_tokenizer.bin" + "tiny_shakespeare_train.bin" + "tiny_shakespeare_val.bin" + ) + + local f + for f in "${files[@]}"; do + download_file "${base}/${f}?download=true" "${GPT2_DIR}/${f}" + done + + echo + echo "GPT-2 ready:" + echo " weights: ${GPT2_DIR}/gpt2_124M.bin" + echo " tokenizer: ${GPT2_DIR}/gpt2_tokenizer.bin" + echo " train: ${GPT2_DIR}/tiny_shakespeare_train.bin" + echo " val: ${GPT2_DIR}/tiny_shakespeare_val.bin" +} + +ensure_llama_python() { + need_cmd "${PYTHON}" + + local venv="${CACHE_DIR}/llama3-venv" + LLAMA_PY="${venv}/bin/python" + local py="${LLAMA_PY}" + + if [[ ! -x "${py}" ]]; then + log "Creating local Python environment for LLaMA3 preparation" + "${PYTHON}" -m venv "${venv}" + fi + + if ! "${py}" - <<'PY' >/dev/null 2>&1 +import numpy +import huggingface_hub +import socksio +import transformers +PY + then + log "Installing LLaMA3 preparation dependencies into ${venv}" + "${py}" -m pip install --upgrade pip + "${py}" -m pip install \ + "numpy>=1.24" \ + "huggingface_hub>=0.24" \ + "socksio>=1.0" \ + "transformers>=4.43" + fi + +} + +prepare_llama3() { + log "Preparing LLaMA 3.2 1B dataset / weights" + + local LLAMA_PY="" + ensure_llama_python + local py="${LLAMA_PY}" + + local tiny_txt="${CACHE_DIR}/tiny_shakespeare.txt" + download_file \ + "https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt" \ + "${tiny_txt}" + + # The model repository is gated. The Python helper accepts either HF_TOKEN or + # the token saved by `hf auth login`. + TINY_SHAKESPEARE_TXT="${tiny_txt}" \ + LLAMA3_OUTPUT_DIR="${LLAMA3_DIR}" \ + LLAMA3_CACHE_DIR="${CACHE_DIR}/llama3-hf" \ + SKIP_LLAMA3_WEIGHTS="${SKIP_LLAMA3_WEIGHTS}" \ + FORCE="${FORCE}" \ + "${py}" "${SCRIPT_DIR}/prepare_llama3_assets.py" + + echo + echo "LLaMA3 ready:" + if [[ "${SKIP_LLAMA3_WEIGHTS}" != "1" ]]; then + echo " weights: ${LLAMA3_DIR}/llama3.2_1B_fp32.bin" + fi + echo " train: ${LLAMA3_DIR}/tiny_shakespeare_train.bin" + echo " val: ${LLAMA3_DIR}/tiny_shakespeare_val.bin" +} + +prepare_mnist() { + log "Preparing MNIST IDX dataset" + + # TorchVision's public MNIST mirror. + local base="https://ossci-datasets.s3.amazonaws.com/mnist" + local files=( + "train-images-idx3-ubyte" + "train-labels-idx1-ubyte" + "t10k-images-idx3-ubyte" + "t10k-labels-idx1-ubyte" + ) + + local f + for f in "${files[@]}"; do + local gz="${MNIST_DIR}/${f}.gz" + local dst="${MNIST_DIR}/${f}" + + download_file "${base}/${f}.gz" "${gz}" + + if [[ ! -s "${dst}" || "${FORCE}" == "1" ]]; then + echo "extract: ${gz}" + gzip -dc "${gz}" > "${dst}.part" + mv "${dst}.part" "${dst}" + else + echo "skip existing: ${dst}" + fi + done + + echo + echo "MNIST ready:" + echo " dataset: ${MNIST_DIR}" +} + +need_cmd curl +need_cmd gzip + +case "${TARGET}" in + gpt2) + prepare_gpt2 + ;; + llama3) + prepare_llama3 + ;; + mnist) + prepare_mnist + ;; + all) + prepare_gpt2 + prepare_llama3 + prepare_mnist + ;; +esac + +log "Done" diff --git a/scripts/assets/prepare_llama3_assets.py b/scripts/assets/prepare_llama3_assets.py new file mode 100755 index 00000000..7487f9d7 --- /dev/null +++ b/scripts/assets/prepare_llama3_assets.py @@ -0,0 +1,394 @@ +#!/usr/bin/env python3 + +import json +import mmap +import os +import struct +from pathlib import Path + +import numpy as np +from huggingface_hub import get_token, snapshot_download +from transformers import AutoTokenizer + +MODEL_ID = "meta-llama/Llama-3.2-1B" + +out_dir = Path(os.environ["LLAMA3_OUTPUT_DIR"]) +cache_dir = Path(os.environ["LLAMA3_CACHE_DIR"]) +tiny_path = Path(os.environ["TINY_SHAKESPEARE_TXT"]) +force = os.environ.get("FORCE", "0") == "1" +skip_weights = os.environ.get("SKIP_LLAMA3_WEIGHTS", "0") == "1" + +out_dir.mkdir(parents=True, exist_ok=True) +cache_dir.mkdir(parents=True, exist_ok=True) + +token = os.environ.get("HF_TOKEN") or get_token() +if not token: + raise SystemExit( + f"\nLLaMA3 preparation needs access to {MODEL_ID}.\n" + "1) Accept the model license on Hugging Face.\n" + "2) Run `hf auth login` or export HF_TOKEN=hf_xxx.\n" + ) + +allow_patterns = [ + "config.json", + "generation_config.json", + "tokenizer.json", + "tokenizer_config.json", + "special_tokens_map.json", + "added_tokens.json", + "*.model", +] +if not skip_weights: + allow_patterns.extend([ + "model.safetensors", + "model-*.safetensors", + "model.safetensors.index.json", + ]) + +print(f"[llama3] downloading/reusing Hugging Face files for {MODEL_ID}") +model_dir = Path(snapshot_download( + repo_id=MODEL_ID, + token=token, + cache_dir=str(cache_dir), + allow_patterns=allow_patterns, +)) + +# --------------------------------------------------------------------------- +# TinyShakespeare -> InfiniTrain / llm.c LLaMA-3 data format +# header: 256 int32 = 1024 bytes +# [0] magic = 20240801 +# [1] version = 7 +# [2] ntokens +# payload: uint32 token ids +# --------------------------------------------------------------------------- + +def write_datafile(path: Path, toks): + if path.exists() and path.stat().st_size > 0 and not force: + print(f"[llama3] skip existing: {path}") + return + + header = np.zeros(256, dtype=" {path}") + +tokenizer = AutoTokenizer.from_pretrained( + model_dir, + local_files_only=True, + token=token, + use_fast=True, +) + +text = tiny_path.read_text(encoding="utf-8") +sections = text.split("\n\n") + +bos = tokenizer.bos_token_id +if bos is None: + probe = tokenizer.encode("") + if not probe: + raise RuntimeError("could not determine LLaMA3 BOS token") + bos = probe[0] + +def encode_no_special(s: str): + # Match llm.c's tinyshakespeare.py behavior as closely as current + # transformers versions permit. + try: + return tokenizer.encode( + s, + add_special_tokens=False, + verbose=False, + split_special_tokens=True, + ) + except TypeError: + return tokenizer.encode(s, add_special_tokens=False) + +tokens = [] +for i, section in enumerate(sections): + tokens.append(int(bos)) + padded = section + "\n\n" if i != len(sections) - 1 else section + tokens.extend(int(x) for x in encode_no_special(padded)) + +val_tokens = tokens[:32768] +train_tokens = tokens[32768:] + +write_datafile(out_dir / "tiny_shakespeare_val.bin", val_tokens) +write_datafile(out_dir / "tiny_shakespeare_train.bin", train_tokens) + +if skip_weights: + print("[llama3] SKIP_LLAMA3_WEIGHTS=1: dataset prepared; weight conversion skipped") + raise SystemExit(0) + +# --------------------------------------------------------------------------- +# Hugging Face safetensors -> InfiniTrain LLaMA3 LLMC FP32 format +# +# InfiniTrain's current loader expects: +# magic = 20240803 +# version = 3 (FP32) +# +# Followed by weights in the same order as llm.c train_llama3.py: +# wte +# all ln_1 +# all packed QKV +# all attention output projections +# all ln_2 +# all up projections +# all gate projections +# all down projections +# final norm +# lm_head +# +# This implementation parses safetensors directly and streams large matrices, +# so it does not need to instantiate the model or load all tensors into memory. +# --------------------------------------------------------------------------- + +weight_out = out_dir / "llama3.2_1B_fp32.bin" +if weight_out.exists() and weight_out.stat().st_size > 0 and not force: + print(f"[llama3] skip existing: {weight_out}") + raise SystemExit(0) + +config = json.loads((model_dir / "config.json").read_text()) + +hidden = int(config["hidden_size"]) +n_layer = int(config["num_hidden_layers"]) +n_head = int(config["num_attention_heads"]) +n_kv_head = int(config["num_key_value_heads"]) +vocab = int(config["vocab_size"]) +intermediate = int(config["intermediate_size"]) +norm_eps = float(config.get("rms_norm_eps", 1e-5)) +rope_theta = float(config.get("rope_theta", 500000.0)) + +# InfiniTrain's current TinyShakespeare reader caps LLaMA-3 sequence length at +# 8192, matching the llm.c reference config used by this loader. +block_size = 8192 + +if (hidden, intermediate) != (2048, 8192): + raise RuntimeError( + "unexpected LLaMA 3.2 1B dimensions: " + f"hidden_size={hidden}, intermediate_size={intermediate}" + ) + +ffn_dim_multiplier = 1.5 +multiple_of = 256 + +# Match the current InfiniTrain-Test 3.2 1B baseline, which trains with an +# 8192-token context and writes use_scaled_rope=0. InfiniTrain does not yet +# implement the extended-context LLaMA RoPE scaling path. +use_scaled_rope = 0 +max_gen_bs = 4 + +# Sanity checks for the requested model. +if n_head % n_kv_head != 0: + raise RuntimeError("num_attention_heads must be divisible by num_key_value_heads") +if hidden % n_head != 0: + raise RuntimeError("hidden_size must be divisible by num_attention_heads") + +# Build key -> safetensors shard mapping. +index_path = model_dir / "model.safetensors.index.json" +if index_path.exists(): + index = json.loads(index_path.read_text()) + weight_map = dict(index["weight_map"]) +else: + weight_map = {} + +class SafeTensorShard: + _dtype_map = { + "F32": np.dtype("