chore: import upstream snapshot with attribution

scripts/chat.py

@@ -0,0 +1,65 @@
+"""
+Chat / inference CLI for any stage checkpoint (base, SFT, DPO, PPO, GRPO).
+
+Model dimensions are read from the checkpoint, so you only pass the path. Use the chat
+template for instruction-tuned models, or --raw for base-model continuation.
+
+One-shot:
+    PYTHONPATH=. python scripts/chat.py --ckpt /ephemeral/ckpts/sft.pt --prompt "What is 13 + 29?"
+    PYTHONPATH=. python scripts/chat.py --ckpt /ephemeral/ckpts/grpo.pt --prompt "..." --greedy
+    PYTHONPATH=. python scripts/chat.py --ckpt /ephemeral/ckpts/base_pretrained.pt --raw --prompt "Once upon a time"
+Interactive REPL (no --prompt):
+    PYTHONPATH=. python scripts/chat.py --ckpt /ephemeral/ckpts/sft.pt
+"""
+
+from __future__ import annotations
+
+import argparse
+
+import torch
+
+from src.post_training.inference import generate_reply, load_model_from_ckpt
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--ckpt", required=True)
+    p.add_argument("--prompt", default=None, help="one-shot prompt; omit for interactive REPL")
+    p.add_argument("--system", default=None, help="optional system message (chat mode)")
+    p.add_argument("--raw", action="store_true", help="base-model continuation (no chat template)")
+    p.add_argument("--max_new_tokens", type=int, default=256)
+    p.add_argument("--temperature", type=float, default=0.8)
+    p.add_argument("--top_p", type=float, default=0.95)
+    p.add_argument("--top_k", type=int, default=None)
+    p.add_argument("--greedy", action="store_true", help="deterministic argmax decoding")
+    p.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
+    args = p.parse_args()
+
+    model = load_model_from_ckpt(args.ckpt, args.device)
+    n = sum(p.numel() for p in model.parameters())
+    print(f"loaded {args.ckpt} ({n/1e6:.0f}M params) on {args.device} | "
+          f"mode={'raw' if args.raw else 'chat'} {'greedy' if args.greedy else f'T={args.temperature} top_p={args.top_p}'}")
+
+    def reply(text):
+        return generate_reply(model, text, device=args.device, system=args.system, raw=args.raw,
+                              max_new_tokens=args.max_new_tokens, temperature=args.temperature,
+                              top_p=args.top_p if args.top_p < 1 else None, top_k=args.top_k, greedy=args.greedy)
+
+    if args.prompt is not None:
+        print(reply(args.prompt))
+        return
+
+    print("Interactive chat (Ctrl-D / 'exit' to quit).")
+    while True:
+        try:
+            text = input("\nyou> ").strip()
+        except EOFError:
+            break
+        if text in ("exit", "quit"):
+            break
+        if text:
+            print("bot>", reply(text))
+
+
+if __name__ == "__main__":
+    main()

scripts/data_download.py

