OPT HellaSwag

HellaSwag inference using a custom OPT (Open Pre-trained Transformer) implementation with fairscale tensor parallelism.

The model architecture is implemented from scratch — no HuggingFace model abstractions at inference time. Weights are converted from HuggingFace format into sharded checkpoints compatible with fairscale model parallel.

• Python 3.10+
• CUDA-capable GPU(s)
• PyTorch 2.0+

pip install -r requirements.txt

bash data-download.sh

The conversion script downloads HuggingFace OPT weights, remaps parameter names, splits tensors for model parallelism, and saves consolidated checkpoints.

OPT-6.7B (single GPU):

python convert_weights.py \
    --hf_model_path facebook/opt-6.7b \
    --output_dir pretrained/opt-6.7b \
    --num_shards 1

OPT-30B (4 GPUs):

python convert_weights.py \
    --hf_model_path facebook/opt-30b \
    --output_dir pretrained/opt-30b \
    --num_shards 4

Setting a HuggingFace token (huggingface-cli login) before downloading is recommended for faster download speeds.

OPT-6.7B (single GPU):

torchrun --nproc_per_node=1 run_hellaswag.py \
    --ckpt_dir pretrained/opt-6.7b \
    --tokenizer_path pretrained/opt-6.7b \
    --batch_size 4 \
    --max_samples 20

OPT-30B (4 GPUs):

torchrun --nproc_per_node=4 run_hellaswag.py \
    --ckpt_dir pretrained/opt-30b \
    --tokenizer_path pretrained/opt-30b \
    --batch_size 4

Set --max_samples N to limit the number of HellaSwag examples (omit or set to 0 for the full validation set of ~10k samples).

For each example the four endings are scored by the summed log-likelihood of their tokens conditioned on the context, and the highest-scoring ending is compared against the gold label. Both acc (raw) and acc_norm (log-likelihood normalized by continuation length — the standard HellaSwag headline metric) are printed and saved. --batch_size is the number of examples per step; each expands internally to 4 scored sequences, so the model runs with max_batch_size = batch_size * 4.

Results (overall acc/acc_norm plus per-example predictions and log-probs) are saved to results_{batch_size}.json.

opt_hellaswag/
├── opt/
│   ├── model.py          # OPT Transformer (LayerNorm, learned pos embeddings, ReLU FFN)
│   ├── generation.py     # Log-likelihood scoring + autoregressive generation
│   └── tokenizer.py      # GPT-2 BPE tokenizer wrapper
├── dataset/
│   └── hellaswag.py      # HellaSwag dataset loader (query, endings, gold label)
├── convert_weights.py    # HuggingFace → consolidated checkpoint converter
├── run_hellaswag.py      # Main inference entry point (torchrun)
├── download.sh           # Shortcut: download + convert in one step
└── data-download.sh      # Download HellaSwag JSONL files

If running on a SLURM cluster:

srun -N 1 -p <partition> --gres=gpu:1 \
    torchrun --nproc_per_node=1 run_hellaswag.py \
    --ckpt_dir pretrained/opt-6.7b \
    --tokenizer_path pretrained/opt-6.7b \
    --batch_size 4 --max_samples 20