Struct Gpt2 

pub struct Gpt2 {
    pub config: Gpt2Config,
    /* private fields */
}

Fields

config: Gpt2Config

Implementations

impl Gpt2

pub fn from_safetensors( path: impl AsRef<Path>, config: Gpt2Config, device: &Device, ) -> Result<Self>

Load HF GPT-2 weights. path is the .safetensors file. Weight names may or may not carry the “transformer.” prefix; both are accepted. The LM head is tied to wte (the checkpoint has no separate tensor).

pub fn from_safetensors_bytes( bytes: &[u8], config: Gpt2Config, device: &Device, ) -> Result<Self>

Load HF GPT-2 weights from in-memory .safetensors bytes — the primary form; on wasm the bytes arrive from an HTTP fetch. wte never touches the device un-chunked: it is uploaded row-chunked straight from the file buffer (roadmap v4, pitfall 2).

pub fn new_cache(&self) -> Result<KvCache>

Empty KV cache sized for this model’s full context.

pub fn forward(&self, ids: &[u32]) -> Result<Tensor>

Full logits [t, vocab] — used by verification tests. The LM head is wte-tied: logits = h @ wte^T.

pub fn logits_last(&self, ids: &[u32]) -> Result<Vec<f32>>

Logits for the last position only, recomputing the full context (no cache) — the reference decode path used by verification tests.

pub fn logits_step(&self, ids: &[u32], cache: &mut KvCache) -> Result<Vec<f32>>

Last-position logits for ids continuing from cache (incremental decode: pass the full prompt once, then one token at a time).

pub async fn logits_step_async( &self, ids: &[u32], cache: &mut KvCache, ) -> Result<Vec<f32>>

Async form of Gpt2::logits_step — identical math; the readback is awaited so it works on wasm32 (roadmap v4, pitfall 14).

pub fn hidden_step(&self, ids: &[u32], cache: &mut KvCache) -> Result<Vec<f32>>

The last position’s hidden state after the final LayerNorm — [n_embd], exactly the vector Gpt2::logits_step hands to the head.

This is the only stage at which two models’ activations can be merged and still decode: it is what wte^T multiplies, so models that share a (frozen) wte share a basis here and nowhere else. Anything deeper — per-head Q/K/V, the post-GELU MLP hidden — is private to one model.

pub async fn hidden_step_async( &self, ids: &[u32], cache: &mut KvCache, ) -> Result<Vec<f32>>

Async form of Gpt2::hidden_step — identical math, awaited readback so it works on wasm32.

pub fn logits_from_hidden(&self, h: &[f32]) -> Result<Vec<f32>>

Logits for a hidden state supplied from outside: h @ wte^T, the same wte-tied head every other decode path applies. Lets a vector this model did not produce — a merge of several models’ Gpt2::hidden_step outputs, say — be decoded to a distribution over the vocabulary.

pub async fn logits_from_hidden_async(&self, h: &[f32]) -> Result<Vec<f32>>

Async form of Gpt2::logits_from_hidden.

pub fn wte_host(&self) -> Result<Vec<f32>>

The token embedding table as host floats, chunks reassembled. Exists so a council can verify its experts really do share one wte — the invariant that makes merging their hidden states mean anything.

pub async fn logits_step_attn_async( &self, ids: &[u32], cache: &mut KvCache, ) -> Result<(Vec<f32>, Vec<AttnStep>)>

Gpt2::logits_step_async plus every block’s attention probabilities for this step, in layer order.

The cheap probe: exactly Gpt2::logits_step_trace_async with no detail layers and no top-n, so the two can never disagree about what the model attended with.

pub async fn logits_step_trace_async( &self, ids: &[u32], cache: &mut KvCache, detail_layers: usize, top_n: usize, ) -> Result<(Vec<f32>, StepTrace)>

Gpt2::logits_step_async plus a full calculation trace for this step: the embedding, every block’s attention, the whole forward pass through the first detail_layers blocks (both LayerNorms, Q/K/V, pre-softmax scores, the concatenated heads and their projection, both residual adds, post-GELU MLP), the final LayerNorm, and the top_n most likely next tokens with their probabilities.

The logits are identical to logits_step_async — the probe only reads tensors the step computes anyway. The whole step is enqueued first, and every captured tensor comes back in a single batched readback: one at a time they cost ~2.7x the decode itself, since the price is the round trip rather than the bytes.

detail_layers is the cost knob. At 0 this is the attention-only probe and nothing but the post-softmax tensors is read. Each detail layer adds q_len * (13 * n_embd + n_head * kv_len) floats, so a long prefill is the only expensive call — a decode step is q_len == 1.

pub async fn surprisal_async(&self, ids: &[u32]) -> Result<Surprisal>

How surprised the model was by text that is already there.

This is reading, not writing: for each position the model is asked what it expected before seeing the character that actually followed, and the answer is scored against what did. It is a teacher-forced scoring pass, so the whole sequence costs one forward pass rather than t decode steps — the model reads a paragraph in the time it would take to generate one character.

bits[i] is -log2 p(ids[i] | ids[..i]): 0 means “entirely expected”, and log2(vocab_size) is what a uniform guess would score. bits[0] is 0 — nothing precedes the first token, so nothing about it can be a surprise.

top[i] and top_p[i] are the token the model would have picked at that position and how sure it was, which is what makes a surprise legible: “it expected e here” says more than a number.

