A kernel-centric path to real-time video generation on Trainium

來源原標題

A kernel-centric path to real-time video generation on Trainium

機構

Amazon

原始連結

https://www.amazon.science/blog/a-kernel-centric-path-to-real-time-video-generation-on-trainium

摘要原文

In 2018, Jürgen Schmidhuber — who in 1990 was the first person to propose world models as a machine learning concept — published a paper with David Ha. They wrote about “a predictive world model” which could “extract useful representations of space and time” and use that to train agents to drive, among other things. The growth in the capability and real-world applications of world models in the eight years since that seminal paper is staggering. The emergence of massive training datasets, paired with variational autoencoders, and both diffusion and autoregressive transformers enable today’s world models to convert text, images, audio, and video inputs into latent space that is used to build and maintain breathtaking, immersive world models which can predict state changes and respond to actions. These models have vital and expanding roles in robotics, transport, climate modeling, video game development, scientific simulation, and design prototyping. The sharp increase in the power and potential of world models has been accompanied by an increased need for hardware and infrastructure capable of supporting them. That is a challenge that Bryce Schmidtchen, co-founder and chief technology officer of Reactor, is well acquainted with. “It's about real time, it's about low latency, and it's about doing that as efficiently as you can at scale,” Schmidtchen observed. Reactor is a platform that allows developers, designers, and researchers to deploy, use, and scale real-time interactive AI models. “Efficiency means everything from how you schedule the inference on the given chip, in our case Trainium, to how you think about maximally bin packing every forward pass of the model.” The challenge of efficiency is made even more acute by Reactor’s global presence. “We think about GPU clusters in terms of regions — we have hundreds all over the world. The latent that turns into a pixel that comes off the GPU needs to be able to get to the network without hopping around through Kubernetes. And you need to do this in a way where it doesn't matter if it's a well-supported cluster or bare metal in a closet. “And then,” Schmidtchen continued, “you need to tie this to world-class networking and media streaming that supports different codecs, different resolutions, and allows that to connect to APIs and SDKs across different languages that can support all different types of applications.” The global reach of AWS already helps Reactor to achieve many of its goals. “Everything they have at the inference layer, the level of scale that we're able to achieve in different regions is what makes a platform of this scale possible and reliable,” Schmidtchen observed. Now, a recent collaboration between Reactor and the Amazon Neuron Science team has forged a path to even greater efficiency for world models. This is a look at the optimization strategy those teams pursued, and how that laid the foundation for future models to achieve real-time capability on Trainium. The rise of autoregressive diffusion models “The Neuron Science team explores new techniques for generative AI model enablement optimization,” explained Jun Wu, a principal applied scientist on the Neuron team. “This might be a new model architecture or new algorithm for optimization or a new way to generate and optimize the code for running those models on Trainium. We identify opportunities to build a prototype and make it feasible to be ported to the production pipeline.” The team spotted one such opportunity around the usage of diffusion models in video generation. “We noticed an evolution from generating shorter, fixed-length videos to infinite-length or dynamic-length videos,” Wu explained. “Generating high-quality video frames at those lengths gave rise to the usage of autoregressive diffusion models.” Those types of models, which combine the sequential next-token prediction found in large language models with the iteration and refinement of diffusion models, are essential for users who want to generate video where they can navigate the generated environment. “Autoregressive diffusion means a user keystroke can be absorbed as input to the model and, conditioned on the previously generated video frame, the model can correctly decide the next move or the next scene,” Wu observed. “Most of the interactive video generation models we are seeing today are using this.” Wu and his team, Mason Fu and Lingfan Yu, both senior applied scientists, saw a chance to optimize how those models are deployed for real-time video generation on Trainium, and saw the opportunity to do this in collaboration with Reactor. “We had been working on real-time video and interactive video generation for over 18 months at that point,” Schmidtchen noted. The teams considered various video generation techniques and aligned on utilizing Rolling Forcing, citing both its ability to consistently generate high-quality 30-second videos and its relative size. The hard part with these models isn't quality, it's that they have to run in real time. Unlike traditional video generation, which renders a full clip offline and returns it later, models like Rolling Forcing are streaming. Each frame is generated and immediately consumed, shown to a user or fed back in as the next input. That is how developers and customers actually use them, and a frame that arrives late breaks the experience, because generation has to stay ahead of the playback timeline. “High-quality generation diffusion models pose the challenge of a very long sequence, which requires a lot of memory consumption,” Wu explained. “Rolling Forcing is relatively small, but its sequence length is very large.” That combination of a small model, a very long sequence, and a hard frame-rate floor is what makes real time so demanding. Rolling Forcing's ability to generate 16 frames per second, the standard for video playback, meant it also met latency requirements. This is where Trainium adds value, bringing the performance and memory to sustain that long-sequence workload at speed and deliver real-time generation above 16 fps rather than merely producing good frames eventually. So the teams set about enabling Rolling Forcing, as a proxy for autoregressive diffusion video generation on Trainium. A developer-friendly approach The Neuron Science and Reactor teams adopted a kernel-centric, bottom-up methodology aimed at making life easier on developers. They focused on three challenges — dynamic shapes, unusual memory access patterns, and heavy cache management — that real-time video generation poses for generic compilers. Each of those challenges recurs on every forward pass, so what might be an insignificant delay on other workloads can compound into latency failures in real-time video generation. For example, the challenges posed by dynamic shapes are partly rooted in the shifting nature of scenes in real-time video generation: Imagine a shot showing a pitcher, alone on the mound, that pans out to show the other players and then thousands of fans in the stands, all within seconds. Real-time video generation also involves frames which are generated and evicted in a continuous sliding-window process. “This means attention lengths vary across forward passes, and there are two distinct passes per window: denoising, then cache cleanup,” Wu said. “All of that makes static compilation hard.” In addition to sliding KV cache copies, video generation workloads contain operations — rotary position embeddings (RoPE) and attention transposes — whose memory access patterns are workload-specific, making them challenging for any general-purpose compiler to fully optimize. For example, in video generation RoPE must contend with three axes (height, width, and time) rather than the single position it accounts for in LLMs. “The 3D rotary embedding interleaves odd and even elements along the innermost dimension, producing many tiny data transfers when compiled generically,” Wu explained.