1#include "models.h"
2
3template <bool iswa>
4llm_build_olmo2<iswa>::llm_build_olmo2(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) {
5 const int64_t n_embd_head = hparams.n_embd_head_v;
6
7 GGML_ASSERT(n_embd_head == hparams.n_embd_head_k);
8 GGML_ASSERT(n_embd_head == hparams.n_rot);
9
10 ggml_tensor * cur;
11 ggml_tensor * inpL;
12
13 inpL = build_inp_embd(model.tok_embd);
14
15 // inp_pos - contains the positions
16 ggml_tensor * inp_pos = build_inp_pos();
17
18 using inp_attn_type = std::conditional_t<iswa, llm_graph_input_attn_kv_iswa, llm_graph_input_attn_kv>;
19 inp_attn_type * inp_attn = nullptr;
20
21 if constexpr (iswa) {
22 inp_attn = build_attn_inp_kv_iswa();
23 } else {
24 inp_attn = build_attn_inp_kv();
25 }
26 ggml_tensor * inp_out_ids = build_inp_out_ids();
27
28 for (int il = 0; il < n_layer; ++il) {
29 ggml_tensor * inpSA = inpL;
30
31 cur = inpL;
32
33 // self_attention
34 {
35 // compute Q and K and RoPE them
36 ggml_tensor * Qcur = build_lora_mm(model.layers[il].wq, cur);
37 cb(Qcur, "Qcur", il);
38
39 ggml_tensor * Kcur = build_lora_mm(model.layers[il].wk, cur);
40 cb(Kcur, "Kcur", il);
41
42 ggml_tensor * Vcur = build_lora_mm(model.layers[il].wv, cur);
43 cb(Vcur, "Vcur", il);
44
45 Qcur = build_norm(Qcur, model.layers[il].attn_q_norm, NULL,
46 LLM_NORM_RMS, il);
47 cb(Qcur, "Qcur_normed", il);
48
49 Kcur = build_norm(Kcur, model.layers[il].attn_k_norm, NULL,
50 LLM_NORM_RMS, il);
51 cb(Kcur, "Kcur_normed", il);
52
53 Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head, n_head, n_tokens);
54 Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens);
55 Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens);
56
57 const bool is_swa = hparams.is_swa(il);
58
59 if (is_swa) {
60 // For sliding window layers, Olmo3 use regular rope with no yarn rope scaling.
61 // This is achieved here by setting freq_scale and attn_factor to 1.
62 // We also set ext_factor to 0 to avoid a few unnecessary computations.
63 Qcur = ggml_rope_ext(
64 ctx0, Qcur, inp_pos, nullptr,
65 n_rot, rope_type, n_ctx_orig, freq_base, 1.0,
66 0.0, 1.0, beta_fast, beta_slow
67 );
68
69 Kcur = ggml_rope_ext(
70 ctx0, Kcur, inp_pos, nullptr,
71 n_rot, rope_type, n_ctx_orig, freq_base, 1.0,
72 0.0, 1.0, beta_fast, beta_slow
73 );
74 } else {
75 Qcur = ggml_rope_ext(
76 ctx0, Qcur, inp_pos, nullptr,
77 n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
78 ext_factor, attn_factor, beta_fast, beta_slow
79 );
80
81 Kcur = ggml_rope_ext(
82 ctx0, Kcur, inp_pos, nullptr,
83 n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
84 ext_factor, attn_factor, beta_fast, beta_slow
85 );
86 }
87 cb(Qcur, "Qcur", il);
88 cb(Kcur, "Kcur", il);
89 cb(Vcur, "Vcur", il);
90
91 cur = build_attn(inp_attn,
92 model.layers[il].wo, NULL,
93 Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, 1.0f/sqrtf(float(n_embd_head)), il);
94 }
95 if (il == n_layer - 1 && inp_out_ids) {
96 cur = ggml_get_rows(ctx0, cur, inp_out_ids);
97 inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
98 }
99 cur = build_norm(cur,
100 model.layers[il].attn_post_norm, NULL,
101 LLM_NORM_RMS, il);
102 cb(cur, "attn_post_norm", il);
103
104 ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);
105 cb(ffn_inp, "ffn_inp", il);
106
107 // feed-forward network
108 cur = build_ffn(ffn_inp,
109 model.layers[il].ffn_up, NULL, NULL,
110 model.layers[il].ffn_gate, NULL, NULL,
111 model.layers[il].ffn_down, NULL, NULL,
112 NULL,
113 LLM_FFN_SILU, LLM_FFN_PAR, il);
114 cb(cur, "ffn_out", il);
115
116 cur = build_norm(cur,
117 model.layers[il].ffn_post_norm, NULL,
118 LLM_NORM_RMS, -1);
119 cb(cur, "ffn_post_norm", -1);
120
121 cur = ggml_add(ctx0, cur, ffn_inp);
122 cb(cur, "ffn_out", il);
123
124 cur = build_cvec(cur, il);
125 cb(cur, "l_out", il);
126
127 // input for next layer
128 inpL = cur;
129 }
130 cur = inpL;
131
132 cur = build_norm(cur,
133 model.output_norm, NULL,
134 LLM_NORM_RMS, -1);
135
136 cb(cur, "result_norm", -1);
137 res->t_embd = cur;
138
139 // lm_head
140 cur = build_lora_mm(model.output, cur);
141
142 cb(cur, "result_output", -1);
143 res->t_logits = cur;
144
145 ggml_build_forward_expand(gf, cur);
146}
147
148// Explicit template instantiations
149template struct llm_build_olmo2<false>;
150template struct llm_build_olmo2<true>;