Post-training GPT Neo 1B with Rust & PyTorch Pt. 1 | ML Engineering

 Hi All 

Today we have an example of post-training a small model (GPT Neo 1B) in pure Rust. This example focuses of the training loop and model updates, assuming you've already loaded the model weights.

Key Areas:

1.    Model Loading: We'll use the rust-bert crate for model loading. (supports GPT-Neo)

2.    Training Loop: Pure Rust implementation with tch-rs (Torch bindings) for tensor operations.

3.    Optimizer: AdamW optimizer for training.

 

View full code example below:

The following is a block-by-block breakdown of the Rust post-training code above.

 

1. Dependencies and Imports

use tch::{nn, Device, Tensor, Kind};  // Torch bindings for Rust
use rust_bert::pipelines::common::ModelType;  // Model types (GPT-Neo/Llama)
use rust_bert::pipelines::text_generation::TextGenerationModel;  // Model pipeline
use anyhow::Result;  // Error handling

*    tch - Rust binding for PyTorch

*     rust-bert - High-level Rust library for HuggingFace models.

*    anyhow - Simplifies error handling (like Pythons try/except).

 2. Main Function & Device Setup

fn main() -> Result<()> {
    // Set device (CPU/GPU)
    let device = Device::cuda_if_available();

*    Device::cuda_if_available();

        *    Automatically uses GPU if available (like torch.device("cuda") in Python).

        *    Falls back to CPU if no GPU is detected.

3.  Load Pre-trained Model

    // Load pre-trained model (GPT-Neo or Llama)
    let model = TextGenerationModel::new(
        ModelType::GptNeo,  // or ModelType::Llama
        None,  // Use default model
        None,  // Use default tokenizer
        device,
    )?;

 *    TextGenerationModel::new():

        *    Loads a HuggingFace-style model (GPT Neo 1B)

        *    ModelType::GPTNeo: Specifies the model architecture.

        *     None: Uses default weights/tokenizer (downloads if not cached)

        *    device: Moves the model to specified device (GPU/CPU)

4. Example Training Data

    // Example training data (tokenized)
    let input_ids = Tensor::of_slice(&[1, 2, 3, 4]).to(device);  // Batch of token IDs
    let labels = Tensor::of_slice(&[5, 6, 7, 8]).to(device);     // Target tokens

 *    Tensor::of_slice(&[1, 2, 3, 4]) :

         *    Creates a 1D "Tensor" in Python (actually a static array or a vector, mathematically tensors are n-dimentional structures )

        *    to_device: Moves the Tensor to CPU/GPU

*    input_ids: Input tokens (eg. [1, 2, 3, 4]).

*    labels: Target tokens for supervised learning (eg. [5, 6, 7, 8]). 

5. Model Parameters and Optimizer

    // Model parameters (simplified)
    let vs = nn::VarStore::new(device);  // Stores model weights
    let mut opt = nn::Adam::default().build(&vs, 1e-3)?;  // AdamW with lr=1e-3

*    nn::VarStore :

        *    Stores model weights (like PyTorch's nn.Parameter)

        *    new(device): Initializes on the specified device.

*    nn:Adam::default().build(&vs, 1e-3) :

        *    Creates Adam optimizer (like torch.optim.Adam)

        *    1e-3: Learning rate (adjustable)

        *    &vs: Links the optimizer to the VarStore. 

6. Training Loop

    // Training loop
    for epoch in 0..10 {
        // Forward pass
        let output = model.forward_t(&input_ids, None, None, false)?;
        let logits = output.logits;

*    model.forward_t():

        *    Performs a forward pass (like model(input_ids) in PyTorch).

        *    &input_ids: Input tensor (again tensor in a GPU context not a math one, this is a static array in 1D).

        *    None: Optional arguments (eg. attention mask, past key values etc.).

        *    false: train flag (set to true for dropout during training).

7. Loss Calculation

        // Compute loss (cross-entropy)
        let loss = logits.cross_entropy_for_logits(&labels);

 *    cross_entropy_for_logits():

        *    Computes cross entropy loss (like torch.nn.functional.cross_entropy).

        *    labels: Target tokens for comparison.

8. Backward Pass & Optimization

        // Backward pass
        opt.backward_step(&loss);

 *    opt.backward_step(&loss); :

        *    Computes gradients (loss.backward() in PyTorch)

        *    Updates weights (opt.step() in PyTorch).

9. Logging & Epoch Completion

        println!("Epoch {}: Loss = {:?}", epoch, f64::from(loss));
    }

*     println! : Logs the loss for each epoch.

*    f64::from(loss) : Converts the tensor loss to a f64 for printing.

That's it. Tch-rs provides Rust bindings for tensor operations, autograd, nn layers and optimizers. 

The full source with ops code and docs is available here. That's all for now, tune in for part two where we'll dive into production grade improvements and get it ready for deployment (K8s, IaC and GKE), track that code here.

 

Later

-    Ed