candle-roberta

A from-scratch port of HuggingFace transformers' RoBERTa to Candle 0.11, with a 17-file educational tutorial and a Python parity check that verifies all 768 output dimensions match transformers to within float32 rounding noise.

$ cargo run
sentence:     "The quick brown fox jumps over the lazy dog."
input_ids:    [0, 133, 2119, 6219, 23602, 13855, 81, 5, 22414, 2335, 4, 2]
sequence_output shape: [1, 12, 768]
mean_pooled shape:      [1, 768]
normalized shape:       [1, 768]
L2 norm:                 1.000000

Full 768-dim parity check:
  dims compared: 768
  mismatches (>1e-4): 0
  max abs diff:      9.03e-8
  PASS — all 768 dims match within atol=1e-4

What this is

A learn-by-porting exercise: take the current transformers RobertaModel (roberta/modeling_roberta.py), translate it to Rust using Candle's native building blocks, and verify the output is numerically identical to PyTorch.

The accompanying tutorial (tutorial/) walks through every module — embeddings, self-attention, encoder, pooler — with the PyTorch source extracted inline, the math explained, and the Rust port annotated line-by-line.

What it does

string → tokenizer → RobertaModel → mean pooling → L2 normalize → 768-dim sentence embedding

• Downloads FacebookAI/roberta-base from the HuggingFace Hub (config, tokenizer, safetensors)
• Tokenizes a sentence with the HuggingFace tokenizers crate
• Runs a full 12-layer forward pass producing per-token contextual embeddings
• Applies mask-weighted mean pooling and L2 normalization for a sentence-level vector

Tokenization

RoBERTa uses byte-level Byte-Pair Encoding (BPE) — the same tokenization scheme as GPT-2. It is not WordPiece (BERT's tokenizer) and not SentencePiece (LLaMA's tokenizer). The vocabulary is 50,265 tokens.

How BPE works (briefly)

1. Pre-tokenization: the input text is split into words/whitespace-delimited chunks using a byte-level pre-tokenizer. Every byte is representable (no <unk> for non-ASCII), which is why it's called "byte-level."
2. BPE merges: starting from individual bytes, the most frequent adjacent pairs are merged iteratively during training. The result is a vocabulary of byte-pair tokens of varying length (from single bytes to whole common words like " the").
3. Encoding: at inference time, the pre-tokenized chunks are greedily merged using the learned merge rules until no more merges apply. Each resulting token maps to an integer id via the vocabulary.

RoBERTa's special tokens

What we use

We use the Rust tokenizers crate — the same library that powers HuggingFace's Python RobertaTokenizer. The Python RobertaTokenizer class (in roberta/tokenization_roberta.py) is itself just a thin wrapper around this Rust crate:

# from roberta/tokenization_roberta.py, line 143
self._tokenizer = Tokenizer(BPE(vocab=..., merges=...))
self._tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=...)
self._tokenizer.post_processor = processors.RobertaProcessing(sep=..., cls=...)

On the Rust side we load the serialized tokenizer directly:

let tokenizer = Tokenizer::from_file(tokenizer_filename)?;   // tokenizer.json from the Hub
let encoded = tokenizer.encode("Hello, world!", true)?;      // true = add_special_tokens (<s>...</s>)
let input_ids: Vec<u32> = encoded.get_ids().to_vec();
let attention_mask: Vec<u32> = encoded.get_attention_mask().to_vec();

The tokenizer.json file on the Hub is the pre-serialized BPE model + byte-level pre-tokenizer + RoBERTa post-processor (which adds <s> and </s>). Loading it with Tokenizer::from_file gives us the exact same tokenization as RobertaTokenizer in Python — no reimplementation needed.

Project layout

candle-roberta/
├── Cargo.toml
├── src/
│   ├── main.rs            # end-to-end: tokenize → forward → mean pool → L2 normalize → parity check
│   ├── lib.rs              # pub mod models; pub mod utils;
│   ├── models/
│   │   ├── mod.rs          # pub mod roberta;
│   │   └── roberta.rs      # the entire model (config, sublayers, RobertaModel, tests) — ~720 lines
│   └── utils.rs            # build_roberta_model_and_tokenizer (hf-hub 1.0 API), round_to_decimal_places
├── tutorial/               # 17 numbered markdown files, ~3,800 lines total
│   ├── 00_GETTING_STARTED.md
│   ├── 01_PYTORCH_CANDLE_PARALLELS.md
│   ├── 02_ROBERTA_OVERVIEW.md
│   ├── 03_BUILDING_BLOCKS.md
│   ├── 04_ROBERTA_CONFIG.md
│   ├── 05_RobertaEmbeddings.md
│   ├── 06_RobertaSelfAttention.md
│   ├── 07_RobertaSelfOutput.md
│   ├── 08_RobertaAttention.md
│   ├── 09_RobertaIntermediate.md
│   ├── 10_RobertaOutput.md
│   ├── 11_RobertaLayer.md
│   ├── 12_RobertaEncoder.md
│   ├── 13_RobertaPooler.md
│   ├── 14_RobertaModel.md
│   ├── 15_LOADING_AND_INFERENCE.md
│   └── 16_DEBUGGING_AND_TESTS.md
└── parity/                 # Python reference implementation
    ├── pyproject.toml
    ├── verify_parity.py    # runs the same pipeline in PyTorch, exports embedding for comparison
    └── pytorch_embedding.json  # pre-generated 768-dim reference embedding (17KB, no Python needed)

