Wesley Kambale
ML Engineer, GDE for AINone
I'm a machine learning engineer and Google Developer Expert for AI, adept at crafting production-ready ML systems that provide impactful solutions in the African market.
Abstract
The future of AI in Africa requires building models that natively understand low-resource languages, rather than relying on translation wrappers around English-centric models. Training foundational generative models from scratch requires immense computational control. This tutorial takes attendees into the trenches of modern machine learning engineering using the JAX AI Stack. We will build "Pearl-Chat," a localized Large Language Model, entirely in JAX using the new object-oriented Flax NNX API. The tutorial will unpack the end-to-end pipeline required for pre-training. We will start by integrating a custom BPE tokenizer for Luganda using Google's Grain library for deterministic, high-throughput data loading. The live demonstration will focus on the architectural mechanics that make JAX uniquely powerful for this task. Attendees will see a side-by-side comparison of PyTorch’s implicit state management versus JAX’s explicit state ledgers. Finally, we will visualize jax.sharding in action, showing exactly how the model’s parameters and training data are split across a multi-core device mesh to achieve extreme distributed performance.