@@ -0,0 +1,81 @@
+import os
+import argparse
+import requests
+from tqdm import tqdm
+from typing import List
+
+# Base URL for the dataset files
+BASE_URL = "https://huggingface.co/datasets/monology/pile-uncopyrighted/resolve/main"
+VAL_URL = f"{BASE_URL}/val.jsonl.zst"  # URL for the validation dataset
+TRAIN_URLS = [f"{BASE_URL}/train/{i:02d}.jsonl.zst" for i in range(65)]  # URLs for 65 training files (adjust the range if needed)
+
+def download_file(url: str, file_name: str) -> None:
+    """
+    Downloads a file from the given URL and saves it with the specified file name.
+    Displays a progress bar using tqdm.
+    
+    Args:
+        url (str): The URL of the file to download.
+        file_name (str): The local path where the file will be saved.
+    """
+    print(f"Downloading: {file_name}...")
+    response = requests.get(url, stream=True)  # Stream the file content
+    total_size = int(response.headers.get('content-length', 0))  # Get total file size if available
+    block_size = 1024  # Size of each block for the progress bar
+    with open(file_name, 'wb') as f:  # Open file for writing in binary mode
+        for chunk in tqdm(response.iter_content(block_size), total=total_size // block_size, desc="Downloading", leave=True):
+            f.write(chunk)  # Write each chunk to the file
+
+def download_dataset(val_url: str, train_urls: List[str], val_dir: str, train_dir: str, max_train_files: int) -> None:
+    """
+    Manages downloading of the dataset, including both validation and training files.
+    
+    Args:
+        val_url (str): URL for the validation dataset.
+        train_urls (list): List of URLs for the training dataset files.
+        val_dir (str): Directory where the validation file will be stored.
+        train_dir (str): Directory where the training files will be stored.
+        max_train_files (int): Maximum number of training files to download.
+    """
+    # Define the path for the validation file
+    val_file_path = os.path.join(val_dir, "val.jsonl.zst")
+    if not os.path.exists(val_file_path):  # Check if the validation file already exists
+        print(f"Validation file not found. Downloading from {val_url}...")
+        download_file(val_url, val_file_path)  # Download the validation file
+    else:
+        print("Validation data already present. Skipping download.")
+
+    # Loop through the training file URLs and download if not already present
+    for idx, url in enumerate(train_urls[:max_train_files]):  # Limit to max_train_files
+        file_name = f"{idx:02d}.jsonl.zst"  # Format file name (e.g., 00.jsonl.zst)
+        file_path = os.path.join(train_dir, file_name)  # Construct the full file path
+        if not os.path.exists(file_path):  # Check if the file already exists
+            print(f"Training file {file_name} not found. Downloading...")
+            download_file(url, file_path)  # Download the training file
+        else:
+            print(f"Training file {file_name} already present. Skipping download.")
+
+def main() -> None:
+    """
+    Main function to parse arguments and orchestrate the dataset download process.
+    """
+    # Parse command-line arguments using argparse
+    parser = argparse.ArgumentParser(description="Download PILE dataset.")  # Description of the script
+    parser.add_argument('--train_max', type=int, default=1, help="Max number of training files to download.")  # Max training files
+    parser.add_argument('--train_dir', default="data/train", help="Directory for storing training data.")  # Training directory
+    parser.add_argument('--val_dir', default="data/val", help="Directory for storing validation data.")  # Validation directory
+
+    args = parser.parse_args()  # Parse the arguments provided by the user
+
+    # Ensure directories for training and validation data exist
+    os.makedirs(args.train_dir, exist_ok=True)  # Create training directory if it doesn't exist
+    os.makedirs(args.val_dir, exist_ok=True)  # Create validation directory if it doesn't exist
+
+    # Start downloading the dataset
+    download_dataset(VAL_URL, TRAIN_URLS, args.val_dir, args.train_dir, args.train_max)
+
+    print("Dataset downloaded successfully.")  # Indicate successful download
+
+if __name__ == "__main__":
+    # Entry point of the script
+    main()

scripts/data_preprocess.py

@@ -0,0 +1,116 @@
+import os
+import json
+import zstandard as zstd
+import tiktoken
+import h5py
+from tqdm import tqdm
+import argparse
+from typing import Optional
+
+def process_files(input_dir: str, output_file: str, tokenizer_name: str, max_data: Optional[int] = None) -> None:
+    """
+    Process a specified number of lines from each .jsonl.zst file in the input directory
+    and save encoded tokens to an HDF5 file.
+
+    Args:
+        input_dir (str): Directory containing input .jsonl.zst files.
+        output_file (str): Path to the output HDF5 file.
+        tokenizer_name (str): Name of the tiktoken tokenizer to use (e.g., 'r50k_base').
+        max_data (int, optional): Maximum number of lines to process from each file.
+                                  If None, process all lines.
+    """
+    # Print processing strategy based on max_data
+    if max_data is not None:
+        print(f"You have chosen max_data = {max_data}. Processing only the top {max_data} JSON objects from each file.")
+    else:
+        print("Processing all available JSON objects from each file.")
+
+    # Load the tokenizer using the provided tokenizer name
+    enc = tiktoken.get_encoding(tokenizer_name)
+
+    # Create an HDF5 file for output
+    with h5py.File(output_file, 'w') as out_f:
+        # Initialize the dataset for storing tokenized data
+        dataset = out_f.create_dataset('tokens', (0,), maxshape=(None,), dtype='i')
+        start_index = 0  # Track the starting index for the next batch of tokens
+
+        # Process each .jsonl.zst file in the input directory
+        for filename in sorted(os.listdir(input_dir)):
+            if filename.endswith(".jsonl.zst"):  # Only process .jsonl.zst files
+                in_file = os.path.join(input_dir, filename)
+                print(f"Processing: {in_file}")
+
+                processed_lines = 0  # Counter for processed lines in the current file
+
+                # Open the compressed .jsonl.zst file for reading
+                with zstd.open(in_file, 'rt', encoding='utf-8') as in_f:
+                    # Iterate over each line in the file
+                    for line in tqdm(in_f, desc=f"Processing {filename}", total=max_data if max_data is not None else None):
+                        try:
+                            # Parse the line as JSON
+                            data = json.loads(line)
+                            text = data.get('text')  # Extract the 'text' field from the JSON object
+
+                            if text:
+                                # Tokenize the text and append an end-of-text token
+                                encoded = enc.encode(text + "<|endoftext|>", allowed_special={'<|endoftext|>'})
+                                encoded_len = len(encoded)
+
+                                # Resize the dataset to accommodate new tokens
+                                end_index = start_index + encoded_len
+                                dataset.resize(dataset.shape[0] + encoded_len, axis=0)
+
+                                # Store the encoded tokens in the dataset
+                                dataset[start_index:end_index] = encoded
+                                start_index = end_index  # Update the start index
+                            else:
+                                # Warn if 'text' key is missing in the JSON object
+                                print(f"Warning: 'text' key missing in line from {filename}")
+                        except json.JSONDecodeError:
+                            # Handle JSON decoding errors
+                            print(f"Warning: Could not decode JSON from line in {filename}")
+                        except Exception as e:
+                            # Handle any other errors
+                            print(f"An error occurred while processing line in {filename}: {e}")
+
+                        processed_lines += 1
+                        # Stop processing if max_data limit is reached
+                        if max_data is not None and processed_lines >= max_data:
+                            break
+
+def main():
+    """
+    Main function to parse arguments, validate directories, and process files.
+    """
+    # Parse command-line arguments
+    parser = argparse.ArgumentParser(description="Preprocess PILE dataset files and save tokens to HDF5.")
+    parser.add_argument("--train_dir", type=str, default="data/train", help="Directory containing training .jsonl.zst files.")
+    parser.add_argument("--val_dir", type=str, default="data/val", help="Directory containing validation .jsonl.zst files.")
+    parser.add_argument("--out_train_file", type=str, default="data/train/pile_train.h5", help="Path to the output training HDF5 file.")
+    parser.add_argument("--out_val_file", type=str, default="data/val/pile_dev.h5", help="Path to the output validation HDF5 file.")
+    parser.add_argument("--tokenizer_name", type=str, default="r50k_base", help="Name of the tiktoken tokenizer to use.")
+    parser.add_argument("--max_data", type=int, default=1000, help="Maximum number of json objects to process from each file in both train and val datasets (default: 1000).")
+
+    args = parser.parse_args()
+
+    # Validate the existence of the training and validation directories
+    if not os.path.isdir(args.train_dir):
+        print(f"Error: Training directory not found: {args.train_dir}")
+        return
+    if not os.path.isdir(args.val_dir):
+        print(f"Error: Validation directory not found: {args.val_dir}")
+        return
+
+    # Process training data
+    print("Starting training data preprocessing...")
+    process_files(args.train_dir, args.out_train_file, args.tokenizer_name, args.max_data)
+    print("Training data preprocessing complete.")
+
+    # Process validation data
+    print("Starting validation data preprocessing...")
+    process_files(args.val_dir, args.out_val_file, args.tokenizer_name, args.max_data)
+    print("Validation data preprocessing complete.")
+
+# Entry point of the script
+if __name__ == "__main__":
+    main()

scripts/eval_post_training.py

@@ -0,0 +1,87 @@
+"""
+Evaluate any stage checkpoint on GSM8K (greedy) and optionally dump sample generations.
+Use it to build the headline "GSM8K accuracy across stages" table:
+
+    for s in base_pretrained sft dpo ppo grpo; do
+      PYTHONPATH=. python scripts/eval_post_training.py --ckpt /ephemeral/ckpts/$s.pt \
+        --label $s --limit 200 --append /ephemeral/logs/stage_table.jsonl
+    done
+    PYTHONPATH=. python scripts/eval_post_training.py --table /ephemeral/logs/stage_table.jsonl
+
+Model dimensions are read from the checkpoint's stored ``cfg`` so you don't have to repeat
+them. Reward checkpoints (which have a reward head, not an LM head only) still load because
+we keep just the backbone keys for generation.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+
+import torch
+
+from src.models.transformer import Transformer
+from src.post_training.evaluation import gsm8k_accuracy, load_gsm8k_eval
+
+
+def model_from_ckpt(ckpt_path: str, device: str, overrides: dict | None = None) -> Transformer:
+    ck = torch.load(ckpt_path, map_location="cpu", weights_only=False)
+    cfg = ck.get("cfg", {}) or {}
+    cfg = {**cfg, **(overrides or {})}
+    model = Transformer(
+        n_head=cfg.get("n_head", 16), n_embed=cfg.get("n_embed", 1024),
+        context_length=cfg.get("context_length", 1024), vocab_size=cfg.get("vocab_size", 50304),
+        N_BLOCKS=cfg.get("n_blocks", 24),
+    )
+    state = ck["model_state_dict"] if "model_state_dict" in ck else ck
+    state = {k.removeprefix("module.").removeprefix("transformer."): v for k, v in state.items()}
+    backbone_keys = set(model.state_dict().keys())
+    filtered = {k: v for k, v in state.items() if k in backbone_keys}
+    model.load_state_dict(filtered, strict=False)
+    return model.to(device).eval()
+
+
+def print_table(path: str):
+    rows = [json.loads(l) for l in open(path) if l.strip()]
+    print(f"\n{'stage':<18}{'GSM8K acc':>10}{'n':>8}")
+    print("-" * 36)
+    for r in rows:
+        print(f"{r['label']:<18}{r['accuracy']*100:>9.1f}%{r['n']:>8}")
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--ckpt")
+    p.add_argument("--label", default="model")
+    p.add_argument("--limit", type=int, default=200)
+    p.add_argument("--split", default="test")
+    p.add_argument("--max_new_tokens", type=int, default=300)
+    p.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
+    p.add_argument("--samples", type=int, default=3)
+    p.add_argument("--append", default=None, help="append the result row to this JSONL")
+    p.add_argument("--table", default=None, help="just print a stage table from this JSONL and exit")
+    args = p.parse_args()
+
+    if args.table:
+        print_table(args.table)
+        return
+
+    model = model_from_ckpt(args.ckpt, args.device)
+    qa = load_gsm8k_eval(args.split, limit=args.limit)
+    res = gsm8k_accuracy(model, qa, device=args.device, max_new_tokens=args.max_new_tokens,
+                         greedy=True, return_samples=args.samples)
+    print(f"[{args.label}] GSM8K {args.split} accuracy: {res['accuracy']*100:.1f}%  ({res['correct']}/{res['n']})")
+    for s in res["samples"]:
+        print(f"\n  Q: {s['q'][:120]}\n  gold={s['gold']} correct={s['correct']}\n  A: {s['response'][:300]}")
+
+    if args.append:
+        os.makedirs(os.path.dirname(args.append) or ".", exist_ok=True)
+        with open(args.append, "a") as f:
+            f.write(json.dumps({"label": args.label, "accuracy": res["accuracy"],
+                                "correct": res["correct"], "n": res["n"]}) + "\n")
+        print(f"\nappended -> {args.append}")
+
+
+if __name__ == "__main__":
+    main()

scripts/generate_text.py

@@ -0,0 +1,68 @@
+import torch
+import tiktoken
+import argparse
+from config.config import default_config as config
+from src.models.transformer import Transformer  # Assuming your Transformer class is in this module
+
+
+def load_checkpoint(model_path: str, device: str):
+    try:
+        return torch.load(model_path, map_location=torch.device(device), weights_only=False)
+    except TypeError:
+        return torch.load(model_path, map_location=torch.device(device))
+
+def generate_text(model_path: str, input_text: str, max_new_tokens: int = 100, device: str = 'cuda') -> str:
+    """
+    Generates text using a pre-trained Transformer model.
+
+    Args:
+        model_path (str): Path to the saved model checkpoint.
+        input_text (str): The initial text to start generation from.
+        max_new_tokens (int): The maximum number of new tokens to generate.
+        device (str): 'cuda' or 'cpu', the device to run the model on.
+
+    Returns:
+        str: The generated text.
+    """
+    # Load the model checkpoint
+    checkpoint = load_checkpoint(model_path, device)
+
+    # Initialize the model using the configuration from config.py
+    model = Transformer(
+        n_head=config['n_head'],
+        n_embed=config['n_embed'],
+        context_length=config['context_length'],
+        vocab_size=config['vocab_size'],
+        N_BLOCKS=config['n_blocks']
+    )
+    model.load_state_dict(checkpoint['model_state_dict'])
+    model.eval().to(device)
+
+    # Load the tokenizer
+    enc = tiktoken.get_encoding("r50k_base")
+
+    start_ids = enc.encode_ordinary(input_text)
+    context = torch.tensor(start_ids, dtype=torch.long, device=device).unsqueeze(0)
+
+    # Generation process
+    with torch.no_grad():
+        generated_tokens = model.generate(context, max_new_tokens=max_new_tokens)[0].tolist()
+
+    # Decode the generated tokens
+    output_text = enc.decode(generated_tokens)
+
+    return output_text
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Generate text using a pre-trained Transformer model.")
+    parser.add_argument('--model_path', type=str, help='Path to the saved model checkpoint.')
+    parser.add_argument('--input_text', type=str, help='The initial text to start generation from.')
+    parser.add_argument('--max_new_tokens', type=int, default=100, help='Maximum number of new tokens to generate.')
+
+    args = parser.parse_args()
+
+    generated = generate_text(args.model_path, args.input_text, args.max_new_tokens, config['device'])
+    print(f"Generated text:\n{generated}")
+
+if __name__ == "__main__":
+    main()

scripts/prepare_preference_data.py

@@ -0,0 +1,112 @@
+"""
+Build preference data for the reward model and DPO from real public datasets:
+  - Anthropic/hh-rlhf                       (helpful/harmless human preferences)
+  - HuggingFaceH4/ultrafeedback_binarized   (LLM-judged preference pairs)
+
+Emits JSONL of ``{"prompt", "chosen", "rejected"}`` (train + held-out test). The held-out
+test split is what reward-model preference accuracy is measured on.
+
+Example:
+    PYTHONPATH=. HF_HOME=/ephemeral/hf_cache python scripts/prepare_preference_data.py \
+        --source both --max_per_source 40000 --out_dir /ephemeral/data
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+
+os.environ.setdefault("HF_HOME", "/ephemeral/hf_cache")
+
+_ASSISTANT_MARKER = "\n\nAssistant:"
+
+
+def _split_hh(text: str) -> tuple[str, str] | None:
+    """Split an HH-RLHF conversation string into (prompt_context, final_response)."""
+    idx = text.rfind(_ASSISTANT_MARKER)
+    if idx == -1:
+        return None
+    prompt = text[:idx].strip()
+    response = text[idx + len(_ASSISTANT_MARKER):].strip()
+    if not prompt or not response:
+        return None
+    return prompt, response
+
+
+def from_hh(max_n: int, split: str) -> list[dict]:
+    from datasets import load_dataset
+
+    ds = load_dataset("Anthropic/hh-rlhf", split=split)
+    if max_n:
+        ds = ds.select(range(min(max_n, len(ds))))
+    out = []
+    for ex in ds:
+        c = _split_hh(ex["chosen"]); r = _split_hh(ex["rejected"])
+        if not c or not r:
+            continue
+        prompt, chosen = c
+        _, rejected = r
+        if chosen == rejected:
+            continue
+        out.append({"prompt": prompt, "chosen": chosen, "rejected": rejected})
+    return out
+
+
+def from_ultrafeedback(max_n: int, split: str) -> list[dict]:
+    from datasets import load_dataset
+
+    hf_split = "train_prefs" if split == "train" else "test_prefs"
+    ds = load_dataset("HuggingFaceH4/ultrafeedback_binarized", split=hf_split)
+    if max_n:
+        ds = ds.select(range(min(max_n, len(ds))))
+    out = []
+    for ex in ds:
+        prompt = ex["prompt"].strip()
+        chosen = ex["chosen"][-1]["content"].strip()
+        rejected = ex["rejected"][-1]["content"].strip()
+        if not prompt or not chosen or not rejected or chosen == rejected:
+            continue
+        out.append({"prompt": prompt, "chosen": chosen, "rejected": rejected})
+    return out
+
+
+def collect(source: str, max_n: int, split: str) -> list[dict]:
+    rows: list[dict] = []
+    if source in ("hh", "both"):
+        print(f"Loading Anthropic/hh-rlhf [{split}] ...")
+        rows += from_hh(max_n, split)
+    if source in ("ultrafeedback", "both"):
+        print(f"Loading ultrafeedback_binarized [{split}] ...")
+        try:
+            rows += from_ultrafeedback(max_n, split)
+        except Exception as e:  # noqa: BLE001
+            print(f"  (skipping ultrafeedback: {e})")
+    return rows
+
+
+def write_jsonl(rows: list[dict], path: str) -> None:
+    os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
+    with open(path, "w") as f:
+        for r in rows:
+            f.write(json.dumps(r) + "\n")
+    print(f"  wrote {len(rows)} pairs -> {path}")
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--source", choices=["hh", "ultrafeedback", "both"], default="both")
+    p.add_argument("--max_per_source", type=int, default=40000)
+    p.add_argument("--out_dir", default="/ephemeral/data")
+    args = p.parse_args()
+
+    train = collect(args.source, args.max_per_source, "train")
+    test = collect(args.source, max(2000, args.max_per_source // 20), "test")
+    import random
+    random.Random(0).shuffle(train)
+    write_jsonl(train, os.path.join(args.out_dir, "preferences.jsonl"))
+    write_jsonl(test, os.path.join(args.out_dir, "preferences_test.jsonl"))
+
+
+if __name__ == "__main__":
+    main()

scripts/prepare_pretrain_data.py

@@ -0,0 +1,139 @@
+"""
+Download + tokenize Pile-uncopyrighted shards into a single flat-token HDF5 for
+pretraining. Faster than the original ``data_preprocess.py`` (which resizes the HDF5 per
+document): here we stream-decompress, batch-tokenize with tiktoken, and write tokens to
+the HDF5 in large chunks.
+
+Writes to /ephemeral by default (the 1.5TB disk).
+
+Examples:
+    # dev split from the Pile validation file
+    PYTHONPATH=. python scripts/prepare_pretrain_data.py --split val \
+        --out /ephemeral/data/pile_dev.h5
+    # one training shard
+    PYTHONPATH=. python scripts/prepare_pretrain_data.py --split train --num_shards 1 \
+        --out /ephemeral/data/pile_train.h5
+"""
+
+from __future__ import annotations
+
+import argparse
+import io
+import json
+import os
+
+import h5py
+import numpy as np
+import requests
+import tiktoken
+import zstandard as zstd
+from tqdm import tqdm
+
+BASE_URL = "https://huggingface.co/datasets/monology/pile-uncopyrighted/resolve/main"
+EOT_ID = 50256
+WRITE_CHUNK = 8_000_000  # flush tokens to HDF5 in ~8M-token chunks
+ENC_BATCH = 1024          # documents per tiktoken batch-encode
+
+
+def shard_urls(split: str, num_shards: int) -> list[str]:
+    if split == "val":
+        return [f"{BASE_URL}/val.jsonl.zst"]
+    return [f"{BASE_URL}/train/{i:02d}.jsonl.zst" for i in range(num_shards)]
+
+
+def download(url: str, dest: str) -> str:
+    if os.path.exists(dest):
+        print(f"  cached: {dest}")
+        return dest
+    os.makedirs(os.path.dirname(dest), exist_ok=True)
+    print(f"  downloading {url}")
+    with requests.get(url, stream=True) as r:
+        r.raise_for_status()
+        total = int(r.headers.get("content-length", 0))
+        with open(dest + ".part", "wb") as f:
+            for chunk in tqdm(r.iter_content(1 << 20), total=total >> 20, unit="MB", desc="dl"):
+                f.write(chunk)
+    os.replace(dest + ".part", dest)
+    return dest
+
+
+def iter_texts(zst_path: str):
+    dctx = zstd.ZstdDecompressor()
+    with open(zst_path, "rb") as fh:
+        reader = dctx.stream_reader(fh)
+        for line in io.TextIOWrapper(reader, encoding="utf-8"):
+            line = line.strip()
+            if not line:
+                continue
+            try:
+                txt = json.loads(line).get("text")
+            except json.JSONDecodeError:
+                continue
+            if txt:
+                yield txt
+
+
+def tokenize_to_h5(zst_paths: list[str], out_path: str, max_tokens: int | None) -> int:
+    enc = tiktoken.get_encoding("r50k_base")
+    os.makedirs(os.path.dirname(out_path) or ".", exist_ok=True)
+    total = 0
+    buf: list[int] = []
+    with h5py.File(out_path, "w") as f:
+        dset = f.create_dataset("tokens", (0,), maxshape=(None,), dtype="i4", chunks=(WRITE_CHUNK,))
+
+        def flush():
+            nonlocal total, buf
+            if not buf:
+                return
+            arr = np.asarray(buf, dtype=np.int32)
+            dset.resize(total + arr.size, axis=0)
+            dset[total: total + arr.size] = arr
+            total += arr.size
+            buf = []
+
+        for zp in zst_paths:
+            print(f"  tokenizing {zp}")
+            docs: list[str] = []
+            pbar = tqdm(iter_texts(zp), unit="doc", desc="tok")
+            for txt in pbar:
+                docs.append(txt)
+                if len(docs) >= ENC_BATCH:
+                    for ids in enc.encode_ordinary_batch(docs):
+                        buf.extend(ids)
+                        buf.append(EOT_ID)
+                    docs = []
+                    if len(buf) >= WRITE_CHUNK:
+                        flush()
+                        pbar.set_postfix(tokens=f"{total/1e6:.1f}M")
+                    if max_tokens and total >= max_tokens:
+                        break
+            if docs:
+                for ids in enc.encode_ordinary_batch(docs):
+                    buf.extend(ids)
+                    buf.append(EOT_ID)
+            flush()
+            if max_tokens and total >= max_tokens:
+                break
+    print(f"  wrote {total:,} tokens -> {out_path}")
+    return total
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--split", choices=["train", "val"], required=True)
+    p.add_argument("--num_shards", type=int, default=1, help="train shards to use")
+    p.add_argument("--raw_dir", default="/ephemeral/data/pile_raw")
+    p.add_argument("--out", required=True)
+    p.add_argument("--max_tokens", type=int, default=None, help="stop after this many tokens")
+    args = p.parse_args()
+
+    urls = shard_urls(args.split, args.num_shards)
+    local = []
+    for u in urls:
+        name = "val.jsonl.zst" if args.split == "val" else os.path.basename(u)
+        local.append(download(u, os.path.join(args.raw_dir, name)))
+    tokenize_to_h5(local, args.out, args.max_tokens)
+
+
+if __name__ == "__main__":
+    main()

scripts/prepare_rl_prompts.py

@@ -0,0 +1,76 @@
+"""
+Build RL prompt sets ({"prompt", "gold"}) for PPO/GRPO:
+  - GSM8K train -> rl_prompts_train.jsonl   (the RL training prompts)
+  - GSM8K test  -> rl_prompts_test.jsonl    (held-out benchmark for eval)
+  - a programmatic arithmetic set -> arithmetic_prompts.jsonl  (RL curriculum warm-up,
+    where even a weak model gets some non-zero reward so RL has signal to start from)
+
+Example:
+    PYTHONPATH=. HF_HOME=/ephemeral/hf_cache python scripts/prepare_rl_prompts.py --out_dir /ephemeral/data
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import random
+
+os.environ.setdefault("HF_HOME", "/ephemeral/hf_cache")
+
+from src.post_training.rewards import gsm8k_gold_answer
+
+
+def write_jsonl(rows, path):
+    os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
+    with open(path, "w") as f:
+        for r in rows:
+            f.write(json.dumps(r) + "\n")
+    print(f"  wrote {len(rows)} prompts -> {path}")
+
+
+def gsm8k_prompts(split: str, limit: int | None):
+    from datasets import load_dataset
+
+    ds = load_dataset("openai/gsm8k", "main", split=split)
+    if limit:
+        ds = ds.select(range(min(limit, len(ds))))
+    rows = []
+    for ex in ds:
+        gold = gsm8k_gold_answer(ex["answer"])
+        if gold is not None:
+            rows.append({"prompt": ex["question"].strip(), "gold": gold})
+    return rows
+
+
+def arithmetic_prompts(n: int, max_val: int, seed: int):
+    rng = random.Random(seed)
+    ops = [("+", lambda a, b: a + b), ("-", lambda a, b: a - b), ("*", lambda a, b: a * b)]
+    rows = []
+    for _ in range(n):
+        a, b = rng.randint(0, max_val), rng.randint(0, max_val)
+        sym, fn = rng.choice(ops)
+        rows.append({"prompt": f"What is {a} {sym} {b}?", "gold": float(fn(a, b))})
+    return rows
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--out_dir", default="/ephemeral/data")
+    p.add_argument("--train_limit", type=int, default=None)
+    p.add_argument("--test_limit", type=int, default=500)
+    p.add_argument("--arith_n", type=int, default=5000)
+    p.add_argument("--arith_max", type=int, default=20)
+    args = p.parse_args()
+
+    print("Loading GSM8K train ...")
+    write_jsonl(gsm8k_prompts("train", args.train_limit), os.path.join(args.out_dir, "rl_prompts_train.jsonl"))
+    print("Loading GSM8K test ...")
+    write_jsonl(gsm8k_prompts("test", args.test_limit), os.path.join(args.out_dir, "rl_prompts_test.jsonl"))
+    print("Generating arithmetic curriculum ...")
+    write_jsonl(arithmetic_prompts(args.arith_n, args.arith_max, seed=0),
+                os.path.join(args.out_dir, "arithmetic_prompts.jsonl"))
+
+
+if __name__ == "__main__":
+    main()

scripts/prepare_sft_data.py

@@ -0,0 +1,127 @@
+"""
+Build the SFT dataset from real public instruction data + GSM8K, render it through the
+chat template (masking prompt tokens), pack to fixed-length rows, and write a packed
+HDF5 (``tokens`` + ``loss_mask``) plus a held-out dev file.
+
+Datasets (downloaded via HuggingFace ``datasets`` to /ephemeral/hf_cache):
+  - tatsu-lab/alpaca                (general instruction following)
+  - databricks/databricks-dolly-15k (general instruction following)
+  - openai/gsm8k (main, train)      (math, reformatted into <think>/<answer> so the model
+                                     learns the reasoning format the RL verifier rewards)
+
+Example:
+    PYTHONPATH=. HF_HOME=/ephemeral/hf_cache python scripts/prepare_sft_data.py \
+        --context_length 1024 --out_dir /ephemeral/data
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import re
+
+import h5py
+import numpy as np
+
+from src.post_training.chat_template import encode_chat, ANSWER_OPEN, ANSWER_CLOSE, THINK_OPEN, THINK_CLOSE
+from src.post_training.sft import pack_examples
+
+os.environ.setdefault("HF_HOME", "/ephemeral/hf_cache")
+_CALC_RE = re.compile(r"<<[^>]*>>")          # GSM8K calculator annotations
+_HASH_RE = re.compile(r"####\s*(.+)\s*$")
+
+
+def gsm8k_to_messages(question: str, answer: str) -> list[dict]:
+    """Reformat a GSM8K (question, answer) into chat messages whose assistant turn uses
+    the <think>...</think><answer>N</answer> structure."""
+    answer = _CALC_RE.sub("", answer).strip()
+    m = _HASH_RE.search(answer)
+    final = m.group(1).strip() if m else answer
+    reasoning = _HASH_RE.sub("", answer).strip()
+    completion = f"{THINK_OPEN}{reasoning}{THINK_CLOSE}{ANSWER_OPEN}{final}{ANSWER_CLOSE}"
+    return [{"role": "user", "content": question.strip()},
+            {"role": "assistant", "content": completion}]
+
+
+def alpaca_to_messages(ex: dict) -> list[dict]:
+    instr = ex["instruction"].strip()
+    inp = (ex.get("input") or "").strip()
+    user = f"{instr}\n\n{inp}" if inp else instr
+    return [{"role": "user", "content": user}, {"role": "assistant", "content": ex["output"].strip()}]
+
+
+def dolly_to_messages(ex: dict) -> list[dict]:
+    instr = ex["instruction"].strip()
+    ctx = (ex.get("context") or "").strip()
+    user = f"{instr}\n\n{ctx}" if ctx else instr
+    return [{"role": "user", "content": user}, {"role": "assistant", "content": ex["response"].strip()}]
+
+
+def collect_examples(context_length: int, limit_per_set: int | None) -> list[tuple[list[int], list[int]]]:
+    from datasets import load_dataset
+
+    examples: list[tuple[list[int], list[int]]] = []
+    n_kept = n_skipped = 0
+
+    def add(messages):
+        nonlocal n_kept, n_skipped
+        ids, mask = encode_chat(messages)
+        if len(ids) <= context_length and sum(mask) > 0:
+            examples.append((ids, mask))
+            n_kept += 1
+        else:
+            n_skipped += 1
+
+    print("Loading tatsu-lab/alpaca ...")
+    alpaca = load_dataset("tatsu-lab/alpaca", split="train")
+    for ex in (alpaca.select(range(min(limit_per_set, len(alpaca)))) if limit_per_set else alpaca):
+        add(alpaca_to_messages(ex))
+
+    print("Loading databricks/databricks-dolly-15k ...")
+    try:
+        dolly = load_dataset("databricks/databricks-dolly-15k", split="train")
+        for ex in (dolly.select(range(min(limit_per_set, len(dolly)))) if limit_per_set else dolly):
+            add(dolly_to_messages(ex))
+    except Exception as e:  # noqa: BLE001
+        print(f"  (skipping dolly: {e})")
+
+    print("Loading openai/gsm8k (main/train) ...")
+    gsm = load_dataset("openai/gsm8k", "main", split="train")
+    for ex in (gsm.select(range(min(limit_per_set, len(gsm)))) if limit_per_set else gsm):
+        add(gsm8k_to_messages(ex["question"], ex["answer"]))
+
+    print(f"Collected {n_kept} examples (skipped {n_skipped} that exceeded context_length).")
+    return examples
+
+
+def write_packed(examples, context_length: int, out_path: str) -> int:
+    tokens, masks = pack_examples(examples, context_length)
+    os.makedirs(os.path.dirname(out_path) or ".", exist_ok=True)
+    with h5py.File(out_path, "w") as f:
+        f.create_dataset("tokens", data=tokens)
+        f.create_dataset("loss_mask", data=masks)
+    print(f"  wrote {tokens.shape[0]} packed rows x {context_length} -> {out_path}")
+    return tokens.shape[0]
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--context_length", type=int, default=1024)
+    p.add_argument("--out_dir", default="/ephemeral/data")
+    p.add_argument("--dev_frac", type=float, default=0.02)
+    p.add_argument("--limit_per_set", type=int, default=None, help="cap examples per dataset (debug)")
+    p.add_argument("--seed", type=int, default=42)
+    args = p.parse_args()
+
+    examples = collect_examples(args.context_length, args.limit_per_set)
+    rng = np.random.default_rng(args.seed)
+    rng.shuffle(examples)
+    n_dev = max(1, int(len(examples) * args.dev_frac))
+    dev, train = examples[:n_dev], examples[n_dev:]
+
+    write_packed(train, args.context_length, os.path.join(args.out_dir, "sft_packed.h5"))
+    write_packed(dev, args.context_length, os.path.join(args.out_dir, "sft_dev_packed.h5"))
+
+
+if __name__ == "__main__":
+    main()

scripts/pretrain_base.py

@@ -0,0 +1,158 @@
+"""
+Pretrain the mid-size (~400M) base model from scratch on the Pile HDF5 corpus.
+
+This is the shared starting checkpoint for every post-training stage. It upgrades the
+original ``train_transformer.py`` recipe with the things needed to actually train a
+mid-size model on 2x H100: DistributedDataParallel, bf16 autocast, gradient accumulation,
+a cosine LR schedule with warmup, weight-decay param groups, and periodic checkpointing.
+The original ``train_transformer.py`` is left untouched.
+
+Single GPU:
+    PYTHONPATH=. python scripts/pretrain_base.py
+Both GPUs:
+    PYTHONPATH=. torchrun --standalone --nproc_per_node=2 scripts/pretrain_base.py
+
+Override any config field from the CLI, e.g. ``--batch_size 16 --train_steps 50000``.
+"""
+
+from __future__ import annotations
+
+import os
+import time
+
+import numpy as np
+import torch
+
+from config.post_training_config import PretrainConfig
+from data_loader.data_loader import get_batch_iterator
+from src.post_training.cli import parse_config_with_json
+from src.post_training.distributed import ddp_setup, ddp_wrap, cleanup, reduce_scalar
+from src.post_training.logging_utils import MetricsLogger
+from src.post_training.optim import configure_optimizer, cosine_lr
+from src.post_training.utils import (
+    amp_autocast, build_model_from_config, save_stage_ckpt, set_seed, unwrap,
+)
+
+
+@torch.no_grad()
+def estimate_loss(model, cfg, ctx, iters: int) -> dict[str, float]:
+    model.eval()
+    out = {}
+    for split, path in [("train", cfg.train_path), ("dev", cfg.dev_path)]:
+        if not os.path.exists(path):
+            continue
+        it = get_batch_iterator(path, cfg.batch_size, cfg.context_length, device=ctx.device)
+        losses = torch.zeros(iters)
+        for k in range(iters):
+            xb, yb = next(it)
+            with amp_autocast(cfg.amp_dtype, ctx.device):
+                _, loss = model(xb, yb)
+            losses[k] = loss.item()
+        out[split] = losses.mean().item()
+    model.train()
+    return out
+
+
+def main():
+    cfg, extras = parse_config_with_json(
+        PretrainConfig, "configs/pretrain.json",
+        extra={"--resume": dict(type=str, default=None, help="checkpoint to resume from")})
+    resume = extras.resume
+    ctx = ddp_setup(cfg.device)
+    # Different data shuffle per rank (the loader shuffles via numpy global RNG).
+    set_seed(cfg.seed + ctx.rank)
+
+    model = build_model_from_config(cfg).to(ctx.device)
+    start_step = 0
+    if resume and os.path.exists(resume):
+        ck = torch.load(resume, map_location="cpu", weights_only=False)
+        unwrap(model).load_state_dict(ck["model_state_dict"])
+        start_step = ck.get("step", 0)
+        if ctx.is_main:
+            print(f"Resumed from {resume} at step {start_step}")
+
+    if cfg.compile:
+        model = torch.compile(model)
+    model = ddp_wrap(model, ctx)
+
+    optimizer = configure_optimizer(unwrap(model), cfg.lr, cfg.weight_decay)
+    if resume and os.path.exists(resume):
+        ck = torch.load(resume, map_location="cpu", weights_only=False)
+        if ck.get("optimizer_state_dict"):
+            optimizer.load_state_dict(ck["optimizer_state_dict"])
+
+    logger = None
+    if ctx.is_main:
+        n_params = sum(p.numel() for p in unwrap(model).parameters())
+        print(f"Model parameters: {n_params:,} (~{n_params/1e6:.0f}M) | world_size={ctx.world_size}")
+        print(f"Effective batch = {cfg.batch_size}*{cfg.grad_accum}*{ctx.world_size} "
+              f"= {cfg.batch_size*cfg.grad_accum*ctx.world_size} seqs/step")
+        logger = MetricsLogger("pretrain", cfg.log_dir, use_wandb=cfg.use_wandb,
+                               wandb_project=cfg.wandb_project, config=vars(cfg).copy() if hasattr(cfg, "__dict__") else None)
+
+    batch_iter = get_batch_iterator(cfg.train_path, cfg.batch_size, cfg.context_length, device=ctx.device)
+    tokens_per_step = cfg.batch_size * cfg.context_length * cfg.grad_accum * ctx.world_size
+
+    model.train()
+    t0 = time.perf_counter()
+    for step in range(start_step, cfg.train_steps):
+        lr = cosine_lr(step, warmup_steps=cfg.warmup_steps, max_steps=cfg.train_steps,
+                       lr=cfg.lr, min_lr=cfg.min_lr)
+        for g in optimizer.param_groups:
+            g["lr"] = lr
+
+        optimizer.zero_grad(set_to_none=True)
+        accum_loss = 0.0
+        for micro in range(cfg.grad_accum):
+            xb, yb = next(batch_iter)
+            # Only sync grads on the last micro-step (DDP optimization).
+            sync = (micro == cfg.grad_accum - 1) or not ctx.enabled
+            cm = model.no_sync() if (ctx.enabled and not sync) else _nullcm()
+            with cm, amp_autocast(cfg.amp_dtype, ctx.device):
+                _, loss = model(xb, yb)
+                loss = loss / cfg.grad_accum
+            loss.backward()
+            accum_loss += loss.item()
+
+        torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.grad_clip)
+        optimizer.step()
+
+        if ctx.is_main and step % 20 == 0:
+            dt = time.perf_counter() - t0
+            tok_s = tokens_per_step * 20 / dt if step > start_step else 0.0
+            t0 = time.perf_counter()
+            print(f"step {step} | loss {accum_loss:.4f} | lr {lr:.2e} | {tok_s:,.0f} tok/s")
+            if logger:
+                logger.log(step, {"train_loss": accum_loss, "lr": lr, "tok_per_s": tok_s})
+
+        if step > start_step and step % cfg.eval_steps == 0:
+            ev = estimate_loss(model, cfg, ctx, cfg.eval_iters)
+            ev = {k: reduce_scalar(v, ctx) for k, v in ev.items()}
+            if ctx.is_main:
+                print(f"  [eval] step {step} | " + " | ".join(f"{k} {v:.4f}" for k, v in ev.items()))
+                if logger:
+                    logger.log(step, {f"eval_{k}": v for k, v in ev.items()})
+
+        if ctx.is_main and step > start_step and step % cfg.save_every == 0:
+            save_stage_ckpt(cfg.out_ckpt, model, optimizer, stage="pretrain",
+                            cfg=cfg, step=step, metrics={"train_loss": accum_loss})
+            print(f"  saved checkpoint -> {cfg.out_ckpt} (step {step})")
+
+    if ctx.is_main:
+        save_stage_ckpt(cfg.out_ckpt, model, optimizer, stage="pretrain", cfg=cfg,
+                        step=cfg.train_steps, metrics={"train_loss": accum_loss})
+        print(f"Done. Final checkpoint -> {cfg.out_ckpt}")
+        if logger:
+            logger.close()
+    cleanup(ctx)
+
+
+import contextlib
+
+
+def _nullcm():
+    return contextlib.nullcontext()
+
+
+if __name__ == "__main__":
+    main()

scripts/render_diagrams.sh

@@ -0,0 +1,29 @@
+#!/usr/bin/env bash
+# Render the hand-drawn, colour-coded diagram sources (docs/diagrams/src/*.mmd) to SVGs
+# (docs/diagrams/*.svg) that are embedded as images in the docs. We pre-render because
+# GitHub's live Mermaid does not reliably support `look: handDrawn`; embedding the SVG makes
+# the hand-drawn look show everywhere.
+#
+# One-time setup:
+#   sudo apt-get install -y nodejs                      # Node >= 18
+#   sudo npm install -g @mermaid-js/mermaid-cli         # provides `mmdc`
+#   # a Chrome/Chromium for headless rendering (e.g. google-chrome-stable)
+#
+# Usage (from repo root):
+#   bash scripts/render_diagrams.sh
+set -euo pipefail
+cd "$(dirname "$0")/.."
+
+SRC=docs/diagrams/src
+OUT=docs/diagrams
+CHROME="${CHROME:-/usr/bin/google-chrome-stable}"
+PP=$(mktemp)
+echo "{\"executablePath\":\"$CHROME\",\"args\":[\"--no-sandbox\",\"--disable-gpu\",\"--disable-dev-shm-usage\"]}" > "$PP"
+
+for m in "$SRC"/*.mmd; do
+  base=$(basename "$m" .mmd)
+  mmdc -p "$PP" -i "$m" -o "$OUT/$base.png" -b white -s 2   # PNG @2x: renders in every viewer (GitHub, VS Code preview)
+  echo "rendered $OUT/$base.png"
+done
+rm -f "$PP"
+echo "Done. Edit a .mmd in $SRC, re-run this script, and the embedded image updates."

scripts/run_posttraining.sh

@@ -0,0 +1,49 @@
+#!/usr/bin/env bash
+# Turnkey post-training pipeline: SFT -> Reward Model -> {DPO, PPO} -> GRPO -> eval table.
+# Assumes the base model is already pretrained (scripts/pretrain_base.py ->
+# /ephemeral/ckpts/base_pretrained.pt) and the datasets are prepared (scripts/prepare_*).
+#
+# Usage (from repo root):
+#   bash scripts/run_posttraining.sh            # use both GPUs (torchrun)
+#   NPROC=1 bash scripts/run_posttraining.sh    # single GPU
+#
+# Each stage writes a checkpoint to /ephemeral/ckpts and metrics JSONL to /ephemeral/logs.
+set -euo pipefail
+
+cd "$(dirname "$0")/.."
+export PYTHONPATH=. HF_HOME=/ephemeral/hf_cache PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
+PY=/ephemeral/venv/bin/python
+NPROC=${NPROC:-2}
+
+run() {  # run a training script single- or multi-GPU
+  if [ "$NPROC" -gt 1 ]; then
+    /ephemeral/venv/bin/torchrun --standalone --nproc_per_node="$NPROC" "$@"
+  else
+    $PY "$@"
+  fi
+}
+
+echo "############ 1/5  SFT ############"
+run scripts/train_sft.py
+
+echo "############ 2/5  Reward Model ############"
+run scripts/train_reward.py
+
+echo "############ 3/5  DPO ############"
+run scripts/train_dpo.py --loss_type dpo
+
+echo "############ 4/5  PPO (GSM8K, verifier reward) ############"
+run scripts/train_ppo.py --reward_source verifier
+
+echo "############ 5/5  GRPO (arithmetic curriculum -> GSM8K) ############"
+run scripts/train_grpo.py
+
+echo "############ Eval: GSM8K accuracy across stages ############"
+TABLE=/ephemeral/logs/stage_table.jsonl
+rm -f "$TABLE"
+for s in base_pretrained sft dpo ppo grpo; do
+  [ -f "/ephemeral/ckpts/$s.pt" ] && \
+    $PY scripts/eval_post_training.py --ckpt "/ephemeral/ckpts/$s.pt" --label "$s" --limit 200 --append "$TABLE"
+done
+$PY scripts/eval_post_training.py --table "$TABLE"
+echo "Done. Metrics in /ephemeral/logs, checkpoints in /ephemeral/ckpts."

scripts/train_dpo.py

@@ -0,0 +1,144 @@
+"""
+Direct Preference Optimization (and ORPO / KTO variants) on preference pairs.
+
+The policy is initialized from the SFT checkpoint; a frozen deep copy of it serves as the
+DPO/KTO reference (ORPO is reference-free). Reports implicit-reward accuracy on held-out
+preferences and GSM8K dev accuracy.
+
+    PYTHONPATH=. python scripts/train_dpo.py --loss_type dpo --beta 0.1
+    PYTHONPATH=. torchrun --standalone --nproc_per_node=2 scripts/train_dpo.py
+"""
+
+from __future__ import annotations
+
+import time
+
+import torch
+
+from config.post_training_config import DPOConfig
+from data_loader.preference_dataset import get_preference_iterator
+from src.post_training.cli import parse_config_with_json
+from src.post_training.distributed import ddp_setup, ddp_wrap, cleanup, reduce_scalar
+from src.post_training.dpo import dpo_loss, orpo_loss, kto_loss, implicit_accuracy
+from src.post_training.logging_utils import MetricsLogger
+from src.post_training.optim import configure_optimizer, cosine_lr
+from src.post_training.rollout import sequence_logprobs
+from src.post_training.utils import (
+    amp_autocast, load_backbone_from_ckpt, make_frozen_copy, save_stage_ckpt, set_seed, unwrap,
+)
+
+TEST_PATH = "/ephemeral/data/preferences_test.jsonl"
+
+
+def _logps(model, ids, mask, requires_grad):
+    return sequence_logprobs(model, ids, mask, requires_grad=requires_grad)
+
+
+def _compute_losses(policy, ref, batch, cfg, ctx):
+    B = batch["chosen_ids"].size(0)
+    ids = torch.cat([batch["chosen_ids"], batch["rejected_ids"]], dim=0)
+    mask = torch.cat([batch["chosen_mask"], batch["rejected_mask"]], dim=0)
+
+    with amp_autocast(cfg.amp_dtype, ctx.device):
+        psum, pn = _logps(policy, ids, mask, requires_grad=True)
+    pc, pr, ncn, nrn = psum[:B], psum[B:], pn[:B], pn[B:]
+
+    if cfg.loss_type == "orpo":
+        return orpo_loss(pc, pr, ncn, nrn, orpo_lambda=cfg.orpo_lambda)
+
+    with torch.no_grad(), amp_autocast(cfg.amp_dtype, ctx.device):
+        rsum, _ = _logps(ref, ids, mask, requires_grad=False)
+    rc, rr = rsum[:B], rsum[B:]
+    if cfg.loss_type == "kto":
+        return kto_loss(pc, pr, rc, rr, beta=cfg.beta)
+    return dpo_loss(pc, pr, rc, rr, beta=cfg.beta)
+
+
+@torch.no_grad()
+def eval_implicit_acc(policy, ref, cfg, ctx, max_batches: int = 100) -> tuple[float, float]:
+    policy.eval()
+    it = get_preference_iterator(TEST_PATH, cfg.batch_size, cfg.max_len, device=ctx.device,
+                                 rank=ctx.rank, world_size=ctx.world_size, shuffle=False, infinite=False)
+    acc, marg, n = 0.0, 0.0, 0
+    for batch in it:
+        loss, cr, rr = _compute_losses(policy, ref, batch, cfg, ctx)
+        acc += implicit_accuracy(cr, rr).item()
+        marg += (cr - rr).mean().item()
+        n += 1
+        if n >= max_batches:
+            break
+    policy.train()
+    return acc / max(1, n), marg / max(1, n)
+
+
+def main():
+    cfg, _ = parse_config_with_json(DPOConfig, "configs/dpo.json")
+    ctx = ddp_setup(cfg.device)
+    set_seed(cfg.seed + ctx.rank)
+
+    policy = load_backbone_from_ckpt(cfg, cfg.sft_ckpt, ctx.device)
+    ref = make_frozen_copy(policy, device=ctx.device) if cfg.loss_type != "orpo" else None
+    policy = ddp_wrap(policy, ctx)
+    optimizer = configure_optimizer(unwrap(policy), cfg.lr, cfg.weight_decay)
+
+    with open(cfg.pref_path) as f:
+        n_rows = sum(1 for line in f if line.strip())
+    total_steps = max(1, (n_rows // (cfg.batch_size * ctx.world_size)) * cfg.epochs)
+
+    logger = None
+    if ctx.is_main:
+        print(f"DPO[{cfg.loss_type}] from {cfg.sft_ckpt} | {n_rows} pairs | total_steps={total_steps} | beta={cfg.beta}")
+        logger = MetricsLogger(f"dpo_{cfg.loss_type}", cfg.log_dir, use_wandb=cfg.use_wandb, wandb_project=cfg.wandb_project)
+
+    train_it = get_preference_iterator(cfg.pref_path, cfg.batch_size, cfg.max_len, device=ctx.device,
+                                       rank=ctx.rank, world_size=ctx.world_size, shuffle=True, infinite=True)
+
+    policy.train()
+    t0 = time.perf_counter()
+    for step in range(total_steps):
+        lr = cosine_lr(step, warmup_steps=cfg.warmup_steps, max_steps=total_steps, lr=cfg.lr, min_lr=cfg.lr * 0.1)
+        for g in optimizer.param_groups:
+            g["lr"] = lr
+
+        batch = next(train_it)
+        loss, cr, rr = _compute_losses(policy, ref, batch, cfg, ctx)
+        optimizer.zero_grad(set_to_none=True)
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(policy.parameters(), cfg.grad_clip)
+        optimizer.step()
+
+        if ctx.is_main and step % 20 == 0:
+            acc = implicit_accuracy(cr, rr).item()
+            dt = time.perf_counter() - t0; t0 = time.perf_counter()
+            print(f"step {step}/{total_steps} | loss {loss.item():.4f} | acc {acc:.3f} | "
+                  f"r_chosen {cr.mean().item():.3f} r_rejected {rr.mean().item():.3f} | {dt:.1f}s/20")
+            if logger:
+                logger.log(step, {"train_loss": loss.item(), "train_acc": acc,
+                                  "r_chosen": cr.mean().item(), "r_rejected": rr.mean().item(), "lr": lr})
+
+        if step > 0 and step % cfg.eval_steps == 0:
+            acc, marg = eval_implicit_acc(policy, ref, cfg, ctx)
+            acc, marg = reduce_scalar(acc, ctx), reduce_scalar(marg, ctx)
+            if ctx.is_main:
+                print(f"  [eval] step {step} | test_acc {acc:.3f} | margin {marg:.3f}")
+                if logger:
+                    logger.log(step, {"test_acc": acc, "test_margin": marg})
+
+        if ctx.is_main and step > 0 and step % cfg.save_every == 0:
+            save_stage_ckpt(cfg.out_ckpt, policy, optimizer, stage=f"dpo_{cfg.loss_type}", cfg=cfg, step=step,
+                            metrics={"train_loss": loss.item()})
+
+    if ctx.is_main:
+        # Unwrap for the final eval: other ranks are already at cleanup(), so a collective on
+        # the DDP-wrapped policy here would hang (NCCL timeout). The periodic eval runs on all ranks.
+        acc, marg = eval_implicit_acc(unwrap(policy), ref, cfg, ctx)
+        save_stage_ckpt(cfg.out_ckpt, policy, optimizer, stage=f"dpo_{cfg.loss_type}", cfg=cfg, step=total_steps,
+                        metrics={"test_acc": acc, "test_margin": marg})
+        print(f"Done DPO[{cfg.loss_type}]. test_acc {acc:.3f} margin {marg:.3f} -> {cfg.out_ckpt}")
+        if logger:
+            logger.close()
+    cleanup(ctx)
+
+
+if __name__ == "__main__":
+    main()

scripts/train_grpo.py

@@ -0,0 +1,139 @@
+"""
+GRPO / RLVR on GSM8K (DeepSeek-R1 style), from scratch -- no critic, group-relative
+advantages, verifiable reward.
+
+Per iteration: for each prompt sample a group of G completions, score each with the GSM8K
+verifier, compute group-relative advantages, and update with a token-level clipped
+surrogate + KL-to-reference penalty. An arithmetic warm-up curriculum runs first so the
+policy gets non-zero reward variance before facing full GSM8K.
+
+    PYTHONPATH=. python scripts/train_grpo.py
+    PYTHONPATH=. torchrun --standalone --nproc_per_node=2 scripts/train_grpo.py
+"""
+
+from __future__ import annotations
+
+import time
+
+import torch
+
+from config.post_training_config import GRPOConfig
+from data_loader.prompt_dataset import get_prompt_iterator
+from src.post_training.cli import parse_config_with_json
+from src.post_training.chat_template import decode, encode_prompt
+from src.post_training.distributed import ddp_setup, ddp_wrap, cleanup, reduce_scalar
+from src.post_training.evaluation import gsm8k_accuracy, load_gsm8k_eval
+from src.post_training.grpo import group_advantages, grpo_loss
+from src.post_training.logging_utils import MetricsLogger
+from src.post_training.optim import configure_optimizer
+from src.post_training.rewards import reward_gsm8k
+from src.post_training.rollout import compute_logprobs, rollout_prompts
+from src.post_training.utils import (
+    amp_autocast, load_backbone_from_ckpt, make_frozen_copy, save_stage_ckpt, set_seed, unwrap,
+)
+
+
+def seq_lengths_from_mask(response_mask, prompt_lens):
+    N, T = response_mask.shape
+    pos = torch.arange(T, device=response_mask.device)
+    last = torch.where(response_mask, pos[None, :], torch.full_like(response_mask, -1, dtype=torch.long)).max(dim=1).values
+    return torch.where(last >= 0, last + 1, prompt_lens)
+
+
+def main():
+    cfg, _ = parse_config_with_json(GRPOConfig, "configs/grpo.json")
+    ctx = ddp_setup(cfg.device)
+    set_seed(cfg.seed + ctx.rank)
+
+    policy = load_backbone_from_ckpt(cfg, cfg.sft_ckpt, ctx.device)
+    ref = make_frozen_copy(policy, device=ctx.device)
+    policy_ddp = ddp_wrap(policy, ctx)
+    optimizer = configure_optimizer(unwrap(policy_ddp), cfg.lr, weight_decay=0.0)
+
+    eval_set = None
+    logger = MetricsLogger("grpo", cfg.log_dir, use_wandb=cfg.use_wandb, wandb_project=cfg.wandb_project) if ctx.is_main else None
+    if ctx.is_main:
+        print(f"GRPO from {cfg.sft_ckpt} | group_size={cfg.group_size} | world={ctx.world_size}")
+
+    warm_it = get_prompt_iterator(cfg.curriculum_path, cfg.prompts_per_iter, rank=ctx.rank,
+                                  world_size=ctx.world_size, seed=cfg.seed)
+    main_it = get_prompt_iterator(cfg.prompt_path, cfg.prompts_per_iter, rank=ctx.rank,
+                                  world_size=ctx.world_size, seed=cfg.seed)
+
+    G = cfg.group_size
+    for it in range(cfg.iterations):
+        rows = next(warm_it if it < cfg.curriculum_iters else main_it)
+        # Replicate each prompt G times, group-contiguously.
+        base_prompts = [encode_prompt([{"role": "user", "content": r["prompt"]}]) for r in rows]
+        prompts = [p for p in base_prompts for _ in range(G)]
+        golds = [r.get("gold") for r in rows for _ in range(G)]
+
+        policy.eval()
+        with amp_autocast(cfg.amp_dtype, ctx.device):
+            seqs, rmask, plens = rollout_prompts(policy, prompts, cfg.rollout_len, device=ctx.device,
+                                                 temperature=cfg.temperature, top_p=cfg.top_p if cfg.top_p < 1 else None)
+        resp = rmask[:, 1:]
+
+        seq_lens = seq_lengths_from_mask(rmask, plens)
+        responses = [decode(seqs[i, plens[i]:seq_lens[i]].tolist()) for i in range(len(prompts))]
+        rewards = torch.tensor([reward_gsm8k(responses[i], golds[i]) for i in range(len(prompts))],
+                               device=ctx.device, dtype=torch.float32)
+        adv = group_advantages(rewards, G)
+
+        with torch.no_grad(), amp_autocast(cfg.amp_dtype, ctx.device):
+            old_logp, _ = compute_logprobs(policy, seqs, rmask, temperature=cfg.temperature, requires_grad=False)
+            ref_logp, _ = compute_logprobs(ref, seqs, rmask, temperature=cfg.temperature, requires_grad=False)
+        old_logp, ref_logp = old_logp.float(), ref_logp.float()
+
+        policy.train()
+        N = seqs.size(0)
+        agg = {"loss": 0.0, "kl": 0.0, "clipfrac": 0.0, "n": 0}
+        for _ in range(cfg.grpo_epochs):
+            perm = torch.randperm(N, device=ctx.device)
+            for s in range(0, N, max(1, G)):  # minibatch ~ one group's worth
+                mb = perm[s:s + max(1, G)]
+                with amp_autocast(cfg.amp_dtype, ctx.device):
+                    new_logp, _ = compute_logprobs(policy_ddp, seqs[mb], rmask[mb], temperature=cfg.temperature, requires_grad=True)
+                loss, st = grpo_loss(new_logp.float(), old_logp[mb], ref_logp[mb], adv[mb], resp[mb],
+                                     clip=cfg.clip, kl_coef=cfg.kl_coef)
+                optimizer.zero_grad(set_to_none=True)
+                loss.backward()
+                torch.nn.utils.clip_grad_norm_(policy_ddp.parameters(), cfg.grad_clip)
+                optimizer.step()
+                agg["loss"] += loss.item(); agg["kl"] += st["kl"]; agg["clipfrac"] += st["clipfrac"]; agg["n"] += 1
+
+        mean_reward = reduce_scalar(rewards.mean().item(), ctx)
+        # Fraction of groups with non-zero reward spread (informative groups).
+        grp_std = rewards.view(-1, G).std(dim=1)
+        informative = reduce_scalar((grp_std > 1e-6).float().mean().item(), ctx)
+        resp_len = reduce_scalar(resp.float().sum(1).mean().item(), ctx)
+        n = max(1, agg["n"])
+        if ctx.is_main and it % 5 == 0:
+            phase = "warmup" if it < cfg.curriculum_iters else "gsm8k"
+            print(f"iter {it}[{phase}] | reward {mean_reward:.3f} | informative {informative:.2f} | "
+                  f"loss {agg['loss']/n:.4f} | KL {agg['kl']/n:.4f} | clipfrac {agg['clipfrac']/n:.3f} | resp_len {resp_len:.0f}")
+            if logger:
+                logger.log(it, {"reward": mean_reward, "informative_groups": informative,
+                                "loss": agg["loss"]/n, "kl": agg["kl"]/n, "resp_len": resp_len})
+
+        if ctx.is_main and it > 0 and it % cfg.eval_every == 0:
+            if eval_set is None:
+                eval_set = load_gsm8k_eval("test", limit=200)
+            res = gsm8k_accuracy(unwrap(policy_ddp), eval_set, device=ctx.device, max_new_tokens=cfg.rollout_len)
+            print(f"  [eval] iter {it} | GSM8K test acc {res['accuracy']:.3f} ({res['correct']}/{res['n']})")
+            if logger:
+                logger.log(it, {"gsm8k_acc": res["accuracy"]})
+        if ctx.is_main and it > 0 and it % cfg.save_every == 0:
+            save_stage_ckpt(cfg.out_ckpt, unwrap(policy_ddp), optimizer, stage="grpo", cfg=cfg, step=it,
+                            metrics={"reward": mean_reward})
+
+    if ctx.is_main:
+        save_stage_ckpt(cfg.out_ckpt, unwrap(policy_ddp), optimizer, stage="grpo", cfg=cfg, step=cfg.iterations, metrics={})
+        print(f"Done GRPO -> {cfg.out_ckpt}")
+        if logger:
+            logger.close()
+    cleanup(ctx)
+
+
+if __name__ == "__main__":
+    main()

scripts/train_ppo.py

@@ -0,0 +1,167 @@
+"""
+PPO RLHF on GSM8K (the classic InstructGPT recipe), from scratch.
+
+Per iteration: roll out completions with the current policy, score them (verifiable GSM8K
+reward, or a trained reward model), add a per-token KL-to-reference penalty, compute GAE
+advantages with the shared value head, then run several clipped-surrogate update epochs.
+Reports mean reward, KL, value loss, clip fraction, and held-out GSM8K accuracy.
+
+    PYTHONPATH=. python scripts/train_ppo.py --reward_source verifier
+    PYTHONPATH=. torchrun --standalone --nproc_per_node=2 scripts/train_ppo.py
+"""
+
+from __future__ import annotations
+
+import time
+
+import torch
+import torch.nn.functional as F
+
+from config.post_training_config import PPOConfig
+from data_loader.prompt_dataset import get_prompt_iterator
+from src.post_training.cli import parse_config_with_json
+from src.post_training.chat_template import EOT_ID, decode, encode_prompt
+from src.post_training.distributed import ddp_setup, ddp_wrap, cleanup, reduce_scalar
+from src.post_training.evaluation import gsm8k_accuracy, load_gsm8k_eval
+from src.post_training.logging_utils import MetricsLogger
+from src.post_training.optim import configure_optimizer
+from src.post_training.ppo import compute_gae, whiten, ppo_policy_loss, ppo_value_loss, approx_kl
+from src.post_training.reward_model import load_reward_model
+from src.post_training.rewards import reward_gsm8k
+from src.post_training.rollout import compute_logprobs, rollout_prompts
+from src.post_training.utils import (
+    amp_autocast, load_backbone_from_ckpt, make_frozen_copy, masked_mean, save_stage_ckpt, set_seed, unwrap,
+)
+from src.post_training.value_head import TransformerWithValueHead
+
+
+def actor_logp_values(actor, seqs, temperature):
+    """One forward through the actor-critic -> (logp, values) in the action frame (B,T-1)."""
+    logits, values = actor(seqs)
+    logits = logits[:, :-1, :]
+    logp_all = F.log_softmax(logits.float() / max(temperature, 1e-6), dim=-1)
+    logp = logp_all.gather(-1, seqs[:, 1:, None]).squeeze(-1)
+    return logp, values[:, :-1]
+
+
+def seq_lengths_from_mask(response_mask, prompt_lens):
+    """Real token count per row = last response position + 1 (or prompt_len if no response)."""
+    N, T = response_mask.shape
+    pos = torch.arange(T, device=response_mask.device)
+    last = torch.where(response_mask, pos[None, :], torch.full_like(response_mask, -1, dtype=torch.long)).max(dim=1).values
+    return torch.where(last >= 0, last + 1, prompt_lens)
+
+
+def main():
+    cfg, _ = parse_config_with_json(PPOConfig, "configs/ppo.json")
+    ctx = ddp_setup(cfg.device)
+    set_seed(cfg.seed + ctx.rank)
+
+    backbone = load_backbone_from_ckpt(cfg, cfg.sft_ckpt, ctx.device)
+    ref = make_frozen_copy(backbone, device=ctx.device)
+    actor = TransformerWithValueHead(backbone).to(ctx.device)
+    actor_ddp = ddp_wrap(actor, ctx)
+    optimizer = configure_optimizer(unwrap(actor_ddp), cfg.lr, weight_decay=0.0)
+
+    rm = load_reward_model(cfg, cfg.reward_ckpt, ctx.device) if cfg.reward_source == "rm" else None
+
+    eval_set = None  # loaded lazily on first eval to keep startup/smoke runs offline
+    logger = MetricsLogger("ppo", cfg.log_dir, use_wandb=cfg.use_wandb, wandb_project=cfg.wandb_project) if ctx.is_main else None
+    if ctx.is_main:
+        print(f"PPO from {cfg.sft_ckpt} | reward={cfg.reward_source} | world={ctx.world_size}")
+
+    prompt_it = get_prompt_iterator(cfg.prompt_path, cfg.prompts_per_iter, rank=ctx.rank,
+                                    world_size=ctx.world_size, seed=cfg.seed)
+
+    for it in range(cfg.iterations):
+        rows = next(prompt_it)
+        prompts = [encode_prompt([{"role": "user", "content": r["prompt"]}]) for r in rows]
+        golds = [r.get("gold") for r in rows]
+
+        # --- rollout ---
+        actor.eval()
+        with amp_autocast(cfg.amp_dtype, ctx.device):
+            seqs, rmask, plens = rollout_prompts(actor, prompts, cfg.rollout_len, device=ctx.device,
+                                                 temperature=cfg.temperature, top_p=cfg.top_p if cfg.top_p < 1 else None)
+        resp = rmask[:, 1:]                       # action-frame response mask (N, T-1)
+
+        # --- score ---
+        seq_lens = seq_lengths_from_mask(rmask, plens)
+        responses = [decode(seqs[i, plens[i]:seq_lens[i]].tolist()) for i in range(len(rows))]
+        if cfg.reward_source == "rm":
+            with torch.no_grad(), amp_autocast(cfg.amp_dtype, ctx.device):
+                task_r = rm(seqs, seq_lengths=seq_lens).float().tolist()
+        else:
+            task_r = [reward_gsm8k(responses[i], golds[i]) for i in range(len(rows))]
+
+        # --- per-token rewards: KL penalty everywhere + task reward at last response token ---
+        with torch.no_grad(), amp_autocast(cfg.amp_dtype, ctx.device):
+            old_logp, old_values = actor_logp_values(actor, seqs, cfg.temperature)
+            ref_logp, _ = compute_logprobs(ref, seqs, rmask, temperature=cfg.temperature, requires_grad=False)
+        old_logp, old_values, ref_logp = old_logp.float(), old_values.float(), ref_logp.float()
+
+        rewards = -cfg.kl_coef * (old_logp - ref_logp) * resp.float()
+        last_idx = seq_lens - 2                    # action index of the last response token
+        last_idx = last_idx.clamp(min=0)
+        task_t = torch.tensor(task_r, device=ctx.device, dtype=torch.float32)
+        rewards[torch.arange(len(rows), device=ctx.device), last_idx] += task_t
+
+        values_next = torch.cat([old_values[:, 1:], torch.zeros_like(old_values[:, :1])], dim=1)
+        adv, returns = compute_gae(rewards, old_values, values_next, resp, gamma=cfg.gamma, lam=cfg.gae_lambda)
+        adv = whiten(adv, resp)
+
+        # --- clipped PPO update epochs over minibatches of rollout rows ---
+        actor.train()
+        N = seqs.size(0)
+        stats = {"policy_loss": 0.0, "value_loss": 0.0, "clipfrac": 0.0, "kl": 0.0, "n": 0}
+        for _ in range(cfg.ppo_epochs):
+            perm = torch.randperm(N, device=ctx.device)
+            for s in range(0, N, cfg.minibatch_size):
+                mb = perm[s:s + cfg.minibatch_size]
+                with amp_autocast(cfg.amp_dtype, ctx.device):
+                    new_logp, new_values = actor_logp_values(actor_ddp, seqs[mb], cfg.temperature)
+                m = resp[mb]
+                p_loss, clipf = ppo_policy_loss(new_logp.float(), old_logp[mb], adv[mb], m, clip=cfg.clip)
+                v_loss = ppo_value_loss(new_values.float(), old_values[mb], returns[mb], m, vf_clip=cfg.vf_clip)
+                loss = p_loss + cfg.vf_coef * v_loss
+                optimizer.zero_grad(set_to_none=True)
+                loss.backward()
+                torch.nn.utils.clip_grad_norm_(actor_ddp.parameters(), cfg.grad_clip)
+                optimizer.step()
+                stats["policy_loss"] += p_loss.item(); stats["value_loss"] += v_loss.item()
+                stats["clipfrac"] += clipf.item()
+                stats["kl"] += approx_kl(new_logp.float(), old_logp[mb], m).item(); stats["n"] += 1
+
+        mean_reward = reduce_scalar(float(sum(task_r) / max(1, len(task_r))), ctx)
+        kl_ref = reduce_scalar(masked_mean(old_logp - ref_logp, resp).item(), ctx)
+        resp_len = reduce_scalar(resp.float().sum(1).mean().item(), ctx)
+        n = max(1, stats["n"])
+        if ctx.is_main and it % 5 == 0:
+            print(f"iter {it} | reward {mean_reward:.3f} | KL_ref {kl_ref:.3f} | "
+                  f"ploss {stats['policy_loss']/n:.4f} | vloss {stats['value_loss']/n:.4f} | "
+                  f"clipfrac {stats['clipfrac']/n:.3f} | resp_len {resp_len:.0f}")
+            if logger:
+                logger.log(it, {"reward": mean_reward, "kl_ref": kl_ref, "policy_loss": stats["policy_loss"]/n,
+                                "value_loss": stats["value_loss"]/n, "clipfrac": stats["clipfrac"]/n, "resp_len": resp_len})
+
+        if ctx.is_main and it > 0 and it % cfg.eval_every == 0:
+            if eval_set is None:
+                eval_set = load_gsm8k_eval("test", limit=200)
+            res = gsm8k_accuracy(unwrap(actor_ddp), eval_set, device=ctx.device, max_new_tokens=cfg.rollout_len)
+            print(f"  [eval] iter {it} | GSM8K test acc {res['accuracy']:.3f} ({res['correct']}/{res['n']})")
+            if logger:
+                logger.log(it, {"gsm8k_acc": res["accuracy"]})
+        if ctx.is_main and it > 0 and it % cfg.save_every == 0:
+            save_stage_ckpt(cfg.out_ckpt, unwrap(actor_ddp).transformer, optimizer, stage="ppo", cfg=cfg, step=it,
+                            metrics={"reward": mean_reward})
+
+    if ctx.is_main:
+        save_stage_ckpt(cfg.out_ckpt, unwrap(actor_ddp).transformer, optimizer, stage="ppo", cfg=cfg, step=cfg.iterations, metrics={})
+        print(f"Done PPO -> {cfg.out_ckpt}")
+        if logger:
+            logger.close()
+    cleanup(ctx)
+
+
+if __name__ == "__main__":
+    main()

scripts/train_reward.py

@@ -0,0 +1,132 @@
+"""
+Train the reward model on preference pairs with the Bradley-Terry loss.
+
+Initializes the reward backbone from the SFT checkpoint, adds a scalar reward head, and
+trains so chosen responses score above rejected ones. Reports held-out preference accuracy.
+
+Single GPU:
+    PYTHONPATH=. python scripts/train_reward.py
+Both GPUs:
+    PYTHONPATH=. torchrun --standalone --nproc_per_node=2 scripts/train_reward.py
+"""
+
+from __future__ import annotations
+
+import time
+
+import torch
+
+from config.post_training_config import RewardConfig
+from data_loader.preference_dataset import get_preference_iterator
+from src.post_training.cli import parse_config_with_json
+from src.post_training.distributed import ddp_setup, ddp_wrap, cleanup, reduce_scalar
+from src.post_training.logging_utils import MetricsLogger
+from src.post_training.optim import configure_optimizer, cosine_lr
+from src.post_training.reward_model import RewardModel
+from src.post_training.reward_train import bradley_terry_loss, preference_accuracy, reward_margin
+from src.post_training.utils import amp_autocast, load_backbone_from_ckpt, save_stage_ckpt, set_seed, unwrap
+
+TEST_PATH = "/ephemeral/data/preferences_test.jsonl"
+
+
+def _pair_rewards(rm, batch, cfg, ctx):
+    """Forward chosen+rejected in one pass; return (chosen_rewards, rejected_rewards)."""
+    B = batch["chosen_ids"].size(0)
+    ids = torch.cat([batch["chosen_ids"], batch["rejected_ids"]], dim=0)
+    lens = torch.cat([batch["chosen_len"], batch["rejected_len"]], dim=0)
+    with amp_autocast(cfg.amp_dtype, ctx.device):
+        rewards = rm(ids, seq_lengths=lens).float()
+    return rewards[:B], rewards[B:]
+
+
+@torch.no_grad()
+def eval_accuracy(rm, cfg, ctx, max_batches: int = 100) -> tuple[float, float]:
+    rm.eval()
+    it = get_preference_iterator(TEST_PATH, cfg.batch_size, cfg.max_len, device=ctx.device,
+                                 rank=ctx.rank, world_size=ctx.world_size, shuffle=False, infinite=False)
+    acc, marg, n = 0.0, 0.0, 0
+    for batch in it:
+        cr, rr = _pair_rewards(rm, batch, cfg, ctx)
+        acc += preference_accuracy(cr, rr).item()
+        marg += reward_margin(cr, rr).item()
+        n += 1
+        if n >= max_batches:
+            break
+    rm.train()
+    return acc / max(1, n), marg / max(1, n)
+
+
+def main():
+    cfg, _ = parse_config_with_json(RewardConfig, "configs/reward.json")
+    ctx = ddp_setup(cfg.device)
+    set_seed(cfg.seed + ctx.rank)
+
+    backbone = load_backbone_from_ckpt(cfg, cfg.sft_ckpt, ctx.device)
+    rm = RewardModel(backbone).to(ctx.device)
+    # find_unused_parameters=True: the reward model uses the backbone's forward_hidden + a
+    # reward head and never its lm_head, so lm_head params get no gradient. Without this flag
+    # DDP errors on the first backward.
+    rm = ddp_wrap(rm, ctx, find_unused_parameters=True)
+    optimizer = configure_optimizer(unwrap(rm), cfg.lr, cfg.weight_decay)
+
+    import json
+    with open(cfg.pref_path) as f:
+        n_rows = sum(1 for line in f if line.strip())
+    total_steps = max(1, (n_rows // (cfg.batch_size * ctx.world_size)) * cfg.epochs)
+
+    logger = None
+    if ctx.is_main:
+        print(f"Reward model from {cfg.sft_ckpt} | {n_rows} pairs | total_steps={total_steps}")
+        logger = MetricsLogger("reward", cfg.log_dir, use_wandb=cfg.use_wandb, wandb_project=cfg.wandb_project)
+
+    train_it = get_preference_iterator(cfg.pref_path, cfg.batch_size, cfg.max_len, device=ctx.device,
+                                       rank=ctx.rank, world_size=ctx.world_size, shuffle=True, infinite=True)
+
+    rm.train()
+    t0 = time.perf_counter()
+    for step in range(total_steps):
+        lr = cosine_lr(step, warmup_steps=cfg.warmup_steps, max_steps=total_steps, lr=cfg.lr, min_lr=cfg.lr * 0.1)
+        for g in optimizer.param_groups:
+            g["lr"] = lr
+
+        batch = next(train_it)
+        cr, rr = _pair_rewards(rm, batch, cfg, ctx)
+        loss = bradley_terry_loss(cr, rr)
+        optimizer.zero_grad(set_to_none=True)
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(rm.parameters(), cfg.grad_clip)
+        optimizer.step()
+
+        if ctx.is_main and step % 20 == 0:
+            acc = preference_accuracy(cr, rr).item()
+            dt = time.perf_counter() - t0; t0 = time.perf_counter()
+            print(f"step {step}/{total_steps} | loss {loss.item():.4f} | train_acc {acc:.3f} | lr {lr:.2e} | {dt:.1f}s/20")
+            if logger:
+                logger.log(step, {"train_loss": loss.item(), "train_acc": acc, "lr": lr})
+
+        if step > 0 and step % cfg.eval_steps == 0:
+            acc, marg = eval_accuracy(rm, cfg, ctx)
+            acc, marg = reduce_scalar(acc, ctx), reduce_scalar(marg, ctx)
+            if ctx.is_main:
+                print(f"  [eval] step {step} | test_acc {acc:.3f} | margin {marg:.3f}")
+                if logger:
+                    logger.log(step, {"test_acc": acc, "test_margin": marg})
+
+        if ctx.is_main and step > 0 and step % cfg.save_every == 0:
+            save_stage_ckpt(cfg.out_ckpt, rm, optimizer, stage="reward", cfg=cfg, step=step,
+                            metrics={"train_loss": loss.item()})
+
+    if ctx.is_main:
+        # Unwrap for the final eval: other ranks are already at cleanup(), so a collective on
+        # the DDP-wrapped model here would hang (NCCL timeout). The periodic eval runs on all ranks.
+        acc, marg = eval_accuracy(unwrap(rm), cfg, ctx)
+        save_stage_ckpt(cfg.out_ckpt, rm, optimizer, stage="reward", cfg=cfg, step=total_steps,
+                        metrics={"test_acc": acc, "test_margin": marg})
+        print(f"Done RM. test_acc {acc:.3f} margin {marg:.3f} -> {cfg.out_ckpt}")
+        if logger:
+            logger.close()
+    cleanup(ctx)
+
+
+if __name__ == "__main__":
+    main()

scripts/train_sft.py

@@ -0,0 +1,127 @@
+"""
+Supervised Fine-Tuning of the pretrained base on packed instruction data.
+
+Loads the base checkpoint, trains with the prompt-masked SFT loss, periodically reports
+masked dev loss, and saves an SFT checkpoint. DDP + bf16, single code path for 1 or N GPUs.
+
+Single GPU:
+    PYTHONPATH=. python scripts/train_sft.py
+Both GPUs:
+    PYTHONPATH=. torchrun --standalone --nproc_per_node=2 scripts/train_sft.py
+"""
+
+from __future__ import annotations
+
+import contextlib
+import math
+import time
+
+import torch
+
+from config.post_training_config import SFTConfig
+from data_loader.sft_dataset import get_sft_batch_iterator
+from src.post_training.cli import parse_config_with_json
+from src.post_training.distributed import ddp_setup, ddp_wrap, cleanup, reduce_scalar
+from src.post_training.logging_utils import MetricsLogger
+from src.post_training.optim import configure_optimizer, cosine_lr
+from src.post_training.sft import sft_loss
+from src.post_training.utils import amp_autocast, load_backbone_from_ckpt, save_stage_ckpt, set_seed, unwrap
+
+DEV_PATH = "/ephemeral/data/sft_dev_packed.h5"
+
+
+@torch.no_grad()
+def eval_dev(model, cfg, ctx, dev_path: str, max_batches: int = 50) -> float:
+    model.eval()
+    it = get_sft_batch_iterator(dev_path, cfg.batch_size, device=ctx.device,
+                                rank=ctx.rank, world_size=ctx.world_size, shuffle=False, infinite=False)
+    total, n = 0.0, 0
+    for tokens, mask, _ in it:
+        with amp_autocast(cfg.amp_dtype, ctx.device):
+            logits, _ = model(tokens)
+            loss = sft_loss(logits, tokens, mask)
+        total += loss.item(); n += 1
+        if n >= max_batches:
+            break
+    model.train()
+    return total / max(1, n)
+
+
+def main():
+    cfg, _ = parse_config_with_json(SFTConfig, "configs/sft.json")
+    ctx = ddp_setup(cfg.device)
+    set_seed(cfg.seed + ctx.rank)
+
+    model = load_backbone_from_ckpt(cfg, cfg.pretrained_ckpt, ctx.device)
+    if cfg.compile:
+        model = torch.compile(model)
+    model = ddp_wrap(model, ctx)
+    optimizer = configure_optimizer(unwrap(model), cfg.lr, cfg.weight_decay)
+
+    logger = None
+    if ctx.is_main:
+        print(f"SFT from {cfg.pretrained_ckpt} | world_size={ctx.world_size}")
+        logger = MetricsLogger("sft", cfg.log_dir, use_wandb=cfg.use_wandb, wandb_project=cfg.wandb_project)
+
+    train_it = get_sft_batch_iterator(cfg.data_path, cfg.batch_size, device=ctx.device,
+                                      rank=ctx.rank, world_size=ctx.world_size, shuffle=True, infinite=True)
+
+    # Estimate total steps for the cosine schedule from dataset size.
+    import h5py
+    with h5py.File(cfg.data_path, "r") as f:
+        n_rows = f["tokens"].shape[0]
+    steps_per_epoch = max(1, n_rows // (cfg.batch_size * ctx.world_size))
+    total_steps = cfg.max_steps if cfg.max_steps > 0 else steps_per_epoch * cfg.epochs
+    if ctx.is_main:
+        print(f"{n_rows} packed rows | ~{steps_per_epoch} steps/epoch | total_steps={total_steps}")
+
+    model.train()
+    t0 = time.perf_counter()
+    for step in range(total_steps):
+        lr = cosine_lr(step, warmup_steps=cfg.warmup_steps, max_steps=total_steps, lr=cfg.lr, min_lr=cfg.min_lr)
+        for g in optimizer.param_groups:
+            g["lr"] = lr
+
+        tokens, mask, epoch = next(train_it)
+        if epoch >= cfg.epochs and cfg.max_steps <= 0:
+            break
+        optimizer.zero_grad(set_to_none=True)
+        with amp_autocast(cfg.amp_dtype, ctx.device):
+            logits, _ = model(tokens)
+            loss = sft_loss(logits, tokens, mask)
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.grad_clip)
+        optimizer.step()
+
+        if ctx.is_main and step % 20 == 0:
+            dt = time.perf_counter() - t0; t0 = time.perf_counter()
+            print(f"step {step}/{total_steps} | loss {loss.item():.4f} | ppl {math.exp(min(20, loss.item())):.2f} | lr {lr:.2e} | {dt:.1f}s/20")
+            if logger:
+                logger.log(step, {"train_loss": loss.item(), "lr": lr})
+
+        if step > 0 and step % cfg.eval_steps == 0:
+            dev = reduce_scalar(eval_dev(model, cfg, ctx, DEV_PATH), ctx)
+            if ctx.is_main:
+                print(f"  [eval] step {step} | dev_loss {dev:.4f} | dev_ppl {math.exp(min(20, dev)):.2f}")
+                if logger:
+                    logger.log(step, {"dev_loss": dev})
+
+        if ctx.is_main and step > 0 and step % cfg.save_every == 0:
+            save_stage_ckpt(cfg.out_ckpt, model, optimizer, stage="sft", cfg=cfg, step=step,
+                            metrics={"train_loss": loss.item()})
+
+    if ctx.is_main:
+        # Use the unwrapped model for the final eval: the other ranks have already reached
+        # cleanup(), so calling the DDP-wrapped model here would launch a collective with no
+        # peer and hang (NCCL timeout). The periodic eval above runs on all ranks, so it is fine.
+        dev = eval_dev(unwrap(model), cfg, ctx, DEV_PATH)
+        save_stage_ckpt(cfg.out_ckpt, model, optimizer, stage="sft", cfg=cfg, step=total_steps,
+                        metrics={"dev_loss": dev})
+        print(f"Done SFT. dev_loss {dev:.4f} -> {cfg.out_ckpt}")
+        if logger:
+            logger.close()
+    cleanup(ctx)
+
+
+if __name__ == "__main__":
+    main()

scripts/train_transformer.py

@@ -0,0 +1,640 @@
+from __future__ import annotations
+
+import argparse
+import contextlib
+import os
+import re
+import sys
+import tempfile
+import time
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple
+
+import numpy as np
+import torch
+from tqdm import tqdm
+
+ROOT_DIR = Path(__file__).resolve().parents[1]
+if str(ROOT_DIR) not in sys.path:
+    sys.path.insert(0, str(ROOT_DIR))
+
+from config.config import default_config as config
+from src.models.transformer import Transformer
+
+
+# --- Runtime Diagnostics Helpers ---
+
+def bytes_to_gib(num_bytes: int) -> float:
+    """Convert a byte count to gibibytes for human-readable memory reports."""
+    return num_bytes / (1024 ** 3)
+
+
+def get_device_report(device: str) -> str:
+    """
+    Build a short report describing the runtime environment: PyTorch/CUDA
+    versions and, when running on a GPU, its name, capability, and total VRAM.
+    This makes it easy to collect comparable training reports across machines.
+    """
+    lines = [
+        f"PyTorch version: {torch.__version__}",
+        f"Configured device: {device}",
+        f"CUDA available: {torch.cuda.is_available()}",
+        f"CUDA version: {torch.version.cuda}",
+    ]
+
+    if device.startswith('cuda') and torch.cuda.is_available():
+        device_index = torch.cuda.current_device()
+        props = torch.cuda.get_device_properties(device_index)
+        total_vram_gib = bytes_to_gib(props.total_memory)
+        lines.extend([
+            f"GPU name: {torch.cuda.get_device_name(device_index)}",
+            f"GPU capability: {props.major}.{props.minor}",
+            f"Total VRAM: {total_vram_gib:.2f} GiB",
+        ])
+    else:
+        lines.append("GPU name: N/A (running without CUDA)")
+
+    return "\n".join(lines)
+
+
+def get_peak_memory_report(device: str) -> str:
+    """Report peak GPU memory (allocated/reserved) since the last reset, or N/A on CPU."""
+    if device.startswith('cuda') and torch.cuda.is_available():
+        peak_allocated = bytes_to_gib(torch.cuda.max_memory_allocated())
+        peak_reserved = bytes_to_gib(torch.cuda.max_memory_reserved())
+        return (
+            f"Peak VRAM allocated: {peak_allocated:.2f} GiB | "
+            f"Peak VRAM reserved: {peak_reserved:.2f} GiB"
+        )
+    return "Peak VRAM allocated: N/A | Peak VRAM reserved: N/A"
+
+
+def estimate_memory_budget(num_params: int, device: str, use_amp: bool) -> str:
+    """
+    Print a rough training VRAM budget so users can predict OOM before launching.
+
+    AdamW keeps fp32 weights + grads + two moment buffers (~16 bytes/param). This is an
+    estimate of optimizer/parameter state only (activations depend on batch/context and
+    are reduced a lot by gradient checkpointing). CUDA-only; returns N/A otherwise.
+    """
+    if not (device.startswith("cuda") and torch.cuda.is_available()):
+        return "VRAM budget: N/A (no CUDA device)"
+    # weights(4) + grad(4) + Adam m(4) + Adam v(4); AMP adds bf16/fp16 copies but keeps
+    # the fp32 master state, so ~16 B/param is a reasonable floor either way.
+    state_gib = bytes_to_gib(num_params * 16)
+    props = torch.cuda.get_device_properties(torch.cuda.current_device())
+    total_gib = bytes_to_gib(props.total_memory)
+    note = " (+ activations; reduce with --grad-checkpointing / --grad-accum)"
+    return (
+        f"VRAM budget: ~{state_gib:.2f} GiB params+optimizer state vs {total_gib:.2f} GiB "
+        f"total on {torch.cuda.get_device_name()}{note}"
+    )
+
+
+# --- Checkpoint Helpers ---
+
+CHECKPOINT_RE = re.compile(r"checkpoint_step_(\d+)\.pt$")
+
+
+def load_checkpoint_file(path: str, device: str) -> Dict[str, Any]:
+    """Load a checkpoint while supporting both newer and older PyTorch versions."""
+    try:
+        return torch.load(path, map_location=torch.device(device), weights_only=False)
+    except TypeError:
+        return torch.load(path, map_location=torch.device(device))
+
+
+def default_checkpoint_dir(out_path: str) -> str:
+    """Return a checkpoint directory tied to the configured final model path."""
+    model_path = Path(out_path)
+    return str(model_path.with_suffix("")) + "_checkpoints"
+
+
+def checkpoint_path(checkpoint_dir: str, step: int) -> str:
+    """Build a stable checkpoint path for the last completed training step."""
+    return os.path.join(checkpoint_dir, f"checkpoint_step_{step:08d}.pt")
+
+
+def checkpoint_step(path: str) -> int:
+    """Extract the step number from a checkpoint filename."""
+    match = CHECKPOINT_RE.search(os.path.basename(path))
+    if not match:
+        return -1
+    return int(match.group(1))
+
+
+def list_checkpoints(checkpoint_dir: str) -> List[str]:
+    """Return periodic checkpoints sorted by training step."""
+    if not os.path.isdir(checkpoint_dir):
+        return []
+    paths = [
+        os.path.join(checkpoint_dir, name)
+        for name in os.listdir(checkpoint_dir)
+        if CHECKPOINT_RE.search(name)
+    ]
+    return sorted(paths, key=checkpoint_step)
+
+
+def resolve_resume_path(resume: Optional[str], checkpoint_dir: str) -> Optional[str]:
+    """
+    Resolve a resume argument.
+
+    ``--resume`` with no value uses the latest periodic checkpoint in checkpoint_dir.
+    ``--resume path/to/file.pt`` loads that exact checkpoint.
+    """
+    if resume is None:
+        return None
+    if resume == "latest":
+        checkpoints = list_checkpoints(checkpoint_dir)
+        if not checkpoints:
+            raise FileNotFoundError(f"No checkpoints found in {checkpoint_dir}")
+        return checkpoints[-1]
+    return resume
+
+
+def current_lr(optimizer: torch.optim.Optimizer) -> float:
+    """Read the learning rate from the first optimizer parameter group."""
+    return float(optimizer.param_groups[0]["lr"])
+
+
+def lr_for_step(train_config: Dict[str, Any], step: int) -> float:
+    """Return the learning rate that should be active at a given step."""
+    if step > train_config['t_lr_decay_step']:
+        return float(train_config['t_lr_decayed'])
+    return float(train_config['t_lr'])
+
+
+def set_optimizer_lr(optimizer: torch.optim.Optimizer, lr: float) -> None:
+    """Set all optimizer parameter groups to the same learning rate."""
+    for group in optimizer.param_groups:
+        group["lr"] = lr
+
+
+def save_training_checkpoint(
+    path: str,
+    model: Transformer,
+    optimizer: torch.optim.Optimizer,
+    train_config: Dict[str, Any],
+    losses: List[float],
+    *,
+    step: int,
+    train_loss: Optional[float] = None,
+    dev_loss: Optional[float] = None,
+    is_final: bool = False,
+) -> None:
+    """
+    Save model, optimizer, loss history, and LR schedule metadata.
+
+    ``step`` is the last completed zero-based training step, so resume starts at
+    ``step + 1``.
+    """
+    os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
+    payload = {
+        'model_state_dict': model.state_dict(),
+        'optimizer_state_dict': optimizer.state_dict(),
+        'losses': losses,
+        'train_loss': train_loss,
+        'dev_loss': dev_loss,
+        'step': step,
+        'last_completed_step': step,
+        'steps': step + 1,
+        'is_final': is_final,
+        'config': dict(train_config),
+        'device': train_config['device'],
+        'pytorch_version': torch.__version__,
+        'cuda_version': torch.version.cuda,
+        'lr_state': {
+            'current_lr': current_lr(optimizer),
+            'initial_lr': train_config['t_lr'],
+            'decayed_lr': train_config['t_lr_decayed'],
+            'decay_step': train_config['t_lr_decay_step'],
+        },
+    }
+    target_dir = os.path.dirname(path) or "."
+    with tempfile.NamedTemporaryFile(
+        dir=target_dir,
+        prefix=f".{os.path.basename(path)}.",
+        suffix=".tmp",
+        delete=False,
+    ) as tmp_file:
+        tmp_path = tmp_file.name
+    try:
+        torch.save(payload, tmp_path)
+        os.replace(tmp_path, path)
+    except Exception:
+        if os.path.exists(tmp_path):
+            os.remove(tmp_path)
+        raise
+
+
+def restore_training_checkpoint(
+    path: str,
+    model: Transformer,
+    optimizer: torch.optim.Optimizer,
+    train_config: Dict[str, Any],
+    device: str,
+) -> Tuple[int, List[float]]:
+    """
+    Restore model/optimizer state and return ``(next_step, losses)``.
+
+    Older checkpoints did not have ``last_completed_step``. For those, ``steps``
+    is treated as the number of completed optimizer steps.
+    """
+    checkpoint = load_checkpoint_file(path, device)
+    model.load_state_dict(checkpoint['model_state_dict'])
+
+    optimizer_state = checkpoint.get('optimizer_state_dict')
+    if optimizer_state:
+        optimizer.load_state_dict(optimizer_state)
+
+    if 'last_completed_step' in checkpoint:
+        last_completed_step = int(checkpoint['last_completed_step'])
+        next_step = last_completed_step + 1
+    else:
+        next_step = int(checkpoint.get('steps', 0))
+        last_completed_step = next_step - 1
+
+    if not optimizer_state:
+        set_optimizer_lr(optimizer, lr_for_step(train_config, next_step))
+
+    losses = [float(loss) for loss in checkpoint.get('losses', [])]
+    print(
+        f"Resumed from {path}. "
+        f"Last completed step: {last_completed_step}. Next step: {next_step}."
+    )
+    return next_step, losses
+
+
+def prune_old_checkpoints(checkpoint_dir: str, keep_last: int) -> None:
+    """Keep only the most recent N periodic checkpoints when requested."""
+    if keep_last <= 0:
+        return
+    checkpoints = list_checkpoints(checkpoint_dir)
+    for old_path in checkpoints[:-keep_last]:
+        os.remove(old_path)
+
+
+def unique_output_path(out_path: str) -> str:
+    """Avoid overwriting an existing final model checkpoint."""
+    modified_model_out_path = out_path
+    save_tries = 0
+    while os.path.exists(modified_model_out_path):
+        save_tries += 1
+        model_out_name = os.path.splitext(out_path)[0]
+        modified_model_out_path = model_out_name + f"_{save_tries}" + ".pt"
+    return modified_model_out_path
+
+
+def as_float(value: Any) -> Optional[float]:
+    """Convert scalar tensors/numbers to plain floats for checkpoint metadata."""
+    if value is None:
+        return None
+    if hasattr(value, "item"):
+        return float(value.item())
+    return float(value)
+
+
+# --- Training / Evaluation ---
+
+@torch.no_grad()
+def estimate_loss(model: Transformer, train_config: Dict[str, Any], steps: int) -> Dict[str, float]:
+    """
+    Evaluate the model on training and development datasets and calculate average loss.
+
+    Args:
+        model (Transformer): The model being trained.
+        train_config (dict): Training configuration values.
+        steps (int): Number of steps to evaluate.
+
+    Returns:
+        dict: Dictionary containing average losses for 'train' and 'dev' splits.
+    """
+    out = {}
+    model.eval()  # Set the model to evaluation mode.
+    from data_loader.data_loader import get_batch_iterator
+
+    for split in ['train', 'dev']:
+        # Select the appropriate data path for the current split.
+        data_path = train_config['train_path'] if split == 'train' else train_config['dev_path']
+
+        # Create a batch iterator for evaluation.
+        batch_iterator_eval = get_batch_iterator(
+            data_path,
+            train_config['t_batch_size'],
+            train_config['t_context_length'],
+            device=train_config['device'],
+        )
+
+        # Track loss values for each evaluation step.
+        losses_eval = []
+        for _ in range(steps):
+            try:
+                # Fetch a batch and calculate the loss.
+                xb, yb = next(batch_iterator_eval)
+                _, loss = model(xb, yb)
+                losses_eval.append(float(loss.item()))
+            except StopIteration:
+                # Handle the case where the data iterator ends early.
+                print(f"Warning: Iterator for {split} ended early.")
+                break
+
+        # Compute the mean loss for the current split.
+        out[split] = float(np.mean(losses_eval)) if losses_eval else float("nan")
+
+    model.train()  # Restore the model to training mode.
+    return out
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Train the Transformer model from scratch.")
+    parser.add_argument(
+        "--resume",
+        nargs="?",
+        const="latest",
+        default=None,
+        help=(
+            "Resume from a checkpoint path. Pass --resume with no value, or "
+            "--resume latest, to use the newest checkpoint in the checkpoint directory."
+        ),
+    )
+    parser.add_argument(
+        "--checkpoint-every",
+        type=int,
+        default=None,
+        help="Save a periodic checkpoint every N completed steps. 0 disables periodic checkpoints.",
+    )
+    parser.add_argument(
+        "--checkpoint-dir",
+        type=str,
+        default=None,
+        help="Directory for periodic checkpoints. Defaults to a directory next to t_out_path.",
+    )
+    parser.add_argument(
+        "--keep-last",
+        type=int,
+        default=None,
+        help="Keep only the most recent N periodic checkpoints. 0 keeps all.",
+    )
+    # --- Memory-optimisation flags (opt-in; all default to the config values, which are OFF) ---
+    parser.add_argument(
+        "--amp",
+        dest="amp",
+        action="store_true",
+        default=None,
+        help="Enable bf16/fp16 mixed-precision autocast (CUDA only; ignored on CPU).",
+    )
+    parser.add_argument(
+        "--amp-dtype",
+        type=str,
+        choices=["bf16", "fp16"],
+        default=None,
+        help="Autocast dtype when --amp is set: bf16 (default, no GradScaler) or fp16.",
+    )
+    parser.add_argument(
+        "--grad-checkpointing",
+        dest="grad_checkpointing",
+        action="store_true",
+        default=None,
+        help="Recompute transformer-block activations in backward to save VRAM.",
+    )
+    parser.add_argument(
+        "--grad-accum",
+        type=int,
+        default=None,
+        help="Accumulate gradients over N micro-batches per optimizer step (effective batch xN).",
+    )
+    parser.add_argument(
+        "--report-memory",
+        dest="report_memory",
+        action="store_true",
+        default=None,
+        help="Print a rough VRAM budget (params + optimizer state) before training (CUDA only).",
+    )
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    train_config = dict(config)
+    checkpoint_every = (
+        args.checkpoint_every
+        if args.checkpoint_every is not None
+        else train_config.get('t_checkpoint_steps', 0)
+    )
+    keep_last = (
+        args.keep_last
+        if args.keep_last is not None
+        else train_config.get('t_keep_last_checkpoints', 0)
+    )
+    checkpoint_dir = (
+        args.checkpoint_dir
+        or train_config.get('t_checkpoint_dir')
+        or default_checkpoint_dir(train_config['t_out_path'])
+    )
+
+    # --- Resolve memory-optimisation options (CLI overrides config; all default OFF) ---
+    use_amp = args.amp if args.amp is not None else bool(train_config.get('use_amp', False))
+    amp_dtype_name = args.amp_dtype or train_config.get('amp_dtype', 'bf16')
+    use_grad_ckpt = (
+        args.grad_checkpointing if args.grad_checkpointing is not None
+        else bool(train_config.get('use_gradient_checkpointing', False))
+    )
+    grad_accum = max(1, args.grad_accum if args.grad_accum is not None
+                     else int(train_config.get('grad_accum_steps', 1)))
+    report_memory = (
+        args.report_memory if args.report_memory is not None
+        else bool(train_config.get('report_memory_budget', False))
+    )
+
+    device_is_cuda = train_config['device'].startswith('cuda') and torch.cuda.is_available()
+    if use_amp and not device_is_cuda:
+        print("[mem-opt] --amp requested but no CUDA device available; disabling AMP.")
+        use_amp = False
+    amp_dtype = torch.bfloat16 if amp_dtype_name == 'bf16' else torch.float16
+
+    def autocast_ctx():
+        if use_amp:
+            return torch.autocast(device_type='cuda', dtype=amp_dtype)
+        return contextlib.nullcontext()
+
+    # GradScaler is only needed for fp16; bf16 has enough range. A disabled scaler is a no-op,
+    # so the scale/unscale_/step/update calls below work unchanged for bf16 and CPU.
+    use_scaler = use_amp and amp_dtype == torch.float16
+    scaler = torch.amp.GradScaler('cuda', enabled=use_scaler)
+
+    # --- Initialize the Model and Print Parameters ---
+
+    # Print runtime/device diagnostics and reset GPU peak-memory stats before training.
+    print(get_device_report(train_config['device']))
+    if train_config['device'].startswith('cuda') and torch.cuda.is_available():
+        torch.cuda.reset_peak_memory_stats()
+
+    model = Transformer(
+        n_head=train_config['n_head'],
+        n_embed=train_config['n_embed'],
+        context_length=train_config['context_length'],
+        vocab_size=train_config['vocab_size'],
+        N_BLOCKS=train_config['n_blocks'],
+    ).to(train_config['device'])
+
+    # Print the total number of parameters.
+    total_params = sum(p.numel() for p in model.parameters())
+    print(f"Total number of parameters in the model: {total_params:,}")
+
+    # Apply opt-in memory optimisations.
+    model.gradient_checkpointing = use_grad_ckpt
+    if report_memory:
+        print(estimate_memory_budget(total_params, train_config['device'], use_amp))
+    if use_amp or use_grad_ckpt or grad_accum > 1:
+        print(
+            f"[mem-opt] amp={use_amp}"
+            f"{'(' + amp_dtype_name + ')' if use_amp else ''} "
+            f"grad_checkpointing={use_grad_ckpt} grad_accum={grad_accum}"
+        )
+
+    # --- Optimizer Setup and Loss Tracking ---
+
+    # Set up the AdamW optimizer with the specified learning rate.
+    optimizer = torch.optim.AdamW(model.parameters(), lr=train_config['t_lr'])
+
+    # List to track loss values during training.
+    losses: List[float] = []
+    start_step = 0
+    last_completed_step = -1
+    resume_path = resolve_resume_path(args.resume, checkpoint_dir)
+    if resume_path is not None:
+        start_step, losses = restore_training_checkpoint(
+            resume_path,
+            model,
+            optimizer,
+            train_config,
+            train_config['device'],
+        )
+        last_completed_step = start_step - 1
+
+    # Define a window size for averaging recent losses in the training loop.
+    avg_window = 64
+
+    # --- Training Loop ---
+
+    from data_loader.data_loader import get_batch_iterator
+
+    # Create a batch iterator for the training data.
+    batch_iterator = get_batch_iterator(
+        train_config['train_path'],
+        train_config['t_batch_size'],
+        train_config['t_context_length'],
+        device=train_config['device'],
+    )
+
+    # Number of tokens processed per optimizer step (batch * context * grad_accum), for throughput.
+    tokens_per_step = train_config['t_batch_size'] * train_config['t_context_length'] * grad_accum
+    last_eval_time = time.perf_counter()
+    latest_train_loss = None
+    latest_dev_loss = None
+
+    # Create a progress bar to monitor training progress.
+    pbar = tqdm(range(start_step, train_config['t_train_steps']))
+    for step in pbar:
+        try:
+            # Start the step timer.
+            step_start_time = time.perf_counter()
+
+            # Accumulate gradients over `grad_accum` micro-batches (==1 => original behaviour).
+            optimizer.zero_grad(set_to_none=True)
+            step_loss = 0.0
+            for _ in range(grad_accum):
+                xb, yb = next(batch_iterator)
+                with autocast_ctx():
+                    _, loss = model(xb, yb)
+                    # Scale so the accumulated gradient equals the full-batch mean gradient.
+                    loss = loss / grad_accum
+                scaler.scale(loss).backward()
+                step_loss += float(loss.item())
+
+            # Record the (accumulated) loss for tracking.
+            losses.append(step_loss)
+            pbar.set_description(f"Train loss: {np.mean(losses[-avg_window:]):.4f}")
+
+            # Clip gradients to prevent exploding gradients (unscale first for fp16 AMP).
+            scaler.unscale_(optimizer)
+            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
+
+            scaler.step(optimizer)
+            scaler.update()
+            last_completed_step = step
+
+            # Measure step time and instantaneous throughput for diagnostics.
+            step_time = time.perf_counter() - step_start_time
+            tokens_per_second = tokens_per_step / step_time if step_time > 0 else float('inf')
+
+            # Periodically evaluate the model on training and development data.
+            if step % train_config['t_eval_steps'] == 0:
+                evaluation_losses = estimate_loss(model, train_config, train_config['t_eval_iters'])
+                latest_train_loss = evaluation_losses['train']
+                latest_dev_loss = evaluation_losses['dev']
+                # Report timing/throughput for the most recent step and wall-time since last eval.
+                now = time.perf_counter()
+                elapsed_since_eval = now - last_eval_time
+                last_eval_time = now
+                print(
+                    f"Step: {step}, Train loss: {latest_train_loss:.4f}, Dev loss: {latest_dev_loss:.4f}, "
+                    f"Step time: {step_time:.3f}s, Throughput: {tokens_per_second:.2f} tokens/s, "
+                    f"Elapsed since last eval: {elapsed_since_eval:.2f}s"
+                )
+                print(get_peak_memory_report(train_config['device']))
+
+            # Decay the learning rate at the specified step.
+            if step == train_config['t_lr_decay_step']:
+                print('Decaying learning rate')
+                set_optimizer_lr(optimizer, train_config['t_lr_decayed'])
+
+            if checkpoint_every and checkpoint_every > 0 and (step + 1) % checkpoint_every == 0:
+                path = checkpoint_path(checkpoint_dir, step)
+                save_training_checkpoint(
+                    path,
+                    model,
+                    optimizer,
+                    train_config,
+                    losses,
+                    step=step,
+                    train_loss=as_float(latest_train_loss),
+                    dev_loss=as_float(latest_dev_loss),
+                )
+                prune_old_checkpoints(checkpoint_dir, int(keep_last or 0))
+                print(f"Saved checkpoint to {path}")
+        except StopIteration:
+            # Handle the case where the training data iterator ends early.
+            print("Training data iterator finished early.")
+            break
+
+    # --- Save Model and Final Evaluation ---
+
+    # Perform a final evaluation of the model on training and development datasets.
+    evaluation_losses = estimate_loss(model, train_config, 200)
+    train_loss = evaluation_losses['train']
+    dev_loss = evaluation_losses['dev']
+
+    final_step = max(last_completed_step, start_step - 1)
+    modified_model_out_path = unique_output_path(train_config['t_out_path'])
+
+    # Save the model's state dictionary, optimizer state, and training metadata
+    # (including the runtime device / PyTorch / CUDA versions for reproducibility).
+    save_training_checkpoint(
+        modified_model_out_path,
+        model,
+        optimizer,
+        train_config,
+        losses,
+        step=final_step,
+        train_loss=train_loss,
+        dev_loss=dev_loss,
+        is_final=True,
+    )
+    print(f"Saved model to {modified_model_out_path}")
+    print(get_peak_memory_report(train_config['device']))
+    print(f"Finished training. Train loss: {train_loss:.4f}, Dev loss: {dev_loss:.4f}")
+
+
+if __name__ == "__main__":
+    main()