rlx_bert/
flow.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3//
4// This program is free software: you can redistribute it and/or modify
5// it under the terms of the GNU General Public License as published by
6// the Free Software Foundation, version 3.
7//
8// This program is distributed in the hope that it will be useful,
9// but WITHOUT ANY WARRANTY; without even the implied warranty of
10// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
11// GNU General Public License for more details.
12//
13// You should have received a copy of the GNU General Public License
14// along with this program. If not, see <https://www.gnu.org/licenses/>.
15
16//! Tier-0 BERT encoder flow — native [`ModelFlow`] assembly.
17
18use anyhow::Result;
19use rlx_flow::{BertQkvStyle, BuiltModel, CompileProfile, ModelFlow};
20use rlx_ir::{DType, Shape};
21use std::borrow::Cow;
22
23use rlx_core::config::BertConfig;
24use rlx_core::flow_util::WeightMapSource;
25use rlx_core::weight_map::WeightMap;
26
27/// Fully-qualified weight-key naming for a BERT checkpoint.
28///
29/// Owns the detected `bert.` / `""` prefix and the Q/K/V layout, and hands out
30/// the embedding keys the flow needs. Centralizing the HF names here keeps a
31/// rename to a single edit (and makes detection unit-testable); returning a
32/// [`Cow`] avoids allocating in the common no-prefix (bare-encoder) case.
33#[derive(Debug, Clone)]
34struct BertKeys {
35    prefix: &'static str,
36    qkv_style: BertQkvStyle,
37}
38
39impl BertKeys {
40    // Weight-name suffixes (relative to the optional `bert.` prefix), each
41    // defined once and shared by `detect` + the accessors below.
42    const WORD_EMBEDDINGS: &'static str = "embeddings.word_embeddings.weight";
43    const POSITION_EMBEDDINGS: &'static str = "embeddings.position_embeddings.weight";
44    const TOKEN_TYPE_EMBEDDINGS: &'static str = "embeddings.token_type_embeddings.weight";
45    const LN_WEIGHT: &'static str = "embeddings.LayerNorm.weight";
46    const LN_BIAS: &'static str = "embeddings.LayerNorm.bias";
47    /// Presence of this key (under the detected prefix) selects the BERT
48    /// (vs MPNet) Q/K/V layout.
49    const QKV_PROBE: &'static str = "encoder.layer.0.attention.self.query.weight";
50
51    /// Detect the `bert.` prefix and Q/K/V weight layout from the checkpoint.
52    fn detect(weights: &WeightMap) -> Self {
53        let prefix = if weights.has(&format!("bert.{}", Self::WORD_EMBEDDINGS)) {
54            "bert."
55        } else {
56            ""
57        };
58        let qkv_style = if weights.has(&format!("{prefix}{}", Self::QKV_PROBE)) {
59            BertQkvStyle::Bert
60        } else {
61            BertQkvStyle::Mpnet
62        };
63        Self { prefix, qkv_style }
64    }
65
66    /// `prefix` + `suffix`, borrowing the static literal when there's no prefix.
67    fn k(&self, suffix: &'static str) -> Cow<'static, str> {
68        if self.prefix.is_empty() {
69            Cow::Borrowed(suffix)
70        } else {
71            Cow::Owned(format!("{}{suffix}", self.prefix))
72        }
73    }
74
75    fn word_embeddings(&self) -> Cow<'static, str> {
76        self.k(Self::WORD_EMBEDDINGS)
77    }
78    fn position_embeddings(&self) -> Cow<'static, str> {
79        self.k(Self::POSITION_EMBEDDINGS)
80    }
81    fn token_type_embeddings(&self) -> Cow<'static, str> {
82        self.k(Self::TOKEN_TYPE_EMBEDDINGS)
83    }
84    fn embeddings_ln_weight(&self) -> Cow<'static, str> {
85        self.k(Self::LN_WEIGHT)
86    }
87    fn embeddings_ln_bias(&self) -> Cow<'static, str> {
88        self.k(Self::LN_BIAS)
89    }
90
91    /// Prefix passed to `repeat_bert_layers` (no trailing dot): `"bert"` or `""`.
92    fn layer_prefix(&self) -> &'static str {
93        self.prefix.trim_end_matches('.')
94    }
95}
96
97#[derive(Debug, Clone)]
98pub struct BertFlow<'a> {
99    cfg: &'a BertConfig,
100    batch: usize,
101    seq: usize,
102    profile: CompileProfile,
103}
104
105impl<'a> BertFlow<'a> {
106    pub fn new(cfg: &'a BertConfig, batch: usize, seq: usize) -> Self {
107        Self {
108            cfg,
109            batch,
110            seq,
111            profile: CompileProfile::encoder(),
112        }
113    }
114
115    pub fn with_profile(mut self, profile: CompileProfile) -> Self {
116        self.profile = profile;
117        self
118    }
119
120    pub fn build(self, weights: &mut WeightMap) -> Result<BuiltModel> {
121        let keys = BertKeys::detect(weights);
122
123        let f = DType::F32;
124        let eps = self.cfg.layer_norm_eps as f32;
125
126        let flow = ModelFlow::new("bert")
127            .with_profile(self.profile)
128            .input("input_ids", Shape::new(&[self.batch, self.seq], f))
129            .input("attention_mask", Shape::new(&[self.batch, self.seq], f))
130            .input("token_type_ids", Shape::new(&[self.batch, self.seq], f))
131            .input("position_ids", Shape::new(&[self.batch, self.seq], f))
132            .embed(keys.word_embeddings())
133            .gather_add("position_ids", keys.position_embeddings())
134            .gather_add("token_type_ids", keys.token_type_embeddings())
135            .layer_norm(keys.embeddings_ln_weight(), keys.embeddings_ln_bias(), eps)
136            .repeat_bert_layers(
137                self.cfg.num_hidden_layers,
138                keys.layer_prefix(),
139                keys.qkv_style,
140                self.cfg.hidden_size,
141                self.cfg.num_attention_heads,
142                eps,
143            )
144            .output("hidden_states");
145
146        flow.build(&mut WeightMapSource(weights))
147    }
148}
149
150pub fn build_bert_built(
151    cfg: &BertConfig,
152    weights: &mut WeightMap,
153    batch: usize,
154    seq: usize,
155) -> Result<BuiltModel> {
156    BertFlow::new(cfg, batch, seq).build(weights)
157}
158
159#[cfg(test)]
160mod tests {
161    use super::*;
162    use crate::test_support::{tiny_bert_weights, weights_with_keys};
163    use rlx_core::config::BertConfig;
164
165    #[test]
166    fn bert_keys_detects_prefix_and_qkv_style() {
167        // Bare encoder, BERT-style Q/K/V.
168        let k = BertKeys::detect(&weights_with_keys(&[
169            "embeddings.word_embeddings.weight",
170            "encoder.layer.0.attention.self.query.weight",
171        ]));
172        assert_eq!(k.layer_prefix(), "");
173        assert_eq!(
174            k.word_embeddings().as_ref(),
175            "embeddings.word_embeddings.weight"
176        );
177        assert!(matches!(k.qkv_style, BertQkvStyle::Bert));
178
179        // `bert.`-prefixed, no BERT-style query key → Mpnet layout. The prefix
180        // is applied to every key — checked structurally, not by re-pinning a
181        // second full key literal.
182        let k2 = BertKeys::detect(&weights_with_keys(&[
183            "bert.embeddings.word_embeddings.weight",
184        ]));
185        assert_eq!(k2.layer_prefix(), "bert");
186        assert!(k2.embeddings_ln_bias().starts_with("bert."));
187        assert!(matches!(k2.qkv_style, BertQkvStyle::Mpnet));
188    }
189
190    #[test]
191    fn bert_flow_builds() {
192        let cfg = BertConfig {
193            vocab_size: 32,
194            hidden_size: 16,
195            num_hidden_layers: 1,
196            num_attention_heads: 4,
197            intermediate_size: 32,
198            max_position_embeddings: 32,
199            type_vocab_size: 2,
200            layer_norm_eps: 1e-12,
201            hidden_act: "gelu".into(),
202        };
203        let mut wm = tiny_bert_weights(&cfg);
204        let built = BertFlow::new(&cfg, 1, 4).build(&mut wm).unwrap();
205        assert!(built.into_hir().unwrap().len() > 10);
206    }
207}
rlx_bert/flow.rs

rlx_bert/
flow.rs
