Candle Tutorial - Convert Pytorch Models to Candle

Candle is an ML framework written in rust that takes advantage of the speed and memory safety Rust provides for writing machine workloads. It can be used as a drop in replacement for ML frameworks like PyTorch, it also has python bindings so you can use it from python...

This repo provides some guide for converting pytorch models from the transformers library to Candle by directly translating the pytorch code to Candle ...

❗️❗️: To make the code easily understandable, I have annotated each line of the Rust/Candle code with the equivalent PyTorch code. Tutorial Structure:

• Getting Started

  • 0. Important things to note
  • 1. Start a new rust project
  • 2. Install Candle & Other Packages

• Parallels between Pytorch and Candle

  • Tensors
  • Tensor Operations

• Translating a PyTorch Transformer Model into Candle

  • RoBERTa
    • a. Writing Building Blocks
    • b. Roberta Config
    • c. RobertaEmbeddings
    • d. RobertaSelfAttention
    • e. RobertaSelfOutput
    • f. RobertaIntermediate
    • g. RobertaOutput
    • h. RobertaLayer
    • i. RobertaEncoder
    • j. RobertaModel

• Debugging/Testing the model

Getting Started:

0. Important things to note

• When Porting an already trained checkpoint to Candle, there's a bunch of PyTorch code that are not relevant and they are mostly included for handling different scenarios in training. It's definitely beneficial to know which functions to bypass if the conversion effort is mostly geared towards loading an already trained model.

• Python Built in Method: Unlike Python where we have built-in methods like __call__ that allow us to use a class as a method and __init__ for initializing a class, In rust we have to explicitly define methods like Class::new() to initialize a class and Class::forward() to perform a forward pass. This is going to be a recurrent theme in most of the code shown below.

• It is important to write unit tests after writing most or every module to ensure that input and output shapes in Candle are consistent with the same module in pytorch.

• In PyTorch, we can initialize module weights by creating a class method _init_weights but in candle it becomes a design decision, you can initialize a tensor using the shape of your weights/bias (e.g. ) and hold it in a VarBuilder which then used to initialize the tensors in each module.

1. Start a new rust project

The command below will create a new rust project called candle-roberta in the current directory with a Cargo.toml file and a src directory with a main.rs file in it.

$ cargo new candle-roberta

2. Install Candle & Other Packages

You can follow the instructions here to install candle or you can use the command below to install candle directly from github.

For this tutorial, we would be using the candle-core and candle-nn crates. candle-core provides the core functionality of the candle framework. It provides an implementation the basic blocks for building neural networks and also integrations with different backends like Cuda, MKL, CPU etc, while candle-nn provides a high level API for building neural networks.

- cargo add --git https://github.com/huggingface/candle.git candle-core  # install candle-core
- cargo add --git https://github.com/huggingface/candle.git candle-nn # install candle-nn

Other frameworks we would need for this tutorial are:

• anyhow for error handling ==> cargo add anyhow
• serde for serialization ==> cargo add serde
• serde_json for json serialization ==> cargo add serde_json
• hf-hub for integrating with the huggingface hub ==> cargo add hf-hub
• tokenizers for tokenizing text ==> cargo add tokenizers

3. Parallels between Pytorch and Candle

To convert a pytorch model to candle, it is important understand the parallels between the two frameworks.

• Candle is a rust framework, so it is statically typed, while pytorch is a python framework, so it is dynamically typed. This means that you need to be explicit about the types of your variables in candle, while in pytorch, you don't need to be explicit about the types of your variables.

Tensors

The examples shows below can be found here;

