From be9e436a3cd1757980893970f4bc14bfa28336e6 Mon Sep 17 00:00:00 2001 From: xbill Date: Sun, 12 Jul 2026 12:23:41 -0400 Subject: [PATCH] Add Gemma-4 E2B text inference sample (torch-neuronx) --- ...trained_gemma4_e2b_inference_on_inf2.ipynb | 410 ++++++++++++++++++ 1 file changed, 410 insertions(+) create mode 100644 torch-neuronx/inference/hf_pretrained_gemma4_e2b_inference_on_inf2.ipynb diff --git a/torch-neuronx/inference/hf_pretrained_gemma4_e2b_inference_on_inf2.ipynb b/torch-neuronx/inference/hf_pretrained_gemma4_e2b_inference_on_inf2.ipynb new file mode 100644 index 0000000..827d2c7 --- /dev/null +++ b/torch-neuronx/inference/hf_pretrained_gemma4_e2b_inference_on_inf2.ipynb @@ -0,0 +1,410 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "004214ba", + "metadata": {}, + "source": [ + "# Run Gemma-4 (E2B) text generation on AWS Inferentia2 with `torch-neuronx`\n", + "\n", + "This notebook compiles and runs [`google/gemma-4-E2B-it`](https://huggingface.co/google/gemma-4-E2B-it)\n", + "on a **single NeuronCore** of an Inferentia2 instance (`inf2.xlarge`), producing greedy output that is\n", + "**token-for-token identical to the CPU reference**.\n", + "\n", + "Gemma-4's smaller models use **Per-Layer Embeddings (PLE)** and a **per-layer KV-sharing** attention\n", + "pattern that don't trace through the vendor stack directly. This sample uses a small, explicit recipe:\n", + "\n", + "* **Two graphs** — a *prefill* graph (whole prompt) and a *single-token decode* graph — traced with\n", + " `torch_neuronx.trace`.\n", + "* **KV cache as graph I/O** — a fixed `MAX`-length KV buffer is passed in and out of the decode graph and\n", + " updated in place with a one-hot masked write, so the cache stays on-device across steps with no host\n", + " round-trip of the graph itself.\n", + "* **Host-side embeddings** — token embeddings and the per-layer inputs are computed on the host and passed\n", + " in as `inputs_embeds` / `per_layer_inputs`, keeping the traced graph free of the PLE lookup.\n", + "* **Eager attention + `tanh` GELU + logit softcap** — matched to the reference so device == CPU.\n", + "\n", + "**Requirements:** an `inf2.xlarge` (or larger) with the Neuron SDK runtime, and access to the gated Gemma-4\n", + "weights on Hugging Face." + ] + }, + { + "cell_type": "markdown", + "id": "7339cdd3", + "metadata": {}, + "source": [ + "## 1. Install dependencies\n", + "\n", + "Pinned to the versions this sample was validated against (Neuron SDK 2.23)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "693eb9c5", + "metadata": {}, + "outputs": [], + "source": [ + "# fmt: off\n", + "%pip install --quiet --extra-index-url https://pip.repos.neuron.amazonaws.com \\\n", + " \"torch-neuronx==2.8.*\" \"neuronx-cc==2.*\" \"transformers==5.13.0\" \"huggingface_hub>=0.34\"\n", + "# fmt: on" + ] + }, + { + "cell_type": "markdown", + "id": "fa7dec1a", + "metadata": {}, + "source": [ + "## 2. Authenticate to Hugging Face\n", + "\n", + "Gemma-4 is a **gated** model — accept the license on the\n", + "[model page](https://huggingface.co/google/gemma-4-E2B-it) first, then log in with a token that has access." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "0a3fd114", + "metadata": {}, + "outputs": [], + "source": [ + "from huggingface_hub import notebook_login\n", + "notebook_login()" + ] + }, + { + "cell_type": "markdown", + "id": "5b82ce11", + "metadata": {}, + "source": [ + "## 3. Configuration" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "66b3f18f", + "metadata": {}, + "outputs": [], + "source": [ + "import os, time, torch\n", + "torch.manual_seed(0)\n", + "\n", + "MODEL_ID = \"google/gemma-4-E2B-it\"\n", + "MAX = 128 # max total sequence length (KV buffer length)\n", + "BUCKET = 32 # prefill bucket: prompts are right-padded to this length\n", + "NEG = torch.finfo(torch.float32).min" + ] + }, + { + "cell_type": "markdown", + "id": "1a1615ee", + "metadata": {}, + "source": [ + "## 4. Load the model\n", + "\n", + "We load in **bf16** so the on-device weight constants fit a single 16 GB core. Gemma-4 uses a `tanh`\n", + "approximation of GELU and a final-logit softcap; we set the activation explicitly and apply the softcap\n", + "in the graph so the device output matches the reference exactly." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1f8afc03", + "metadata": { + "lines_to_next_cell": 1 + }, + "outputs": [], + "source": [ + "from transformers import AutoTokenizer, Gemma4ForConditionalGeneration, DynamicCache\n", + "\n", + "tok = AutoTokenizer.from_pretrained(MODEL_ID)\n", + "m = Gemma4ForConditionalGeneration.from_pretrained(\n", + " MODEL_ID, torch_dtype=torch.bfloat16, attn_implementation=\"eager\").eval()\n", + "lang = m.model.language_model\n", + "lm_head = m.lm_head\n", + "cfg = lang.config\n", + "SW = cfg.sliding_window\n", + "WDT = m.dtype # bf16 — host KV buffers / one-hot must match this\n", + "softcap = getattr(m.config.text_config, \"final_logit_softcapping\", None)\n", + "\n", + "class GeluTanh(torch.nn.Module):\n", + " def forward(self, x):\n", + " return 0.5 * x * (1.0 + torch.tanh(0.7978845608028654 * (x + 0.044715 * x * x * x)))\n", + "for mod in lang.modules():\n", + " if hasattr(mod, \"act_fn\"):\n", + " mod.act_fn = GeluTanh()\n", + "\n", + "# Discover the layers that own their K/V (non KV-shared) and their (n_kv_heads, head_dim).\n", + "NONSHARED, LINFO = [], {}\n", + "for i, lyr in enumerate(lang.layers[:cfg.num_hidden_layers]):\n", + " a = lyr.self_attn\n", + " if not a.is_kv_shared_layer:\n", + " hd = a.head_dim\n", + " NONSHARED.append(i)\n", + " LINFO[i] = (a.k_proj.out_features // hd, hd)\n", + "print(\"non-shared KV layers:\", NONSHARED)\n", + "\n", + "def softcap_logits(lg):\n", + " return softcap * torch.tanh(lg / softcap) if softcap else lg" + ] + }, + { + "cell_type": "markdown", + "id": "b750ecd8", + "metadata": {}, + "source": [ + "## 5. Define the prefill and decode graphs\n", + "\n", + "The **prefill** graph runs the whole (padded) prompt and returns the logits plus the K/V for every\n", + "non-shared layer. The **decode** graph runs a single new token against a fixed `MAX`-length KV buffer,\n", + "writing the new K/V into the buffer with a one-hot mask and returning the updated buffers." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "dce6bdd1", + "metadata": { + "lines_to_next_cell": 1 + }, + "outputs": [], + "source": [ + "class PreWrap(torch.nn.Module):\n", + " def __init__(s): super().__init__(); s.lang = lang; s.head = lm_head\n", + " def forward(s, ie, am, ple):\n", + " cache = DynamicCache()\n", + " out = s.lang(inputs_embeds=ie, per_layer_inputs=ple, attention_mask=am,\n", + " use_cache=True, past_key_values=cache)\n", + " lg = softcap_logits(s.head(out.last_hidden_state))\n", + " ks = [cache.layers[i].keys for i in NONSHARED]\n", + " vs = [cache.layers[i].values for i in NONSHARED]\n", + " return (lg, ks, vs)\n", + "\n", + "class StaticKV:\n", + " \"\"\"A KV cache backed by fixed MAX-length buffers; update() writes the new token via a one-hot mask.\"\"\"\n", + " is_compileable = False\n", + " def __init__(s, key_bufs, val_bufs, onehot):\n", + " s.key = {i: key_bufs[j] for j, i in enumerate(NONSHARED)}\n", + " s.val = {i: val_bufs[j] for j, i in enumerate(NONSHARED)}\n", + " s.oh = onehot # [1,1,MAX,1] float\n", + " def update(s, k, v, idx, *a, **kw): # k,v: [1,nkv,1,hd]\n", + " s.key[idx] = s.key[idx] * (1.0 - s.oh) + k * s.oh\n", + " s.val[idx] = s.val[idx] * (1.0 - s.oh) + v * s.oh\n", + " return s.key[idx], s.val[idx]\n", + " def get_seq_length(s, *a, **k): return 0 # position_ids are supplied explicitly\n", + " def export(s):\n", + " return [s.key[i] for i in NONSHARED], [s.val[i] for i in NONSHARED]\n", + "\n", + "class DecWrap(torch.nn.Module):\n", + " def __init__(s): super().__init__(); s.lang = lang; s.head = lm_head\n", + " def forward(s, ie, ple, position_ids, onehot, full_mask, slide_mask, key_bufs, val_bufs):\n", + " cache = StaticKV(key_bufs, val_bufs, onehot)\n", + " masks = {\"full_attention\": full_mask, \"sliding_attention\": slide_mask}\n", + " out = s.lang(inputs_embeds=ie, per_layer_inputs=ple, position_ids=position_ids,\n", + " attention_mask=masks, use_cache=True, past_key_values=cache)\n", + " lg = softcap_logits(s.head(out.last_hidden_state))\n", + " ks, vs = cache.export()\n", + " return (lg, ks, vs)\n", + "\n", + "pre = PreWrap().eval()\n", + "dec = DecWrap().eval()\n", + "\n", + "# Host helpers: embeddings (incl. per-layer inputs) and all position-dependent tensors.\n", + "def embed_ids(id_list):\n", + " ids = torch.tensor([id_list])\n", + " with torch.no_grad():\n", + " ie = lang.embed_tokens(ids)\n", + " ple = lang.get_per_layer_inputs(ids, ie)\n", + " return ie, ple\n", + "\n", + "def host_pos_tensors(pos):\n", + " ar = torch.arange(MAX)\n", + " onehot = (ar == pos).view(1, 1, MAX, 1).to(WDT)\n", + " valid = ar <= pos\n", + " full = torch.where(valid, 0.0, NEG).view(1, 1, 1, MAX)\n", + " slide = torch.where(valid & (ar > pos - SW), 0.0, NEG).view(1, 1, 1, MAX)\n", + " return torch.tensor([[pos]], dtype=torch.long), onehot, full, slide\n", + "\n", + "ec = m.generation_config.eos_token_id\n", + "EOS = set(ec) if isinstance(ec, (list, tuple)) else {ec}" + ] + }, + { + "cell_type": "markdown", + "id": "ad8d955a", + "metadata": {}, + "source": [ + "## 6. Greedy generation loop (works with either the eager modules or the traced graphs)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "aa9a217b", + "metadata": {}, + "outputs": [], + "source": [ + "def run_greedy(pre_fn, dec_fn, prompt, maxnew=30):\n", + " n0 = len(prompt)\n", + " pad = prompt + [0] * (BUCKET - n0)\n", + " ie, ple = embed_ids(pad)\n", + " am = torch.tensor([[1] * n0 + [0] * (BUCKET - n0)])\n", + " with torch.no_grad():\n", + " lg, ks, vs = pre_fn(ie, am, ple)\n", + " first = int(lg[0, n0 - 1].argmax())\n", + "\n", + " key_bufs = [torch.zeros(1, LINFO[i][0], MAX, LINFO[i][1], dtype=WDT) for i in NONSHARED]\n", + " val_bufs = [torch.zeros(1, LINFO[i][0], MAX, LINFO[i][1], dtype=WDT) for i in NONSHARED]\n", + " for j in range(len(NONSHARED)):\n", + " key_bufs[j][:, :, :n0, :] = ks[j][:, :, :n0, :]\n", + " val_bufs[j][:, :, :n0, :] = vs[j][:, :, :n0, :]\n", + "\n", + " seq, cur = [first], n0\n", + " for _ in range(maxnew):\n", + " if seq[-1] in EOS:\n", + " break\n", + " ie1, ple1 = embed_ids([seq[-1]])\n", + " position_ids, onehot, full_mask, slide_mask = host_pos_tensors(cur)\n", + " with torch.no_grad():\n", + " lg1, key_bufs, val_bufs = dec_fn(ie1, ple1, position_ids, onehot,\n", + " full_mask, slide_mask, key_bufs, val_bufs)\n", + " seq.append(int(lg1[0, 0].argmax()))\n", + " cur += 1\n", + " return seq\n", + "\n", + "msgs = [{\"role\": \"user\", \"content\": \"What is the capital of France?\"}]\n", + "enc = tok.apply_chat_template(msgs, add_generation_prompt=True, return_tensors=\"pt\", return_dict=True)\n", + "prompt = enc[\"input_ids\"][0].tolist()\n", + "assert len(prompt) <= BUCKET, \"prompt longer than BUCKET; increase BUCKET (and re-trace)\"" + ] + }, + { + "cell_type": "markdown", + "id": "ca2d370e", + "metadata": {}, + "source": [ + "## 7. Compile the two graphs for Inferentia\n", + "\n", + "`torch_neuronx.trace` compiles each graph to a NEFF. Compilation takes a few minutes the first time." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "366bd1db", + "metadata": {}, + "outputs": [], + "source": [ + "import torch_neuronx\n", + "\n", + "ie, ple = embed_ids(prompt + [0] * (BUCKET - len(prompt)))\n", + "am = torch.tensor([[1] * len(prompt) + [0] * (BUCKET - len(prompt))])\n", + "CARGS = [\"--model-type\", \"transformer\", \"--auto-cast\", \"all\", \"--auto-cast-type\", \"bf16\"]\n", + "\n", + "t = time.time()\n", + "pre_neff = torch_neuronx.trace(pre, (ie, am, ple), compiler_args=CARGS)\n", + "print(f\"prefill compiled in {time.time()-t:.0f}s\")\n", + "\n", + "key_bufs = [torch.zeros(1, LINFO[i][0], MAX, LINFO[i][1], dtype=WDT) for i in NONSHARED]\n", + "val_bufs = [torch.zeros(1, LINFO[i][0], MAX, LINFO[i][1], dtype=WDT) for i in NONSHARED]\n", + "ie1, ple1 = embed_ids([prompt[-1]])\n", + "position_ids, onehot, full_mask, slide_mask = host_pos_tensors(len(prompt))\n", + "t = time.time()\n", + "dec_neff = torch_neuronx.trace(\n", + " dec, (ie1, ple1, position_ids, onehot, full_mask, slide_mask, key_bufs, val_bufs),\n", + " compiler_args=CARGS)\n", + "print(f\"decode compiled in {time.time()-t:.0f}s\")\n", + "\n", + "# Optionally persist the NEFFs so you can reload without recompiling:\n", + "# torch.jit.save(pre_neff, \"gemma4_e2b_prefill.pt\")\n", + "# torch.jit.save(dec_neff, \"gemma4_e2b_decode.pt\")" + ] + }, + { + "cell_type": "markdown", + "id": "fb925906", + "metadata": {}, + "source": [ + "## 8. Validate: device output == CPU reference" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "77627c9f", + "metadata": { + "lines_to_next_cell": 1 + }, + "outputs": [], + "source": [ + "cpu_seq = run_greedy(pre, dec, prompt)\n", + "dev_seq = run_greedy(pre_neff, dec_neff, prompt)\n", + "\n", + "print(\"CPU: \", repr(tok.decode([s for s in cpu_seq if s not in EOS], skip_special_tokens=True)))\n", + "print(\"Device:\", repr(tok.decode([s for s in dev_seq if s not in EOS], skip_special_tokens=True)))\n", + "print(\"SEQ_MATCH:\", cpu_seq == dev_seq)\n", + "assert cpu_seq == dev_seq, \"device output diverged from CPU reference\"" + ] + }, + { + "cell_type": "markdown", + "id": "bd2ea1d4", + "metadata": {}, + "source": [ + "## 9. Generate and measure throughput" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1f7c126c", + "metadata": {}, + "outputs": [], + "source": [ + "def generate(question, maxnew=80):\n", + " enc = tok.apply_chat_template([{\"role\": \"user\", \"content\": question}],\n", + " add_generation_prompt=True, return_tensors=\"pt\", return_dict=True)\n", + " p = enc[\"input_ids\"][0].tolist()\n", + " t = time.time()\n", + " seq = run_greedy(pre_neff, dec_neff, p, maxnew=maxnew)\n", + " dt = time.time() - t\n", + " print(tok.decode([s for s in seq if s not in EOS], skip_special_tokens=True))\n", + " print(f\"\\n[{len(seq)} tokens in {dt:.1f}s = {len(seq)/dt:.1f} tok/s]\")\n", + "\n", + "generate(\"In one sentence, what is AWS Inferentia?\")" + ] + }, + { + "cell_type": "markdown", + "id": "c18f8b59", + "metadata": {}, + "source": [ + "## Notes / troubleshooting (Gemma-4 specifics)\n", + "\n", + "* **Per-Layer Embeddings (PLE).** The smaller Gemma-4 models add per-layer inputs. We compute\n", + " `embed_tokens` + `get_per_layer_inputs` on the host and pass them in, so the traced graph never does the\n", + " PLE lookup. (The larger encoder-free `gemma4_unified` models have no PLE.)\n", + "* **KV-sharing as graph I/O.** Passing the fixed `MAX`-length KV buffers in and out of the decode graph is\n", + " what lets the per-layer KV-sharing pattern trace cleanly — the sharing shows up as ordinary graph data\n", + " dependencies. The one-hot masked write (`buf*(1-oh) + kv*oh`) avoids any dynamic index op in the graph.\n", + "* **bf16 is required for a single core.** In fp32 the on-device weight constants (~15 GB for E2B's larger\n", + " sibling) exceed a 16 GB core; bf16 halves them. `--auto-cast bf16` matches the host bf16 embeddings.\n", + "* **Eager attention + `tanh` GELU + softcap.** Set `attn_implementation=\"eager\"`, replace `act_fn` with the\n", + " `tanh` GELU, and apply the final-logit softcap in the graph — all three are needed for exact CPU parity.\n", + "* **Bigger models need tensor parallelism.** E2B fits one core. The 4B/12B variants exceed 16 GB even in\n", + " bf16 and require a TP=2 build across both NeuronCores of an `inf2.8xlarge`." + ] + } + ], + "metadata": { + "jupytext": { + "cell_metadata_filter": "-all", + "main_language": "python", + "notebook_metadata_filter": "-all" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +}