scaling sparse ngram embedding tables needs new infra
Hyperstition came out with a nanogpt speedrun world record this week that was pretty historic - you can now train nanogpt in 39.9 seconds. very cool!
I wanted to unpack the ngram optimizations made here.
What is n-gram
n-gram is a sequence of n consecutive tokens
a probabilistic model that can estimate the probability of a word n given the n-1 previous words
example - if the sequence is cat sat on the mat
1-gram: cat, sat, on, the, mat
2-gram: cat sat, sat on, on the, the mat
3-gram: cat sat on, sat on the, on the mat
What is a Bigram & Trigram table?
bigram table is a bigram embedding look up table, that stores embedding vectors for each pair
so the model can now look up the token embedding table or the bigram embedding table
this way the model can vector representations for commonly occurring pairs
I love machine learning
Ilovemachinelearning
Ilove
"I love"
lovemachine
"love machine"
machinelearning
"machine learning"
Bigram embedding look up table
Bigram pair
Embedding vector
("I", "love")
[0.4, -0.3, ...]
("love", "machine")
[0.9, 0.1, ...]
("machine", "learning")
[-0.2, 0.8, ...]
What was wrong with the original bigram table:
the previous speedrun record also had a bigram table, but the optimizer state for each row occupied ~6KB of memory
for adam optimizer that stores the first and second moment for every param, this can quickly take up a lot of memory if each row / pair in the bigram table has to store both moments
to fix this, instead of storing the entire second moment vector, they store only a scalar and a timestamp per row. so the 2nd moment decay can be recalculated using the timestamp.
Hashing
but with vocab size of ~50K tokens, there can be over 2.5B possible pairs, and this would take up a lot of memory
the solution is to map the pairs into a finite number of rows using hashing
the previous record on nanogpt had very few rows (377,280 rows) and caused lot of collisions due to this.
Adding a Trigram table
now that they saved some memory space from storing all optimizer states, with this new space, they increased the size of the table from 377280 to 84.6million rows
also added trigram channels
trigram lookup is useful as it can memorize more precise short-context patterns
What this means
they've scaled the emb table from 640M to 65B for a 124M parameter model
and that's 500 times the size of the model
the speedrun benchmark caps the model at 124M active params per token, but there's no cap on the total params that can be stored
so they store an enormous sparse embedding table as long as each token only uses 124M
we can have huge / sparse embedding weights for all the memory related weights responsible for retrieval
and the transformer model's parameters for performing FLOPs computation and reasoning
Do we need new infra??
to scale the 500x embedding table for frontier model size we need new infra
for ex. if the model was 100B params, and the embedding table is only around ~12x the size like what ds v4.1 flash does, it's still possible to shard the memory weights across a big enough GPU cluster.
but if its 500x the size of the 100B model, that's a lot of params to store / shard and standard FSDP wouldn't work
new infra might look like distributed database systems