AnyMoE: Build an MoE model from anything, quickly

[EricLBuehler]

This PR implements AnyMoE, a method to build a flexible MoE model from any combination of fine-tuned expert models. Please see the paper for reference, although currently there are a few implementation-level differences.

The technique implemented here is similar to the gating method found in models such as Cephalo Vision 3x8b beta or Mixtral: it uses a gating layer to select experts in the MLP layer. This PR implements a built-in pretraining strategy with the option to save the resulting gating layer weights, powered by Candle's autograd infrastructure.

By allowing users to create an MoE model quickly by using the expert weights and a small pretraining dataset, AnyMoE lowers the barrier to entry for users who want to create and deploy tailored MoE models for their application.

Status

• AnyMoEPipeline
  • Add most of it
  • Pretraining infrastructure
  • Connect to the rest of mistral.rs
    • Python API
    • CLI
    • Server

[github-actions]

Code Metrics Report

  ===============================================================================
 Language            Files        Lines         Code     Comments       Blanks
===============================================================================
 Dockerfile              1           34           25            0            9
 Happy                   1          442          369            0           73
 JSON                    9           21           21            0            0
 Python                 33         1274         1089           37          148
 TOML                   16          445          403            2           40
-------------------------------------------------------------------------------
 Jupyter Notebooks       1            0            0            0            0
 |- Markdown             1           60           30           22            8
 |- Python               1           96           87            1            8
 (Total)                            156          117           23           16
-------------------------------------------------------------------------------
 Markdown               19         1373            0         1028          345
 |- BASH                 5          101           98            0            3
 |- Python               5           98           88            0           10
 |- Rust                 3          151          135            6           10
 (Total)                           1723          321         1034          368
-------------------------------------------------------------------------------
 Rust                  122        38838        35152          695         2991
 |- Markdown            70          682           13          631           38
 (Total)                          39520        35165         1326         3029
===============================================================================
 Total                 203        42427        37059         1762         3606
===============================================================================
  

[EricLBuehler]

cargo run --release --features cuda -- -i toml -f toml-selectors/anymoe.toml