Pure JAX on G5g: Serving Gemma 4 on Graviton and a T4G
This article provides a step by step deployment guide for serving Google's Gemma 4 on an AWS EC2 G5g instance using pure JAX. The code is here: github.com/xbill9/gemma4-dev What is this project trying to Do? This project aims to serve a modern open model on the cheapest whole CUDA GPU AWS will rent you, and to measure honestly what that costs. Aren't You Using The Wrong GPU? Probably! The T4G is a Turing chip from 2018. It has no bfloat16 and no fp8. But it is cheap, it is available when nothing else is, and it is attached to a Graviton2 host — which makes G5g the rare hardware axis that almost nothing in the ML ecosystem targets: aarch64 and CUDA together . So let's give pure JAX a shot on G5g! AWS EC2 G5g G5g instances pair an AWS Graviton2 (64-bit Arm) processor with NVIDIA T4G Tensor Core GPUs. At g5g.xlarge they are the cheapest EC2 instance carrying a whole NVIDIA GPU , and the only Arm-based GPU family AWS offers. Two GPU instances are cheaper per hour and neither can serve this model (us-east-1, Linux, on-demand, checked against the Pricing API on 2026-08-28): g6f.large at $0.2020 is genuinely NVIDIA and genuinely CUDA — but it is one eighth of a GPU with 3 GB , and the weights alone are 6.155 GB. The first g6f that fits is g6f.4xlarge at $0.9500, which is 1.7x this rig's g5g.2xlarge . g4ad.xlarge at $0.3785 carries an AMD Radeon Pro V520 — no CUDA at any price. Among whole NVIDIA GPUs, G5g is the floor: g5g.xlarge at $0.4200, and the next one up is g4dn.xlarge at $0.5260. More information is available here: https://aws.amazon.com/ec2/instance-types/g5g/ The default in this rig is g5g.2xlarge — 1 GPU, 8 vCPU, 16 GiB RAM. Note- the T4G reports 15,360 MiB of device memory, not the nominal 16 GB. Budget against the measured number. Gemma 4 Gemma is Google's family of open models built from the same research as Gemini. This rig serves google/gemma-4-E2B-it , the instruction-tuned reference release. JAX JAX is Google's array computing library — NumPy semantics, c