catgrad_llm/models/
llama.rs

1// Llama-3 model description
2
3use super::utils::{Cache, Config, ModelBuilder};
4use catgrad::backend::cpu::eval::Builder;
5use catgrad::core::nn::layers::*;
6use catgrad::core::{Dtype, NdArrayType, Shape, Var};
7
8pub struct Model;
9
10impl ModelBuilder for Model {
11    fn build(
12        &self,
13        builder: &Builder,
14        config: &Config,
15        cache: &mut Cache,
16        pos: usize,
17        x: Var,
18    ) -> Var {
19        let tokens = x.label.shape.0[1];
20        let emb = Model::embeddings(builder, config, x);
21
22        let mut result = emb;
23
24        for i in 0..config.num_hidden_layers {
25            result = Model::layer(
26                builder,
27                i,
28                config,
29                cache,
30                pos,
31                &format!("model.layers.{i}"),
32                result,
33            );
34        }
35
36        result = rmsnorm(builder, config.rms_norm_eps, "model.norm", result);
37
38        // Get the logits for the last token only
39        if tokens > 1 {
40            result = narrow(builder, 1, tokens - 1, 1, result);
41        }
42
43        // Add lm_head if weight tying is used
44        if config.tie_word_embeddings {
45            result = linear_no_bias(
46                builder,
47                config.hidden_size,
48                config.vocab_size,
49                "model.embed_tokens",
50                result,
51            );
52        }
53        result
54    }
55}
56
57impl Model {
58    pub fn embeddings(builder: &Builder, config: &Config, x: Var) -> Var {
59        let t = NdArrayType::new(
60            Shape(vec![config.vocab_size, config.hidden_size]),
61            Dtype::F32,
62        );
63        let weights = parameter(builder, t, "model.embed_tokens.weight".to_string());
64        embedding(builder, x, weights)
65    }
66
67    pub fn attention(
68        builder: &Builder,
69        layer_id: usize,
70        config: &Config,
71        cache: &mut Cache,
72        pos: usize,
73        name: &str,
74        x: Var,
75    ) -> Var {
76        let dim = config.hidden_size;
77        let num_heads = config.num_attention_heads;
78        let num_kv_heads = config.num_key_value_heads;
79        let rep = num_heads / num_kv_heads;
80        let head_dim = config.hidden_size / num_heads;
81        let b = x.label.shape.0[0];
82        let s = x.label.shape.0[1];
83
84        let q = linear_no_bias(builder, dim, dim, &format!("{name}.q_proj"), x.clone());
85        let k = linear_no_bias(
86            builder,
87            dim,
88            dim / rep,
89            &format!("{name}.k_proj"),
90            x.clone(),
91        );
92        let v = linear_no_bias(builder, dim, dim / rep, &format!("{name}.v_proj"), x);
93
94        let q = reshape(builder, Shape(vec![b, s, num_heads, head_dim]), q);
95        let k = reshape(builder, Shape(vec![b, s, num_kv_heads, head_dim]), k);
96        let v = reshape(builder, Shape(vec![b, s, num_kv_heads, head_dim]), v);
97
98        let q = transpose(builder, 1, 2, q);
99        let k = transpose(builder, 1, 2, k);
100        let v = transpose(builder, 1, 2, v);
101
102        let q = apply_rope_embedding(builder, pos, cache.cos.clone(), cache.sin.clone(), q);
103        let k = apply_rope_embedding(builder, pos, cache.cos.clone(), cache.sin.clone(), k);
104
105        let (k, v) = cache.update_kv_cache(builder, layer_id, k, v);
106
107        let k = repeat_kv(builder, rep, k);
108        let v = repeat_kv(builder, rep, v);
109
110        let tk = transpose(builder, 2, 3, k);
111        let attn = mat_mul(builder, q, tk);
112        let denom = constant(builder, attn.label.clone(), f32::sqrt(head_dim as f32));
113        let attn = attn / denom;
114
115        let mask = causal_mask(builder, s);
116        let mask = expand(builder, attn.label.shape.clone(), mask);
117        let attn = attn + mask;
118
119        let attn = softmax(builder, attn);
120        let attn = mat_mul(builder, attn, v);
121        let x = transpose(builder, 1, 2, attn);
122        let x = reshape(builder, Shape(vec![b, s, dim]), x);
123
124        linear_no_bias(builder, dim, dim, &format!("{name}.o_proj"), x)
125    }
126
127    pub fn mlp(builder: &Builder, config: &Config, name: &str, x: Var) -> Var {
128        let gated = linear_no_bias(
129            builder,
130            config.hidden_size,
131            config.intermediate_size,
132            &format!("{name}.gate_proj"),
133            x.clone(),
134        );
135        let up = linear_no_bias(
136            builder,
137            config.hidden_size,
138            config.intermediate_size,
139            &format!("{name}.up_proj"),
140            x,
141        );
142        let x = silu(builder, gated) * up; // SwiGLU
143
144        linear_no_bias(
145            builder,
146            config.intermediate_size,
147            config.hidden_size,
148            &format!("{name}.down_proj"),
149            x,
150        )
151    }
152
153    pub fn layer(
154        builder: &Builder,
155        layer_id: usize,
156        config: &Config,
157        cache: &mut Cache,
158        pos: usize,
159        name: &str,
160        x: Var,
161    ) -> Var {
162        let res = x.clone();
163        let x = rmsnorm(
164            builder,
165            config.rms_norm_eps,
166            &format!("{name}.input_layernorm"),
167            x,
168        );
169        let x = Model::attention(
170            builder,
171            layer_id,
172            config,
173            cache,
174            pos,
175            &format!("{name}.self_attn"),
176            x,
177        );
178        let x = res + x;
179        let res = x.clone();
180        let x = rmsnorm(
181            builder,
182            config.rms_norm_eps,
183            &format!("{name}.post_attention_layernorm"),
184            x,
185        );
186        let x = Model::mlp(builder, config, &format!("{name}.mlp"), x);
187        x + res
188    }
189}
catgrad_llm/models/llama.rs

catgrad_llm/models/
llama.rs
