On building Lumos - foundation models for personalization
The devil in the details
(Read the Lumos technical report and the official blog post.)
Lumos is an opinionated, large transformer training framework I wrote at Dream11. It was a culmination of everything I wanted to do differently in ML. My core tenets while building it-
Process over outcomes - systems and infra over modelling
People who build better systems with a software mindset produce better results faster and make fewer mistakes. In this day and age of democratisation of architectures and techniques, I don’t think brilliant people who do EDA, training, and evaluation in a single notebook, no matter how smart, will make it.
Lumos had 17k lines of code. Only ~250 of them were PyTorch.
Reproducibility over flexibility - create recipes, not results
Everything should be in the code or the config. That is the documentation. I should not have to talk to you to understand how you sampled your data and tweaked your loss to get the results you presented. I should be able to look at your code, run it, and see for myself. If it’s not EASILY reproducible, it did not happen.
Reproducibility was built into lumos such that people didn’t have to think about it(going back to systems), using extensive logging, enforcing config schemas, and storing and consuming metadata with artifacts.
We had about half a dozen researchers running training jobss and I was sure that if any of them were promising, we would be able to reproduce, combine, and deploy into other training runs without any overhead.This also helped us automate inference. Since all the recipes for feature creation were always logged, promoting a training run to the production model was a single click.
Go big or go home - build for scale from day one
The biggest lesson that can be read from 70 years of AI research is that general methods that leverage computation are ultimately the most effective, and by a large margin. The ultimate reason for this is Moore's law, or rather its generalization of continued exponentially falling cost per unit of computation. Most AI research has been conducted as if the computation available to the agent were constant (in which case leveraging human knowledge would be one of the only ways to improve performance) but, over a slightly longer time than a typical research project, massively more computation inevitably becomes available. Seeking an improvement that makes a difference in the shorter term, researchers seek to leverage their human knowledge of the domain, but the only thing that matters in the long run is the leveraging of computation.
- The bitter lesson by Richard SuttonIn the era of scaling laws, scale is the MOAT. The labs are not winning because they have some architectural secret sauce. It’s because they can scale.
Building small models to get confidence and scaling up later does not work in this paradigm. So I built for scale from day one. I used Spark and Ray, where I could have gotten away with pandas. I used DDP, where I could have used a single GPU and normal PyTorch to begin with.
Until I could train a 200m model with at least 10 terabytes of data, I did not train a single model. By the end, we ran thousands of training runs, each worth ~1.5 TFLOPs, using a simple 50-line bash script.Begin with the end in mind- eval before dev
Evals are the craze rn, but they always were the first thing I’d do for any ML project. Once you know what you are looking for, you can put your headphones on and ruthlessly optimise. I’ve been in so many meetings where people build models for months only to realise that the metric they were optimising didn’t make sense. Or that they did not look at another important metric. I spent a lot of time standardising and systematising the evaluation of models.
Before training models, I built an internal Kaggle using MLFlow that automated evals for 150+ internal tasks, removing ambiguity and providing a public leaderboard for each task.
Apart from the core tenets, I made several interesting discoveries. You can also read the Lumos technical report and the official blog post.
On memory management
Managing GPU memory is hard. Some will have you believe that tracking memory requirements is easy. 4 x the number of params for grads and Adam states. But what really kills you is the activations.
A very succinct example in the attention matrix, which is the size of seq_len^2. When increasing num_heads for your encoder from 1 to 8, you go from one attention matrix to eight attention matrices with the same number of params. This explodes the peak_mem requirement without increasing the number of parameters.
The only sane way to figure out the batch size that works for a given architecture is to run one batch through the model to see if the GPU does not go OOM when you run a forward and a backward pass.
This makes hyperparameter tuning a nightmare. Always pick one batch size that would work across your crazy experiments and the best learning rate for that batch size. Then keep those fixed for as long as possible.
Throughput management
Remember when I said managing GPU memory is hard. Scratch that. Memory management looks like a piece of cake when you consider managing throughput. The goal of throughput management is to ensure your GPUs are always busy doing floating point operations(or FLOPS).
Today, our training loop is linked to a data loading path that is slower than the compute. Accio(our propritary internal laoder based on delta tables) retrieves one large Delta partition at a time, processes it through Spark, converts it to pandas, and then iterates over the rows. The GPUs wait while the CPU loads each new block into memory. We experience latency spikes when a new partition arrives, RAM expands to fit the partition, and 90% of total time is spent loading and shaping data. Parallelism depends on batch shape and DDP rank calculations; increasing the number of GPUs does not speed up the input but just increases the wait time. Going beyond a single node requires code changes, increases the chances of errors, and offers no hiding spot for I/O delays.With Ray, the shape of the system changes. Data becomes a Ray Dataset. Direct parquet reads, distributed by default, processed with vectorized batch transforms instead of row-wise applies. More importantly, loading infra can be decoupled from training infra. We can scale the CPUs independently of GPUs to increase parallelism here since this is the bottleneck.
- Me in internal docs
You can, of course, go further. Within the GPU, you have an HBM and a tensor core. The battle within the GPU is similar to the battle outside. We need to make sure the tensor cores are always busy and not waiting on data to be loaded from the HBM. This is the problem that stuff like flash attention and other kernel fusing techniques solve.
Data version control
DVC is a cool idea, but I never saw anyone really do it. Lumos was a chance for me to know if it’s possible. Not only was it possible - it was possibly the highest ROI feature I invested in.
Only made possible by using delta tables for all my data. I abused the versioning and user_metadata features of the delta table to make a full fledged feature and config store out of a single table. I don't think people at Databricks will forgive me for this. When I showed this to one of my colleagues, he was repulsed, and maybe rightly so.
Data was not just version-controlled; we embedded the *config used to generate it* into the Delta table’s metadata. Every training run was tagged with metadata that could track the entire lineage from raw data to feature creation, model code, model config, and training runs. Without this, the experimentation cannot scale.
def get_version_metadata(table_name, version):
# Fetch specific version using Delta Time Travel
history_df = spark.sql(f”DESCRIBE HISTORY {table_name}”)
version_data = history_df.filter(f”version = {version}”).collect()
# Extract the JSON config stored in userMetadata
metadata = json.loads(version_data[0][”userMetadata”])
return metadataThe dark arts of architecture
We wanted to use the backbone of the transformer but we were not afraid to switch out stuff that made sense for modelling language, but not for modelling user behavior. We learned our own positioned embeddings - something that noone is doing now. We did not use causal masking because we were not going to do autoregressive generation. We used an encoder-decoder architecture - just like in the original “Attention is all you need paper”- in a world where all LLMs are decoder-only.
Inference is another ballgame
The priorities change for inference. Things became easier for us because we did not care about latency at inference time - just throughout. We were running 100 million inferences once a day. We were able to get away with using Spark to distribute the load in mini-batches to workers who all had their own small GPUs. It is a relief not needing any GPU-to-GPU communication and synchronization.
dataclasses ftw
Dataclasses are amazing, and I find every excuse to use them in Python. Everything as a data class in lumos - training config, feature definitions - you name it. I’ll probably have a blog post just on data classes, so I’ll stop here.
…and a bunch of other stuff that I will keep adding