Add ALBERT masked language model

README.md

@@ -7,6 +7,7 @@ examples constructed using the [Burn] deep learning framework.
 
 | Model         | Description                   | Repository                            |
 | ------------- | ----------------------------- | ------------------------------------- |
+| [ALBERT]      | Masked language model         | [albert-burn](albert-burn/)           |
 | [Llama]       | Large language models         | [llama-burn](llama-burn/)             |
 | [MiniLM]      | Sentence embeddings           | [minilm-burn](minilm-burn/)           |
 | [MobileNetV2] | Mobile image classification   | [mobilenetv2-burn](mobilenetv2-burn/) |
@@ -51,6 +52,7 @@ respective repositories for specific license information.
 
 <!-- Official Models -->
 
+[ALBERT]: https://arxiv.org/abs/1909.11942
 [Llama]: https://github.com/meta-llama/llama3
 [MiniLM]: https://huggingface.co/sentence-transformers/all-MiniLM-L12-v2
 [MobileNetV2]: https://arxiv.org/abs/1801.04381

albert-burn/Cargo.toml

@@ -0,0 +1,46 @@
+[package]
+authors = ["Dilshod Tadjibaev (@antimora)"]
+license = "MIT OR Apache-2.0"
+name = "albert-burn"
+version = "0.1.0"
+edition = "2024"
+description = "ALBERT masked language model with Burn"
+
+[features]
+default = ["pretrained"]
+pretrained = ["burn/network", "dep:dirs", "dep:hf-hub"]
+
+# Backend selection
+ndarray = ["burn/ndarray"]
+tch-cpu = ["burn/tch"]
+tch-gpu = ["burn/tch"]
+wgpu = ["burn/wgpu"]
+cuda = ["burn/cuda"]
+
+[dependencies]
+burn = { version = "0.21.0", git = "https://github.com/tracel-ai/burn.git", rev = "5e5557ce52e9c434fe735217f91c4f4d48d3d151", default-features = false, features = ["std"] }
+burn-store = { version = "0.21.0", git = "https://github.com/tracel-ai/burn.git", rev = "5e5557ce52e9c434fe735217f91c4f4d48d3d151", features = ["std", "safetensors"] }
+
+# Tokenizer
+tokenizers = { version = "0.19.1", default-features = false, features = ["onig"] }
+
+# HuggingFace model download
+hf-hub = { version = "0.4.3", optional = true }
+dirs = { version = "6.0.0", optional = true }
+
+# Serialization
+serde = { version = "1.0", default-features = false, features = ["derive", "alloc"] }
+serde_json = "1.0"
+
+[dev-dependencies]
+clap = { version = "4.5", features = ["derive"] }
+divan = "0.1"
+
+[[bench]]
+name = "inference"
+harness = false
+required-features = ["pretrained", "ndarray"]
+
+[[example]]
+name = "inference"
+required-features = ["pretrained", "ndarray"]

albert-burn/README.md

