# xuefuzhao/openmoe architecture

> Architecture diagram of the GitHub repository xuefuzhao/openmoe, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-09-28). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/xuefuzhao/openmoe
- Repository: https://github.com/xuefuzhao/openmoe

## Overview

OpenMoE is an open-source Mixture-of-Experts language-model project: it releases model checkpoints and supports training, inference, and evaluation. The principal runtime path is a user prompt through the multi-device inference script to a loaded checkpoint and generated text; the sampled MoE layer implements sparse expert routing and distributed expert computation. Evaluation retrieves and normalizes BIG-bench scores before plotting comparisons, while TriQA plotting uses embedded data. Training is documented and has a launcher, but its implementation and wiring are not shown in the sampled source.

Main entry points: [`script/inference_on_multi_devices.py`](https://github.com/xuefuzhao/openmoe/blob/main/script/inference_on_multi_devices.py) (Multi-device inference), [`analysis/colossalai_replace/layer.py`](https://github.com/xuefuzhao/openmoe/blob/main/analysis/colossalai_replace/layer.py) (Sparse expert layer), [`eval/result_retrieval_bigbenchlite.py`](https://github.com/xuefuzhao/openmoe/blob/main/eval/result_retrieval_bigbenchlite.py) (Score retrieval), [`eval/triqa_plot.py`](https://github.com/xuefuzhao/openmoe/blob/main/eval/triqa_plot.py) (TriQA comparison plot)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_execution["Model execution"]
  node_inference["Multi-device inference"]
  node_checkpoint["Model checkpoint"]
  node_hub(("Hugging Face Hub"))
  node_tokenizer["Tokenizer"]
  node_model["Loaded MoE model"]
  node_sparsemlp["Sparse expert layer<br/>[layer.py]"]
  node_colossalai["ColossalAI MoE ops"]
  node_generated["Generated text"]
end

subgraph group_evaluation["Evaluation and analysis"]
  node_bigbench["BIG-bench results"]
  node_retrieval["Score retrieval"]
  node_csv["Retrieved score CSV"]
  node_bigbenchplot["BIG-bench plots<br/>[plot_bigbench.py]"]
  node_triqa["TriQA comparison plot<br/>[triqa_plot.py]"]
end

subgraph group_training["Training"]
  node_pretrain["Pretraining launcher<br/>[run_pretrain.sh]"]
  node_finetune["Evaluation launcher<br/>[run_eval.sh]"]
end

node_user(("User"))

node_user -->|"provides prompt"| node_inference
node_inference -->|"downloads checkpoint"| node_hub
node_hub -->|"returns weights"| node_checkpoint
node_inference -->|"loads weights"| node_checkpoint
node_inference -->|"loads tokenizer"| node_tokenizer
node_inference -->|"dispatches inference"| node_model
node_model -->|"uses layer"| node_sparsemlp
node_sparsemlp -->|"calls MoE ops"| node_colossalai
node_inference -->|"encodes prompt"| node_tokenizer
node_inference -->|"returns generation"| node_generated
node_bigbench -->|"supplies scores"| node_retrieval
node_retrieval -->|"writes scores"| node_csv
node_csv -->|"reads scores"| node_bigbenchplot
node_triqa -->|"plots embedded data"| node_triqa

click node_inference "https://github.com/xuefuzhao/openmoe/blob/main/script/inference_on_multi_devices.py"
click node_tokenizer "https://github.com/xuefuzhao/openmoe/blob/main/script/inference_on_multi_devices.py"
click node_model "https://github.com/xuefuzhao/openmoe/blob/main/script/inference_on_multi_devices.py"
click node_sparsemlp "https://github.com/xuefuzhao/openmoe/blob/main/analysis/colossalai_replace/layer.py"
click node_retrieval "https://github.com/xuefuzhao/openmoe/blob/main/eval/result_retrieval_bigbenchlite.py"
click node_bigbenchplot "https://github.com/xuefuzhao/openmoe/blob/main/eval/plot_bigbench.py"
click node_triqa "https://github.com/xuefuzhao/openmoe/blob/main/eval/triqa_plot.py"
click node_pretrain "https://github.com/xuefuzhao/openmoe/blob/main/script/run_pretrain.sh"
click node_finetune "https://github.com/xuefuzhao/openmoe/blob/main/script/run_eval.sh"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_inference,node_checkpoint,node_hub,node_tokenizer,node_model,node_sparsemlp,node_colossalai,node_generated,node_user toneBlue
class node_bigbench,node_retrieval,node_csv,node_bigbenchplot,node_triqa toneAmber
class node_pretrain,node_finetune toneMint
```

## Components

### Model execution

- [Multi-device inference](https://github.com/xuefuzhao/openmoe/blob/main/script/inference_on_multi_devices.py): `script/inference_on_multi_devices.py`
- **Model checkpoint**
- **Hugging Face Hub**
- [Tokenizer](https://github.com/xuefuzhao/openmoe/blob/main/script/inference_on_multi_devices.py): `script/inference_on_multi_devices.py`
- [Loaded MoE model](https://github.com/xuefuzhao/openmoe/blob/main/script/inference_on_multi_devices.py): `script/inference_on_multi_devices.py`
- [Sparse expert layer](https://github.com/xuefuzhao/openmoe/blob/main/analysis/colossalai_replace/layer.py): `analysis/colossalai_replace/layer.py`
- **ColossalAI MoE ops**
- **Generated text**

### Evaluation and analysis

- **BIG-bench results**
- [Score retrieval](https://github.com/xuefuzhao/openmoe/blob/main/eval/result_retrieval_bigbenchlite.py): `eval/result_retrieval_bigbenchlite.py`
- **Retrieved score CSV**
- [BIG-bench plots](https://github.com/xuefuzhao/openmoe/blob/main/eval/plot_bigbench.py): `eval/plot_bigbench.py`
- [TriQA comparison plot](https://github.com/xuefuzhao/openmoe/blob/main/eval/triqa_plot.py): `eval/triqa_plot.py`

### Training

- [Pretraining launcher](https://github.com/xuefuzhao/openmoe/blob/main/script/run_pretrain.sh): `script/run_pretrain.sh`
- [Evaluation launcher](https://github.com/xuefuzhao/openmoe/blob/main/script/run_eval.sh): `script/run_eval.sh`

### Other components

- **User**

## Connections

- User → Multi-device inference: provides prompt
- Multi-device inference → Hugging Face Hub: downloads checkpoint
- Hugging Face Hub → Model checkpoint: returns weights
- Multi-device inference → Model checkpoint: loads weights
- Multi-device inference → Tokenizer: loads tokenizer
- Multi-device inference → Loaded MoE model: dispatches inference
- Loaded MoE model → Sparse expert layer: uses layer
- Sparse expert layer → ColossalAI MoE ops: calls MoE ops
- Multi-device inference → Tokenizer: encodes prompt
- Multi-device inference → Generated text: returns generation
- BIG-bench results → Score retrieval: supplies scores
- Score retrieval → Retrieved score CSV: writes scores
- Retrieved score CSV → BIG-bench plots: reads scores
- TriQA comparison plot → TriQA comparison plot: plots embedded data