Skip to content

PyTorch BERT-base

Overview

This tutorial shows how to compile and run inference with the HuggingFace BERT-base model for masked language modeling using both the RBLN Python API and PyTorch's torch.compile() API.

Setup & Installation

Before you begin, ensure that your system environment is properly configured and that all required packages are installed. This includes:

  • System Requirements:
  • Package Requirements:
  • Installation Command:
    1
2
3
4
    pip install --extra-index-url https://download.pytorch.org/whl/cpu torch transformers numpy
pip install \
  --extra-index-url https://pypi.rbln.ai/simple \
  rebel-compiler==0.12.0

Info

  • rebel-compiler requires an RBLN Portal account.
  • The commands above are intended for a default pip install on Debian-based Linux such as Ubuntu. For all other configurations, refer to the Installation Guide for the supported install matrix and the applicable commands.

Using RBLN Python API

Model Compilation

Load the pre-trained BERT-base model from HuggingFace, set it to evaluation mode, compile it using the RBLN Compiler, and save the compiled model to local storage.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
import torch  
from transformers import BertForMaskedLM  
import rebel  # RBLN Compiler  

# Instantiate the BERT-base model  
bert_model = BertForMaskedLM.from_pretrained('bert-base-uncased', return_dict=False)  
bert_model.eval()  

# Compile the model  
MAX_SEQ_LEN = 128  
input_info = [  
    ('input_ids', [1, MAX_SEQ_LEN], 'int64'),  
    ('attention_mask', [1, MAX_SEQ_LEN], 'int64'),  
    ('token_type_ids', [1, MAX_SEQ_LEN], 'int64'),  
]  
compiled_model = rebel.compile_from_torch(  
    bert_model,  
    input_info,  
)  
compiled_model.save('bert_base.rbln')

Model Inference

Prepare the input using BertTokenizer, load the compiled model with RBLN Runtime, run inference, and display the predicted word for the masked token.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
import torch  
from transformers import BertTokenizer  
import rebel  # RBLN Runtime  

MAX_SEQ_LEN = 128  
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')  
text = 'the color of rose is [MASK].'  
inputs = tokenizer(text, return_tensors='pt', padding='max_length', max_length=MAX_SEQ_LEN)  

# Load the compiled model  
module = rebel.Runtime('bert_base.rbln', tensor_type='pt')  

# Run inference  
output = module(**inputs)  

# Display results  
mask_token_index = (inputs.input_ids == tokenizer.mask_token_id)[0].nonzero(as_tuple=True)[0]  
predicted_token_id = output[0, mask_token_index].argmax(axis=-1)  
print(f'Predicted token: {tokenizer.decode(predicted_token_id)}')

Example Output: 

1
Predicted token: white

Using the torch.compile() API

Model and Input Preparation

Load the BERT-base model and prepare the input using BertTokenizer.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
import torch  
from transformers import BertForMaskedLM, BertTokenizer  

# Instantiate the model  
bert_model = BertForMaskedLM.from_pretrained('bert-base-uncased')  
bert_model.eval()  

MAX_SEQ_LEN = 128  
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')  
text = 'The color of a rose is [MASK].'  
inputs = tokenizer(text, return_tensors='pt', padding='max_length', max_length=MAX_SEQ_LEN)

Compilation and Inference

Compile the model using torch.compile() with the RBLN backend, trigger compilation on the first forward pass, and run inference on the input.

Note

To use the RBLN backend with torch.compile(), you must specify backend='rbln'. This requires import rebel before compilation.

Warning

If the RBLN backend fails to compile the graph, torch.compile() falls back to eager execution on the CPU and only logs a warning. Add 'mode': 'strict' to options to raise an error instead.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
import rebel  # Need to use torch dynamo's "rbln" backend.

compiled_model = torch.compile(  
    bert_model,  
    backend='rbln',  
    options={'cache_dir': './.cache'},  
    dynamic=False  
)  

# Trigger compilation with the first forward pass  
compiled_model(**inputs)  

# Run inference  
logits = compiled_model(**inputs).logits  

# Get the mask token index  
mask_token_index = (inputs.input_ids == tokenizer.mask_token_id)[0].nonzero(as_tuple=True)[0]  

# Get the predicted token ID  
predicted_token_id = logits[0, mask_token_index].argmax(axis=-1)  

# Print the predicted word  
print(f'Predicted word: {tokenizer.decode(predicted_token_id)}')

Example Output: 

1
Predicted word: white

References