@@ -0,0 +1,126 @@
+# ALBERT-Burn
+
+ALBERT (A Lite BERT) masked language model implementation in Rust using
+[Burn](https://github.com/tracel-ai/burn).
+
+ALBERT uses factorized embedding parameterization and cross-layer parameter sharing to reduce model
+size while maintaining performance.
+
+Supports all v2 variants from HuggingFace:
+
+| Variant          | Hidden Size | Parameters | HuggingFace                                                          |
+| ---------------- | ----------- | ---------- | -------------------------------------------------------------------- |
+| BaseV2 (default) | 768         | ~12M       | [albert-base-v2](https://huggingface.co/albert/albert-base-v2)       |
+| LargeV2          | 1,024       | ~18M       | [albert-large-v2](https://huggingface.co/albert/albert-large-v2)     |
+| XLargeV2         | 2,048       | ~60M       | [albert-xlarge-v2](https://huggingface.co/albert/albert-xlarge-v2)   |
+| XXLargeV2        | 4,096       | ~235M      | [albert-xxlarge-v2](https://huggingface.co/albert/albert-xxlarge-v2) |
+
+## Usage
+
+```rust
+use burn::backend::ndarray::NdArray;
+use albert_burn::{AlbertMaskedLM, AlbertVariant, tokenize_batch};
+
+type B = NdArray<f32>;
+
+fn main() -> Result<(), Box<dyn std::error::Error>> {
+    let device = Default::default();
+
+    // Load pretrained model and tokenizer (downloads from HuggingFace)
+    let (model, tokenizer) = AlbertMaskedLM::<B>::pretrained(&device, AlbertVariant::BaseV2, None)?;
+
+    // Tokenize input with [MASK] token
+    let sentence = "The capital of France is [MASK].";
+    let (input_ids, attention_mask) = tokenize_batch::<B>(&tokenizer, &[sentence], &device);
+
+    // Forward pass returns logits over vocabulary
+    let logits = model.forward(input_ids, attention_mask, None);
+
+    Ok(())
+}
+```
+
+## Features
+
+- `pretrained` - Enables model download utilities (default)
+- `ndarray` - NdArray backend
+
+Backend features:
+
+- `wgpu` - WebGPU backend
+- `cuda` - CUDA backend
+- `tch-cpu` - LibTorch CPU backend
+- `tch-gpu` - LibTorch GPU backend
+
+## Example
+
+Run the fill-mask inference example:
+
+```bash
+cargo run --example inference --features "pretrained,ndarray" --release
+```
+
+Specify a variant:
+
+```bash
+cargo run --example inference --features "pretrained,ndarray" --release -- xxlarge
+```
+
+### Results by variant
+
+Prompt: `"The capital of France is [MASK]."`
+
+**BaseV2** (12M params):
+
+| Rank | Token    | Logit |
+| ---- | -------- | ----- |
+| 1    | reims    | 16.35 |
+| 2    | toulouse | 16.17 |
+| 3    | paris    | 15.89 |
+| 4    | amiens   | 15.66 |
+| 5    | cannes   | 15.62 |
+
+**LargeV2** (18M params):
+
+| Rank | Token      | Logit |
+| ---- | ---------- | ----- |
+| 1    | paris      | 14.41 |
+| 2    | strasbourg | 12.26 |
+| 3    | lyon       | 11.82 |
+| 4    | brest      | 11.62 |
+| 5    | cannes     | 11.58 |
+
+**XLargeV2** (60M params):
+
+| Rank | Token      | Logit |
+| ---- | ---------- | ----- |
+| 1    | paris      | 16.82 |
+| 2    | lyon       | 16.06 |
+| 3    | strasbourg | 15.86 |
+| 4    | toulouse   | 15.02 |
+| 5    | grenoble   | 13.91 |
+
+**XXLargeV2** (235M params):
+
+| Rank | Token      | Logit |
+| ---- | ---------- | ----- |
+| 1    | paris      | 20.15 |
+| 2    | reims      | 17.17 |
+| 3    | marseille  | 17.02 |
+| 4    | versailles | 17.01 |
+| 5    | nantes     | 16.96 |
+
+## Testing
+
+Integration tests (requires model download):
+
+```bash
+cargo test --features "pretrained,ndarray" -- --ignored
+```
+
+Tests verify logit values, top-5 predictions, statistics, and per-position L2 norms across 3
+sentences against Python HuggingFace reference at 5e-4 relative tolerance.
+
+## License
+
+MIT OR Apache-2.0

albert-burn/benches/inference.rs

@@ -0,0 +1,135 @@
+//! Benchmark ALBERT BaseV2 inference (forward pass) across backends.
+//!
+//! Run:
+//! ```bash
+//! cargo bench --bench inference -p albert-burn
+//! ```
+
+use albert_burn::{AlbertMaskedLM, AlbertVariant, tokenize_batch};
+use burn::prelude::*;
+use divan::{AllocProfiler, Bencher};
+use std::cell::RefCell;
+
+#[global_allocator]
+static ALLOC: AllocProfiler = AllocProfiler::system();
+
+// Backend type aliases
+type NdArrayBackend = burn::backend::ndarray::NdArray<f32>;
+
+#[cfg(feature = "wgpu")]
+type WgpuBackend = burn::backend::wgpu::Wgpu;
+
+#[cfg(feature = "cuda")]
+type CudaBackend = burn::backend::cuda::Cuda<f32, i32>;
+
+#[cfg(feature = "tch-cpu")]
+type TchBackend = burn::backend::libtorch::LibTorch<f32>;
+
+// Shared model + inputs, initialized once in main()
+thread_local! {
+    static NDARRAY_STATE: RefCell<Option<BenchState<NdArrayBackend>>> = const { RefCell::new(None) };
+
+    #[cfg(feature = "wgpu")]
+    static WGPU_STATE: RefCell<Option<BenchState<WgpuBackend>>> = const { RefCell::new(None) };
+
+    #[cfg(feature = "cuda")]
+    static CUDA_STATE: RefCell<Option<BenchState<CudaBackend>>> = const { RefCell::new(None) };
+
+    #[cfg(feature = "tch-cpu")]
+    static TCH_STATE: RefCell<Option<BenchState<TchBackend>>> = const { RefCell::new(None) };
+}
+
+struct BenchState<B: Backend> {
+    model: AlbertMaskedLM<B>,
+    input_ids: Tensor<B, 2, Int>,
+    attention_mask: Tensor<B, 2>,
+}
+
+/// Run a forward pass to warm up the backend (triggers shader compilation, etc.).
+fn warmup<B: Backend>(state: &BenchState<B>) {
+    let _ = state.model.forward(
+        state.input_ids.clone(),
+        state.attention_mask.clone(),
+        None,
+    );
+}
+
+fn init_state<B: Backend>(device: &B::Device) -> BenchState<B> {
+    let (model, tokenizer) =
+        AlbertMaskedLM::<B>::pretrained(device, AlbertVariant::BaseV2, None)
+            .expect("Failed to load pretrained model");
+
+    let sentence = "The capital of France is [MASK].";
+    let (input_ids, attention_mask) = tokenize_batch::<B>(&tokenizer, &[sentence], device);
+
+    BenchState {
+        model,
+        input_ids,
+        attention_mask,
+    }
+}
+
+fn main() {
+    println!("Loading ALBERT BaseV2 for benchmarking...");
+
+    // Initialize all enabled backends
+    NDARRAY_STATE.with(|s| *s.borrow_mut() = Some(init_state(&Default::default())));
+
+    #[cfg(feature = "wgpu")]
+    WGPU_STATE.with(|s| {
+        let state = init_state(&Default::default());
+        warmup(&state);
+        *s.borrow_mut() = Some(state);
+    });
+
+    #[cfg(feature = "cuda")]
+    CUDA_STATE.with(|s| {
+        let state = init_state(&Default::default());
+        warmup(&state);
+        *s.borrow_mut() = Some(state);
+    });
+
+    #[cfg(feature = "tch-cpu")]
+    TCH_STATE.with(|s| *s.borrow_mut() = Some(init_state(&Default::default())));
+
+    println!("Model loaded. Running benchmarks...\n");
+
+    divan::main();
+}
+
+macro_rules! bench_backend {
+    ($backend:ty, $state:ident, $mod_name:ident, $backend_name:literal) => {
+        #[divan::bench_group(name = $backend_name, sample_count = 10)]
+        mod $mod_name {
+            use super::*;
+
+            #[divan::bench]
+            fn forward_pass(bencher: Bencher) {
+                bencher.bench(|| {
+                    $state.with(|s| {
+                        let s = s.borrow();
+                        let s = s.as_ref().expect("state not initialized");
+                        // Benchmark only the raw forward pass, excluding argmax/top-k
+                        // post-processing so results are comparable across backends.
+                        s.model.forward(
+                            s.input_ids.clone(),
+                            s.attention_mask.clone(),
+                            None,
+                        )
+                    })
+                });
+            }
+        }
+    };
+}
+
+bench_backend!(NdArrayBackend, NDARRAY_STATE, ndarray_backend, "NdArray (CPU)");
+
+#[cfg(feature = "wgpu")]
+bench_backend!(WgpuBackend, WGPU_STATE, wgpu_backend, "WGPU (GPU)");
+
+#[cfg(feature = "cuda")]
+bench_backend!(CudaBackend, CUDA_STATE, cuda_backend, "CUDA (NVIDIA GPU)");
+
+#[cfg(feature = "tch-cpu")]
+bench_backend!(TchBackend, TCH_STATE, tch_backend, "LibTorch (CPU)");