• Initializing a Tensor: Tensors can be directly created from an array in both frameworks

  • Pytorch: in pytorch the data type is automatically inffereed from the data;

    import torch
    from typing import List
    
    data: List = [1, 2, 3]
    tensor = torch.tensor(data)
    print(tensor)
    
    nested_data = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
    nested_tensor = torch.tensor(nested_data)
    print(nested_tensor)
    

  • Candle: in candle, the data type needs to be explicitly specified;

    use candle_core::{DType, Device, Tensor};
    use anyhow::Result;
    
    let data: [u32; 3] = [1u32, 2, 3];
    let tensor = Tensor::new(&data, &Device::Cpu)?;
    println!("tensor: {:?}", tensor.to_vec1::<u32>()?);
    
    let nested_data: [[u32; 3]; 3] = [[1u32, 2, 3], [4, 5, 6], [7, 8, 9]];
    let nested_tensor = Tensor::new(&nested_data, &Device::Cpu)?;
    println!("nested_tensor: {:?}", nested_tensor.to_vec2::<u32>()?);
    

• Creating a tensor from another tensor

  • Pytorch: in pytorch, the data type is automatically inferred from the data;

    zero_tensor = torch.zeros_like(tensor)
    ones_tensor = torch.ones_like(tensor)
    random_tensor = torch.rand_like(tensor)
    

  • Candle: in candle, the data type needs to be explicitly specified;

    let data: [u32; 3] = [1u32, 2, 3];
    let tensor = Tensor::new(&data, &Device::Cpu)?;
    
    let zero_tensor = tensor.zeros_like()?;
    println!("zero_tensor: {:?}", zero_tensor.to_vec1::<u32>()?);
    
    let ones_tensor = tensor.ones_like()?;
    println!("ones_tensor: {:?}", ones_tensor.to_vec1::<u32>()?);
    
    let random_tensor = tensor.rand_like(0.0, 1.0)?;
    println!("random_tensor: {:?}", random_tensor.to_vec1::<f64>()?);
    

• Checking tensor dimensions:

  • PyTorch

    print(tensor.shape)
    print(tensor.size())
    

  • Candle

    // 1 dimensional tensor
    println!("tensor shape: {:?}", tensor.shape().dims()); 
    // 2 dimensional tensor
    println!("tensor shape: {:?}", tensor.shape().dims2()); 
    // 3 dimensional tensor
    println!("tensor shape: {:?}", tensor.shape().dims3()); 
    

Tensor Operations:

Performing tensor operations is pretty similar across both frameworks
Some examples can be found here:: [Candle CheatSheet](https://github.com/huggingface/candle/blob/main/README.md#how-to-use)

3. Translating a PyTorch Transformer Model into Candle

Here's the fun part! In this section we are going to take a look at translating models from the transformers library to candle. We would be using the RoBERTa and XLM-Roberta model for this tutorial.

We would be translating the Pytorch Source Code into Candle Code and then load the pretrained checkpoint into Rust and compare the output from both frameworks.

Note ❗️❗️: To make the code easily understandable, I have annotated each line of the Rust/Candle code with the equivalent PyTorch code.

3.1. RoBERTa

RoBERTa is a variant of the BERT model. Although both models have different pretraining approaches, structurally both models are very similar and the major difference between both models is that in the RoBERTa layer, Position numbers begin at padding_idx+1, While in BERT, Position numbers begin at 0.

Following the transformers PyTorch implementation, RoBERTa Model can be divided into the 2 main parts (embeddings and encoder):

RobertaModel(
  (embeddings): RobertaEmbeddings(
    (word_embeddings): Embedding(50265, 768, padding_idx=1)
    (position_embeddings): Embedding(514, 768, padding_idx=1)
    (token_type_embeddings): Embedding(1, 768)
    (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
    (dropout): Dropout(p=0.1, inplace=False)
  )
  (encoder): RobertaEncoder(
    (layer): ModuleList(
      (0-11): 12 x RobertaLayer(
        (attention): RobertaAttention(
          (self): RobertaSelfAttention(
            (query): Linear(in_features=768, out_features=768, bias=True)
            (key): Linear(in_features=768, out_features=768, bias=True)
            (value): Linear(in_features=768, out_features=768, bias=True)
            (dropout): Dropout(p=0.1, inplace=False)
          )
          (output): RobertaSelfOutput(
            (dense): Linear(in_features=768, out_features=768, bias=True)
            (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
            (dropout): Dropout(p=0.1, inplace=False)
          )
        )
        (intermediate): RobertaIn