Mixture of Thoughts (MoT)

A framework for orchestrating multiple Large Language Models (LLMs) using sparse top-k routing. MoT dynamically selects and combines outputs from multiple expert models based on input characteristics, enabling efficient and specialized model deployment.

The Mixture of Thoughts framework implements a multi-expert system where different language models act as specialized experts. A learned router dynamically selects the most appropriate experts for each input, combining their outputs through attention mechanisms and stack-based layer partitioning.

• We use Sparse Top-K Routing. You can plug and play any router as you please.
• Stack-Based Layer Partitioning: Divides model layers into stacks for fine-grained expert interaction
• Cross-Expert Attention: Enables information flow between selected experts
• Distributed Training: Supports multi-GPU training with DDP (Distributed Data Parallel)
• Multiple Benchmark Support: Compatible with MMLU, GSM8K, CMMLU, ARC-Challenge, and HumanEval datasets

• Python 3.8 or higher
• CUDA-capable GPU (recommended)
• PyTorch 2.0.0 or higher

1. Clone the repository:

git clone <repository-url>
cd mot

2. Install dependencies:

pip install -r requirements.txt

3. Download expert models (optional, will download automatically during first use):

python download_experts.py --config configs/routerdc_mot_config.json

mot/
├── mixture_of_thoughts.py    # Core MoT framework implementation
├── training.py               # Training utilities and loss functions
├── train_ddp.py             # Distributed training script
├── dataset_loaders.py       # Dataset loading utilities
├── utils.py                 # Helper functions and utilities
├── example.py               # Usage examples
├── configs/                 # Configuration files
│   └── routerdc_mot_config.json
├── experiments/             # Experiment results and checkpoints
├── logs/                    # Training logs
└── requirements.txt         # Package dependencies

Run a simple inference example:

python example.py --demo inference

python train_ddp.py --config configs/routerdc_mot_config.json

./run_experiments.sh --gpus 0,1,2,3 --exp-name mot_experiment

Or using torchrun directly:

torchrun --nproc_per_node=4 train_ddp.py --config configs/routerdc_mot_config.json

The framework is configured through JSON files. Key configuration parameters include:

{
  "experiment_name": "mot_experiment",
  "expert_models": [
    "model_name_1",
    "model_name_2"
  ],
  "mot_config": {
    "num_stacks": 4,
    "top_k": 3,
    "shared_dim": 768,
    "router_hidden_dim": 256,
    "interaction_heads": 8
  },
  "training": {
    "batch_size": 8,
    "learning_rate": 1e-4,
    "num_epochs": 10,
    "gradient_accumulation_steps": 4
  }
}

Basic usage example:

from mixture_of_thoughts import MixtureOfThoughts, MoTConfig
from transformers import AutoModelForCausalLM, AutoTokenizer

# Load expert models
expert_models = [
    AutoModelForCausalLM.from_pretrained("gpt2"),
    AutoModelForCausalLM.from_pretrained("distilgpt2")
]
tokenizers = [
    AutoTokenizer.from_pretrained("gpt2"),
    AutoTokenizer.from_pretrained("distilgpt2")
]

# Configure MoT
config = MoTConfig(
    num_stacks=4,
    top_k=2,
    shared_dim=768
)

# Initialize framework
mot_model = MixtureOfThoughts(
    expert_models=expert_models,
    tokenizers=tokenizers,
    config=config
)

# Run inference
input_ids = tokenizers[0]("Hello, world!", return_tensors="pt").input_ids
outputs = mot_model(input_ids=input_ids)

The framework includes loaders for the following benchmark datasets:

• MMLU: Massive Multitask Language Understanding
• GSM8K: Grade School Math 8K
• CMMLU: Chinese Massive Multitask Language Understanding
• ARC-Challenge: AI2 Reasoning Challenge
• HumanEval: Code generation benchmark

./run_experiments.sh [OPTIONS]

Options:
  -g, --gpus          GPU IDs to use (e.g., '0,1,2,3')
  -n, --num-gpus      Number of GPUs to use
  -c, --config        Path to configuration file
  -e, --exp-name      Experiment name for logging
  -w, --wandb-mode    WandB mode: online, offline, disabled
  -r, --resume        Path to checkpoint to resume from

python download_experts.py --config configs/routerdc_mot_config.json --max-workers 4

The MoT framework consists of several key components:

1. Router Network: A learnable MLP that assigns scores to each expert based on input embeddings
2. Expert Models: Pre-trained language models that serve as specialized experts
3. Stack Partitioning: Divides each expert into Q stacks of layers
4. Interaction Layers: Cross-attention mechanisms between selected experts
5. Output Aggregation: Combines expert outputs using learned weights

The framework uses multiple loss components:

• Primary task loss (e.g., language modeling)
• Router entropy regularization
• Load balancing loss for expert utilization
• Auxiliary expert-specific losses

• Models are cached in ~/.cache/huggingface/hub by default
• Supports 8-bit quantization for memory efficiency
• Implements gradient checkpointing for large models
• Uses mixed precision training (fp16/bf16) when available

• Training logs are saved to logs/ directory
• Experiment results and checkpoints stored in experiments/
• Supports Weights & Biases (wandb) integration for experiment tracking
• Real-time training metrics displayed during training

pytest tests/

black .
isort .
flake8 .

This project is licensed under the MIT License. See LICENSE file for details.

1. Out of Memory: Reduce batch size or enable gradient checkpointing
2. Model Download Failures: Check network connection and HuggingFace hub access
3. DDP Training Issues: Ensure all GPUs are visible and NCCL is properly installed

Enable debug logging:

import logging
logging.basicConfig(level=logging.DEBUG)

For questions or issues, please open an issue on the repository.