The main Transformer model. This class combines token and position embeddings with a sequence of Transformer blocks and a final linear l

Wrap a :class:`Transformer` and add a scalar reward head (no ``lm_head`` used).

Wrap a :class:`Transformer` and add a scalar value head for PPO. ``forward(idx)`` returns ``(logits, values)`` where ``logits`` are the policy lo

A single Transformer block. This block consists of a multi-head attention layer followed by an MLP, with layer normalization and residua

A simple Multi-Layer Perceptron with one hidden layer. This module is used within the Transformer block for feed-forward processing. It

Multi-Head Attention module. This module combines multiple attention heads in parallel. The outputs of each head are concatenated and pa

A single attention head. This module calculates attention scores and applies them to the values. It includes key, query, and value proje

Result of :func:`generate_with_logprobs`. Attributes: sequences: (B, P+G) prompt tokens followed by generated tokens. respons

Pretrain the mid base model from scratch on the Pile HDF5 (mix in task text late).