Velocidade e Python ao mesmo tempo: como usar CUDA Graphs para executar código Python rapidamente em deep learning

Nos últimos anos, a velocidade das GPUs aumentou de forma explosiva, e os métodos de otimização para cargas de trabalho de deep learning também estão mudando. O PyTorch também vem adicionando recursos de otimização, como torch.compile(), mas em algumas cargas de trabalho, incluindo LLMs, as melhorias ainda estão em andamento.

(Enquanto esperamos melhorias no torch.compile()), encontrei um artigo que apresenta o CUDA Graph como um método de otimização que pode ser aplicado imediatamente, e resolvi traduzi-lo. (⚠️ Atenção: no fim do texto há parte de material promocional da Fireworks.ai, empresa que desenvolve e oferece uma plataforma de inferência para LLMs.)

Neste texto, o CUDA Graph é apresentado na seguinte ordem:

• introdução ao método de otimização existente de overlap entre CPU e GPU

• trechos em que ocorre overhead da CPU

• técnicas para otimizar o overhead da CPU e o CUDA Graph

• apresentação de um caso de aplicação de CUDA Graph ao modelo LLaMA2-7B

• apresentação dos ganhos de desempenho obtidos com CUDA Graph

• apêndice: problemas ao usar torch.compile() no momento atual (PyTorch 2.0.1) e como resolvê-los