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
S
Speaker
Stefano Bosisio
1
talks
MLOps Engineer
NVIDIA
Stefano Bosisio is an accomplished MLOps Engineer with a solid background in Biomedical Engineering, focusing on cellular biology, genetics, and molecular simulations. He earned his PhD in Computational Chemistry from the University of Edinburgh, where he developed a strong foundation in computational methods. After completing his PhD, Stefano transitioned into Data Science, where he began his career as a Data Scientist. His interest in machine learning engineering grew, leading him to specialize in building ML platforms that drive business success. Stefano's expertise bridges the gap between complex scientific research and practical machine learning applications.
Bio from: JAX for LLM Development at NVIDIA & AI in Production at Mollie
Filter by Event / Source
Talks & appearances
1 activities · Newest first