forge/nn/
mod.rs

1//! Minimal neural-network modules — exactly what GPT-2 needs.
2
3use crate::error::Result;
4use crate::ops::{self, MatmulSpec};
5use crate::tensor::Tensor;
6
7/// y = x @ W + b, with W stored [in_features, out_features].
8///
9/// This matches the HF GPT-2 "Conv1D" convention, so checkpoint weights load
10/// without transposition (see roadmap "Known Pitfalls").
11pub struct Linear {
12    pub w: Tensor,
13    pub b: Option<Tensor>,
14}
15
16impl Linear {
17    pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
18        ops::matmul(x, &self.w, self.b.as_ref(), MatmulSpec::default())
19    }
20}
21
22pub struct LayerNorm {
23    pub gamma: Tensor,
24    pub beta: Tensor,
25    pub eps: f32,
26}
27
28impl LayerNorm {
29    pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
30        ops::layernorm(x, &self.gamma, &self.beta, self.eps)
31    }
32}
33
34/// Token + positional embedding table. The token table is stored row-chunked
35/// so no single GPU binding exceeds the device limit (see roadmap pitfalls);
36/// on CPU it is a single chunk.
37pub struct Embedding {
38    /// Row chunks of the [vocab, n_embd] token table.
39    pub wte_chunks: Vec<Tensor>,
40    /// Rows per chunk (the last chunk may be smaller).
41    pub chunk_rows: usize,
42    /// [n_ctx, n_embd]
43    pub wpe: Tensor,
44}
45
46impl Embedding {
47    /// Split a host-side [vocab, n_embd] table into device chunks that each
48    /// fit in a single binding on `device`.
49    pub fn from_host_wte(
50        wte: &[f32],
51        vocab: usize,
52        n_embd: usize,
53        wpe: Tensor,
54        device: &crate::device::Device,
55    ) -> Result<Self> {
56        let row_bytes = n_embd * 4;
57        let max_rows = (device.max_binding_bytes() / row_bytes).min(vocab);
58        let chunk_rows = max_rows.max(1);
59        let mut chunks = Vec::new();
60        let mut start = 0usize;
61        while start < vocab {
62            let rows = chunk_rows.min(vocab - start);
63            chunks.push(Tensor::from_f32(
64                &wte[start * n_embd..(start + rows) * n_embd],
65                [rows, n_embd],
66                device,
67            )?);
68            start += rows;
69        }
70        Ok(Embedding {
71            wte_chunks: chunks,
72            chunk_rows,
73            wpe,
74        })
75    }
76
77    /// ids: u32 tensor `[t]`; `pos` is the absolute position of `ids[0]`.
78    pub fn forward(&self, ids: &Tensor, pos: usize) -> Result<Tensor> {
79        ops::embedding_chunked(ids, &self.wte_chunks, self.chunk_rows, Some(&self.wpe), pos)
80    }
81}
forge/nn/mod.rs

forge/nn/
mod.rs
