Return the underlying module behind a DDP wrapper (or the model itself).

Teacher-forced recomputation of per-token log-probs of ``sequences`` under ``model``. Mirrors the model's training shift: ``logits[:, t]`` p

Return a bf16 autocast context on CUDA when requested, else a no-op context. bf16 needs no GradScaler (unlike fp16), which keeps the training loo

All-reduce a python scalar across ranks (sum or mean). No-op when single-process.

Save a checkpoint in the repo's existing shape (``model_state_dict`` / ``optimizer_state_dict``) plus post-training metadata (``stage``, ``cf

Mean of ``values`` over positions where ``mask`` is truthy (safe if mask empty).

Reward for a GSM8K response: correctness (dominant) + small format bonus, clipped.

Seed python / numpy / torch (incl. CUDA) for reproducible runs.

Autoregressively sample a completion for each prompt and record per-token log-probs. ``prompt_ids`` must be (B, P) with a shared length P (l

Deep-copy a model, put it in eval mode, and disable all gradients. Used for the DPO/PPO/GRPO reference model and the PPO old-policy snapshot.

AdamW with weight decay applied only to >=2D parameters (matrices), not to biases / norms / 1D params. Standard GPT recipe.

Initialize the process group if launched under torchrun; otherwise single-process. Reads ``RANK`` / ``LOCAL_RANK`` / ``WORLD_SIZE`` from the envi

Wrap a model in DDP when running multi-GPU; return it unchanged otherwise. Only the *trainable* model should be wrapped. Reference / old-policy /

Decode token ids back to text (for eval / verifier parsing / sample dumps). Defensive: drops the EOT terminator and any id >= EOT (50256), s

Tokenize a conversation and build an aligned per-token loss mask. The mask is 1 over assistant content (and its terminating EOT) and 0 every

Extract the model's final numeric answer with graceful fallback: 1) inside an ``<answer>...</answer>`` tag, else 2) after a GSM8K-style `

Like :func:`parse_config`, but resolves the config from a JSON file too. Adds two flags on top of the per-field ``--field`` overrides:

Read a field from either a dataclass/object (attr) or a dict.

Convert a byte count to gibibytes for human-readable memory reports.

Generalized Advantage Estimation over the response. Args (all (B, L) in the action frame; ``resp_mask`` bool): rewards: per-acti

Run the backbone and return the final hidden states AFTER the final layer norm. This is exactly the tensor that ``lm_head`` consumes

Creates an iterator for generating batches of data from an HDF5 file. Args: data_path (str): Path to the HDF5 file containing tokeni

Extract the numeric gold answer from a GSM8K ``answer`` field (ends in ``#### N``).

Build a Transformer from ``cfg`` and load backbone weights from a checkpoint saved by the pretraining script or any post-training stage (``mo

Parse the first number out of a string, tolerating ``$`` and thousands commas.

Save model, optimizer, loss history, and LR schedule metadata. ``step`` is the last completed zero-based training step, so resume starts at

Construct a fresh :class:`Transformer` from a config carrying the standard keys ``n_head, n_embed, context_length, vocab_size, n_blocks``. Wo

Build a stable checkpoint path for the last completed training step.

Linear warmup to ``lr`` over ``warmup_steps``, then cosine decay to ``min_lr`` by ``max_steps`` (constant ``min_lr`` afterwards).

Token ids for the prompt form (ends in the assistant header, ready to generate).

Yield collated preference batches (dict of tensors). Rows are sharded across ranks.

Group-relative advantage: ``(r - group_mean) / (group_std + eps)``. ``rewards`` is (num_prompts * group_size,) laid out group-contiguously (

Token-level clipped surrogate + KL penalty. Args: new_logp/old_logp/ref_logp: (B, L) per-token log-probs (policy / sampling / ref).

Compute greedy GSM8K accuracy over ``(question, answer_field)`` pairs. Returns ``{"accuracy", "n", "correct", "samples"}`` where ``samples``

Whether the model's parsed answer matches the gold answer (used for accuracy eval).

Load ``(question, answer_field)`` pairs from GSM8K for evaluation.

Clipped surrogate policy loss. Returns (loss, clip_fraction).

Encode text as ordinary tokens (no special-token handling).

Get logits whether ``model`` returns ``(logits, loss)`` (Transformer) or ``(logits, values)`` (value-head wrapper) or a bare logits tensor.

Generate a response string for each tokenized prompt, returned in input order. Prompts are grouped by identical length (the model has no pad

Infinitely yield lists of ``prompts_per_iter`` rows (this rank's shard).

Return the shared r50k_base encoder (cached so we build it once).

Return periodic checkpoints sorted by training step.

Build a :class:`Transformer` from a checkpoint's stored cfg and load its backbone weights (tolerates ``module.``/``transformer.`` prefixes and rew

Clipped value-function loss (0.5 * max of clipped/unclipped squared error).

Normalize advantages to zero mean / unit std over masked (response) positions.

Map an empty text box back to None for optional (``| None``) string fields.

Numerically stable log(1 - exp(x)) for x < 0.

Forward chosen+rejected in one pass; return (chosen_rewards, rejected_rewards).

Combines token and position embeddings. Args: idx (torch.Tensor): Input token indices. Returns: tor

Split an HH-RLHF conversation string into (prompt_context, final_response).

One forward through the actor-critic -> (logp, values) in the action frame (B,T-1).

Convert scalar tensors/numbers to plain floats for checkpoint metadata.

Fraction of generated tokens equal to the target token (dense, easy to learn).

Downloads a file from the given URL and saves it with the specified file name. Displays a progress bar using tqdm. Args: url

Evaluate the model on training and development datasets and calculate average loss. Args: model (Transformer): The model being train

Given per-token ``values`` (B, T) and ``seq_lengths`` (B,), return the value at the last real token of each row, i.e. ``values[i, seq_lengths

Generates new tokens given a starting sequence. Args: idx (torch.Tensor): Initial sequence of token indices.

Generate a response to ``user_text``. - chat mode (default): wraps the prompt in the chat template (optionally with a ``system`` messa

Report peak GPU memory (allocated/reserved) since the last reset, or N/A on CPU.

Yield ``(tokens, loss_mask, epoch)`` batches from a packed SFT HDF5 file. Rows are sharded across ranks (each rank sees a disjoint stride) s

True iff there is exactly one well-formed ``<answer>...</answer>`` block.

Fraction of pairs where the implicit (DPO/KTO/ORPO) reward prefers chosen.

Per-token unbiased, non-negative KL estimator (Schulman's k3) for KL(policy||ref).

Resolve ``cfg_cls`` from ``base.json`` + the stage JSON + CLI overrides. Args: cfg_cls: the stage dataclass (e.g. ``SFTConfig``).

Fraction of pairs where the model scores the chosen response higher.

Process a specified number of lines from each .jsonl.zst file in the input directory and save encoded tokens to an HDF5 file. Args:

Keep only the most recent N periodic checkpoints when requested.

Resolve a resume argument. ``--resume`` with no value uses the latest periodic checkpoint in checkpoint_dir. ``--resume path/to/file.pt`