Avoid overwriting an existing final model checkpoint.

Prove the GSM8K scoring is correct independent of any model: a response containing the gold scores correct; a wrong number scores incorrect.

Gold answers in the RL prompt files must match the GSM8K dataset's #### answer.

The critical SFT check: the loss_mask must be 1 EXACTLY on assistant-completion tokens and 0 on the prompt / role markers. We decode masked vs unm

Initializes the multi-head attention module. Args: n_head (int): The number of parallel attention heads. n_e

Initializes the Transformer block. Args: n_head (int): The number of attention heads. n_embed (int): The dim

Initializes the MLP module. Args: n_embed (int): The dimensionality of the input embedding.

Initializes the Transformer model. Args: n_head (int): Number of attention heads. n_embed (int): Embedding d

Extract the step number from a checkpoint filename.

Return the content of the first ``<think>...</think>`` block, if any.

Scalar reward per sequence (B,). Args: idx: (B, T) token ids. seq_lengths: (B,) number of real (non-pad) tok

Forward pass through the attention head. Args: x (torch.Tensor): Input tensor of shape (B, T, C). Returns:

Forward pass through the multi-head attention. Args: x (torch.Tensor): Input tensor of shape (B, T, C). Returns

Forward pass through the Transformer block. Args: x (torch.Tensor): Input tensor. Returns: torch.Te

Forward pass through the MLP. Args: x (torch.Tensor): Input tensor of shape (B, T, C), where B is batch size,

Forward pass through the Transformer. Args: idx (torch.Tensor): Input token indices. targets (torch.Tensor,

Forward pass focusing on the embedding and attention parts. Args: x (torch.Tensor): Input tensor. Returns:

Forward pass focusing on the embedding and attention blocks. Args: idx (torch.Tensor): Input token indices. Ret

Return the job_id of a *running* GPU job, if any (for the launch guard).

Return a list of (idx, name, used_mb, total_mb, util) via nvidia-smi, or [].

Start ``argv`` as a detached background job. ``kind`` is 'gpu' or 'cpu' (for the guard).

Per-row (B,) mean of ``values`` (B,T) over masked positions.

Numeric columns worth plotting (drop bookkeeping).

Parse CLI overrides for a config dataclass. Returns ``(cfg, extras)`` where ``cfg`` is an instance of ``cfg_cls`` with any provided ``--

Render a list of ``{"role", "content"}`` messages to the plain-text chat format. This is purely for display/debugging; tokenization goes thr

Pure format reward (no gold): 1.0 if exactly one well-formed answer block else 0.

Per-token predictive entropy ``(B, T-1)`` of ``model`` on ``sequences`` (for an optional PPO entropy bonus / diagnostics).

Return this rank's strided slice of ``range(num_items)`` for data sharding.

Return 'running' | 'finished' | 'failed' | 'none'.

Terminate a running job and all its workers (process group).

Per-token values without computing the (large) vocab logits -- used when scoring rollouts where only V(s_t) is needed.