HeyGen × Google Cloud：Avatar IV を TPU で動かすHeyGen x Google Cloud: Bringing Avatar IV to TPUs

HeyGen, a company that produces AI-generated avatar videos, has ported its Avatar IV model — a video generation system with more than 18 billion parameters — onto Google Cloud's Trillium (v6e) Tensor Processing Units. The work, described on the Google Developers blog, is notable because it demonstrates how a large, production-grade generative video model can be adapted from its original framework to run on Google's custom silicon, and reportedly emerges faster: the team cites a 1.86x speedup for real-time streaming.

Trillium is Google's latest generation of TPU, marketed as v6e, and represents an iteration of accelerators designed specifically for machine learning rather than general-purpose computing. TPUs are built around large matrix-multiplication units and high-bandwidth interconnects, which makes them well suited to the dense linear algebra at the heart of transformer-based models. Running a model across an "eight-chip mesh," as HeyGen did, means the workload is distributed over eight interconnected TPU chips that can exchange data quickly, allowing a model too large or too slow for a single chip to be split across several.

A central part of the effort was bridging software ecosystems. Many contemporary models, including Avatar IV, are developed in PyTorch, while TPUs are most naturally programmed through XLA, the Accelerated Linear Algebra compiler that also underpins JAX and TensorFlow. To connect the two, HeyGen used torchax, a library that lets PyTorch code execute against an XLA backend. This approach appears to let the team preserve much of their existing PyTorch model definition while still compiling down to instructions the TPU can run efficiently, avoiding a full rewrite in a different framework.

To fit an 18-billion-parameter model onto the hardware and keep it fast, HeyGen combined two parallelism strategies. Fully Sharded Data Parallel, or FSDP, splits a model's parameters, gradients, and optimizer state across devices so that no single chip has to hold the entire model in memory. Ulysses sequence parallelism, a technique popularized by the DeepSpeed project, instead partitions work along the sequence dimension. That is especially relevant for video generation, where long sequences — many frames and long attention contexts — can dominate memory and compute costs. Using the two together lets the system shard both the model itself and the sequences it processes across the eight-chip mesh.

The headline result is a 1.86x speedup for real-time streaming, the scenario in which frames must be produced quickly and continuously rather than rendered offline. To reach it, the engineering team appears to have pipelined parts of the workload, overlapping computation and data movement so that chips spend less time idle waiting for inputs. Pipelining is a common optimization on distributed accelerators, and it typically complements sharding strategies like FSDP by keeping the hardware busy across the stages of generation.

The move fits a broader industry pattern.