{"slug": "porting-gemma-4-2b-4b-12b-to-aws-inferentia2", "title": "Porting Gemma-4 (2B / 4B / 12B) to AWS Inferentia2", "summary": "A developer ported Google's Gemma-4 family (2B, 4B, 12B) to AWS Inferentia2, overcoming three architectural obstacles that break the vendor stack. The port achieved 44 tok/s for the 2B model, 33-39 tok/s for the 4B model, and 15 tok/s for the 12B model with token-for-token identical output to CPU. The key challenges involved mixed attention heads, cross-layer KV-sharing, and GQA sharding, which were resolved by tracing the Hugging Face forward pass directly.", "body_md": "*A field report on running Google's Gemma-4 family on AWS Inferentia2 ( inf2), covering the\nthree architectural obstacles that break the vendor stack — **mixed attention heads*\n\n`optimum-neuron`\n\n/ NxD`neuronx-cc`\n\n) limits —| Models |\n`google/gemma-4-E2B-it` , `google/gemma-4-E4B-it` , `google/gemma-4-12B-it`\n|\n| Hardware | AWS Inferentia2 — `inf2.xlarge` (1 chip / 2 cores / 32 GB HBM) and `inf2.8xlarge` (same accelerator, 128 GB host RAM) |\n| Software | Neuron SDK 2.23 · `torch-neuronx` 2.8.0 · `neuronx-cc` 2.23.6484 · `neuronx-distributed` 0.17 · `transformers` 5.13.0 |\n| Result | E2B ~44 tok/s (1 core), E4B ~33–39 tok/s (TP=2), 12B ~15 tok/s (TP=2) — greedy decode token-for-token identical to the CPU reference on all three |\n| Artifacts | HF: `xbill9/gemma-4-{E2B,E4B,12B}-it-inferentia2` · Docker Hub: `xbill9/gemma4-optb{,-e4b,-12b}`\n|\n\nGemma-4 is not a vanilla decoder. Across the family it combines several features that each map\n\ncleanly onto TPU/XLA (where the model was designed) but individually break the AWS inference path:\n\n`tanh`\n\ncap at 30) over a The AWS vendor stack (`optimum-neuron`\n\n+ `neuronx-distributed`\n\n+ the Neuron vLLM backend) has **no\nGemma-4 model class at all**, and its graph builder cannot express KV-sharing. The public Neuron\n\nE2B |\nE4B |\n12B |\n|\n|---|---|---|---|\n| HF class |\n`Gemma4ForConditionalGeneration` (text) |\nsame | `Gemma4UnifiedForConditionalGeneration` |\n| model_type | `gemma4_text` |\n`gemma4_text` |\n`gemma4_unified` (encoder-free multimodal) |\n| Params (effective) | ~5B (2B) | ~8B (4B) | 12B |\n| Per-Layer Embeddings | yes |\nyes |\nno (`hidden_size_per_layer_input=0` ) |\n| Cross-layer KV-sharing |\nyes (15 non-shared layers own KV) |\nyes |\nno (`num_kv_shared_layers=0` , every layer owns KV) |\n| Query / KV heads | GQA | 8 q / 2 kv |\n16 q / 8 kv, global layers nkv=1\n|\n| Attention | sliding + global, `sliding_window=512`\n|\nsliding + global, `sliding_window=512`\n|\nsliding + global, `sliding_window=1024` , `attention_k_eq_v=true`\n|\n| head_dim | — | 256 | 256 (global 512) |\n| Logit softcap | 30 | 30 | 30 |\n| Vocab | 262 144 | 262 144 | 262 144 (tied embeddings) |\n| Fits one 16 GB core? |\nyes (bf16) |\nno → TP=2 |\nno → TP=2 |\n\nThree model sizes, three *different* reasons the vendor path fails, and — as it turns out — three\n\ndistinct meanings of \"mixed heads.\"\n\n\"Mixed heads\" was the single most expensive class of bug in this port. It shows up in three\n\ncompletely different forms as you move up the size ladder.\n\nOn E2B/E4B a layer's attention is flagged `self_attn.is_kv_shared_layer`\n\n. Shared layers do **not**\n\nrun their own `k_proj`\n\n/`v_proj`\n\n; they attend against a KV tensor produced by an earlier \"owner\"\n\nlayer. On TPU this is a free graph view. In AWS's `neuronx-distributed`\n\n(NxD) `ModelBuilder`\n\n, which\n\nwants a static, per-layer weight→buffer mapping, it **cannot be represented** — there is no place to\n\nsay \"this layer's K/V is that layer's K/V, computed once, live.\"\n\nThe fix (**\"Option B\"**, see §3) is to stop fighting the framework and trace the Hugging Face\n\nforward pass directly. KV-sharing then falls out as an ordinary live data dependency in the traced\n\ngraph. The port only has to enumerate which layers actually own a cache:\n\n```\n# discover the layers that write KV; shared layers never touch the cache\nNONSHARED = []\nfor i, lyr in enumerate(lang.layers[:cfg.num_hidden_layers]):\n    a = lyr.self_attn\n    if not a.is_kv_shared_layer:                 # <- the crux\n        NONSHARED.append(i)\n        LINFO[i] = (a.k_proj.out_features // a.head_dim, a.head_dim)   # (n_kv, head_dim)\n```\n\nOnly the `NONSHARED`\n\nlayers (15 of them on E2B) allocate a KV buffer; the shared layers read\n\nthrough. This is why the KV cache is far smaller than a naïve \"one buffer per layer\" implementation\n\nwould allocate — and why it fits a single core in bf16.\n\nE4B does not fit one 16 GB core (§4.3), so it must be sharded across both NeuronCores (TP=2). GQA\n\nnow bites: E4B has **8 query heads and 2 KV heads**. Column-sharding the query projection across 2\n\nranks is trivial (4 q-heads/rank). The KV projection has only **2** heads, and the subtle failure is\n\nthat `repeat_kv`\n\ninside attention broadcasts each KV head to a *group* of query heads — so you cannot\n\nindependently halve the KV heads and the query heads without scrambling **which query attends which\nKV**.\n\nThe rule that makes it correct: shard KV to `nkv // TP`\n\nheads per rank **only when it divides**, and\n\ncrucially **keep num_key_value_groups unchanged** so\n\n`repeat_kv`\n\nstill maps each rank's 4 query\n\n``` python\ndef _kv_rank_width(nkv):\n    return nkv // TP if nkv % TP == 0 else nkv     # 2 // 2 = 1 KV head / rank on E4B\n\n# ... per layer:\nnkv = a.k_proj.out_features // a.head_dim\nif nkv % TP == 0:                                  # divisible: shard k/v, keep groups\n    a.k_proj = col(a.k_proj)\n    a.v_proj = col(a.v_proj)\n# else: see 2.3\n```\n\nNaïvely also dividing `num_key_value_groups`\n\nproduces the classic \"4-vs-8 `repeat_kv`\n\nmismatch\":\n\nplausible-looking but wrong output. The KV cache is then kept **device-resident and sliced per rank**\n\n(head *r* → rank *r*), because NxD's `forward()`\n\nonly returns rank 0's KV (§5.2).\n\nThe 12B is where \"mixed heads\" becomes structural. `gemma4_unified`\n\ninterleaves two attention types\n\nwith **different KV-head counts**:\n\n`nkv = 8`\n\n→ divisible by TP=2 → shard to 4 heads/rank (the §2.2 path).`num_global_key_value_heads = 1`\n\n→ You cannot column-shard a single KV head across two ranks. The working answer is to **leave the\nglobal layers' K/V replicated** (a plain\n\n`nn.Linear`\n\n, which NxD loads identically on every rank) and`repeat_kv`\n\nmatches the \n\n```\na = lyr.self_attn\nhd = a.head_dim\nnq = a.q_proj.out_features // hd            # capture BEFORE replacing q_proj (see gotcha)\na.q_proj = col(a.q_proj); a.o_proj = row(a.o_proj)\nnkv = a.k_proj.out_features // hd\nif nkv % TP == 0:                            # sliding layers: nkv=8 -> 4/rank, keep groups\n    a.k_proj = col(a.k_proj); a.v_proj = col(a.v_proj)\nelse:                                        # global layers: nkv=1, indivisible\n    # leave k/v REPLICATED; shrink groups to the per-rank sharded q-head count\n    a.num_key_value_groups = (nq // TP) // nkv\n```\n\nTwo gotchas made this a multi-hour fight (commit `d70dc94`\n\n):\n\n`nq = q_proj.out_features // head_dim`\n\n`q_proj`\n\nwith a `ColumnParallelLinear`\n\n, whose `.out_features`\n\nreports the `AttributeError`\n\n/wrong groups.`else`\n\nbranch had literally never been exercised by E4B (all E4B layers are\n`nkv=2`\n\n, divisible), so it was dead, untested code until the 12B walked into it.**Takeaway:** \"mixed heads\" is not one problem. It is (a) mixed KV *sharing* across layers, (b) mixed\n\nquery/KV *counts* within a layer (GQA), and (c) mixed attention *types* with different KV counts per\n\ntype. Each needs a different sharding rule, and the naïve \"just divide everything by TP\" is wrong in\n\nall three.\n\n`optimum-neuron`\n\n/ NxD couldn't do it\nThe port deliberately does **not** use the AWS vendor inference path. Here is why each vendor layer\n\nfailed, and what replaced it.\n\n** optimum-neuron has no Gemma-4.** There is simply no\n\n`gemma4_text`\n\n/ `gemma4_unified`\n\nmodel class`optimum-neuron`\n\n. The Neuron vLLM backend dispatches through `optimum-neuron`\n\n, so vLLM has nothing**NxD ModelBuilder cannot represent KV-sharing.** NxD wants a static graph where each layer maps\n\n**The vendor endpoint served gibberish.** The public Neuron vLLM/NxD deployment we began with came up\n\n\"healthy\" and produced fluent-looking but semantically empty text — the worst failure mode, because\n\nit *looks* like it works.\n\n**An automated port falsely reported success.** A framework auto-port pass claimed **\"100% PASS\"**.\n\nOn inspection its golden reference had been built from a **PLE-stripped** checkpoint, so both sides of\n\nthe comparison were wrong in the same way. A green check on a broken oracle is worse than a red one.\n\n**The replacement — \"Option B.\"** Instead of teaching NxD about Gemma-4, `torch_neuronx.trace()`\n\n(for\n\nsingle-core E2B) and `neuronx_distributed.trace.parallel_model_trace`\n\n/ `ModelBuilder`\n\n(for TP E4B/12B)\n\nare pointed **straight at the Hugging Face transformers 5.13 text forward pass**. KV-sharing,\n\nThere is still a TP-launch wrinkle: `parallel_model_trace`\n\nspawns rank workers, and\n\n`transformers`\n\n5.13's FX utilities plus the spawn need a `transformers.utils.fx`\n\nshim and a\n\n`_TP_CHILD`\n\nenvironment sentinel to guard against infinite re-spawn.\n\n`neuronx-cc`\n\n) wall\nOnce the graph is traceable, `neuronx-cc`\n\nimposes its own hard limits. The recurring one is **SBUF**,\n\nthe on-chip state buffer (SRAM), capped at **196 608 B per partition**. Several Gemma-4 ops blow\n\nstraight through it, each with a distinct fix.\n\nThe Per-Layer Embedding table is `262144 × hidden`\n\n. Materialised on device it trips the compiler and\n\nby itself over-commits the 16 GB core. **Fix:** keep the PLE (and word) embedding **on the host**,\n\ngather on CPU, and feed the result in **as activations**. The device graph never sees the table.\n\nGemma-4 applies `softcap * tanh(logits / softcap)`\n\nover the full 262 144-token vocabulary. In fp32\n\nthat `tanh`\n\nis a custom-call whose working set is `128 × 524288`\n\nbytes — **524 288 > 196 608 B/partition**:\n\n```\n[NCC_INLA001] Allocated memory out of bound (128x524288) 524288 vs 196608 B/partition\n```\n\n**Fix (commit 328ca88):**\n\n`argmax(softcap(x)) == argmax(x)`\n\n— greedy decode is completely unaffected. The device returns `temperature > 0`\n\n). Correct and free.`tanh`\n\ncap does fit and is kept in-graph.)`sliding_window=1024`\n\n(12B)\nDropping softcap was necessary but not sufficient — the real overflow on the 12B was the **fused\nattention custom-call**, because the 12B's\n\n`sliding_window = 1024`\n\nis `a0a2a75`\n\n):`cfg._attn_implementation = \"eager\"`\n\non both the config and its `text_config`\n\n). Eager tiles attention`sliding_window=512`\n\nthe fused pathA single Option-B neff for E4B materialises **~15.4 GB of fp32 model constants**; a second neff\n\n(prefill + decode) on the same core tips past the 16 GB NeuronCore budget →\n\n`NRT_RESOURCE / status=4 Allocation Failure`\n\n. Two things are needed:\n\n`MB_WDTYPE=bf16`\n\n): NxD's `shard_children`\n\ncasts the checkpoint to the layer dtype,\nhalving the neff. E2B is the exception: in bf16 it fits one core, which is why it stays single-core at ~44 tok/s.\n\n`neuronx-cc`\n\nis happiest with plain arithmetic. Several things are written out by hand rather than\n\nleft to library helpers or dynamic control flow:\n\n`0.5*x*(1+tanh(0.7978845608*(x + 0.044715*x^3)))`\n\n.`DynamicCache`\n\nat trace time`buf*(1-oh) + k*oh`\n\n— pure arithmetic, trace-safe, no scatter op or\ndata-dependent indexing.`layer_scalar`\n\nis a Each Gemma-4 layer does `hidden_states *= self.layer_scalar`\n\n. `layer_scalar`\n\nis a registered\n\n**buffer** (default `1.0`\n\n, real value ≈ 0.06). NxD's weight-sharding (`shard_children`\n\n/\n\n`get_sharded_checkpoint`\n\n) loads **parameters only, never buffers** — so without intervention every\n\nlayer over-scales ~16×, compounding across all layers into `cos ≈ 0`\n\ngarbage. **Fix (commit\ne785f6d):** read the per-layer\n\n`layer_scalar`\n\nfrom the checkpoint and copy it into the module by`neuronx-cc`\n\nruns on **CPU** — no NeuronCore needed to compile, so you can build neffs on any box.\n\nBut a TP=2 compile runs `neuronx-cc`\n\nfor **both ranks concurrently** and peaks **past 128 GB host RAM**;\n\nadd ≥55 GB of swap before compiling (the resulting neff runs fine without swap).\n\nCorrectness and SBUF are only half the story; throughput comes from keeping the KV cache **on the\ndevice** and never round-tripping it through the host. Three designs evolved:\n\nTwo neffs share one static KV buffer: **prefill** (padded prompt ≤ `KV_BUCKET`\n\n, returns the 15\n\nnon-shared layers' K/V) and **decode** (single-token forward against a fixed `KV_MAX`\n\nbuffer, one-hot\n\nmasked write). KV tensors are graph inputs/outputs. ~44 tok/s, ~0.06 s prefill after a one-time\n\n~100 s neff load.\n\nUnder TP a neff spans **both** cores, so you cannot park prefill on core 0 and decode on core 1. Two\n\nsub-designs handle prefill:\n\n`tp_alias`\n\n`nn.Parameter`\n\ns aliased as graph I/O`input_output_aliases`\n\n), updated in place each step, so the cache never leaves the cores. Each\nrank's KV is seeded with head `ModelBuilder`\n\n)\n\n```\n# KV parameters aliased as graph I/O so the cache is never round-tripped through the host\naliases = {}\nfor j in range(NK): aliases[w.kbuf[j]] = 1 + j\nfor j in range(NK): aliases[w.vbuf[j]] = 1 + NK + j\n```\n\nPer-token cost is essentially **flat across context length** because the cache is device-resident.\n\nThe traced executor is saved with `torch.jit.save(model, path)`\n\nand reloaded with\n\n`torch.jit.load`\n\n+ `nxd_model.initialize_with_saved_weights(torch.tensor(start_rank))`\n\n. (The executor's\n\nown `.save`\n\nis TorchScript's, not NxD's — a subtle trap.)\n\nAll three ports match the CPU float reference **token-for-token** on greedy decode\n\n(`SEQ_MATCH True`\n\n), e.g. *\"The capital of France is **Paris**.\"*\n\n| Model | Build | Hardware | TP | Context | First token | Decode | Host RAM |\n|---|---|---|---|---|---|---|---|\nE2B |\ntwo-graph static KV |\n`inf2.8xlarge` / `inf2.xlarge` (slim) |\n1 core | 512 / 128 | ~0.06 s | ~44 tok/s |\n~6 GB (slim) |\nE4B |\n`tp_alias` |\n`inf2.8xlarge` |\n2 | 512 / 128 · 2048 / 512 | ~1.4–1.6 s | ~33 tok/s | — |\nE4B |\ndevice-prefill (bf16) |\n`inf2.8xlarge` / `inf2.xlarge` (slim) |\n2 | 512 / 128 | ~0.11–0.16 s |\n~36–39 tok/s |\n~8 GB (slim) |\n12B |\ndevice-prefill (bf16) |\n`inf2.8xlarge` / `inf2.xlarge` (slim) |\n2 | 256 / 64 | ~0.1 s |\n~15 tok/s |\n~8 GB (slim) |\n\nThe 12B's lower decode rate is inherent (dense 12B, all weights active per token), not a port defect;\n\nper-token cost matches the 8xlarge full server. Notably, **all three run on a single inf2.xlarge**\n\n`tokenizer.json`\n\nmasquerades as a device bug.`hf download`\n\nonce left\n`tokenizer.json`\n\nabsent; `GemmaTokenizer`\n\nloaded with no vocab and mapped `<unk>`\n\n. The model then faithfully emitted garbage (unused high-id tokens), which looked exactly like\na device/reload/precision failure and consumed hours. The decisive diagnostic was a `tok(\"hello world\").input_ids`\n\nbefore\nsuspecting the compiler.`inf2.xlarge`\n\n(16 GB) needs swap.`inf2.xlarge`\n\n.`inf2.8xlarge`\n\n(700 GB root); the push only uploads the tiny server layer since the base is already on Hub.The three obstacles above were the entry point. The port also produced a set of findings that\n\ngeneralize beyond Gemma-4 and were, in aggregate, worth more than any single fix.\n\nEvery time this port produced garbage, the NeuronCore was **innocent**. The causes, in order of how\n\noften they bit: a broken/missing **tokenizer** (§7), a mis-restored **weight reload** (§8.2), a wrong\n\n**buffer** (`layer_scalar`\n\n, §4.6), and a mismatched **driver/SDK** version (a precompiled neff can\n\n*mis-execute into garbage rather than error* on the wrong runtime). All four are **cheaper to rule out\nthan a device bug**, and all four\n\nThe fastest oracle turned out to be a **CPU-reference run on the same box**: load the model in bf16\n\nwith `from_pretrained`\n\n(≈5 s off the page cache once the weights are downloaded) and run one forward.\n\nIf the CPU reference produces the *same* garbage as the device, the accelerator is exonerated and the\n\nbug is upstream (tokenizer, input, weights). This single technique collapsed a multi-hour \"device\n\nbug\" into a two-minute tokenizer fix. **Reach for the CPU oracle before profiling the neff.**\n\nThe most dangerous illusion in the whole project was a green `SEQ_MATCH`\n\n. Correctness was validated\n\n**in-process** on the freshly traced model (`ModelBuilder.trace(initialize_model_weights=True)`\n\n) — but\n\nthe server loads a **saved** model in a **fresh process** and calls\n\n`initialize_with_saved_weights()`\n\n. Those are different code paths, and the in-process pass never\n\nexercised the one the server actually uses. (`torch.jit.load`\n\n*without* the init call fails outright —\n\n\"This model is not initialized\" — so the init step is load-bearing, not cosmetic.)\n\nCombined with the auto-port's **false \"100% PASS\"** against a PLE-stripped golden (§3), the lesson is\n\nblunt: **a passing test against the wrong oracle, or on the wrong execution path, is worse than a\nfailing one.** Validate (a) the exact artifact you will ship, (b) in the exact way you will load it,\n\nE2B is marketed as \"2B\" and E4B as \"4B\" — the MatFormer/PLE *effective* parameter counts. But the\n\n**device footprint is the full parameter count** (~5B and ~8B), and that is what has to fit 16 GB. This\n\nis exactly why E4B, the \"4B\" model, does **not** fit one core while E2B, the \"2B\" model, does.\n\n**Plan capacity from real parameters, never the effective headline.**\n\nThe second half of the trap: **bf16 halves the on-disk neff but not the on-device constant count.**\n\nThe intuition \"just use lower precision to fit\" is wrong here — E4B's ~15.4 GB of resident fp32\n\nconstants stay ~15.4 GB of *slots* regardless of dtype tricks at the boundaries, and a second neff\n\nstill tips past 16 GB. **Tensor parallelism, not precision, is the lever that actually fits the model\non the core** (§4.4).\n\nThe same cross-layer KV-sharing that NxD can't represent (§2.1, §3) is what makes E2B fit one core in\n\nthe first place. Only the **15 non-shared layers allocate a KV buffer**; the rest read through. The\n\ndevice-resident cache is therefore far smaller than a naïve one-buffer-per-layer implementation, which\n\nis a large part of why the memory budget closes. An architectural feature that breaks the vendor\n\nabstraction can still be a **net win** once you stop fighting it.\n\nTwo of the compiler wins came from *math*, not from the compiler:\n\n`argmax`\n\n— greedy decode is identical whether or not it runs. Moving it host-side (only needed for\nsampling) freed the SBUF it was overflowing (§4.2). Generalize: `layer_scalar`\n\n(§4.6). When a model multiplies by a learned scalar that\nlives in a buffer, that scalar will be wrong after any param-only load — a bug with no error message,\nonly degraded output.`inf2.xlarge`\n\nand `inf2.8xlarge`\n\ncarry the **identical 2-NeuronCore / 32 GB-HBM accelerator**; they\n\ndiffer only in **host** vCPU (4 vs 32) and RAM (16 vs 128 GB). Since Gemma-4's transformer runs\n\nentirely on the cores, the only thing standing between the ~4× cheaper box and full performance is\n\n**host memory** — solved by \"slim\" servers that keep just the embedding table on the CPU and drop the\n\ntransformer layers (they live in the neff). Result: **all three models serve on a single\ninf2.xlarge**, and because throughput barely drops, the cheap box is\n\nTwo prefill strategies coexist and trade off cleanly, so the \"right\" one depends on the workload:\n\nhost-seed (`tp_alias` ) |\ndevice-prefill (`ModelBuilder` ) |\n|\n|---|---|---|\n| First token | ~1.4–1.6 s (CPU prefill) | ~0.1–0.16 s |\n| Decode | up to ~60 tok/s (E4B slim) | ~36–39 tok/s (E4B) |\n| Best for | long generations | chat / first-token-bound |\n\nBecause the KV cache is **device-resident** in both, per-token latency is essentially **flat across\ncontext length** — there is no host round-trip that grows with the sequence. The knob to turn is\n\n**Hugging Face** (recipe + compiled neffs + Dockerfiles + model cards):\n\n`xbill9/gemma-4-E2B-it-inferentia2`\n\n`xbill9/gemma-4-E4B-it-inferentia2`\n\n`xbill9/gemma-4-12B-it-inferentia2`\n\n**Docker Hub** (prebuilt, run with `--device /dev/neuron0 --ipc=host`\n\n):\n\n`xbill9/gemma4-optb`\n\n— E2B: `latest`\n\n/`512-128`\n\n, `slim`\n\n, `tp2-slim`\n\n, `tp2-2048`\n\n`xbill9/gemma4-optb-e4b`\n\n— E4B: `latest`\n\n/`512-128`\n\n, `tp2-devprefill-512`\n\n, `slim-devprefill`\n\n`xbill9/gemma4-optb-12b`\n\n— 12B: `tp2-devprefill-256`\n\n, `slim-devprefill`\n\n**Key source** (branch `gemma4-inf2-nxd-kvshare`\n\n): `optb_kv.py`\n\n(E2B two-graph), `tp_alias_trace.py`\n\n(E4B TP+alias), `tp_mb.py`\n\n(E4B/12B device-prefill `ModelBuilder`\n\n), `optb_server_*.py`\n\n(full / slim /\n\ndevice-prefill HTTP servers with OpenAI routes). Fix history: `e785f6d`\n\n(layer_scalar), `d70dc94`\n\n(12B global-layer sharding), `328ca88`\n\n(drop on-device softcap), `a0a2a75`\n\n(force eager attention),\n\n`78967ac`\n\n(bf16 weights).\n\n`gemma4_unified`\n\naudio (640 samples/40 ms) and image (48×48 patch) projection\npaths exist but are not yet wired to the device graph.Natural next steps: static batching, a 12B 512-context recompile, a `tp_alias`\n\n-style 12B decode path\n\n(the E4B data suggests ~2× throughput headroom), and wiring the 12B multimodal projections.\n\n*Base models © Google, Apache-2.0. The compiled neffs embed the base weights in bf16 and are\nredistributed as Apache-2.0 derivatives with attribution. The Option-B recipe, TP sharding rules,\ndevice-resident KV design, and servers are Apache-2.0. Not affiliated with or endorsed by Google or\nAWS; provided for research/testing.*", "url": "https://wpnews.pro/news/porting-gemma-4-2b-4b-12b-to-aws-inferentia2", "canonical_source": "https://dev.to/gde/porting-gemma-4-2b-4b-12b-to-aws-inferentia2-2jnf", "published_at": "2026-07-13 13:36:03+00:00", "updated_at": "2026-07-13 13:46:33.285772+00:00", "lang": "en", "topics": ["artificial-intelligence", "large-language-models", "ai-infrastructure", "developer-tools"], "entities": ["Google", "Gemma-4", "AWS Inferentia2", "Hugging Face", "Neuron SDK", "optimum-neuron", "neuronx-distributed"], "alternates": {"html": "https://wpnews.pro/news/porting-gemma-4-2b-4b-12b-to-aws-inferentia2", "markdown": "https://wpnews.pro/news/porting-gemma-4-2b-4b-12b-to-aws-inferentia2.md", "text": "https://wpnews.pro/news/porting-gemma-4-2b-4b-12b-to-aws-inferentia2.txt", "jsonld": "https://wpnews.pro/news/porting-gemma-4-2b-4b-12b-to-aws-inferentia2.jsonld"}}