A team at AllenAI has released a complete post-training pipeline for the Open Instruct framework, specifically adapted to run on a single consumer GPU with 16 GB of memory. The method combines Supervised Fine-Tuning, Direct Preference Optimization, and Reinforcement Learning with Verifiable Rewards using GRPO, while replacing heavy distributed components like Ray and DeepSpeed with standard PyTorch and Hugging Face libraries.
Setting up the environment
The process begins by cloning the Open Instruct repository and installing a specific set of lightweight dependencies. This includes packages for model acceleration, logging, and symbolic math verification.
import os, sys, subprocess, textwrap, json, math, random, re, ast, types, dataclasses, gc, contextlib
REPO_URL = "https://github.com/allenai/open-instruct.git"
REPO_DIR = "/content/open-instruct" if os.path.isdir("/content") else "./open-instruct"
PIP_PKGS = [
"peft", "accelerate",
"ray", "wandb", "beaker-py",
"langdetect==1.0.9", "immutabledict==1.2.0", "nltk",
"absl-py", "sympy", "antlr4-python3-runtime==4.11",
"tiktoken",
]
def sh(*args):
print("$", " ".join(args))
subprocess.run(args, check=False)
def setup():
sh(sys.executable, "-m", "pip", "install", "-q", *PIP_PKGS)
if not os.path.isdir(REPO_DIR):
sh("git", "clone", "--depth", "1", REPO_URL, REPO_DIR)
if REPO_DIR not in sys.path:
sys.path.insert(0, REPO_DIR)
os.environ.setdefault("WANDB_MODE", "disabled")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
os.environ.setdefault("RAY_DISABLE_IMPORT_WARNING", "1")
setup()
import numpy as np
import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader
from datasets import load_dataset, Dataset
from transformers import AutoModelForCausalLM, DataCollatorForSeq2Seq, get_cosine_schedule_with_warmup
from peft import LoraConfig, get_peft_model
DEV = "cuda" if torch.cuda.is_available() else "cpu"
try:
_bf16 = DEV == "cuda" and torch.cuda.is_bf16_supported(including_emulation=False)
except TypeError:
_bf16 = DEV == "cuda" and torch.cuda.get_device_properties(0).major >= 8
AMP_DTYPE = torch.bfloat16 if _bf16 else torch.float16
USE_SCALER = AMP_DTYPE is torch.float16
print(f"device={DEV} autocast dtype={AMP_DTYPE} gpu={torch.cuda.get_device_name(0) if DEV=='cuda' else '-'}")
def oi_load(relpath, names, ns=None):
src = open(os.path.join(REPO_DIR, relpath)).read()
tree = ast.parse(src)
found = {n.name: n for n in tree.body
if isinstance(n, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)) and n.name in names}
missing = set(names) - set(found)
if missing:
raise KeyError(f"{relpath}: could not find {missing} (upstream may have renamed them)")
ns = {} if ns is None else dict(ns)
ns.update({"torch": torch, "F": F, "np": np, "enum": __import__("enum"),
"dataclasses": dataclasses, "math": math, "os": os})
future = ast.parse("from __future__ import annotations").body
mod = ast.Module(body=future + [found[n] for n in names], type_ignores=[])
exec(compile(ast.fix_missing_locations(mod), f"<open_instruct:{relpath}>", "exec"), ns)
return {n: ns[n] for n in names}
_dpo = oi_load("open_instruct/dpo_utils.py", ["dpo_loss", "_get_batch_logps"])
_pf = oi_load("open_instruct/padding_free_collator.py", ["calculate_per_token_logps"])
_rl = oi_load("open_instruct/rl_utils.py", ["masked_mean"])
_mu = oi_load("open_instruct/model_utils.py", ["estimate_kl"])
_grpo = oi_load("open_instruct/grpo_utils.py", ["GRPOLossType", "compute_grpo_loss"],
ns={"model_utils": types.SimpleNamespace(**_mu)})
dpo_loss = _dpo["dpo_loss"]
get_batch_logps = _dpo["_get_batch_logps"]
per_token_logps_fn = _pf["calculate_per_token_logps"]
masked_mean = _rl["masked_mean"]
compute_grpo_loss = _grpo["compute_grpo_loss"]
GRPOLossType = _grpo["GRPOLossType"]
print("lifted from repo:", [f.__name__ for f in (dpo_loss, get_batch_logps, per_token_logps_fn,
masked_mean, compute_grpo_loss)])
from open_instruct.dataset_transformation import (
CHAT_TEMPLATES, TokenizerConfig,
sft_tulu_tokenize_and_truncate_v1, sft_tulu_filter_v1,
preference_tulu_tokenize_and_truncate_v1_2,
rlvr_tokenize_v1, visualize_token_role,
)
from open_instruct.ground_truth_utils import GSM8KVerifier, MathVerifier, IFEvalVerifierOldThe script detects the available GPU precision mode, selecting either FP16 or BF16 for automatic casting based on hardware capabilities. It then extracts specific loss and utility functions directly from the repository files without importing the full distributed training stack.
Configuration and tokenization
A central configuration class manages the model selection, dataset sizes, learning rates, and batch settings for every stage of the training pipeline. The default model is set to Qwen/Qwen2.5-0.5B-Instruct.
@dataclasses.dataclass
class CFG:
model: str = "Qwen/Qwen2.5-0.5B-Instruct"
max_seq_len: int = 640
seed: int = 42
n_sft: int = 192
sft_steps: int = 40
sft_micro_bs: int = 2
sft_accum: int = 4
sft_lr: float = 1e-4
n_dpo: int = 96
dpo_steps: int = 24
dpo_micro_bs: int = 1
dpo_accum: int = 4
dpo_lr: float = 5e-5
dpo_beta: float = 0.1
dpo_norm: bool = True
grpo_iters: int = 6
prompts_per_iter: int = 4
samples_per_prompt: int = 4
grpo_micro_bs: int = 1
grpo_inner_epochs: int = 2
grpo_lr: float = 2e-5
grpo_temperature: float = 1.0
grpo_max_new: int = 200
grpo_kl_beta: float = 0.02
clip_lower: float = 0.2
clip_higher: float = 0.272
kl_estimator: int = 2
adv_norm: str = "centered"
n_eval: int = 24
cfg = CFG()
random.seed(cfg.seed); np.random.seed(cfg.seed); torch.manual_seed(cfg.seed)
tc = TokenizerConfig(tokenizer_name_or_path=cfg.model, chat_template_name=None, use_fast=True)
tok = tc.tokenizer
print(f"\navailable CHAT_TEMPLATES: {list(CHAT_TEMPLATES)[:12]} ... ({len(CHAT_TEMPLATES)} total)")
print(f"pad={tok.pad_token!r}({tok.pad_token_id}) eos={tok.eos_token!r}({tok.eos_token_id})")
_demo = {"messages": [
{"role": "user", "content": "What is 12 * 3?"},
{"role": "assistant", "content": "12 * 3 = 36. The answer is 36."},
{"role": "user", "content": "And minus 6?"},
{"role": "assistant", "content": "36 - 6 = 30. The answer is 30."},
]}
_enc = sft_tulu_tokenize_and_truncate_v1(dict(_demo), tok, cfg.max_seq_len)
print("\n[SFT label masking — colour 0 = masked out of the loss, colour 1 = trained on]")
visualize_token_role(_enc["input_ids"].tolist(), (_enc["labels"] != -100).long().tolist(), tok)
print(f"trainable tokens: {(_enc['labels'] != -100).sum().item()}/{_enc['labels'].numel()}")Initialisation includes setting random seeds for reproducibility and loading the tokenizer with the correct chat template. The code visualises the token roles during Supervised Fine-Tuning, highlighting which assistant tokens are masked out of the loss calculation versus those actively used for training.
Preparing the GSM8K dataset
The workflow loads the GSM8K dataset, which contains grade-school math problems. The system prepares data for three distinct phases: SFT, DPO, and RLVR.
For Supervised Fine-Tuning, the code constructs a dataset where the system prompt instructs the model to reason step by step before providing the final answer. The dataset is tokenised and filtered to ensure only relevant tokens contribute to the loss.
Direct Preference Optimization requires pairs of responses. The script generates a “chosen” example with the correct answer and a “rejected” example with an intentionally incorrect number derived from the correct solution.
gsm = load_dataset("openai/gsm8k", "main")
SYS = "You are a careful math assistant. Reason step by step, then finish with 'The answer is N.'"
def gsm_answer(a):
return a.split("####")[-1].strip().replace(",", "")
def gsm_solution(a):
body = a.split("####")[0].strip()
body = re.sub(r"<<.*?>>", "", body)
return f"{body}\nThe answer is {gsm_answer(a)}."Source Read original →Related articles
