Posts

Advanced Inference: Model portability across GPU backends (Rocm and Nvidia) Pt. 1 [PyTorch] | ML Engineering

Image
    Hi All I'm aiming to solve some tough issues in the ML backend, and heterogeneous backend support is up there as one of the most pressing . Today we'll be training BERT on a Rocm (AMD) backend and running the binary (inference) on a Nvidia backend.  Training a complex model like BERT (250M+ params) on a ROCm backend and then running the trained model on NVIDIA backend is not impossible, but it comes with some challenges and considerations. Here's a breakdown of the key aspects:  Steps to achieve Cross-Backend Training/Inference Option 1: Train on ROCm, Inference on NVIDIA (Same Framework) 1.     Train the model on AMD GPU's using PyTorch/Tensorflow with ROCm. # Example PyTorch training script (ROCm) import torch device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # ... rest of the training code *      Note: ROCm uses cuda device string for compatibility, but maps it to AMD hardware. 2.      Save t...

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

Image
 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 *  ...

Training TinyLlama 1.1B with PyTorch in Google Colab | ML Engineering

Image
Hi All Today we'll be loading a pre-trained model and train it on Colab. Here's a minimal example using Huggingface for transformers and accelerate for memory efficiency. View source below. Let's break down the above code. Block by block. I'm assuming you're using a notebook for this example. However in the full example in my GitHub I've included the Ops files too (Docker and k8s for GKE or AKS). I don't go into those here, I'll have a series of articles tacking MLOps in depth, this is about fine-tuning a pre-trained small LLM in Colab.  Prep: Setup Environment   !pip install -q transformers accelerate datasets peft bitsandbytes   *     transformers provides pre-trained models and training utilities *      peft enables LoRA (low-rank Adaptation), reducing memory usage *        bitsandbytes allows 8-bit precision to shrink model size *        accelerate helps manage GPU memory and multi-GPU setups ...

Advanced Rust ML: Custom Modules with Tch-rs | ML Engineering

Image
  Hi All Another day, another Rust ML backend application. Today we're looking at defining custom modules in Rust in the ML backend. We're going to be creating a custom linear layer in Rust using Tch-rs ( Rust PyTorch bindings ).   View full source below, also available here :   Let's break the above code down. Block by block. 1. Imports use tch::{nn, nn::Module, Tensor}; What it does:  *        tch::nn - Neural network module containing layer definitions. *      nn:Module - The trait that all neural network modules must implement. *      Tensor - The Tensor type used throughout tch-rs   2. Struct Definition   struct CustomLayer { weight: Tensor, bias: Tensor, } What it does: *      weight - A tensor holding a layer's weights (matrix) *      bias -  A tensor holding the layer's bias (vector) This is essentially a linear layer (fully connected laye...

Advanced Rust ML: Loading Pre-trained models with Tch-rs | ML Engineering

Image
  Hi All Today we're looking at some Rust ML with  Tch-rs . Now Tch-rs is a  Rust bindings library for the C++ api of PyTorch. "The goal of the tch crate is to provide some thin wrappers around the C++ PyTorch api (a.k.a. libtorch)" - github . Let's get to it then. Let's demonstrate tch-rs's core capability: creating tensors, moving them to devices, performing operations and getting results - all in Rust while using PyTorch's backend. See introductory example below.  use tch::{Device, Tensor, Kind}; fn main() { let device = Device::cuda_if_available(); let x = Tensor::of_slice(&[1.0, 2.0, 3.0]).to_device(device); let y = Tensor::of_slice(&[4.0, 5.0, 6.0]).to_device(device); let z = x + y; println!("{:?}", z); } Breakdown: *      Imports : Device , Tensor and Kind   from Tch-rs, the core types for device management and memory operations. *        Device Selection : Device::cuda_if_available(); a...

Rust experiments: Wrapping unsafe (malloc/free) in safe abstractions | GPU engineering

Image
    I'll have a full example soon on how to wrap unsafe code (malloc/free like in a C++ Cuda kernel), in safe abstractions using Rust. See snippet below: