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

 

 

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 the model (eg. model.save_pretrained("bert_model")) for HuggingFace or torch.save(model.state_dict(), "model.pt")).

3.    Load and run the model on NVIDIA:


# Load the model on NVIDIA
model = torch.load("model.pt", map_location="cuda")  # or "cpu" if no GPU
model.eval()
# ... run inference

Option 2: Convert to ONNX for Cross-Backend Inference

1.    Train and save the model in PyTorch/TensorFlow.

2.    Convert to ONNX.


# PyTorch to ONNX
torch.onnx.export(model, dummy_input, "model.onnx", input_names=["input"], output_names=["output"])

3.    Run ONNX model on NVIDIA with ONNX Runtime:


import onnxruntime as ort
sess = ort.InferenceSession("model.onnx")
outputs = sess.run(None, {"input": input_data})

 Option 3: Use HuggingFace Transformers (Framework Agnostic)

*    HuggingFace's transformers library abstracts backend differences. Train on ROCm, then load the model on NVIDIA:


from transformers import BertForMaskedLM
model = BertForMaskedLM.from_pretrained("bert_model")  # Load saved model
model.to("cuda")  # Works on NVIDIA

When it might fail 

*    Custom CUDA kernels: If your model uses custom CUDA kernels(eg. for attention), these won't work on ROCm unless you rewrite them for ROCm's HIP.

*    Vendor Specific libraries: If you rely on NVIDIA libraries (eg. cuDNN, TensorRT) during training, you'll need AMD equivalents (eg. rocBLAS, MIOpen).

*    Framework Bugs: ROCm support in PyTorch/TensorFlow is improving but may have edge-case bugs. Test thoroughly.

Performance Considerations

*     ROCm's performance on AMD GPUs may not match NVIDIA's CUDA for all operations (eg. mixed precision, sparse tensors).

*    Benchmark both backends to ensure the model meets your latency/thoughput requirements.

 

Conclusions

Yes you can train a BERT-like model on ROCm (AMD) and run it on NVIDIA GPU's, provided:

1.    You use a framework like PyTorch or TensorFlow that abstracts the backend.

2.    The model is saved in a framework compatible format (eg. .pt, .bin, ONNX).

3.    You avoid hardware-specific optimizations that lock you into one backend.

For most use cases, this works seamlessly. If you hit issues, converting to OONX or using HuggingFace's transformers library can simplify cross-backend deployment. Full source available here. Catch part 2 in this series, where we try model portability with JAX. 

Later

-    Ed

 

 

Comments

Post a Comment

Popular posts from this blog

Fine-Tuning Mistral 7B using QLoRA with PyTorch pt. 1: The Model | ML Engineering

Image
     Hi All Today we're working with a popular and slightly bigger model than our previous example. Mistral 7B is capable of chat and light coding tasks, for older hardware it's a winner for sure.  Here's a complete, runnable example of fine-tuning Mistral 7B using QLoRA with the peft , transformers , and bitsandbytes libraries. This example assumes you're working with a single GPU (eg. an A100 or similar). First install the required packages: pip install -q bitsandbytes datasets accelerate peft transformers trl View full script below, also available here :   Full breakdown of the script above, block-by-block. 1.      Dataset Loading dataset = load_dataset("timdettmers/openassistant-guanaco", split="train") *      Loads a preprocessed instruction-following dataset (Guanco, derived from OpenAssistant). *      split="train" selects the training portion *      The dataset is in a conversational ...
Read more

Deploying LoRA Optimised BERT as a FastApi service on GKE | ML Engineering & MLOps

Image
  Hi All  Alright today on "Bored MLE" we're doing some inference, using a LoRA optimized BERT model on GKE . Now there's a lot more to MLOps than just this (like cluster level observability for example), but for this example I'm keeping it brief.  The full code is available on GitHub with added minikube deployment instructions (I only cover the Cloud deployment here).    I assume familiarity with Python, PyTorch, K8s, FastApi, Minikube and GKE. Let's get on with it.  View the full FastApi service source below, also available here : Let's breakdown the above code, block by block. 0. Install Dependencies :  pip install fastapi uvicorn torch transformers peft accelerate  1. Imports from fastapi import FastAPI, HTTPException from pydantic import BaseModel from typing import List import torch from transformers import AutoModelForSequenceClassification, AutoTokenizer from peft import PeftModel, PeftConfig  *      fastapi : F...
Read more

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...
Read more