The softmax and the gather are done on the host deliberately. The GPU ops that would do them (ops::softmax over [t, vocab], and ops::gather_nll) sit behind the train feature, and for a character-level vocabulary the host loop is free. Note the cost model for a BPE model, though: the readback is t × vocab × 4 bytes, which for GPT-2’s 50257-token vocabulary is ~196 KB per position.

pub fn init_random( config: Gpt2Config, device: &Device, seed: u64, ) -> Result<Gpt2>

Random initialization for training from scratch: N(0, 0.02) weights, residual projections scaled by 1/sqrt(2*n_layer), zero biases, identity LayerNorms.

pub fn param_specs(&self) -> Vec<(String, bool)>

Parameter names (safetensors keys) and weight-decay flags, in the canonical order used by params, params_mut, and loss_grads.

pub fn params(&self) -> Result<Vec<&Tensor>>

Parameters in param_specs order.

pub fn params_mut(&mut self) -> Result<Vec<&mut Tensor>>

Mutable parameters in param_specs order (for the optimizer).

pub fn save_safetensors(&self, path: impl AsRef<Path>) -> Result<()>

Save a checkpoint loadable by from_safetensors. A chunked wte is reassembled host-side first.

pub fn loss(&self, input: &[u32], targets: &[u32]) -> Result<f32>

Available on crate feature train only.

Mean cross-entropy of targets given input, forward only — no tape, no gradients, no dropout. This is the evaluation path: a validation split scored with Gpt2::loss_grads would build and discard a full backward graph per window.

pub fn loss_grads( &self, input: &[u32], targets: &[u32], dropout_p: f32, seed: u32, ) -> Result<(f32, Vec<Tensor>)>

Available on crate feature train only.

One training forward + backward on a single sequence: mean cross-entropy of targets given input, and gradients for every parameter in param_specs order. Batching is achieved by accumulating grads over calls (roadmap v4 batching policy).

pub fn generate( &self, tokenizer: &impl Tokenizer, prompt: &str, max_new_tokens: usize, sampling: Sampling, ) -> Result<String>

Autoregressive generation with KV-cache decode. Returns the full text (prompt + completion).

pub fn generate_streaming( &self, tokenizer: &impl Tokenizer, prompt: &str, max_new_tokens: usize, sampling: Sampling, on_token: impl FnMut(u32, &str), ) -> Result<String>

Sync streaming generation with KV-cache decode. on_token fires exactly once per generated token (never for the prompt) with the token id and the newly decodable text delta — which may be empty when a multi-byte character spans several BPE tokens, since a partial UTF-8 sequence is held back until its trailing bytes arrive.

This is the accurate throughput hook: callback invocations map 1:1 to tokens, unlike Gpt2::generate_async’s text-delta callback.

pub fn generate_streaming_ctl( &self, tokenizer: &impl Tokenizer, prompt: &str, max_new_tokens: usize, sampling: Sampling, on_token: impl FnMut(u32, &str) -> ControlFlow<()>, ) -> Result<String>

Gpt2::generate_streaming with an interruptible callback: returning ControlFlow::Break stops after the current token and returns the text generated so far. Long interactive runs need this — a callback that cannot abort can only be “cancelled” once the whole run finishes.

pub async fn generate_async( &self, tokenizer: &impl Tokenizer, prompt: &str, max_new_tokens: usize, sampling: Sampling, on_text: impl FnMut(&str), ) -> Result<String>

Async autoregressive generation with KV-cache decode — the browser form of Gpt2::generate. on_text receives each newly decoded text fragment (the delta of the full decode, so multi-byte characters that span BPE tokens are never split) so a page can stream output.

pub async fn generate_async_ctl( &self, tokenizer: &impl Tokenizer, prompt: &str, max_new_tokens: usize, sampling: Sampling, on_text: impl FnMut(&str) -> ControlFlow<()>, ) -> Result<String>

Gpt2::generate_async with an interruptible callback — the async counterpart of Gpt2::generate_streaming_ctl. Returning ControlFlow::Break stops after the current token, which is how a page offers a working “stop” button.

pub async fn generate_async_probe( &self, tokenizer: &impl Tokenizer, prompt: &str, max_new_tokens: usize, sampling: Sampling, on_text: impl FnMut(&str) -> ControlFlow<()>, on_attn: Option<impl FnMut(&[AttnStep])>, ) -> Result<String>

Gpt2::generate_async_ctl with an optional attention probe: on_attn, when present, fires once per decode step with every block’s attention probabilities — including the prompt prefill, whose q_len is the prompt length rather than 1.

Opt-in because it costs one readback per block per token. Passing None is exactly the non-probing path; the generated text is identical either way, since the probe reads tensors the step already computed.

pub async fn generate_async_trace( &self, tokenizer: &impl Tokenizer, prompt: &str, max_new_tokens: usize, sampling: Sampling, on_text: impl FnMut(&str) -> ControlFlow<()>, on_trace: Option<impl FnMut(&StepTrace)>, detail_layers: usize, top_n: usize, ) -> Result<String>

Gpt2::generate_async_probe with the full calculation trace: on_trace, when present, fires once per step — including the prompt prefill, whose q_len is the prompt length rather than 1 — with everything Gpt2::logits_step_trace_async captured.

None is byte-for-byte the non-probing path. With Some, the generated text is still identical: the probe reads tensors the step already computed and changes no arithmetic.