Juliacon 2024

Training large neural networks has been a hard challenge in the Julia ecosystem. It requires the infrastructure to scale training to many GPUs simultaneously. In this talk we demonstrate several pieces required to train LLMs such as Llama2 in Julia. Using Dagger.jl, we demonstrate scaling to multiple GPUs during the training of a Flux.jl model.

This approach has several applications; since the underlying implementation is agnostic of the model being trained, it can thus be reused to parallelize training across a wide variety of models to solve current and emerging problems faster than before.

In practice, LLMs aren't trained from scratch. Pre-trained models are typically loaded and then fine tuned on specialised tasks with techniques such as Low Rank Adaptation of LLMs (LoRA). This greatly reduces the training time required to fit a model to become suitable for a specific task.

We will present the various components of the pipeline - from the model and data description to setting up the parallel training job. We will present the scaling in training time achieved, and the size of models we could accommodate with this approach. We will walk through the various challenges encountered during the process, and how we envision scaling of large GPU workloads using the Julia ecosystem in the future.