HeyGen Speeds Up Avatar IV Model with Google Cloud Optimization

HeyGen, an AI video generation platform featuring avatar models, has successfully optimized its Avatar IV model to run on Google Cloud's specialized AI infrastructure. The company worked with Google Cloud's performance optimization team and achieved a 1.86x speedup compared to their first working version.

The Avatar IV pipeline was ported from GPUs to Trillium (v6e) hosts, which consists of eight chips. The model code runs unmodified, and the TPU-specific engineering concentrates on where the hardware differs. Each attention variant dispatches to a Pallas kernel built for its shape.

The team hid the collective's footprint by pipelineing it, allowing XLA to overlap the transfer with other tasks. This resulted in a 5x reduction in the collective's impact on the compute stream without affecting attention time or wire time.