Avaliação de Desempenho e Uso de VRAM em GPUs: Um Estudo Comparativo entre JAX (XLA) e PyTorch no Vision Transformer

• Martony Demes da Silva UFMS

DOI: https://doi.org/10.5753/eradco.2026.33745

Resumo

O dimensionamento de arquiteturas Transformer impõe restrições de recursos computacionais e de memória em GPUs. Este artigo apresenta um estudo comparativo de desempenho entre JAX (XLA) e PyTorch (Eager e Compiled) aplicando um bloco de Atenção Multicabeça de um Vision Transformer em uma GPU NVIDIA Tesla T4. Os resultados empíricos demonstram que o JAX atinge até 3984,74 amostras/s com uma redução observável na alocação de pico de VRAM em relação ao PyTorch Eager, evidenciando sua eficiência em cenários de memória limitada, a despeito do tempo inicial de compilação JIT.

Referências

Bradbury, J., Frostig, R., Hawkins, P., Johnson, M. J., Leary, C., Maclaurin, D., Necula, G., Paszke, A., VanderPlas, J., Wanderman-Milne, S., et al. (2018). JAX: composable transformations of Python+NumPy programs. [link], 1(1):0–1.

Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., et al. (2021). An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations (ICLR).

SBC, editor (2020). Anais da VI Escola Regional de Alto Desempenho do Centro-Oeste (ERAD-CO 2020). Sociedade Brasileira de Computação (SBC).

SBC, editor (2021). Anais da VII Escola Regional de Alto Desempenho do Centro-Oeste (ERAD-CO 2021). Sociedade Brasileira de Computação (SBC).