BERT: PyTorch Reimplementation

A from-scratch PyTorch reimplementation of the BERT-Base, Uncased architecture from the 2018 paper. The implementation is validated against the released google-bert/bert-base-uncased checkpoint.

Implemented

• BERT-Base embeddings, bidirectional self-attention, feed-forward blocks, residual connections, LayerNorm, dropout, and pooler.
• Masked-language-model and next-sentence-prediction heads.
• Tied MLM decoder and input word-embedding weights.
• MLM and NSP cross-entropy losses with the same model structure and state dictionary layout as the reference implementation.

Verification

verify.py uses the pinned Hugging Face checkpoint as a reference and checks:

1. Exact parameter names and tensor shapes with strict state-dictionary loading.
2. WordPiece tokenization, special tokens, sentence pairs, omitted segment IDs, padding masks, and the maximum sequence length.
3. MLM and NSP logits on a deterministic masked sentence pair.
4. The combined loss and its MLM/NSP decomposition.
5. Weight tying between the MLM decoder and input embeddings.

The verification compares the local implementation with the reference model; it does not claim to reproduce the original multi-billion-word pretraining run.

Run verification

uv sync
uv run --extra verify verify.py

Observed verification output from the latest run:

pretraining state_dict: compatible
weight tying: cls.predictions.decoder.weight is bert.embeddings.word_embeddings.weight
tokenizer: compatible ('Hello, world!')
tokens: ['[CLS]', 'hello', ',', 'world', '!', '[SEP]']
input_ids: [101, 7592, 1010, 2088, 999, 102]
raw text: sequence=0.000e+00, pooler=0.000e+00
sentence pair: sequence=0.000e+00, pooler=0.000e+00
omitted token_type_ids: sequence=0.000e+00, pooler=0.000e+00
different padding and bool mask: sequence=0.000e+00, pooler=0.000e+00
sequence length 512: sequence=0.000e+00, pooler=0.000e+00
encoder max sequence error: 0.000e+00
encoder max pooler error:   0.000e+00
forward parity: prediction_logits=0.000e+00, seq_relationship_logits=0.000e+00, loss=0.000e+00
loss decomposition: mlm=1.865043, nsp=0.007077, total=1.872119
model_id: google-bert/bert-base-uncased
model_revision: 86b5e0934494bd15c9632b12f734a8a67f723594
torch: 2.13.0+cpu
transformers: 5.14.1
encoder parameters: 109482240
pretraining model parameters: 110106428

References

• Devlin et al., BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding
• BERT-Base Uncased checkpoint
• Hugging Face BERT implementation