New Record: Untie Value Embeddings (-0.4s, -25 steps)#209
Conversation
* train script * add to record * typing ext * fix record * add commentary --------- Co-authored-by: Your Name <[email protected]>
photomz
changed the title
|
Did you experiment with value embeddings on every other layer, or skipping some layers? I think that is what Karpathy is doing in his Nanochat model |
|
Cool, I'll merge later if the timing/loss holds up. When I consider the loss difference, the time change here is pretty marginal but the step count decrease is significant and interesting. |
@varunneal Thanks for the tip! Glad that Karpathy is interested in Value Embeddings, but I did a light ablation and couldn't replicate his alternate pattern. I suspect our config for each Attn blocks is too custom, or maybe his approach worked at different scales -- can't say for sure.
I think the marginal time change just emphasizes the need for someone to find an efficient sparse gradient update. Most of our params are now in Embeddings, that apart from |
|
It would be interesting to measure how much time one embedding table costs us over the course of training. There's probably some ratio there, like a table needs 10-15 steps to justify itself. Also, on the pattern, there's actually someone over in nanochat who played around and found that VEs on "deeper layers" worked best (here), but was experimenting at our scale. The default / target model for nanochat is a 20-layer / 1280-dim model, though, so it could definitely be a scale thing. Here's some thoughts I had / conjecture on why our particular pattern might be working well: We believe that the first 8 layers of our model are where it's building its compute graph--laying down the logic (via keys and values) that will dictate how it processes future tokens. And then the last three layers are for the actual token prediction, and that's where we have our VEs. So maybe the value embeddings are more 'semantic'? That might also help explain why the same set of VEs are useful on the first two layers? |
|
Merging at 99.0s (-0.3s) based on: step:1575/1575 val_loss:3.2795 train_time:98292ms step_avg:62.41ms sum([98650, 98674, 98718])/3-sum([98292, 98514, 98450])/3 |
4 participants
Updates in PR, building off latest #201
VE[.12...012]withVE[.01...345].Also suggest we close #198 by pinning
typing-extensions==4.15.0for housekeeping.Ideas
The latest idea is inspired by the principle of Bigram/Engram Hash embeddings to allow more uncontextualised information into the architecture. I believe there's further room in adding sparsity to be explored, i.e. Meta's STEM and optimizing the bwd to not communicate unaccessed indices in sparse embeddings.
I tried these ablations:
ve[012...012](1)ve[01…01]ve[.01…01.]Surprisingly, the old finding in New Record: Partition Residual Dims and Symmetric Value Embeddings (-0.7s) #194 that adding VE to the first layer no longer held up -- there was no loss delta at all compared to before. I hypothesize that improvements in between (Partial Key Offset, Bigram Hash) removed the need for this somehow, but I'm not quite sure why.
Unrelated to this VE idea, I tried to optimize the bwd with
torch.Embedding(..., sparse=True), but I'm running into 2 issues: 1) the metadata overhead of keeping count of the accessed indices complicates the Adam step, especially when distributed, and 2) Full CUDA graphs aren't preserved under sparse embeddings in 2.10.Timing
The timings (-25 steps) hold up to the latest PR, and are significant (p<0.001).
The 3 extra embedding weights in H1 added to the average step time, compared to H0. I suggest we time this improvement
time delta = 98.1-97.7 = 0.4s.