The src/models/roberta.rs layout mirrors candle-transformers/src/models/bert.rs — one file per model, sublayers private, only RobertaModel / RobertaConfig / RobertaModelOutput public.

The roberta/ directory contains the reference Python source from transformers (modeling_roberta.py, configuration_roberta.py, modular_roberta.py, tokenization_roberta.py) — not compiled or used at runtime, just there for side-by-side reading while following the tutorial.

Quick start

Prerequisites

• Rust (any recent stable; edition = "2024")
• Internet access on first run (downloads roberta-base, ~500MB, cached by hf-hub)

Run

$ cargo run

That's it. The binary downloads FacebookAI/roberta-base, tokenizes a sentence, runs the forward pass, applies mean pooling + L2 normalization, and prints the resulting 768-dim embedding.

Run the tests

$ cargo test --lib

Four tests: position-id generation (with and without padding), create_position_ids_from_inputs_embeds, and a full-model shape check using VarBuilder::zeros (no download needed).

Verify parity against PyTorch

The repo includes a pre-generated parity/pytorch_embedding.json (17KB) — the 768-dim reference embedding from transformers' RobertaModel on the same input. You can verify parity without setting up Python at all:

$ cargo run

The Rust binary loads pytorch_embedding.json and compares all 768 dimensions. Expected result:

Full 768-dim parity check:
  dims compared: 768
  mismatches (>1e-4): 0
  max abs diff:      9.03e-8
  PASS — all 768 dims match within atol=1e-4

To regenerate the reference yourself (requires Python + uv):

$ cd parity && uv run verify_parity.py   # regenerates pytorch_embedding.json
$ cd .. && cargo run                      # compares against the fresh reference

Switching devices

The device is created in src/main.rs and passed to the loader. To run on Metal (Apple Silicon):

# Cargo.toml
candle-core = { version = "0.11.0", features = ["metal"] }
candle-nn   = { version = "0.11.0", features = ["metal"] }

// src/main.rs
let device = Device::Metal(0);  // instead of Device::Cpu

$ cargo run --features candle-core/metal,candle-nn/metal

Same for CUDA:

candle-core = { version = "0.11.0", features = ["cuda"] }
candle-nn   = { version = "0.11.0", features = ["cuda"] }

let device = Device::Cuda(0);

What was ported

What was intentionally dropped

• past_key_values / KV cache (autoregressive generation)
• is_decoder / cross-attention (encoder-only)
• GradientCheckpointingLayer / apply_chunking_to_forward (training-only)
• ALL_ATTENTION_FUNCTIONS dispatch (we use the eager math)
• _init_weights (we load a pretrained checkpoint)
• Task heads: RobertaForMaskedLM, RobertaForSequenceClassification, RobertaForTokenClassification, RobertaForQuestionAnswering (we stop at embeddings)

Dependencies

The tutorial

The tutorial/ directory is a from-scratch, step-by-step guide to the port. Each file covers one concept or one module, with the PyTorch source extracted inline so you never have to leave the file. No external links to follow — everything is self-contained.

Start at tutorial/00_GETTING_STARTED.md and follow the "Next" links at the bottom of each file. By file 14 you have a working RobertaModel; file 15 wires up the HuggingFace Hub and runs it; file 16 tests it.

Credits

• The original candle-tutorial by @ToluClassics — this project is a modernized, expanded rewrite of that work.
• Candle by HuggingFace.
• RoBERTa by Liu et al. (2019).
• The roberta/ reference source is from transformers (Apache-2.0).

License

Apache License 2.0 — see LICENSE. The roberta/ reference Python source retains its original copyright (The HuggingFace Inc. team, NVIDIA CORPORATION, Google AI Language Team authors), also under Apache-2.0.