What we learned

[ClashLuke]

Whenever a pull request is made, please put a summary of your knowledge in here.

[ClashLuke]

#5 - Mixture of Experts

1. DeepSpeed is fundamentally broken, and we shouldn't get it near our code unless we manually verify their code first. Thus, their claims are not necessarily accurate.
2. MoE's memory operations are incredibly slow. So slow that they're almost as expensive as computing the entire matrix multiplication.

[ClashLuke]

#7 - Parameter Offload

1. DeepSpeed is fundamentally broken, and we shouldn't get it near our code unless we manually verify their code first. Thus, their claims are not necessarily accurate.
2. Synchronous parameter offload is surprisingly trivial, but it is hard to do it asynchronously with caching and prefetching.

This pull request does it synchronously by loading a parameter onto the GPU in a PyTorch autograd function and offloads the gradient and optimizer pass in the backward pass.

[ClashLuke]

#12 - Sharpness Aware Minimization

SAM significantly outperforms its non-SAM counterparts and allows you to find bugs in your code trivially. For example, below is a run that shifts the context window by 1 token at every step like [1, 2, 3] -> [2, 3, 4].

SAM quickly finds the problem and gets to a loss of 1 in no time.

[ClashLuke]

#13 - Custom Kernel
Ignoring that the gradients are slightly wrong, the model is 20% faster and uses 10% less memory when writing both forward and backward pass by hand in C++. This is likely because the normalization used here performed its operations in-place while the baseline doesn't.
To summarize, custom C++ code isn't worth it. Custom CUDA kernels might be, but calling torch from a precompiled C++ kernel doesn't improve performance.

[ClashLuke]

#16 - TorchScript/Trace
A compiled language is likely significantly faster. Thus, a switch to JAX would most likely be reasonable to compile both forward and backward pass in one go. A JAX-based variant is created to validate this hypothesis.

[ClashLuke]

#20 - TorchScript v2
You can script the forward pass, but not the backward or optimiser step. With torch.script.trace, a step takes about 10% less time. Experiments are linked in the PR. Assuming forward, backward, and optimiser all take the same time, the scripted forward pass now consumes nearly 40% less time. As mentioned above, a fully compiled language is likely going to be significantly faster.
First results in HomebrewNLP-Jax show that 350x speedups are possible by using TPUv2 instead of a local GPU (without offload) and Jax.