Introduction

Generative AI powered by LLMs (Large Language Models) has a wide range of applications at Uber, like Uber Eats recommendations and search, customer support chatbots, code development, and SQL query generation. 

To support these applications, Uber leverages open-source models like Meta^® Llama 2 and Mistral AI Mixtral^®, and closed-source models from OpenAI, Google, and other third-party providers. As a leading company in mobility and delivery, Uber also has considerable domain-specific knowledge that can improve LLM performance for these applications. ‌One way Uber incorporates this domain-specific knowledge is through RAG (Retrieval Augmented Generation). 

Uber also explores ways to adapt LLMs to Uber’s knowledge base through continuous pre-training and instruction fine-tuning. For example, for Uber Eats, we found that a model finetuned on Uber’s knowledge of items, dishes, and restaurants could improve the accuracy of item tagging, search queries, and user preference understanding compared to open-source model results. Even further, these finetuned models can achieve similar performance to GPT-4 models while allowing for much more traffic at Uber’s scale. 

AI community support and open-source libraries like transformers, Microsoft DeepSpeed^®, and PyTorch FSDP empower Uber to rapidly build infrastructure to efficiently train and evaluate LLMs. Emerging open-source initiatives like Meta^® Llama 3 llama-recipes, Microsoft LoRA^®, QLoRA^™, and Hugging Face PEFT^™ simplify the fine-tuning lifecycle for LLMs and reduce engineering efforts. Tools like Ray^® and vLLM^™ maximize the throughput of large-scale pre-training, fine-tuning, offline batch prediction, and online serving capabilities for open-source LLMs.

The novel in-house LLM training approach described in this blog ensures Uber’s flexibility and speed in prototyping and developing Generative AI-driven services. We use SOTA (state-of-the-art) open-source models for faster, cheaper, and more secure and scalable experimentation. Optimized in-house LLM training helps Uber maintain cutting-edge technology innovations and passes on the benefits to people who use the Uber app.

Infrastructure Stack

A critical component of LLM training at Uber is the thoroughly tested infrastructure stack that enables rapid experimentation.  

Layer 0: Hardware

Uber’s in-house LLM workflows are scheduled on 2 kinds of computing instances: (1) NVIDIA^® A100 GPU instances in Uber’s on-prem clusters and (2) NVIDIA H100 GPU instances on Google Cloud. Each Uber on-prem A100 host is equipped with 4 A100 GPUs, 600 GB memory, and 3 TB SSD.  On Google Cloud, each host is a 3-highgpu-8g machine type, with 8 H100 GPUs, 1872GB CPU memory, and 6TB SSD. These machines are managed as part of the Crane infrastructure stack.

Layer 1: Orchestration

Computing resources are managed by Kubernetes^® and Ray. Kubernetes is used to schedule training workloads and manage hardware requirements. Ray, along with the KubeRay operator, is used for distributing the workload to the workers. 

Layer 2: Federation

The end-to-end flow of resource management is depicted in Figure 1 below. Our multiple Kubernetes clusters are managed by a federation layer, which schedules workloads based on resource availability. Training jobs are modeled as a Job with multiple Tasks. The JobSpec defines the goal state of the Job and includes information like the instance SKUs, clusters, compute/storage resources, and post-Docker-launch commands to set up tokens and environment variables. 

Training Stack

We fully embraced open source when building our LLM training stack. Specifically, we integrated PyTorch, Ray, Hugging Face, DeepSpeed, and NCLL to enable training LLMs with the Michelangelo platform. 

As shown in Figure 2 below, Ray is at the top of the LLM training stack for coordinating tasks. NCCL is at the bottom level for GPU communication.

Distributed Training Pipeline

To support our in-house implementation, we built an LLM distributed training pipeline that includes host communication, data preparation, distributed model training, and model checkpoint management. Here is how it works:

Training Results

Our exploration involved demonstrating that the in-house Michelangelo platform has the capability and scalability to train any open-source LLM with throughput optimization.    

Training on State-Of-The-Art LLMs

We found that the Michelangelo platform can support the largest open-source LLMs while training under different settings, including full-parameter fine tuning and parameter-efficient fine tuning with LoRA and QLoRA.  

Figure 4 (above) shows the training loss curve for Llama 2 13B and Llama 2 70B with and without LoRA. We found that, even if LoRA and QLoRA use far fewer GPUs and train much faster with fewer GPUs, the loss decreases much less than the full parameter training. Therefore, it’s important to improve the throughput/Model Flops Utilization (MFU) for full parameter fine-tuning of LLMs.

Throughput/MFU Optimization

Scaling the transformer architecture is heavily bottlenecked by the model’s self-attention mechanism, which has quadratic time and memory complexity. To achieve better training throughput, we’ve explored several industry recommended efforts to optimize GPU memory usage: CPU offload and flash attention.

The first training throughput optimization we explored was using DeepSpeed ZeRO-stage-3 CPU Optimizer Offload, which led to at least 34% GPU memory reduction with the same batch size when training Llama 2 70B. This allowed us to increase the batch size by 3-4 times but still keep the same forward and backward speed, so the training throughput increased by 2-3 times.

The second training throughput optimization explored was following Hugging Face’s suggestion to use flash attention. Flash attention is an attention algorithm used to reduce quadratic complexity and scale transformer-based models more efficiently, enabling faster training. With flash attention, we could save 50% of GPU memory with the same batch size. If we maximized the usage of GPU memory, then we could double the batch size while keeping compatible forward and backward speed.

To study training efficiency, we used Hardware Model Flops Utilization (MFU). MFU is the ratio of the observed throughput to the theoretical maximum throughput if the benchmarked hardware setup was operating at peak FLOPS with no memory or communication overhead.

In our benchmark, we used Deepspeed Flops Profiler to obtain the expected FLOPS number. FLOPS per GPU was calculated as: forward and backward FLOPS per GPU/iteration latency. Then we divided it by the device’s ​​peak theoretical performance and obtained our final MFU metric. ‌In all our experiments, we maximized the batch size under different optimization settings so that GPU memory was fully utilized. We did this with the training arguments, setting gradient_accumulation_steps = 1 so that macro_batch_size = micro_batch_size x num_gpus x 1.

Here’s what we found: 

LLM Scorer

To evaluate the raw or finetuned models, we also implemented an offline LLM scorer to predict on large datasets. We used Ray to create a cluster on Kubernetes with multiple instances and each instance with multiple GPUs. This way, we could distribute the data and score in parallel. On each instance, we used inference servers such as vLLM.

In our implementation, we used Ray jobs as the operational foundation. Each Ray job allocates a specified number of CPU and GPU resources, downloads models, and partitions datasets by rank. The Ray `ActorPool` aggregates the outputs from different Ray Actors. Figure 7 below shows the implementation of this LLM scorer. 

Figure 8 below summarizes the performance of batch prediction for Mixtral 8x7b models using 2 A100 and H100 GPUs, where the input token size is 4K and the maximum output tokens are 700. We observed that output tokens per second increased linearly for batch sizes up to 64 on both GPU types. Notably, the H100 achieved a throughput 3 times higher than the A100. This benchmark helps teams make production decisions and plan resource requirements.  

Conclusion

As more and larger open-source models get published, like Llama 3 450B, we’ll improve our LLM training infrastructure to support fine-tuning them. Using these finetuned models will help us improve things like Uber Eats recommendations and search.

Thinking about the broader industry, our journey exploring in-house LLM training has brought us these insights:

Acknowledgements

This major step for GenAI at Uber couldn’t have happened without the many teams at Uber who contributed to it: Michelangelo, Applied AI, Dev Platform, and Platform Core Service.