[WIP] Voltron v0

[winglian]

copied various implementations from around GitHub to get this all hacked together

@haeggee llm-baselines has no LICENSE, so definitely want to check with you on this w using the MoD
@cg123 Tried your implementation of bitnet, but it doesn't seem to work with torch.compile/activation_checkpointing
@kyegomez copied your implementation of MGQA, there was a small bug in the dimensions of the out_proj and your BitLinear also didn't work with torch.compile/activation_checkpointing, also needed to add rotary embeddings to that.

• No Infini-Attention yet. Might be more complexity over the BitNet Attention too, so might have to tackle that once this is working

• BitLinear

• DenseFormer

• Mixture-of-Depth

• Infini-Attention

• BitLinear CUDA/Triton Kernels

[winglian]

I'm pretty sure I screwed up some layernorms somewhere or something else. It doesn't crash atm, but even a 50M parameter model with activation checkpointing uses 31GB of VRAM, and the loss stalls out and doesn't report after about 15 steps.

EDIT: wandb https://wandb.ai/wing-lian/voltronformer?nw=nwuserwinglian

EDIT2: using accelerate launch drops the VRAM use to ~9GB/GPU

[cg123]

Excited to see where this goes!

The main branch on my bitnet repo was a bit stale - when the official code got released I rewrote it a bit to be more in line with what the original authors did but didn't get around to merging it into main. Definitely safer to go with the official code.

[winglian]

Thanks @haeggee for updating the llm-baselines license!

[geronimi73]

🍿 @winglian you ever sleep?