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 t...