Implement fused modules

[casper-hansen]

There are two common ways to fuse layers in Llama/Mistral type of models. Speed and memory are measured on RTX 3090 with TinyLlama 1B.

• MLP: Merge gate_proj and up_proj together
  • Handle splitting SwiGLU weights into original module post-training
• Attention: Merge q, k, v together
  • Handle replacing attention module better, currently only works with sample packing off
  • Fix the loss being much higher

All fusing of layers must happen AFTER the model is loaded in order to load the pretrained weights into the fused modules.

TinyLlama 1.1B - A6000

Conclusion: Fusing MLP can save roughly 27% memory in cache. Fusing attention seems to do nothing for the speed but increases memory with about 1GB.

1. None fused (Main):
   • Memory: 6.186GB (+15.724GB cache, +0.781GB misc)
   • Speed: 0.53 seconds per step
2. MLP fused (PR):
   • Memory: 6.188GB (+11.443GB cache, +0.781GB misc)
   • Speed: 0.51 seconds per step
3. MLP + Attention fused (PR):
   • Memory: 6.618GB (+12.195GB cache, +0.781GB misc)
   • Speed: 0.50 seconds per step

Llama-2-7B - A100

Conclusion: Saves enough memory to load using adamw_torch.

1. None fused (main):

   • adamw_torch Memory: 37.732GB (+39.764GB cache, +1.366GB misc)
   • adamw_torch_fused Memory: 37.732GB (+14.506GB cache, +1.366GB misc)
   • adamw_bnb_8bit Memory: 25.393GB (+14.494GB cache, +1.366GB misc)

2. MLP fused (PR):

   • adamw_torch Memory: 37.732GB (+38.813GB cache, +1.366GB misc)
   • adamw_torch_fused Memory: 37.732GB (+14.647GB cache, +1.366GB misc)
   • adamw_bnb_8bit Memory: 25.269GB (+14.137GB cache, +1.366GB misc)

3. MLP + Attention fused (PR):

   • adamw_bnb_8bit Memory: 31.332GB (+13.752GB cache, +1.366GB misc)
   • adamw_torch Memory: OOM

QLoRA

Currently, it is not compatible with QLoRA. But there is potential to do so. In bitsandbytes, you can import the 4-bit and 8-bit linears and use them instead of nn.Linear.

https://github.com/TimDettmers/bitsandbytes/blob/main/bitsandbytes/nn/modules.py#L258