models/roberta-baseREADME.md3.1 kB
md
# RoBERTa Base
Meta's more thoroughly trained BERT: the same post-norm encoder --
`h = LayerNorm(x + attention(x))`, then `out = LayerNorm(h + ffn(h))`, the
block order `bert-base-uncased` and `all-MiniLM-L6-v2` also use and the
pre-norm ViT and GPT-2 already in Nest do not -- with dynamic masking, no
next-sentence-prediction objective (so `type_vocab_size` is 1: a single
token-type row, looked up like any other embedding so the checkpoint's
tensor loads unmodified), and a byte-level BPE vocabulary of 50265.
RoBERTa's position ids are not zero-based. `padding_idx` (RoBERTa's pad
token, id 1) reserves position row 0 for padding, and the real positions
start at `padding_idx + 1`, so token index i (0-based) reads position row
`i + 2`; the position table therefore has 514 rows for 512 usable
positions. Upstream derives this from the attention mask
(`create_position_ids_from_input_ids`, a running count of non-pad tokens);
with no padding in the sequence -- the only case this encoder's full
attention supports -- that count is just the token's own index, so the
offset collapses to the constant `+2` the source uses.
This checkpoint (`RobertaForMaskedLM`) carries no pooler, only a
masked-language-model head this card does not expose, so there is a single
entry.
Upstream's `hidden_act` is plain GELU, the error-function form, so the
source uses `std.nn.activations::gelu_erf` rather than the tanh
approximation `gelu`; over twelve layers the tanh form would move the hidden
states by 5e-2.
## Loading
```python
from linnet import nest
model = nest.load("roberta-base", backend="torch")
hidden = model(tokens, token_types) # Tensor[B, S; i32], Tensor[B, S; i32] -> Tensor[B, S, 768; f32]
```
`token_types` is the segment-id tensor; pass zeros of the same shape as
`tokens` (RoBERTa's single token-type row means any value looks up the same
embedding, but zero matches upstream's default). Sequences longer than 512
positions are rejected at compile time (`where S + 2 <= MaxPositions`). The
encoder runs full attention over every position: `std.nn.attention::
attention` takes an unbatched `[Q, K]` mask, so a per-sequence padding mask
cannot be expressed here. Pad-free batches (or a batch of one) give exact
results; padded batches attend onto the padding.
## Entries
| Entry | |
| --- | --- |
| `forward<B, S>(tokens, token_types)` | last hidden state, `Tensor[B, S, 768; f32]` |
## Numerics
Validated against `transformers.RobertaModel` (`AutoModel.from_pretrained`,
`numerics="fast"`) on CPU in f32, comparing `last_hidden_state` on
pad-free sentences with RoBERTa's own BPE tokenizer: the maximum absolute
difference is 4.8e-6. The position-id formula above (`i + 2`) was checked
directly against `transformers`' own `create_position_ids_from_input_ids`
and matches on every position.
## Provenance
- Weights: [FacebookAI/roberta-base](https://huggingface.co/FacebookAI/roberta-base), MIT.
- Code: [fairseq/examples/roberta](https://github.com/facebookresearch/fairseq/tree/main/examples/roberta).
- Paper: [RoBERTa: A Robustly Optimized BERT Pretraining Approach](https://arxiv.org/abs/1907.11692).bench.json9.4 kB
json
{
"date": "2026-09-28T18:17:48+00:00",
"environment": {
"python": "3.12.3",
"machine": "x86_64",
"torch": "2.14.0",
"jax": "0.11.2",
"transformers": "5.17.0",
"diffusers": "0.40.0",
"onnxruntime": "1.30.0",
"linnet": "0.1.0",
"gpu": "NVIDIA H100 80GB HBM3",
"gpus": "1",
"driver": "580.126.09"
},
"workload": {
"device": "cuda",
"prompt_tokens": 512,
"new_tokens": 128,
"batch": 1,
"warmup": 3,
"iters": 10,
"seed": 0,
"offload_gib": 8.0,
"serve_requests": 256,
"serve_concurrency": 64,
"serve_prompt_min": 128,
"serve_prompt_max": 512,
"serve_new": 128
},
"methods": [
{
"method": "transformers (eager)",
"kind": "reference",
"metrics": {
"latency_ms": 3.520517610013485,
"throughput_per_s": 10856.98931687499,
"load_s": 15.56502415984869,
"peak_vram_mib": 1190.0
},
"max_abs_diff": 0.0,
"notes": "unpadded batches of exactly 128 tokens (this card takes no padding mask)",
"error": null,
"key": "transformers-eager"
},
{
"method": "transformers (torch.compile)",
"kind": "reference",
"metrics": {
"latency_ms": 1.5702378004789352,
"throughput_per_s": 14513.337342412828,
"load_s": 1.5454451404511929,
"peak_vram_mib": 1170.0
},
"max_abs_diff": 0.125,
"notes": "unpadded batches of exactly 128 tokens (this card takes no padding mask)",
"error": null,
"key": "transformers-compile"
},
{
"method": "Linnet torch (generated source)",
"kind": "linnet",
"metrics": {
"latency_ms": 1.918918453156948,
"throughput_per_s": 13836.24954566742,
"load_s": 1.135751385241747,
"peak_vram_mib": 1136.0
},
"max_abs_diff": 0.15625,
"notes": "unpadded batches of exactly 128 tokens; bf16 on cuda",
"error": null,
"key": "linnet-torch"
},
{
"method": "Linnet torch (CUDA graphs)",
"kind": "linnet",
"metrics": {
"latency_ms": 0.7324917241930962,
"throughput_per_s": 18180.842753032324,
"load_s": 1.1443926822394133,
"peak_vram_mib": 1176.0
},
"max_abs_diff": 0.125,
"notes": "unpadded batches of exactly 128 tokens; bf16 on cuda",
"error": null,
"key": "linnet-cudagraphs"
},
{
"method": "Linnet torch (torch.compile, inductor)",
"kind": "linnet",
"metrics": {
"latency_ms": 1.375451683998108,
"throughput_per_s": 17222.55960361747,
"load_s": 1.1061877142637968,
"peak_vram_mib": 1088.0
},
"max_abs_diff": 0.125,
"notes": "unpadded batches of exactly 128 tokens; bf16 on cuda",
"error": null,
"key": "linnet-inductor"
},
{
"method": "Linnet JAX (XLA, StableHLO)",
"kind": "linnet",
"metrics": {
"latency_ms": 1.1764634400606155,
"throughput_per_s": 12433.080629648419,
"load_s": 0.5444975607097149,
"peak_vram_mib": 784.8,
"driver_vram_mib": 1646.0
},
"max_abs_diff": 0.15625,
"notes": "unpadded batches of exactly 128 tokens; bf16 on cuda; peak memory is JAX's allocator peak; the driver shows its pool, which grows in whole regions",
"error": null,
"key": "linnet-jax"
},
{
"method": "Linnet JAX (XLA, generated source)",
"kind": "linnet",
"metrics": {
"latency_ms": 0.9294366464018822,
"throughput_per_s": 11985.084235622411,
"load_s": 0.5710759032517672,
"peak_vram_mib": 771.2,
"driver_vram_mib": 1646.0
},
"max_abs_diff": 0.21875,
"notes": "unpadded batches of exactly 128 tokens; bf16 on cuda; peak memory is JAX's allocator peak; the driver shows its pool, which grows in whole regions",
"error": null,
"key": "linnet-jax-source"
},
{
"method": "Linnet ONNX f32 -> ONNX Runtime (CUDA)",
"kind": "linnet",
"metrics": {
"latency_ms": 1.5253759920597076,
"throughput_per_s": 5275.919706523964,
"load_s": 0.284932192414999,
"peak_vram_mib": 2702.0
},
"max_abs_diff": 0.08825016021728516,
"notes": "f32; ONNX Runtime's CUDA execution provider; first calls build the sessions; unpadded batches of exactly the sequence length",
"error": null,
"key": "linnet-onnx-f32"
},
{
"method": "Linnet ONNX f32 -> ONNX Runtime (TensorRT)",
"kind": "linnet",
"metrics": {
"latency_ms": 0.7475130259990692,
"throughput_per_s": 9050.218524668815,
"load_s": 0.27336307615041733,
"peak_vram_mib": 3410.0
},
"max_abs_diff": 0.08764839172363281,
"notes": "f32; ONNX Runtime's TensorRT execution provider; first calls build the sessions; unpadded batches of exactly the sequence length",
"error": null,
"key": "linnet-onnx-f32-trt"
},
{
"method": "Linnet ONNX f16 -> ONNX Runtime (CUDA)",
"kind": "linnet",
"metrics": {
"latency_ms": 1.2336215004324913,
"throughput_per_s": 9663.79615647425,
"load_s": 0.2871305625885725,
"peak_vram_mib": 1804.0
},
"max_abs_diff": 0.087890625,
"notes": "f16; ONNX Runtime's CUDA execution provider; first calls build the sessions; unpadded batches of exactly the sequence length",
"error": null,
"key": "linnet-onnx-f16"
},
{
"method": "Linnet ONNX f16 -> ONNX Runtime (TensorRT)",
"kind": "linnet",
"metrics": {
"latency_ms": 0.528271310031414,
"throughput_per_s": 19752.20359641158,
"load_s": 0.2741014324128628,
"peak_vram_mib": 3130.0
},
"max_abs_diff": 0.09765625,
"notes": "f16; ONNX Runtime's TensorRT execution provider; first calls build the sessions; unpadded batches of exactly the sequence length",
"error": null,
"key": "linnet-onnx-f16-trt"
},
{
"method": "Linnet ONNX bf16 -> ONNX Runtime (CUDA)",
"kind": "linnet",
"metrics": {
"latency_ms": 1.230708323419094,
"throughput_per_s": 9497.392652317736,
"load_s": 0.3679192364215851,
"peak_vram_mib": 1804.0
},
"max_abs_diff": 0.125,
"notes": "bf16; ONNX Runtime's CUDA execution provider; first calls build the sessions; unpadded batches of exactly the sequence length",
"error": null,
"key": "linnet-onnx-bf16"
},
{
"method": "Linnet ONNX bf16 -> ONNX Runtime (TensorRT)",
"kind": "linnet",
"metrics": {
"latency_ms": 0.5563749000430107,
"throughput_per_s": 18619.208476031108,
"load_s": 0.2691136598587036,
"peak_vram_mib": 3128.0
},
"max_abs_diff": 0.125,
"notes": "bf16; ONNX Runtime's TensorRT execution provider; first calls build the sessions; unpadded batches of exactly the sequence length",
"error": null,
"key": "linnet-onnx-bf16-trt"
},
{
"method": "KerasHub (JAX)",
"kind": "reference",
"metrics": {
"latency_ms": 4.94545791298151,
"throughput_per_s": 5082.257220972192,
"load_s": 5.992139536887407,
"peak_vram_mib": 567.1,
"driver_vram_mib": 1650.0
},
"max_abs_diff": 0.1171875,
"notes": "bf16, under jax.jit; unpadded batches of exactly 128 tokens; peak memory is JAX's allocator peak; the driver shows its pool, which grows in whole regions",
"error": null,
"key": "keras-hub"
},
{
"method": "transformers -> torch.onnx -> ONNX Runtime",
"kind": "reference",
"metrics": {
"latency_ms": 2.440323121845722,
"throughput_per_s": 4955.115322029787,
"load_s": 18.557513615116477,
"peak_vram_mib": 2874.0
},
"max_abs_diff": 0.08825016021728516,
"notes": "the reference model's own ONNX export, f32 like Linnet's; unpadded batches of exactly 128 tokens",
"error": null,
"key": "onnx-reference"
},
{
"method": "Triton Inference Server (ONNX Runtime backend, torch.onnx export)",
"kind": "reference",
"metrics": {
"latency_ms": 3.7250053137540817,
"throughput_per_s": 566.3469271844568,
"load_s": 10.177207563072443,
"peak_vram_mib": 3584.9
},
"max_abs_diff": 0.08825016021728516,
"notes": "over HTTP; f32; throughput with 4 requests of batch 64 in flight",
"error": null,
"key": "triton-onnx"
},
{
"method": "Triton Inference Server (ONNX Runtime backend, Linnet ONNX)",
"kind": "linnet",
"metrics": {
"latency_ms": 3.321124240756035,
"throughput_per_s": 502.4203527063062,
"load_s": 17.033509828150272,
"peak_vram_mib": 3960.9
},
"max_abs_diff": 0.08825016021728516,
"notes": "over HTTP; f32; throughput with 4 requests of batch 64 in flight",
"error": null,
"key": "triton-linnet-onnx"
},
{
"method": "Triton Inference Server (Python backend, Linnet torch)",
"kind": "linnet",
"metrics": {
"latency_ms": 2.6769787073135376,
"throughput_per_s": 539.6953591739821,
"load_s": 6.0245058834552765,
"peak_vram_mib": 2936.1
},
"max_abs_diff": 0.125,
"notes": "over HTTP; bf16, CUDA graphs; throughput with 4 requests of batch 64 in flight",
"error": null,
"key": "triton-linnet-python"
}
],
"reference": "transformers-eager"
}bindings.json16.4 kB
json
{
"embeddings.word_embeddings": "roberta.embeddings.word_embeddings.weight",
"embeddings.position_embeddings": "roberta.embeddings.position_embeddings.weight",
"embeddings.token_type_embeddings": "roberta.embeddings.token_type_embeddings.weight",
"embeddings.norm.weight": "roberta.embeddings.LayerNorm.weight",
"embeddings.norm.bias": "roberta.embeddings.LayerNorm.bias",
"layers.0.attention.query.weight": "roberta.encoder.layer.0.attention.self.query.weight",
"layers.0.attention.query.bias": "roberta.encoder.layer.0.attention.self.query.bias",
"layers.0.attention.key.weight": "roberta.encoder.layer.0.attention.self.key.weight",
"layers.0.attention.key.bias": "roberta.encoder.layer.0.attention.self.key.bias",
"layers.0.attention.value.weight": "roberta.encoder.layer.0.attention.self.value.weight",
"layers.0.attention.value.bias": "roberta.encoder.layer.0.attention.self.value.bias",
"layers.0.attention.out.weight": "roberta.encoder.layer.0.attention.output.dense.weight",
"layers.0.attention.out.bias": "roberta.encoder.layer.0.attention.output.dense.bias",
"layers.0.attention_norm.weight": "roberta.encoder.layer.0.attention.output.LayerNorm.weight",
"layers.0.attention_norm.bias": "roberta.encoder.layer.0.attention.output.LayerNorm.bias",
"layers.0.up.weight": "roberta.encoder.layer.0.intermediate.dense.weight",
"layers.0.up.bias": "roberta.encoder.layer.0.intermediate.dense.bias",
"layers.0.down.weight": "roberta.encoder.layer.0.output.dense.weight",
"layers.0.down.bias": "roberta.encoder.layer.0.output.dense.bias",
"layers.0.output_norm.weight": "roberta.encoder.layer.0.output.LayerNorm.weight",
"layers.0.output_norm.bias": "roberta.encoder.layer.0.output.LayerNorm.bias",
"layers.1.attention.query.weight": "roberta.encoder.layer.1.attention.self.query.weight",
"layers.1.attention.query.bias": "roberta.encoder.layer.1.attention.self.query.bias",
"layers.1.attention.key.weight": "roberta.encoder.layer.1.attention.self.key.weight",
"layers.1.attention.key.bias": "roberta.encoder.layer.1.attention.self.key.bias",
"layers.1.attention.value.weight": "roberta.encoder.layer.1.attention.self.value.weight",
"layers.1.attention.value.bias": "roberta.encoder.layer.1.attention.self.value.bias",
"layers.1.attention.out.weight": "roberta.encoder.layer.1.attention.output.dense.weight",
"layers.1.attention.out.bias": "roberta.encoder.layer.1.attention.output.dense.bias",
"layers.1.attention_norm.weight": "roberta.encoder.layer.1.attention.output.LayerNorm.weight",
"layers.1.attention_norm.bias": "roberta.encoder.layer.1.attention.output.LayerNorm.bias",
"layers.1.up.weight": "roberta.encoder.layer.1.intermediate.dense.weight",
"layers.1.up.bias": "roberta.encoder.layer.1.intermediate.dense.bias",
"layers.1.down.weight": "roberta.encoder.layer.1.output.dense.weight",
"layers.1.down.bias": "roberta.encoder.layer.1.output.dense.bias",
"layers.1.output_norm.weight": "roberta.encoder.layer.1.output.LayerNorm.weight",
"layers.1.output_norm.bias": "roberta.encoder.layer.1.output.LayerNorm.bias",
"layers.2.attention.query.weight": "roberta.encoder.layer.2.attention.self.query.weight",
"layers.2.attention.query.bias": "roberta.encoder.layer.2.attention.self.query.bias",
"layers.2.attention.key.weight": "roberta.encoder.layer.2.attention.self.key.weight",
"layers.2.attention.key.bias": "roberta.encoder.layer.2.attention.self.key.bias",
"layers.2.attention.value.weight": "roberta.encoder.layer.2.attention.self.value.weight",
"layers.2.attention.value.bias": "roberta.encoder.layer.2.attention.self.value.bias",
"layers.2.attention.out.weight": "roberta.encoder.layer.2.attention.output.dense.weight",
"layers.2.attention.out.bias": "roberta.encoder.layer.2.attention.output.dense.bias",
"layers.2.attention_norm.weight": "roberta.encoder.layer.2.attention.output.LayerNorm.weight",
"layers.2.attention_norm.bias": "roberta.encoder.layer.2.attention.output.LayerNorm.bias",
"layers.2.up.weight": "roberta.encoder.layer.2.intermediate.dense.weight",
"layers.2.up.bias": "roberta.encoder.layer.2.intermediate.dense.bias",
"layers.2.down.weight": "roberta.encoder.layer.2.output.dense.weight",
"layers.2.down.bias": "roberta.encoder.layer.2.output.dense.bias",
"layers.2.output_norm.weight": "roberta.encoder.layer.2.output.LayerNorm.weight",
"layers.2.output_norm.bias": "roberta.encoder.layer.2.output.LayerNorm.bias",
"layers.3.attention.query.weight": "roberta.encoder.layer.3.attention.self.query.weight",
"layers.3.attention.query.bias": "roberta.encoder.layer.3.attention.self.query.bias",
"layers.3.attention.key.weight": "roberta.encoder.layer.3.attention.self.key.weight",
"layers.3.attention.key.bias": "roberta.encoder.layer.3.attention.self.key.bias",
"layers.3.attention.value.weight": "roberta.encoder.layer.3.attention.self.value.weight",
"layers.3.attention.value.bias": "roberta.encoder.layer.3.attention.self.value.bias",
"layers.3.attention.out.weight": "roberta.encoder.layer.3.attention.output.dense.weight",
"layers.3.attention.out.bias": "roberta.encoder.layer.3.attention.output.dense.bias",
"layers.3.attention_norm.weight": "roberta.encoder.layer.3.attention.output.LayerNorm.weight",
"layers.3.attention_norm.bias": "roberta.encoder.layer.3.attention.output.LayerNorm.bias",
"layers.3.up.weight": "roberta.encoder.layer.3.intermediate.dense.weight",
"layers.3.up.bias": "roberta.encoder.layer.3.intermediate.dense.bias",
"layers.3.down.weight": "roberta.encoder.layer.3.output.dense.weight",
"layers.3.down.bias": "roberta.encoder.layer.3.output.dense.bias",
"layers.3.output_norm.weight": "roberta.encoder.layer.3.output.LayerNorm.weight",
"layers.3.output_norm.bias": "roberta.encoder.layer.3.output.LayerNorm.bias",
"layers.4.attention.query.weight": "roberta.encoder.layer.4.attention.self.query.weight",
"layers.4.attention.query.bias": "roberta.encoder.layer.4.attention.self.query.bias",
"layers.4.attention.key.weight": "roberta.encoder.layer.4.attention.self.key.weight",
"layers.4.attention.key.bias": "roberta.encoder.layer.4.attention.self.key.bias",
"layers.4.attention.value.weight": "roberta.encoder.layer.4.attention.self.value.weight",
"layers.4.attention.value.bias": "roberta.encoder.layer.4.attention.self.value.bias",
"layers.4.attention.out.weight": "roberta.encoder.layer.4.attention.output.dense.weight",
"layers.4.attention.out.bias": "roberta.encoder.layer.4.attention.output.dense.bias",
"layers.4.attention_norm.weight": "roberta.encoder.layer.4.attention.output.LayerNorm.weight",
"layers.4.attention_norm.bias": "roberta.encoder.layer.4.attention.output.LayerNorm.bias",
"layers.4.up.weight": "roberta.encoder.layer.4.intermediate.dense.weight",
"layers.4.up.bias": "roberta.encoder.layer.4.intermediate.dense.bias",
"layers.4.down.weight": "roberta.encoder.layer.4.output.dense.weight",
"layers.4.down.bias": "roberta.encoder.layer.4.output.dense.bias",
"layers.4.output_norm.weight": "roberta.encoder.layer.4.output.LayerNorm.weight",
"layers.4.output_norm.bias": "roberta.encoder.layer.4.output.LayerNorm.bias",
"layers.5.attention.query.weight": "roberta.encoder.layer.5.attention.self.query.weight",
"layers.5.attention.query.bias": "roberta.encoder.layer.5.attention.self.query.bias",
"layers.5.attention.key.weight": "roberta.encoder.layer.5.attention.self.key.weight",
"layers.5.attention.key.bias": "roberta.encoder.layer.5.attention.self.key.bias",
"layers.5.attention.value.weight": "roberta.encoder.layer.5.attention.self.value.weight",
"layers.5.attention.value.bias": "roberta.encoder.layer.5.attention.self.value.bias",
"layers.5.attention.out.weight": "roberta.encoder.layer.5.attention.output.dense.weight",
"layers.5.attention.out.bias": "roberta.encoder.layer.5.attention.output.dense.bias",
"layers.5.attention_norm.weight": "roberta.encoder.layer.5.attention.output.LayerNorm.weight",
"layers.5.attention_norm.bias": "roberta.encoder.layer.5.attention.output.LayerNorm.bias",
"layers.5.up.weight": "roberta.encoder.layer.5.intermediate.dense.weight",
"layers.5.up.bias": "roberta.encoder.layer.5.intermediate.dense.bias",
"layers.5.down.weight": "roberta.encoder.layer.5.output.dense.weight",
"layers.5.down.bias": "roberta.encoder.layer.5.output.dense.bias",
"layers.5.output_norm.weight": "roberta.encoder.layer.5.output.LayerNorm.weight",
"layers.5.output_norm.bias": "roberta.encoder.layer.5.output.LayerNorm.bias",
"layers.6.attention.query.weight": "roberta.encoder.layer.6.attention.self.query.weight",
"layers.6.attention.query.bias": "roberta.encoder.layer.6.attention.self.query.bias",
"layers.6.attention.key.weight": "roberta.encoder.layer.6.attention.self.key.weight",
"layers.6.attention.key.bias": "roberta.encoder.layer.6.attention.self.key.bias",
"layers.6.attention.value.weight": "roberta.encoder.layer.6.attention.self.value.weight",
"layers.6.attention.value.bias": "roberta.encoder.layer.6.attention.self.value.bias",
"layers.6.attention.out.weight": "roberta.encoder.layer.6.attention.output.dense.weight",
"layers.6.attention.out.bias": "roberta.encoder.layer.6.attention.output.dense.bias",
"layers.6.attention_norm.weight": "roberta.encoder.layer.6.attention.output.LayerNorm.weight",
"layers.6.attention_norm.bias": "roberta.encoder.layer.6.attention.output.LayerNorm.bias",
"layers.6.up.weight": "roberta.encoder.layer.6.intermediate.dense.weight",
"layers.6.up.bias": "roberta.encoder.layer.6.intermediate.dense.bias",
"layers.6.down.weight": "roberta.encoder.layer.6.output.dense.weight",
"layers.6.down.bias": "roberta.encoder.layer.6.output.dense.bias",
"layers.6.output_norm.weight": "roberta.encoder.layer.6.output.LayerNorm.weight",
"layers.6.output_norm.bias": "roberta.encoder.layer.6.output.LayerNorm.bias",
"layers.7.attention.query.weight": "roberta.encoder.layer.7.attention.self.query.weight",
"layers.7.attention.query.bias": "roberta.encoder.layer.7.attention.self.query.bias",
"layers.7.attention.key.weight": "roberta.encoder.layer.7.attention.self.key.weight",
"layers.7.attention.key.bias": "roberta.encoder.layer.7.attention.self.key.bias",
"layers.7.attention.value.weight": "roberta.encoder.layer.7.attention.self.value.weight",
"layers.7.attention.value.bias": "roberta.encoder.layer.7.attention.self.value.bias",
"layers.7.attention.out.weight": "roberta.encoder.layer.7.attention.output.dense.weight",
"layers.7.attention.out.bias": "roberta.encoder.layer.7.attention.output.dense.bias",
"layers.7.attention_norm.weight": "roberta.encoder.layer.7.attention.output.LayerNorm.weight",
"layers.7.attention_norm.bias": "roberta.encoder.layer.7.attention.output.LayerNorm.bias",
"layers.7.up.weight": "roberta.encoder.layer.7.intermediate.dense.weight",
"layers.7.up.bias": "roberta.encoder.layer.7.intermediate.dense.bias",
"layers.7.down.weight": "roberta.encoder.layer.7.output.dense.weight",
"layers.7.down.bias": "roberta.encoder.layer.7.output.dense.bias",
"layers.7.output_norm.weight": "roberta.encoder.layer.7.output.LayerNorm.weight",
"layers.7.output_norm.bias": "roberta.encoder.layer.7.output.LayerNorm.bias",
"layers.8.attention.query.weight": "roberta.encoder.layer.8.attention.self.query.weight",
"layers.8.attention.query.bias": "roberta.encoder.layer.8.attention.self.query.bias",
"layers.8.attention.key.weight": "roberta.encoder.layer.8.attention.self.key.weight",
"layers.8.attention.key.bias": "roberta.encoder.layer.8.attention.self.key.bias",
"layers.8.attention.value.weight": "roberta.encoder.layer.8.attention.self.value.weight",
"layers.8.attention.value.bias": "roberta.encoder.layer.8.attention.self.value.bias",
"layers.8.attention.out.weight": "roberta.encoder.layer.8.attention.output.dense.weight",
"layers.8.attention.out.bias": "roberta.encoder.layer.8.attention.output.dense.bias",
"layers.8.attention_norm.weight": "roberta.encoder.layer.8.attention.output.LayerNorm.weight",
"layers.8.attention_norm.bias": "roberta.encoder.layer.8.attention.output.LayerNorm.bias",
"layers.8.up.weight": "roberta.encoder.layer.8.intermediate.dense.weight",
"layers.8.up.bias": "roberta.encoder.layer.8.intermediate.dense.bias",
"layers.8.down.weight": "roberta.encoder.layer.8.output.dense.weight",
"layers.8.down.bias": "roberta.encoder.layer.8.output.dense.bias",
"layers.8.output_norm.weight": "roberta.encoder.layer.8.output.LayerNorm.weight",
"layers.8.output_norm.bias": "roberta.encoder.layer.8.output.LayerNorm.bias",
"layers.9.attention.query.weight": "roberta.encoder.layer.9.attention.self.query.weight",
"layers.9.attention.query.bias": "roberta.encoder.layer.9.attention.self.query.bias",
"layers.9.attention.key.weight": "roberta.encoder.layer.9.attention.self.key.weight",
"layers.9.attention.key.bias": "roberta.encoder.layer.9.attention.self.key.bias",
"layers.9.attention.value.weight": "roberta.encoder.layer.9.attention.self.value.weight",
"layers.9.attention.value.bias": "roberta.encoder.layer.9.attention.self.value.bias",
"layers.9.attention.out.weight": "roberta.encoder.layer.9.attention.output.dense.weight",
"layers.9.attention.out.bias": "roberta.encoder.layer.9.attention.output.dense.bias",
"layers.9.attention_norm.weight": "roberta.encoder.layer.9.attention.output.LayerNorm.weight",
"layers.9.attention_norm.bias": "roberta.encoder.layer.9.attention.output.LayerNorm.bias",
"layers.9.up.weight": "roberta.encoder.layer.9.intermediate.dense.weight",
"layers.9.up.bias": "roberta.encoder.layer.9.intermediate.dense.bias",
"layers.9.down.weight": "roberta.encoder.layer.9.output.dense.weight",
"layers.9.down.bias": "roberta.encoder.layer.9.output.dense.bias",
"layers.9.output_norm.weight": "roberta.encoder.layer.9.output.LayerNorm.weight",
"layers.9.output_norm.bias": "roberta.encoder.layer.9.output.LayerNorm.bias",
"layers.10.attention.query.weight": "roberta.encoder.layer.10.attention.self.query.weight",
"layers.10.attention.query.bias": "roberta.encoder.layer.10.attention.self.query.bias",
"layers.10.attention.key.weight": "roberta.encoder.layer.10.attention.self.key.weight",
"layers.10.attention.key.bias": "roberta.encoder.layer.10.attention.self.key.bias",
"layers.10.attention.value.weight": "roberta.encoder.layer.10.attention.self.value.weight",
"layers.10.attention.value.bias": "roberta.encoder.layer.10.attention.self.value.bias",
"layers.10.attention.out.weight": "roberta.encoder.layer.10.attention.output.dense.weight",
"layers.10.attention.out.bias": "roberta.encoder.layer.10.attention.output.dense.bias",
"layers.10.attention_norm.weight": "roberta.encoder.layer.10.attention.output.LayerNorm.weight",
"layers.10.attention_norm.bias": "roberta.encoder.layer.10.attention.output.LayerNorm.bias",
"layers.10.up.weight": "roberta.encoder.layer.10.intermediate.dense.weight",
"layers.10.up.bias": "roberta.encoder.layer.10.intermediate.dense.bias",
"layers.10.down.weight": "roberta.encoder.layer.10.output.dense.weight",
"layers.10.down.bias": "roberta.encoder.layer.10.output.dense.bias",
"layers.10.output_norm.weight": "roberta.encoder.layer.10.output.LayerNorm.weight",
"layers.10.output_norm.bias": "roberta.encoder.layer.10.output.LayerNorm.bias",
"layers.11.attention.query.weight": "roberta.encoder.layer.11.attention.self.query.weight",
"layers.11.attention.query.bias": "roberta.encoder.layer.11.attention.self.query.bias",
"layers.11.attention.key.weight": "roberta.encoder.layer.11.attention.self.key.weight",
"layers.11.attention.key.bias": "roberta.encoder.layer.11.attention.self.key.bias",
"layers.11.attention.value.weight": "roberta.encoder.layer.11.attention.self.value.weight",
"layers.11.attention.value.bias": "roberta.encoder.layer.11.attention.self.value.bias",
"layers.11.attention.out.weight": "roberta.encoder.layer.11.attention.output.dense.weight",
"layers.11.attention.out.bias": "roberta.encoder.layer.11.attention.output.dense.bias",
"layers.11.attention_norm.weight": "roberta.encoder.layer.11.attention.output.LayerNorm.weight",
"layers.11.attention_norm.bias": "roberta.encoder.layer.11.attention.output.LayerNorm.bias",
"layers.11.up.weight": "roberta.encoder.layer.11.intermediate.dense.weight",
"layers.11.up.bias": "roberta.encoder.layer.11.intermediate.dense.bias",
"layers.11.down.weight": "roberta.encoder.layer.11.output.dense.weight",
"layers.11.down.bias": "roberta.encoder.layer.11.output.dense.bias",
"layers.11.output_norm.weight": "roberta.encoder.layer.11.output.LayerNorm.weight",
"layers.11.output_norm.bias": "roberta.encoder.layer.11.output.LayerNorm.bias"
}linnet.toml62 B
toml
[package]
name = "roberta"
version = "0.1.0"
language = "0.1"nest.toml856 B
toml
[model]
name = "roberta-base"
title = "RoBERTa Base"
summary = "Meta's more thoroughly trained BERT: the same post-norm encoder, dynamic masking, no next-sentence-prediction, and a byte-level BPE vocabulary of 50265."
license = "MIT"
family = "bert"
tags = ["fill-mask", "text-encoder", "encoder-only"]
[links]
huggingface = "https://huggingface.co/FacebookAI/roberta-base"
github = "https://github.com/facebookresearch/fairseq/tree/main/examples/roberta"
arxiv = "https://arxiv.org/abs/1907.11692"
[source]
path = "src/lib.linnet"
root = "Model"
entry = "forward"
[generics]
Vocab = 50265
MaxPositions = 514
TypeVocab = 1
D = 768
Heads = 12
Inner = 3072
Layers = 12
T = "f32"
[check]
B = 1
S = 8
[weights]
repo = "FacebookAI/roberta-base"
revision = "e2da8e2f811d1448a5b465c236feacd80ffbac7b"
files = ["model.safetensors"]
bindings = "bindings.json"samples.json1.3 kB
json
{
"kind": "similarity",
"sentences": [
"The cat curled up on the warm windowsill and fell asleep in the sun.",
"Our new kitten refuses to eat anything except wet food.",
"The central bank raised interest rates again to cool inflation.",
"She rebalanced her portfolio after the market's sharp swings.",
"The function threw an exception when the array index ran out of bounds.",
"He refactored the module to remove three copies of the same logic."
],
"matrix": [
[
1.0,
0.970358,
0.950273,
0.961396,
0.954863,
0.956733
],
[
0.970358,
1.0,
0.95579,
0.962669,
0.960474,
0.963584
],
[
0.950273,
0.95579,
1.0,
0.973575,
0.947647,
0.955688
],
[
0.961396,
0.962669,
0.973575,
1.0,
0.956782,
0.969118
],
[
0.954863,
0.960474,
0.947647,
0.956782,
1.0,
0.977434
],
[
0.956733,
0.963584,
0.955688,
0.969118,
0.977434,
1.0
]
],
"caption": "Not a trained sentence embedder: the last hidden state, mean-pooled over tokens and L2-normalized by hand, since this checkpoint has no embedding head."
}src
lib.linnet4.8 kB
linnet
// RoBERTa: the same post-norm BERT encoder as `bert-base-uncased`, but
// with one training-only quirk baked into the checkpoint's position table:
// positions are not zero-based. `padding_idx` (RoBERTa's pad token, id 1)
// reserves position row 0, and the embedding module offsets every real
// position by `padding_idx + 1`, so token index i (0-based) reads position
// row `i + 2`. Upstream computes this from the attention mask
// (`create_position_ids_from_input_ids`, cumulative count of non-pad
// tokens); with no padding in the sequence -- the only case this encoder's
// full attention supports -- that count is just the token's own index, so
// the offset collapses to a constant. `type_vocab_size` is 1: RoBERTa has
// no next-sentence-prediction pretraining, so the token-type table is a
// single row broadcast to every token, kept as an explicit embedding
// lookup so the checkpoint's tensor loads unmodified.
module roberta
use std.nn.activations::{gelu_erf}
use std.nn.attention::{attention}
use std.nn.embedding::{embedding}
use std.nn.linear::{Linear}
use std.nn.norm::{layer_norm}
pub const EPS: f32 = 1e-5
// RoBERTa's padding index is 1; real positions start at `padding_idx + 1`.
const POSITION_OFFSET = 2
pub block LayerNorm<D: Dim, T: Float> {
param weight: Tensor[D; T]
param bias: Tensor[D; T]
pub fn forward<*S: Shape>(x: Tensor[*S, D; T]) -> Tensor[*S, D; T] {
return layer_norm(x, weight, some(bias), EPS)
}
}
pub block Embeddings<Vocab: Dim, MaxPositions: Dim, TypeVocab: Dim, D: Dim, T: Float> {
param word_embeddings: Tensor[Vocab, D; T]
param position_embeddings: Tensor[MaxPositions, D; T]
param token_type_embeddings: Tensor[TypeVocab, D; T]
sub norm: LayerNorm<D, T>
pub fn forward<B: Dim, S: Dim>(
tokens: Tensor[B, S; i32],
token_types: Tensor[B, S; i32],
) -> Tensor[B, S, D; T]
where S + POSITION_OFFSET <= MaxPositions {
// Token i reads position row i + 2, not i; see the module comment.
let positions = iota<i32>(S) + POSITION_OFFSET
let x =
embedding(tokens, word_embeddings) +
embedding(positions, position_embeddings) +
embedding(token_types, token_type_embeddings)
return norm.forward(x)
}
}
pub block Attention<D: Dim, Heads: Dim, T: Float>
where
Heads > 0,
D % Heads == 0
{
sub query: Linear<D, D, T>
sub key: Linear<D, D, T>
sub value: Linear<D, D, T>
sub out: Linear<D, D, T>
// Full attention over every position: the language's attention op takes
// an unbatched [Q, K] mask, so a per-sequence padding mask cannot be
// expressed here. Callers pad-free or accept attention onto padding.
pub fn forward<B: Dim, S: Dim>(x: Tensor[B, S, D; T]) -> Tensor[B, S, D; T] {
let mixed = attention(
heads<B, S, Heads, D / Heads, T>(query.forward(x)),
heads<B, S, Heads, D / Heads, T>(key.forward(x)),
heads<B, S, Heads, D / Heads, T>(value.forward(x)),
rsqrt(cast<f32>(D / Heads)),
none,
)
return out.forward(reshape(permute(mixed, [0, 2, 1, 3]), [B, S, D]))
}
}
fn heads<B: Dim, S: Dim, Heads: Dim, Dh: Dim, T: Float>(
x: Tensor[B, S, Heads * Dh; T],
) -> Tensor[B, Heads, S, Dh; T] {
return permute(reshape(x, [B, S, Heads, Dh]), [0, 2, 1, 3])
}
// Post-norm encoder block: `h = LayerNorm(x + attention(x))`, then
// `out = LayerNorm(h + ffn(h))`.
pub block EncoderLayer<D: Dim, Heads: Dim, Inner: Dim, T: Float>
where
Heads > 0,
D % Heads == 0
{
sub attention: Attention<D, Heads, T>
sub attention_norm: LayerNorm<D, T>
sub up: Linear<D, Inner, T>
sub down: Linear<Inner, D, T>
sub output_norm: LayerNorm<D, T>
pub fn forward<B: Dim, S: Dim>(x: Tensor[B, S, D; T]) -> Tensor[B, S, D; T] {
let attended = attention_norm.forward(x + attention.forward(x))
return output_norm.forward(attended + down.forward(gelu_erf(up.forward(attended))))
}
}
pub block Model<
Vocab: Dim,
MaxPositions: Dim,
TypeVocab: Dim,
D: Dim,
Heads: Dim,
Inner: Dim,
Layers: Dim,
T: Float = f32,
>
where
Heads > 0,
D % Heads == 0
{
sub embeddings: Embeddings<Vocab, MaxPositions, TypeVocab, D, T>
sub layers: [EncoderLayer<D, Heads, Inner, T>; Layers]
// The last hidden state: every token's contextualized representation.
// This checkpoint (`RobertaForMaskedLM`) carries no pooler; only the
// masked-language-model head, which this card does not expose.
pub entry forward<B: Dim, S: Dim>(
tokens: Tensor[B, S; i32],
token_types: Tensor[B, S; i32],
) -> Tensor[B, S, D; T]
where S + POSITION_OFFSET <= MaxPositions {
var x = embeddings.forward(tokens, token_types)
static for layer in layers {
x = layer.forward(x)
}
return x
}
}