Skip to content

Model APIs Overview

The RBLN Optimum library provides APIs for using HuggingFace models on RBLN NPUs. This overview focuses on the API design and usage patterns.

API Design

Naming Convention

RBLN Optimum uses the following naming pattern:

  • Model Classes: Add "RBLN" prefix to original HuggingFace class names 

    1
2
3
4
5
    # Original HuggingFace
from transformers import LlamaForCausalLM

# Optimum RBLN equivalent
from optimum.rbln import RBLNLlamaForCausalLM

  • Configuration Classes: Add "Config" suffix to the RBLN model class 

    1
2
    # Configuration for RBLN model
from optimum.rbln import RBLNLlamaForCausalLMConfig

  • Auto Classes: Add "RBLN" prefix to original Auto Model/Pipeline names

    • AutoModel 

      1
2
3
4
5
      # Original HuggingFace AutoModel
from transformers import AutoModelForCausalLM

# Optimum RBLN equivalent
from optimum.rbln import RBLNAutoModelForCausalLM

    • AutoPipeline 

      1
2
3
4
5
      # Original Diffusers AutoPipeline
from diffusers import AutoPipelineForText2Image

# Optimum RBLN equivalent
from optimum.rbln import RBLNAutoPipelineForText2Image

Static Graph Compilation

The RBLN SDK compiles models into static graphs, so the following parameters must be specified at compile time:

  • Batch size
  • Sequence length (for language models)
  • Number of devices (for distributed execution)
  • Other model-specific parameters

The configuration classes (RBLNModelNameConfig) allow you to specify these static values.

Supported Methods

RBLN Optimum preserves the same API interface as the original HuggingFace models:

  • Language models support .generate() method
  • Vision and other models support both direct call syntax (model(inputs)) and .forward() method
  • Diffusers pipelines support the __call__() method (used with direct call syntax)

Usage Patterns

There are multiple ways to configure and use RBLN models:

Method 1: Using rbln_* parameters

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
from optimum.rbln import RBLNLlamaForCausalLM
from transformers import AutoTokenizer

# Load and compile the model with inline parameters
model = RBLNLlamaForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-chat-hf",
    export=True,  # Set to True for compilation
    rbln_batch_size=1,
    rbln_max_seq_len=512,  # Optional, inferred from model config if not specified
    rbln_num_devices=1,  # Rebel100™ compiles for one NPU
)

# Load tokenizer and prepare inputs
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf")
inputs = tokenizer("Hello, I'm a language model", return_tensors="pt")

# Generate text with the model
outputs = model.generate(**inputs, max_new_tokens=64)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
from optimum.rbln import RBLNLlamaForCausalLM
from transformers import AutoTokenizer

# Load and compile the model with inline parameters
model = RBLNLlamaForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-chat-hf",
    export=True,  # Set to True for compilation
    rbln_batch_size=1,
    rbln_max_seq_len=512,  # Optional, inferred from model config if not specified
    rbln_num_devices=4,  # For multi-device execution
)

# Load tokenizer and prepare inputs
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf")
inputs = tokenizer("Hello, I'm a language model", return_tensors="pt")

# Generate text with the model
outputs = model.generate(**inputs, max_new_tokens=64)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
from optimum.rbln import RBLNLlamaForCausalLM, RBLNLlamaForCausalLMConfig
from transformers import AutoTokenizer

# Define configuration through RBLNConfig object
config = RBLNLlamaForCausalLMConfig(batch_size=1, max_seq_len=4096, num_devices=1)

# Load and compile the model with config object
model = RBLNLlamaForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-chat-hf", export=True, rbln_config=config
)

# Load tokenizer and prepare inputs
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf")
inputs = tokenizer("Hello, I'm a language model", return_tensors="pt")

# Generate text with the model
outputs = model.generate(**inputs, max_new_tokens=64)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
from optimum.rbln import RBLNLlamaForCausalLM, RBLNLlamaForCausalLMConfig
from transformers import AutoTokenizer

# Define configuration through RBLNConfig object
config = RBLNLlamaForCausalLMConfig(batch_size=1, max_seq_len=4096, num_devices=4)

# Load and compile the model with config object
model = RBLNLlamaForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-chat-hf", export=True, rbln_config=config
)

# Load tokenizer and prepare inputs
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf")
inputs = tokenizer("Hello, I'm a language model", return_tensors="pt")

# Generate text with the model
outputs = model.generate(**inputs, max_new_tokens=64)

Method 3: Using Auto Classes

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
from optimum.rbln import RBLNAutoModelForCausalLM
from transformers import AutoTokenizer

# Load and compile the model with inline parameters
model = RBLNAutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-chat-hf",
    export=True,  # Set to True for compilation
    rbln_batch_size=1,
    rbln_max_seq_len=512,  # Optional, inferred from model config if not specified
    rbln_num_devices=1,  # Rebel100™ compiles for one NPU
)

# Load tokenizer and prepare inputs
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf")
inputs = tokenizer("Hello, I'm a language model", return_tensors="pt")

# Generate text with the model
outputs = model.generate(**inputs, max_new_tokens=64)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
from optimum.rbln import RBLNAutoModelForCausalLM
from transformers import AutoTokenizer

# Load and compile the model with inline parameters
model = RBLNAutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-chat-hf",
    export=True,  # Set to True for compilation
    rbln_batch_size=1,
    rbln_max_seq_len=512,  # Optional, inferred from model config if not specified
    rbln_num_devices=4,  # For multi-device execution
)

# Load tokenizer and prepare inputs
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf")
inputs = tokenizer("Hello, I'm a language model", return_tensors="pt")

# Generate text with the model
outputs = model.generate(**inputs, max_new_tokens=64)

Benefits

  • Drop-in Replacement: Replace original HuggingFace imports with RBLN equivalents
  • Same Familiar API: Use the same methods you're already familiar with
  • Fine-grained Control: Configure static parameters for maximum performance

For detailed model support information and hardware compatibility, refer to the RBLN Optimum Overview.