AI Weekly Malaysia

Google open-sourced TorchTPU, a PyTorch version that lets AI apps run on TPUs and move to Google Cloud with a configuration change rather than a rewrite; Amazon released TorchNeuron for Trainium/AWS a day earlier and is submitting PyTorch code to decouple features from Nvidia chips. PyTorch, originally built by Meta in 2016 and moved to the Linux Foundation in 2022, is used by some 85% of AI apps, according to the article. Google says it worked with Meta, Hugging Face, vLLM maker Inferact, and SGLang maker RadixArk on TorchTPU.

Why: If the 'no rewriting' claim holds for your model, TPU and Trainium become a second deployment target and a negotiating lever on GPU cloud pricing; if it does not, you will burn time on portability work. Before committing, ask for benchmarks on your exact model and confirm the vLLM/SGLang/Hugging Face paths you rely on are supported, because the article gives no performance or price numbers.

Hugging Face published a Space that checks whether a given Transformers model runs unmodified on AWS Trainium via torch-neuronx, the native PyTorch backend. Each result row reports the number of Trainium cores the model is sharded over, wall time including compilation, and peak device memory on a trn2 instance, plus the exact run.py script, Neuron stack, and neuronx-cc/Transformers versions used. Coverage is currently limited to text-generation and image-text-to-text Transformers models, with diffusers and sentence-transformers planned; "runs" means the model loads, shards (tensor parallel, or expert parallel for MoE without a TP plan), and generates 32 tokens.

Why: If you are costing out inference for a text-generation or image-text-to-text model, look the model up here before budgeting GPU capacity: the Space gives you Trainium core count, wall time with compilation, and peak memory for a trn2 run, and if it is listed you can skip the first-run compile by syncing the prebuilt graphs with `hf buckets sync hf://buckets/aws-neuron/torch-neuronx-neff-cache/neuronxcc-<version>` using the matching compiler version. Treat a green result as "loads and emits 32 tokens", not as a throughput benchmark — you still need your own numbers before migrating a production workload.

Hugging Face and AWS are shipping TorchNeuron, a PyTorch plugin that makes AWS Trainium a native device (eager execution, torch.compile, FSDP2/DTensor) instead of the lazy-tensor PyTorch/XLA path used by optimum-neuron. Public beta starts Oct 20, 2026 via `pip install torch-neuronx`, and Transformers, Diffusers, Accelerate, PEFT and TRL now run on Trainium directly using `device_map="neuron"`. HF reports 89% of the top 100 most-used text-generation architectures work out of the box, versus 46% with Optimum Neuron v0.4.5, and now recommends Neuron users skip Optimum Neuron and use the libraries directly.

Why: If you have any code built on optimum-neuron, this is a migration signal: HF explicitly recommends moving to the libraries directly, and the architecture coverage gap (89% vs 46%) is the concrete reason. The cost is real though — the sample requires torch.distributed setup, tensor parallelism and a static cache for generation to avoid frequent recompilations, so budget time for that rather than assuming a drop-in swap. For teams comparing Trainium against GPUs on price, the wider out-of-the-box model coverage is the number to weigh, not the announcement itself.

Part 1 of a Hugging Face/Amazon series walks through system-level profiling of TorchNeuron on AWS Trainium. It shows how to set NeuronConfig with DEVICE and RUNTIME modes, capture_enabled_for_nc='0', and correlate_device_time=True, then use NeuronProfiler.export_trace to write trace.json and read key_averages() to compare Self CPU time with Neuron device time. The walkthrough follows one operation through Neuron Explorer from PyTorch dispatch to runtime scheduling to NeuronCore execution, and compares three implementations of the same operation.

Why: If you run PyTorch on Trainium, this gives a concrete profiling recipe: set modes=[DEVICE, RUNTIME] and correlate_device_time=True so key_averages() lines up CPU and Neuron device time, then use that table to decide whether an op is CPU-dispatch-bound or NeuronCore-bound before rewriting it. If you are not on Trainium, there is no action here.