talk-data.com talk-data.com

Description

JAX is a key framework for LLM development, offering composable function transformations and a powerful bridge between low-level compilers and high-level code. To help address the challenges of moving from development to large-scale production, this talk introduces JAX-Toolbox, an open-source project that provides a robust foundation for the LLM development lifecycle. The session covers the CI/CD architecture that provides a stable foundation for JAX-based frameworks, how to build GPU-optimized containers for LLM frameworks such as MaxText and AXLearn to ensure reproducible workflows, and practical methods for deploying frameworks' containers on Kubernetes and SLURM-based clusters.