Mixture of Experts model

[deepsaia]

Thanks @karpathy for the amazing repo.
Translating the inspiration into a working prototype now.

Built a mixture of experts model with somewhat modular architecture:
https://github.com/deepsaia/moe_llama

• Example training run on tiny-shakespeare dataset timed at 6m 11s (on a M1Pro 16GB machine)
  • The response looks totally weird though - partly because of limited data and even the model architecture

Prompt: what say you

Generating...
================================================================================
Generated Text:
--------------------------------------------------------------------------------
what say you at : I shall be the night . QUEEN MARGARET : O , the duke , 
I have be no too . You ' s a man , we have I have thee . LADY LADY CAPULET : 
The father , I , I speak , my son . Second First Murderer : O , he shall I have 
your blood : What have not ' s a a man ' d to me , That have been the time . 
LUCIO : O , this . DUKE VINCENTIO : I I am this . I ' s the heart ! DUKE 
VINCENTIO : Ay , my brother , let me , let you am have the other s life : I 
' Tis ' s , I will be be your body . HENRY BOLINGBROKE : Why , so ? O , 
they ' ll see ' s with me . Why , and me , and I have been it will ; And so 
I think thou do my good your queen : And that you I know with a a man : 
And not , and a son ' d , let : I I am not
--------------------------------------------------------------------------------

Stats: 204 tokens, 27.13 tokens/sec
================================================================================

• Works on CPU, CUDA (NVIDIA), MPS (Apple Silicon)
• For training, we can define the structure via a hocon/json object:

  model {
    dim = 256              # Model dimension
    num_layers = 4         # Transformer layers
    num_heads = 8          # Attention heads
    num_experts = 8        # Number of experts
    top_k = 2              # Experts activated per token
    max_seq_len = 256      # Maximum sequence length
    dropout = 0.1          # Dropout rate
    shared_expert = true   # Include shared expert
    load_balancing_loss_coef = 0.01  # Load balancing weight
  }

• Things to do:
  • Train and test on more LLM-oriented datasets
  • Test different model architectures
  • Build an interactive chat ui
  • Lots of improvements