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.
talk-data.com
Topic
jax-toolbox
1
tagged
Activity Trend
1
peak/qtr
2020-Q1
2026-Q1