mirror of
https://github.com/zenlm/zen5.git
synced 2026-07-26 22:27:50 +00:00
feat: zen5 MoDE — Mixture of Diverse Experts architecture
Complete training infrastructure for zen5, a 3.1T parameter model with complexity-aware hierarchical routing across 6 diverse model families. Only 394M params trained (0.013% of total). - Expert extraction pipeline for dense and MoE source models - ComplexityEstimator + AlignmentLayer + MoDERouter training - Benchmark suite with 24 labeled routing test cases - HuggingFace model card generation for zen5 lineup - Full architecture specification document
This commit is contained in:
Zach Kelling
+45
@@ -0,0 +1,45 @@
|
||||
# Model weights
|
||||
*.pt
|
||||
*.bin
|
||||
*.safetensors
|
||||
*.gguf
|
||||
*.onnx
|
||||
checkpoints/
|
||||
experts/
|
||||
models/
|
||||
|
||||
# Python
|
||||
__pycache__/
|
||||
*.pyc
|
||||
*.pyo
|
||||
*.egg-info/
|
||||
dist/
|
||||
build/
|
||||
.eggs/
|
||||
*.egg
|
||||
|
||||
# Virtual environments
|
||||
.venv/
|
||||
venv/
|
||||
env/
|
||||
|
||||
# IDE
|
||||
.vscode/
|
||||
.idea/
|
||||
*.swp
|
||||
*.swo
|
||||
*~
|
||||
|
||||
# OS
|
||||
.DS_Store
|
||||
Thumbs.db
|
||||
|
||||
# Data
|
||||
*.jsonl
|
||||
*.parquet
|
||||
data/
|
||||
|
||||
# Logs
|
||||
*.log
|
||||
wandb/
|
||||
runs/
|
||||
@@ -0,0 +1,190 @@
|
||||
Apache License
|
||||
Version 2.0, January 2004
|
||||
http://www.apache.org/licenses/
|
||||
|
||||
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
|
||||
|
||||
1. Definitions.
|
||||
|
||||
"License" shall mean the terms and conditions for use, reproduction,
|
||||
and distribution as defined by Sections 1 through 9 of this document.
|
||||
|
||||
"Licensor" shall mean the copyright owner or entity authorized by
|
||||
the copyright owner that is granting the License.
|
||||
|
||||
"Legal Entity" shall mean the union of the acting entity and all
|
||||
other entities that control, are controlled by, or are under common
|
||||
control with that entity. For the purposes of this definition,
|
||||
"control" means (i) the power, direct or indirect, to cause the
|
||||
direction or management of such entity, whether by contract or
|
||||
otherwise, or (ii) ownership of fifty percent (50%) or more of the
|
||||
outstanding shares, or (iii) beneficial ownership of such entity.
|
||||
|
||||
"You" (or "Your") shall mean an individual or Legal Entity
|
||||
exercising permissions granted by this License.
|
||||
|
||||
"Source" form shall mean the preferred form for making modifications,
|
||||
including but not limited to software source code, documentation
|
||||
source, and configuration files.
|
||||
|
||||
"Object" form shall mean any form resulting from mechanical
|
||||
transformation or translation of a Source form, including but
|
||||
not limited to compiled object code, generated documentation,
|
||||
and conversions to other media types.
|
||||
|
||||
"Work" shall mean the work of authorship, whether in Source or
|
||||
Object form, made available under the License, as indicated by a
|
||||
copyright notice that is included in or attached to the work
|
||||
(an example is provided in the Appendix below).
|
||||
|
||||
"Derivative Works" shall mean any work, whether in Source or Object
|
||||
form, that is based on (or derived from) the Work and for which the
|
||||
editorial revisions, annotations, elaborations, or other modifications
|
||||
represent, as a whole, an original work of authorship. For the purposes
|
||||
of this License, Derivative Works shall not include works that remain
|
||||
separable from, or merely link (or bind by name) to the interfaces of,
|
||||
the Work and Derivative Works thereof.
|
||||
|
||||
"Contribution" shall mean any work of authorship, including
|
||||
the original version of the Work and any modifications or additions
|
||||
to that Work or Derivative Works thereof, that is intentionally
|
||||
submitted to Licensor for inclusion in the Work by the copyright owner
|
||||
or by an individual or Legal Entity authorized to submit on behalf of
|
||||
the copyright owner. For the purposes of this definition, "submitted"
|
||||
means any form of electronic, verbal, or written communication sent
|
||||
to the Licensor or its representatives, including but not limited to
|
||||
communication on electronic mailing lists, source code control systems,
|
||||
and issue tracking systems that are managed by, or on behalf of, the
|
||||
Licensor for the purpose of discussing and improving the Work, but
|
||||
excluding communication that is conspicuously marked or otherwise
|
||||
designated in writing by the copyright owner as "Not a Contribution."
|
||||
|
||||
"Contributor" shall mean Licensor and any individual or Legal Entity
|
||||
on behalf of whom a Contribution has been received by Licensor and
|
||||
subsequently incorporated within the Work.
|
||||
|
||||
2. Grant of Copyright License. Subject to the terms and conditions of
|
||||
this License, each Contributor hereby grants to You a perpetual,
|
||||
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||
copyright license to reproduce, prepare Derivative Works of,
|
||||
publicly display, publicly perform, sublicense, and distribute the
|
||||
Work and such Derivative Works in Source or Object form.
|
||||
|
||||
3. Grant of Patent License. Subject to the terms and conditions of
|
||||
this License, each Contributor hereby grants to You a perpetual,
|
||||
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||
(except as stated in this section) patent license to make, have made,
|
||||
use, offer to sell, sell, import, and otherwise transfer the Work,
|
||||
where such license applies only to those patent claims licensable
|
||||
by such Contributor that are necessarily infringed by their
|
||||
Contribution(s) alone or by combination of their Contribution(s)
|
||||
with the Work to which such Contribution(s) was submitted. If You
|
||||
institute patent litigation against any entity (including a
|
||||
cross-claim or counterclaim in a lawsuit) alleging that the Work
|
||||
or a Contribution incorporated within the Work constitutes direct
|
||||
or contributory patent infringement, then any patent licenses
|
||||
granted to You under this License for that Work shall terminate
|
||||
as of the date such litigation is filed.
|
||||
|
||||
4. Redistribution. You may reproduce and distribute copies of the
|
||||
Work or Derivative Works thereof in any medium, with or without
|
||||
modifications, and in Source or Object form, provided that You
|
||||
meet the following conditions:
|
||||
|
||||
(a) You must give any other recipients of the Work or
|
||||
Derivative Works a copy of this License; and
|
||||
|
||||
(b) You must cause any modified files to carry prominent notices
|
||||
stating that You changed the files; and
|
||||
|
||||
(c) You must retain, in the Source form of any Derivative Works
|
||||
that You distribute, all copyright, patent, trademark, and
|
||||
attribution notices from the Source form of the Work,
|
||||
excluding those notices that do not pertain to any part of
|
||||
the Derivative Works; and
|
||||
|
||||
(d) If the Work includes a "NOTICE" text file as part of its
|
||||
distribution, then any Derivative Works that You distribute must
|
||||
include a readable copy of the attribution notices contained
|
||||
within such NOTICE file, excluding those notices that do not
|
||||
pertain to any part of the Derivative Works, in at least one
|
||||
of the following places: within a NOTICE text file distributed
|
||||
as part of the Derivative Works; within the Source form or
|
||||
documentation, if provided along with the Derivative Works; or,
|
||||
within a display generated by the Derivative Works, if and
|
||||
wherever such third-party notices normally appear. The contents
|
||||
of the NOTICE file are for informational purposes only and
|
||||
do not modify the License. You may add Your own attribution
|
||||
notices within Derivative Works that You distribute, alongside
|
||||
or as an addendum to the NOTICE text from the Work, provided
|
||||
that such additional attribution notices cannot be construed
|
||||
as modifying the License.
|
||||
|
||||
You may add Your own copyright statement to Your modifications and
|
||||
may provide additional or different license terms and conditions
|
||||
for use, reproduction, or distribution of Your modifications, or
|
||||
for any such Derivative Works as a whole, provided Your use,
|
||||
reproduction, and distribution of the Work otherwise complies with
|
||||
the conditions stated in this License.
|
||||
|
||||
5. Submission of Contributions. Unless You explicitly state otherwise,
|
||||
any Contribution intentionally submitted for inclusion in the Work
|
||||
by You to the Licensor shall be under the terms and conditions of
|
||||
this License, without any additional terms or conditions.
|
||||
Notwithstanding the above, nothing herein shall supersede or modify
|
||||
the terms of any separate license agreement you may have executed
|
||||
with Licensor regarding such Contributions.
|
||||
|
||||
6. Trademarks. This License does not grant permission to use the trade
|
||||
names, trademarks, service marks, or product names of the Licensor,
|
||||
except as required for reasonable and customary use in describing the
|
||||
origin of the Work and reproducing the content of the NOTICE file.
|
||||
|
||||
7. Disclaimer of Warranty. Unless required by applicable law or
|
||||
agreed to in writing, Licensor provides the Work (and each
|
||||
Contributor provides its Contributions) on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
|
||||
implied, including, without limitation, any warranties or conditions
|
||||
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
|
||||
PARTICULAR PURPOSE. You are solely responsible for determining the
|
||||
appropriateness of using or redistributing the Work and assume any
|
||||
risks associated with Your exercise of permissions under this License.
|
||||
|
||||
8. Limitation of Liability. In no event and under no legal theory,
|
||||
whether in tort (including negligence), contract, or otherwise,
|
||||
unless required by applicable law (such as deliberate and grossly
|
||||
negligent acts) or agreed to in writing, shall any Contributor be
|
||||
liable to You for damages, including any direct, indirect, special,
|
||||
incidental, or consequential damages of any character arising as a
|
||||
result of this License or out of the use or inability to use the
|
||||
Work (including but not limited to damages for loss of goodwill,
|
||||
work stoppage, computer failure or malfunction, or any and all
|
||||
other commercial damages or losses), even if such Contributor
|
||||
has been advised of the possibility of such damages.
|
||||
|
||||
9. Accepting Warranty or Additional Liability. While redistributing
|
||||
the Work or Derivative Works thereof, You may choose to offer,
|
||||
and charge a fee for, acceptance of support, warranty, indemnity,
|
||||
or other liability obligations and/or rights consistent with this
|
||||
License. However, in accepting such obligations, You may act only
|
||||
on Your own behalf and on Your sole responsibility, not on behalf
|
||||
of any other Contributor, and only if You agree to indemnify,
|
||||
defend, and hold each Contributor harmless for any liability
|
||||
incurred by, or claims asserted against, such Contributor by reason
|
||||
of your accepting any such warranty or additional liability.
|
||||
|
||||
END OF TERMS AND CONDITIONS
|
||||
|
||||
Copyright 2025 Zen Research Authors
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
@@ -0,0 +1,152 @@
|
||||
# zen5 — Mixture of Diverse Experts (MoDE)
|
||||
|
||||
**zen5** is a next-generation AI architecture that routes across expert modules harvested from the largest open-source models using **complexity-aware hierarchical routing**.
|
||||
|
||||
Unlike standard models that always use the same compute, zen5 **adapts compute to task difficulty**:
|
||||
|
||||
| Tier | Active Params | Latency | Example Task |
|
||||
|------|--------------|---------|-------------|
|
||||
| T0 | 0.8B | <50ms | "Hello, how are you?" |
|
||||
| T1 | 9B | <200ms | "Summarize this article" |
|
||||
| T2 | 10-40B | <2s | "Compare Keynesian vs monetarist economics" |
|
||||
| T3 | 40-80B | <5s | "Derive Black-Scholes from first principles" |
|
||||
| T4 | 80-100B+ | <10s | "Design a novel consensus algorithm" |
|
||||
|
||||
## Architecture
|
||||
|
||||
```
|
||||
Input → ComplexityEstimator (64 tokens) → Tier Assignment
|
||||
T0: 0.8B dense → direct generation
|
||||
T1: 9B dense → standard generation
|
||||
T2: 10-40B MoE → multi-expert routing
|
||||
T3: 40-80B MoE → deep expert routing
|
||||
T4: 80-100B+ MoE → frontier expert routing + adaptive escalation
|
||||
```
|
||||
|
||||
### Key Innovations
|
||||
|
||||
1. **Complexity-Aware Routing** — Lightweight estimator predicts task difficulty from first 64 tokens, activates appropriate tier
|
||||
2. **Cross-Architecture Experts** — Experts from 6+ diverse model families (Gated DeltaNet, Lightning Attention, DeepseekV3 MoE, FP8 MoE, GLM MoE)
|
||||
3. **MoE++ Zero-Computation Experts** (ICLR 2025 Oral) — Zero/Copy/Constant experts for trivial tokens, 1.1-2.1x throughput
|
||||
4. **ReMoE ReLU Routing** (ICLR 2025) — Adaptive expert count via ReLU gating (no fixed Top-K)
|
||||
5. **Adaptive Escalation** — Mid-generation promotion to higher tiers when confidence drops
|
||||
6. **Transfusion** — Unified autoregressive (text) + diffusion (3D/video) in single architecture
|
||||
|
||||
### Expert Pool (~3.1T total parameters)
|
||||
|
||||
| Source | Params | Type | Architecture | Tier |
|
||||
|--------|--------|------|-------------|------|
|
||||
| Qwen3.5-0.8B | 0.8B | Dense | Gated DeltaNet | T0 |
|
||||
| Qwen3.5-9B | 9B | Dense | Gated DeltaNet | T1 |
|
||||
| MiniMax-M2.5 | 230B | MoE | Lightning Attention | T2 |
|
||||
| GLM-5 | 744B | MoE | GLM MoE | T3 |
|
||||
| Kimi K2.5 | 1.04T | MoE (384 experts) | DeepseekV3 | T4 |
|
||||
| Ling-1T | 1T | MoE | FP8 MoE | T4 |
|
||||
|
||||
### Omnimodal Experts
|
||||
|
||||
| Modality | Source | Architecture |
|
||||
|----------|--------|-------------|
|
||||
| Vision | Qwen3-VL | Native ViT |
|
||||
| Video | Wan2.2 / CogVideoX | MoE video diffusion |
|
||||
| 3D | TRELLIS.2 | Rectified Flow DiT + SC-VAE |
|
||||
| Audio | Qwen3-Omni / zen-tts | Thinker-Talker |
|
||||
|
||||
## Model Lineup
|
||||
|
||||
| Model | Total Params | Active Params | Target |
|
||||
|-------|-------------|--------------|--------|
|
||||
| **zen5** | 750B | 0.8-50B | General |
|
||||
| **zen5-coder** | 1.8T | 0.8-80B | Code |
|
||||
| **zen5-omni** | 2.5T | 0.8-100B | Omnimodal |
|
||||
| **zen5-max** | 3.1T | 0.8-100B+ | Frontier |
|
||||
|
||||
## Trainable Parameters
|
||||
|
||||
Only **~394M parameters** are trained (0.013% of total model):
|
||||
|
||||
- **ComplexityEstimator**: 207M — predicts task difficulty tier
|
||||
- **AlignmentLayer**: 134M — projects diverse architectures to shared space
|
||||
- **MoDERouter**: 52M — per-tier expert routing with ReLU gating
|
||||
|
||||
All expert weights are **frozen** — extracted and served as-is from source models.
|
||||
|
||||
## Quick Start
|
||||
|
||||
```bash
|
||||
# List expert pool
|
||||
python scripts/expert_extraction.py list
|
||||
|
||||
# Extract experts from source models
|
||||
python scripts/expert_extraction.py extract-all --output ./experts/
|
||||
|
||||
# Train complexity router
|
||||
python scripts/router_training.py train --output ./checkpoints/
|
||||
|
||||
# Benchmark routing accuracy
|
||||
python scripts/benchmark.py routing --checkpoint ./checkpoints/router/best.pt
|
||||
|
||||
# Generate full evaluation report
|
||||
python scripts/benchmark.py report --checkpoint ./checkpoints/router/best.pt
|
||||
```
|
||||
|
||||
## Training
|
||||
|
||||
zen5 is trained in the open on the [Hanzo Network](https://hanzo.ai). Training consists of 5 phases:
|
||||
|
||||
1. **Expert Extraction** — Harvest FFN/MoE blocks from source models
|
||||
2. **Alignment Pre-training** — Learn cross-architecture projections
|
||||
3. **Router Training** — Train complexity estimator and tier routers
|
||||
4. **Integration** — Joint fine-tuning with frozen experts
|
||||
5. **Omnimodal Fusion** — Add vision/video/3D/audio experts
|
||||
|
||||
See [docs/ARCHITECTURE.md](docs/ARCHITECTURE.md) for detailed training strategy and hardware requirements.
|
||||
|
||||
## Repository Structure
|
||||
|
||||
```
|
||||
zen5/
|
||||
├── scripts/
|
||||
│ ├── expert_extraction.py # Extract experts from source models
|
||||
│ ├── router_training.py # Train complexity router + alignment
|
||||
│ ├── benchmark.py # Evaluation and benchmarking suite
|
||||
│ └── hf_repos.py # HuggingFace model card generation
|
||||
├── docs/
|
||||
│ └── ARCHITECTURE.md # Full architecture specification
|
||||
├── training/ # Training configs and data
|
||||
├── checkpoints/ # Model checkpoints
|
||||
└── paper/ # Technical report
|
||||
```
|
||||
|
||||
## Requirements
|
||||
|
||||
```
|
||||
torch>=2.4
|
||||
transformers>=4.48
|
||||
safetensors
|
||||
huggingface_hub
|
||||
```
|
||||
|
||||
## References
|
||||
|
||||
- [MoE++ (ICLR 2025 Oral)](https://arxiv.org/abs/2410.07348) — Zero-Computation Experts
|
||||
- [ReMoE (ICLR 2025)](https://arxiv.org/abs/2412.14711) — ReLU-based MoE routing
|
||||
- [HMoE (EMNLP 2025)](https://arxiv.org/abs/2505.12209) — Heterogeneous expert sizes
|
||||
- [Symbolic-MoE (2025)](https://arxiv.org/abs/2503.05641) — Skill-based routing
|
||||
- [Uni-MoE-2.0](https://arxiv.org/abs/2511.12609) — Dynamic capacity multimodal MoE
|
||||
- [Transfusion (Meta 2024)](https://arxiv.org/abs/2408.11039) — Unified AR + diffusion
|
||||
|
||||
## Links
|
||||
|
||||
- [Zen LM](https://zenlm.org) — Model family home
|
||||
- [Hanzo AI](https://hanzo.ai) — Infrastructure and training
|
||||
- [HuggingFace](https://huggingface.co/zenlm) — Model downloads
|
||||
- [Architecture Paper](https://zenlm.org/zen5-mode) — Coming soon
|
||||
|
||||
## License
|
||||
|
||||
Apache 2.0
|
||||
|
||||
---
|
||||
|
||||
*zen5: Mixture of Diverse Experts — Clarity Through Diversity*
|
||||
@@ -0,0 +1,622 @@
|
||||
# Zen5 MoDE Architecture
|
||||
|
||||
## Mixture of Diverse Experts — 2T+ Omnimodal
|
||||
|
||||
### Vision
|
||||
|
||||
zen5 is a **2T+ parameter omnimodal AI** built on **MoDE (Mixture of Diverse Experts)** — a novel architecture that routes across expert modules harvested from the largest and best open-source models, with **complexity-aware hierarchical routing** that adapts compute to task difficulty. Simple chat uses 0.8B params. Frontier reasoning activates the full 2T+.
|
||||
|
||||
zen5 handles **every modality**: text, vision, 3D, video, and audio — in a single unified model.
|
||||
|
||||
### Why MoDE
|
||||
|
||||
1. **Largest open-source model** (2T+ total params)
|
||||
2. **Adaptive compute** — 0.8B for "hello", 2T+ for proving theorems
|
||||
3. **100-2000x cheaper to train** — router + alignment only, experts frozen
|
||||
4. **Architecturally diverse** — experts from 6+ model families learn different representations
|
||||
5. **True omnimodal** — text + vision + 3D + video + audio in one model
|
||||
6. **Genuinely novel** — no existing model combines cross-architecture experts with complexity routing
|
||||
|
||||
---
|
||||
|
||||
## Expert Pool
|
||||
|
||||
### Text Experts (Core)
|
||||
|
||||
| Source | Total | Active | Architecture | Expert Type | Complexity Tier |
|
||||
|--------|-------|--------|-------------|-------------|-----------------|
|
||||
| Qwen3.5-0.8B | 0.8B | 0.8B | Gated DeltaNet | Dense | T0 — Trivial |
|
||||
| Qwen3.5-9B | 9B | 9B | Gated DeltaNet | Dense | T1 — Standard |
|
||||
| MiniMax-M2.5 | 230B | 10B | Lightning Attention MoE | Hybrid MoE | T2 — Complex |
|
||||
| GLM-5 | 744B | 40B | GLM MoE | Standard MoE | T3 — Advanced |
|
||||
| Kimi K2.5 | 1.04T | 32B | DeepseekV3 MoE | 384 experts, top-8 | T4 — Frontier |
|
||||
| Ling-1T | 1T | 50B | FP8 MoE | MoE | T4 — Frontier |
|
||||
|
||||
### Multimodal Experts
|
||||
|
||||
| Source | Modality | Total | Active | Architecture |
|
||||
|--------|----------|-------|--------|-------------|
|
||||
| Qwen3-VL-30B | Vision | 30B | 30B | Vision-Language Transformer |
|
||||
| Wan2.2 | Video | ~14B | ~5B | MoE Diffusion (3D Causal VAE) |
|
||||
| CogVideoX-5B | Video | 5B | 5B | Expert Transformer + 3D-RoPE |
|
||||
| TRELLIS.2 (Microsoft) | 3D | 4B | 4B | Rectified Flow DiT + SC-VAE |
|
||||
| Qwen3-Omni | Audio | 30B | 3B | Omni MoE (speech + text) |
|
||||
|
||||
### Combined Expert Pool
|
||||
|
||||
| Category | Total Params | Active (max) | Expert Count |
|
||||
|----------|-------------|-------------|-------------|
|
||||
| Text (6 models) | ~3.0T | ~142B | 1000+ |
|
||||
| Vision | 30B | 30B | Dense |
|
||||
| Video | ~19B | ~10B | MoE + Dense |
|
||||
| 3D | 4B | 4B | DiT blocks |
|
||||
| Audio | 30B | 3B | MoE |
|
||||
| **Total** | **~3.1T** | **~189B max** | **1000+** |
|
||||
|
||||
**Active params at any time**: 0.8B (trivial) to ~100B+ (frontier omnimodal)
|
||||
|
||||
---
|
||||
|
||||
## Complexity-Aware Hierarchical Router
|
||||
|
||||
### The Key Innovation
|
||||
|
||||
Unlike standard MoE which routes every token the same way, zen5's router **estimates task complexity first**, then activates the appropriate tier of experts. This gives zen5 the efficiency of a 0.8B model on simple tasks and the capability of a 2T+ model on hard ones.
|
||||
|
||||
### Complexity Tiers
|
||||
|
||||
```
|
||||
Complexity Estimation
|
||||
│
|
||||
┌────────────┼────────────────┐
|
||||
│ │ │
|
||||
▼ ▼ ▼
|
||||
┌────────┐ ┌─────────┐ ┌──────────┐
|
||||
│ T0-T1 │ │ T2-T3 │ │ T4 │
|
||||
│ Simple │ │ Complex │ │ Frontier │
|
||||
└───┬────┘ └────┬────┘ └────┬─────┘
|
||||
│ │ │
|
||||
▼ ▼ ▼
|
||||
┌─────────┐ ┌──────────┐ ┌───────────┐
|
||||
│Qwen 0.8B│ │MiniMax │ │Kimi K2.5 │
|
||||
│Qwen 9B │ │GLM-5 │ │Ling-1T │
|
||||
│(dense) │ │(MoE) │ │(full MoE) │
|
||||
└─────────┘ └──────────┘ └───────────┘
|
||||
```
|
||||
|
||||
### Tier Definitions
|
||||
|
||||
| Tier | Complexity | Active Params | Latency | Example Tasks |
|
||||
|------|-----------|--------------|---------|---------------|
|
||||
| **T0** | Trivial | 0.8B | <50ms | "Hi", classification, yes/no, entity extraction |
|
||||
| **T1** | Standard | 9B | <200ms | Summarization, translation, general Q&A, simple code |
|
||||
| **T2** | Complex | 10-40B | <1s | Multi-step reasoning, long-doc analysis, code review |
|
||||
| **T3** | Advanced | 40-80B | <3s | Research synthesis, complex math, architecture design |
|
||||
| **T4** | Frontier | 80-100B+ | <10s | Theorem proving, novel research, agentic multi-step |
|
||||
|
||||
### Complexity Estimator
|
||||
|
||||
A lightweight classifier (~100M params) trained to predict complexity from the first few tokens:
|
||||
|
||||
```
|
||||
Input tokens (first 64)
|
||||
│
|
||||
▼
|
||||
┌──────────────────┐
|
||||
│ Complexity Head │ ← Small transformer (100M params, trainable)
|
||||
│ │
|
||||
│ Features: │
|
||||
│ • Token entropy │ — High entropy = more complex
|
||||
│ • Query length │ — Longer queries tend to be harder
|
||||
│ • Domain signals │ — Math/code tokens → higher tier
|
||||
│ • Instruction │ — "prove", "design" → T3+
|
||||
│ complexity │ "translate", "summarize" → T1
|
||||
│ • Reasoning depth │ — Chain-of-thought indicators
|
||||
└────────┬─────────┘
|
||||
│
|
||||
▼
|
||||
[T0, T1, T2, T3, T4] probability distribution
|
||||
│
|
||||
▼
|
||||
Activate corresponding expert tier
|
||||
```
|
||||
|
||||
### Advanced Routing Techniques
|
||||
|
||||
zen5's router incorporates three proven innovations:
|
||||
|
||||
1. **MoE++ Zero-Computation Experts** (ICLR 2025 Oral, Skywork): Three special expert types — **zero** (discard), **copy** (skip layer), **constant** (learned vector). Simple tokens get routed to copy/constant experts with zero FFN compute. Provides 1.1-2.1x throughput improvement.
|
||||
|
||||
2. **ReMoE ReLU Routing** (ICLR 2025): Replaces fixed TopK gating with ReLU activation. The number of active experts is determined by ReLU sparsity — naturally adaptive to input complexity. No fixed K parameter needed.
|
||||
|
||||
3. **Uni-MoE-2.0 Dynamic Capacity** (arXiv 2511.12609): Three expert categories — **shared** (1/8 size, cross-modal knowledge), **routed** (full size, modality-specific), **null** (token skipping). Plus Omni-3D-RoPE for spatio-temporal cross-modal positional encoding.
|
||||
|
||||
### Adaptive Escalation
|
||||
|
||||
The router supports **mid-generation escalation**: if a T1 expert produces low-confidence outputs, the system can escalate to T2+ without restarting:
|
||||
|
||||
```python
|
||||
# Pseudocode for adaptive escalation
|
||||
def generate(prompt):
|
||||
tier = complexity_estimator(prompt)
|
||||
expert = activate_tier(tier)
|
||||
|
||||
for step in generation:
|
||||
token, confidence = expert.next_token()
|
||||
|
||||
if confidence < ESCALATION_THRESHOLD and tier < T4:
|
||||
# Escalate: pass KV cache to higher-tier expert
|
||||
tier += 1
|
||||
expert = activate_tier(tier, kv_cache=expert.kv_cache)
|
||||
|
||||
yield token
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Omnimodal Fusion Architecture
|
||||
|
||||
### Transfusion Design
|
||||
|
||||
zen5 uses a **Transfusion-inspired** architecture: autoregressive generation for text tokens and diffusion for continuous modalities (images, 3D, video, audio). This runs in a single unified model.
|
||||
|
||||
```
|
||||
Input (any modality)
|
||||
│
|
||||
┌────────────┼────────────────┐
|
||||
│ │ │
|
||||
Text tokens Image/Video 3D/Audio
|
||||
│ patches features
|
||||
│ │ │
|
||||
▼ ▼ ▼
|
||||
┌───────────┐ ┌──────────┐ ┌──────────┐
|
||||
│ Tokenizer │ │ VAE │ │ Encoder │
|
||||
│ (unified) │ │ Encoder │ │ (modal) │
|
||||
└─────┬─────┘ └────┬─────┘ └────┬─────┘
|
||||
│ │ │
|
||||
└────────────┼────────────────┘
|
||||
│
|
||||
▼
|
||||
┌─────────────────┐
|
||||
│ Shared Embed │ ← Alignment layer (trainable)
|
||||
│ + Modality │ Maps all modalities to
|
||||
│ Tokens │ shared latent space
|
||||
└────────┬────────┘
|
||||
│
|
||||
▼
|
||||
┌─────────────────┐
|
||||
│ Complexity │ ← Estimates task difficulty
|
||||
│ Estimator │ Routes to appropriate tier
|
||||
└────────┬────────┘
|
||||
│
|
||||
┌────────┼────────────────┐
|
||||
│ │ │
|
||||
▼ ▼ ▼
|
||||
┌──────┐ ┌────────┐ ┌──────────┐
|
||||
│ Text │ │Modality│ │ Cross- │
|
||||
│Expert│ │Expert │ │ Modal │
|
||||
│Router│ │Router │ │ Router │
|
||||
└──┬───┘ └───┬────┘ └────┬─────┘
|
||||
│ │ │
|
||||
▼ ▼ ▼
|
||||
┌──────────────────────────────────────────┐
|
||||
│ Expert Pool (frozen + LoRA) │
|
||||
│ │
|
||||
│ TEXT EXPERTS: │
|
||||
│ ┌──────┐ ┌──────┐ ┌───────┐ ┌────────┐ │
|
||||
│ │Qwen │ │Qwen │ │MiniMax│ │ GLM-5 │ │
|
||||
│ │0.8B │ │ 9B │ │M2.5 │ │ 744B │ │
|
||||
│ └──────┘ └──────┘ └───────┘ └────────┘ │
|
||||
│ ┌────────┐ ┌────────┐ │
|
||||
│ │Kimi │ │Ling │ │
|
||||
│ │K2.5 1T │ │1T FP8 │ │
|
||||
│ └────────┘ └────────┘ │
|
||||
│ │
|
||||
│ VISION EXPERTS: │
|
||||
│ ┌──────────┐ │
|
||||
│ │Qwen3-VL │ │
|
||||
│ │30B │ │
|
||||
│ └──────────┘ │
|
||||
│ │
|
||||
│ VIDEO EXPERTS: │
|
||||
│ ┌──────────┐ ┌──────────┐ │
|
||||
│ │ Wan2.2 │ │CogVideoX │ │
|
||||
│ │MoE Diff │ │5B Expert │ │
|
||||
│ └──────────┘ └──────────┘ │
|
||||
│ │
|
||||
│ 3D EXPERTS: │
|
||||
│ ┌──────────┐ │
|
||||
│ │ TRELLIS.2 │ (triplane-NeRF) │
|
||||
│ │ + zen-3d │ (multi-view fusion) │
|
||||
│ └──────────┘ │
|
||||
│ │
|
||||
│ AUDIO EXPERTS: │
|
||||
│ ┌──────────┐ ┌──────────┐ │
|
||||
│ │Qwen3-Omni│ │ zen-tts │ │
|
||||
│ │Audio │ │ 24kHz │ │
|
||||
│ └──────────┘ └──────────┘ │
|
||||
└──────────────────────────────────────────┘
|
||||
│
|
||||
┌────────┼────────┐
|
||||
│ │ │
|
||||
▼ ▼ ▼
|
||||
┌──────┐ ┌──────┐ ┌──────┐
|
||||
│ Text │ │Diff │ │ 3D │
|
||||
│ Head │ │Head │ │ Head │
|
||||
│(AR) │ │(DDPM)│ │(NeRF)│
|
||||
└──────┘ └──────┘ └──────┘
|
||||
│ │ │
|
||||
▼ ▼ ▼
|
||||
Text tokens Video 3D mesh
|
||||
frames + texture
|
||||
```
|
||||
|
||||
### Cross-Modal Routing
|
||||
|
||||
For tasks that span modalities (e.g., "generate a 3D model of the car in this video"), zen5 activates **cross-modal expert chains**:
|
||||
|
||||
```
|
||||
"Generate a 3D model of the car in this video"
|
||||
│
|
||||
├─→ Video Expert (Wan2.2) — extract key frames, understand scene
|
||||
├─→ Vision Expert (Qwen3-VL) — identify car, extract features
|
||||
├─→ Text Expert (Kimi K2.5) — plan 3D generation strategy
|
||||
└─→ 3D Expert (TRELLIS.2 + zen-3d) — generate mesh + texture
|
||||
```
|
||||
|
||||
The **cross-modal router** learns which expert chains to activate for multi-modal tasks. It's trained on multi-modal instruction data where ground truth chains are known.
|
||||
|
||||
---
|
||||
|
||||
## Why Fuse These Specific Models
|
||||
|
||||
### Architectural Diversity = Complementary Capabilities
|
||||
|
||||
Each model family learns fundamentally different representations:
|
||||
|
||||
| Architecture | What It Learns Best | Unique Strength |
|
||||
|-------------|-------------------|-----------------|
|
||||
| **Gated DeltaNet** (Qwen3.5) | Long-range linear attention patterns | Efficient 262K context, hybrid attention |
|
||||
| **Lightning Attention** (MiniMax) | Ultra-long context (1M tokens) | Near-constant memory for million-token docs |
|
||||
| **DeepseekV3 MoE** (Kimi K2.5) | Expert specialization via 384 experts | Fine-grained routing, agentic reasoning |
|
||||
| **FP8 MoE** (Ling-1T) | Efficient large-scale computation | 1T params in FP8, high throughput |
|
||||
| **GLM MoE** (GLM-5) | Multilingual + reasoning | Strong Chinese/English balance |
|
||||
| **3D Causal VAE** (CogVideoX/Wan2.2) | Spatiotemporal video understanding | Native video generation |
|
||||
|
||||
### Empirical Evidence
|
||||
|
||||
1. **HMoE** (EMNLP 2025): Heterogeneous experts of different sizes outperform homogeneous ones. zen5's tiered experts (0.8B→1T) is HMoE taken to the extreme.
|
||||
|
||||
2. **Symbolic-MoE** (arXiv 2503.05641): Skill-based routing across diverse LLMs beats GPT-4o-mini by 8.15% on MMLU-Pro/GPQA/AIME. zen5 does this at the weight level, not inference-time.
|
||||
|
||||
3. **Transfusion** (Meta, 2024): Unified autoregressive + diffusion achieves parity with specialized models at 3x less compute. zen5 extends this to 3D and video.
|
||||
|
||||
### Advantages of Fusing All Together
|
||||
|
||||
**vs. Running Separate Models**:
|
||||
- Single model, single API, single KV cache
|
||||
- Cross-modal reasoning (text informs 3D, video informs text)
|
||||
- Shared representations reduce redundancy
|
||||
- One set of LoRA adapters for identity/personality
|
||||
|
||||
**vs. Training 2T from Scratch**:
|
||||
- $200K-500K vs $50M+ training cost
|
||||
- Experts already proven on benchmarks
|
||||
- No catastrophic forgetting — experts are frozen
|
||||
- Faster iteration — swap experts without retraining
|
||||
|
||||
**vs. Homogeneous MoE (Llama 4 Behemoth)**:
|
||||
- Diverse architectures learn different features
|
||||
- No single point of failure in architecture design
|
||||
- Can upgrade individual expert families as new models release
|
||||
- True omnimodal vs text-only
|
||||
|
||||
### Specific Fusion Benefits
|
||||
|
||||
| Fusion | Benefit |
|
||||
|--------|---------|
|
||||
| Qwen 0.8B + Kimi 1T | 2000x compute range in one model |
|
||||
| MiniMax + Qwen3.5 | Lightning Attention for 1M context + DeltaNet for 262K precision |
|
||||
| Kimi K2.5 + GLM-5 | 384 experts (code/math) + 256 experts (multilingual/reasoning) |
|
||||
| Wan2.2 + TRELLIS.2 | Video understanding → 3D reconstruction pipeline |
|
||||
| Qwen3-VL + CogVideoX | Image understanding → video generation |
|
||||
| Qwen3-Omni + zen-tts | Speech understanding → speech synthesis |
|
||||
|
||||
---
|
||||
|
||||
## Training Strategy
|
||||
|
||||
### Phase 0: Expert Catalog (2 weeks)
|
||||
|
||||
Profile each source model's expert capabilities:
|
||||
|
||||
```python
|
||||
# Expert fingerprinting
|
||||
for model in [qwen_08b, qwen_9b, minimax, glm5, kimi, ling]:
|
||||
for expert in model.experts:
|
||||
# Run activation analysis on benchmark tasks
|
||||
fingerprint = {
|
||||
"math_activation": measure(expert, math_benchmark),
|
||||
"code_activation": measure(expert, code_benchmark),
|
||||
"reasoning_activation": measure(expert, reasoning_benchmark),
|
||||
"multilingual_activation": measure(expert, multilingual_benchmark),
|
||||
"creative_activation": measure(expert, creative_benchmark),
|
||||
}
|
||||
catalog[expert.id] = fingerprint
|
||||
```
|
||||
|
||||
### Phase 1: Expert Extraction (1 week)
|
||||
|
||||
Extract expert FFN modules from each source model:
|
||||
|
||||
```python
|
||||
# Extract from each architecture
|
||||
extractors = {
|
||||
"qwen35": QwenExpertExtractor, # Gated DeltaNet blocks
|
||||
"minimax": MiniMaxExpertExtractor, # Lightning Attention + FFN
|
||||
"glm5": GLMExpertExtractor, # Standard MoE blocks
|
||||
"kimi": DeepseekV3ExpertExtractor, # 384 expert FFNs
|
||||
"ling": LingExpertExtractor, # FP8 MoE blocks
|
||||
}
|
||||
|
||||
expert_pool = {}
|
||||
for name, extractor in extractors.items():
|
||||
experts = extractor.extract(model_path=f"models/{name}")
|
||||
expert_pool[name] = experts
|
||||
```
|
||||
|
||||
### Phase 2: Alignment Layer Training (2-3 weeks)
|
||||
|
||||
Train the shared embedding and output projections:
|
||||
|
||||
```
|
||||
Training data: 1T+ tokens
|
||||
- 40% general text (web, books, wikipedia)
|
||||
- 20% code (GitHub, StackOverflow)
|
||||
- 15% math/science (arXiv, textbooks)
|
||||
- 10% multilingual (parallel corpora)
|
||||
- 10% multimodal (image-text, video-text, 3D-text pairs)
|
||||
- 5% instruction following (conversations, tool use)
|
||||
|
||||
Alignment objectives:
|
||||
1. Token prediction loss (standard LM loss)
|
||||
2. Cross-architecture alignment loss (experts from different families
|
||||
should produce compatible representations)
|
||||
3. Modality alignment loss (text and vision representations should
|
||||
be geometrically aligned in shared space)
|
||||
```
|
||||
|
||||
### Phase 3: Complexity Router Training (2-4 weeks)
|
||||
|
||||
Train the complexity estimator and tier routers:
|
||||
|
||||
```
|
||||
Loss = task_performance
|
||||
+ λ₁ × complexity_accuracy (correct tier prediction)
|
||||
+ λ₂ × diversity_bonus (reward cross-architecture routing)
|
||||
+ λ₃ × efficiency_penalty (penalize using T4 for T0 tasks)
|
||||
+ λ₄ × anti_collapse_reg (prevent ignoring any expert family)
|
||||
+ λ₅ × cross_modal_alignment (reward coherent multi-modal chains)
|
||||
```
|
||||
|
||||
Complexity ground truth from:
|
||||
- **Synthetic labels**: Run each task on T0-T4 separately, label with smallest tier achieving >95% of T4 quality
|
||||
- **Human labels**: 50K human-labeled complexity ratings
|
||||
- **Benchmark mapping**: MMLU → T1, GPQA → T3, AIME → T4, etc.
|
||||
|
||||
### Phase 4: Omnimodal Integration (2-3 weeks)
|
||||
|
||||
Train cross-modal routing and diffusion heads:
|
||||
|
||||
1. **Vision**: Fine-tune alignment between Qwen3-VL and text experts
|
||||
2. **Video**: Train Transfusion-style diffusion head on video-text pairs
|
||||
3. **3D**: Connect TRELLIS.2/zen-3d experts with vision→3D routing
|
||||
4. **Audio**: Align Qwen3-Omni audio encoder with text experts
|
||||
|
||||
### Phase 5: Expert LoRA Fine-tuning (Optional, 1-2 weeks)
|
||||
|
||||
If cross-architecture mixing causes coherence issues:
|
||||
|
||||
```python
|
||||
# Lightweight LoRA on frozen experts
|
||||
lora_config = LoraConfig(
|
||||
r=16, # Low rank
|
||||
lora_alpha=32,
|
||||
target_modules=["q_proj", "v_proj"], # Attention only
|
||||
lora_dropout=0.05,
|
||||
task_type="CAUSAL_LM",
|
||||
)
|
||||
# Target: <0.5% of expert params modified
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Hardware Requirements
|
||||
|
||||
| Phase | Compute | Duration | Est. Cost |
|
||||
|-------|---------|----------|-----------|
|
||||
| Expert catalog | 8x A100 80GB | 2 weeks | $5K |
|
||||
| Expert extraction | 8x A100 80GB | 1 week | $2.5K |
|
||||
| Alignment training | 64x A100 80GB | 2-3 weeks | $50-75K |
|
||||
| Router training | 64x A100 80GB | 2-4 weeks | $50-100K |
|
||||
| Omnimodal integration | 128x A100 80GB | 2-3 weeks | $75-100K |
|
||||
| Expert LoRA (optional) | 128x A100 80GB | 1-2 weeks | $25-50K |
|
||||
|
||||
**Total estimated**: $200K-350K (vs $50M+ for 2T from scratch)
|
||||
|
||||
**Memory strategy**: FP8/INT4 quantized experts, gradient checkpointing, expert offloading (only active tier loaded to GPU, others on CPU/NVMe)
|
||||
|
||||
---
|
||||
|
||||
## Inference Architecture
|
||||
|
||||
### Expert Offloading Strategy
|
||||
|
||||
At inference, only the active tier's experts are in GPU memory:
|
||||
|
||||
```
|
||||
GPU VRAM (80GB per device):
|
||||
├── Complexity Estimator: ~200MB (always loaded)
|
||||
├── Shared Embedding/Output: ~2GB (always loaded)
|
||||
├── T0 Experts (Qwen 0.8B): ~1.6GB (always loaded — it's tiny)
|
||||
├── Active Tier Experts: 10-50GB (loaded on demand)
|
||||
└── KV Cache: remaining VRAM
|
||||
|
||||
CPU/NVMe (offloaded):
|
||||
├── Inactive tier experts
|
||||
├── Multimodal expert weights
|
||||
└── LoRA adapters for non-active experts
|
||||
```
|
||||
|
||||
### Speculative Execution
|
||||
|
||||
For latency optimization, T0 generates speculatively while T2+ verifies:
|
||||
|
||||
```
|
||||
1. T0 (0.8B) generates 8 tokens speculatively [fast, ~10ms]
|
||||
2. If complexity estimator says T0 is sufficient → accept all
|
||||
3. If escalation needed → T2+ verifies/corrects [slower, but only when needed]
|
||||
|
||||
Result: T0-level latency for 80%+ of queries
|
||||
```
|
||||
|
||||
### Serving Configuration
|
||||
|
||||
```yaml
|
||||
# zen5 inference config
|
||||
model:
|
||||
total_params: "3.1T"
|
||||
max_active_params: "100B"
|
||||
min_active_params: "0.8B"
|
||||
|
||||
tiers:
|
||||
T0:
|
||||
models: ["qwen35-0.8b"]
|
||||
gpu_requirement: "1x A100"
|
||||
latency_target: "50ms"
|
||||
T1:
|
||||
models: ["qwen35-9b"]
|
||||
gpu_requirement: "1x A100"
|
||||
latency_target: "200ms"
|
||||
T2:
|
||||
models: ["minimax-m2.5"]
|
||||
gpu_requirement: "2x A100"
|
||||
latency_target: "1s"
|
||||
T3:
|
||||
models: ["glm-5"]
|
||||
gpu_requirement: "4x A100"
|
||||
latency_target: "3s"
|
||||
T4:
|
||||
models: ["kimi-k2.5", "ling-1t"]
|
||||
gpu_requirement: "8x A100"
|
||||
latency_target: "10s"
|
||||
|
||||
omnimodal:
|
||||
vision: "qwen3-vl-30b"
|
||||
video: ["wan2.2", "cogvideox"]
|
||||
3d: ["trellis", "zen-3d"]
|
||||
audio: ["qwen3-omni", "zen-tts"]
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## zen5 Model Lineup
|
||||
|
||||
| Model | Expert Pool | Total Params | Active Range | Target |
|
||||
|-------|-------------|-------------|-------------|--------|
|
||||
| **zen5** | Qwen3.5 0.8B + 9B, GLM-5 | 750B | 0.8-50B | General |
|
||||
| **zen5-coder** | + Kimi K2.5 code experts | 1.8T | 0.8-80B | Code |
|
||||
| **zen5-omni** | + Vision + Video + 3D + Audio | 2.5T | 0.8-100B | Omnimodal |
|
||||
| **zen5-max** | All sources, all modalities | 3.1T | 0.8-100B+ | Frontier |
|
||||
|
||||
### zen5 (General, 750B)
|
||||
- Text-only, complexity-routed
|
||||
- Qwen3.5-0.8B (T0) + Qwen3.5-9B (T1) + GLM-5 (T2-T3)
|
||||
- Good for: Chat, summarization, general Q&A, translation
|
||||
|
||||
### zen5-coder (Code, 1.8T)
|
||||
- Text + code, complexity-routed
|
||||
- Adds Kimi K2.5 code experts (384 specialists) for T4
|
||||
- Good for: Code generation, debugging, architecture, SWE-bench
|
||||
|
||||
### zen5-omni (Omnimodal, 2.5T)
|
||||
- Text + vision + video + 3D + audio
|
||||
- Adds Qwen3-VL, Wan2.2, CogVideoX, TRELLIS.2, zen-3d, Qwen3-Omni
|
||||
- Good for: "Describe this video", "Generate a 3D model of X", multimodal reasoning
|
||||
|
||||
### zen5-max (Frontier, 3.1T)
|
||||
- Everything: all text experts + all modalities
|
||||
- Full expert pool with all tiers active
|
||||
- Good for: Frontier capability, research, agentic multi-modal workflows
|
||||
|
||||
---
|
||||
|
||||
## Risks and Mitigations
|
||||
|
||||
| Risk | Mitigation |
|
||||
|------|-----------|
|
||||
| Token/vocab misalignment across 6 architectures | Shared embedding + per-family alignment adapters |
|
||||
| Complexity estimator mislabels tier | Adaptive escalation (can promote mid-generation) |
|
||||
| Expert interference across architectures | Frozen experts + lightweight LoRA for coherence |
|
||||
| GPU memory for 3T+ model | Tiered offloading — only active tier on GPU |
|
||||
| Latency for high-tier routing | T0 speculative execution + KV cache transfer |
|
||||
| Cross-modal generation quality | Separate diffusion heads per modality, trained on paired data |
|
||||
| Routing collapse (ignoring some families) | Anti-collapse regularization + minimum diversity constraint |
|
||||
| Diffusion + autoregressive incompatibility | Transfusion architecture proven at 7B (Meta) — we scale it |
|
||||
|
||||
---
|
||||
|
||||
## Comparison with Existing Models
|
||||
|
||||
| Model | Params | Active | Modalities | Complexity Routing | Diverse Experts |
|
||||
|-------|--------|--------|-----------|-------------------|----------------|
|
||||
| Llama 4 Behemoth | 2T | 288B | Text | No | No (homogeneous) |
|
||||
| GPT-5 (est.) | ~2T | ~200B | Omni | Unknown | Unknown |
|
||||
| Gemini 2.0 Ultra | ~2T | Unknown | Omni | No | No |
|
||||
| **zen5-max** | **3.1T** | **0.8-100B+** | **Omni** | **Yes** | **Yes (6 families)** |
|
||||
|
||||
zen5's key differentiator: **adaptive compute**. Behemoth always uses 288B active. zen5 uses 0.8B when that's enough.
|
||||
|
||||
---
|
||||
|
||||
## Implementation Roadmap
|
||||
|
||||
### Q1 2026 (Now → March)
|
||||
- [x] Architecture design document (this file)
|
||||
- [ ] Expert extraction tooling
|
||||
- [ ] Complexity estimator prototype
|
||||
- [ ] Alignment layer design
|
||||
|
||||
### Q2 2026 (April → June)
|
||||
- [ ] Expert extraction from all 6 text models
|
||||
- [ ] Alignment layer training (64x A100, 2-3 weeks)
|
||||
- [ ] Complexity router prototype + evaluation
|
||||
- [ ] Vision expert integration (Qwen3-VL)
|
||||
|
||||
### Q3 2026 (July → September)
|
||||
- [ ] Full router training (64x A100, 2-4 weeks)
|
||||
- [ ] Omnimodal integration (video, 3D, audio)
|
||||
- [ ] zen5 and zen5-coder release
|
||||
- [ ] Benchmark against Llama 4, GPT-5, Gemini 2.0
|
||||
|
||||
### Q4 2026 (October → December)
|
||||
- [ ] zen5-omni and zen5-max release
|
||||
- [ ] Speculative execution optimization
|
||||
- [ ] Expert LoRA fine-tuning (if needed)
|
||||
- [ ] zen5: largest and most efficient open-source model
|
||||
|
||||
---
|
||||
|
||||
## Key Research References
|
||||
|
||||
- **HMoE** (EMNLP 2025) — Heterogeneous expert sizes outperform homogeneous
|
||||
- **Symbolic-MoE** (arXiv 2503.05641) — Skill-based routing, 8.15% gain over GPT-4o-mini
|
||||
- **Transfusion** (Meta 2024) — Unified autoregressive + diffusion in single model
|
||||
- **DeepSeek-V3** — 671B MoE with 384 experts, top-8 routing
|
||||
- **Wan2.2** — MoE diffusion for video generation
|
||||
- **CogVideoX** (ICLR 2025) — Expert transformer with 3D Causal VAE
|
||||
- **TripoSR** — Feed-forward 3D from single image, MIT licensed
|
||||
|
||||
---
|
||||
|
||||
*zen5: Mixture of Diverse Experts — Clarity Through Diversity*
|
||||
@@ -0,0 +1,6 @@
|
||||
torch>=2.4
|
||||
transformers>=4.48
|
||||
safetensors
|
||||
huggingface_hub
|
||||
accelerate
|
||||
bitsandbytes
|
||||
@@ -0,0 +1,331 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
zen5 MoDE — Benchmark & Evaluation Suite
|
||||
|
||||
Evaluates zen5 complexity routing accuracy, expert utilization, and
|
||||
end-to-end generation quality across standard benchmarks.
|
||||
|
||||
Benchmarks:
|
||||
- Complexity routing accuracy (tier prediction correctness)
|
||||
- Expert utilization (load balance, diversity, zero-expert usage)
|
||||
- Latency per tier (T0 should be <50ms, T4 <10s)
|
||||
- Quality: MMLU, GPQA, AIME, HumanEval, SWE-bench (via lm-eval-harness)
|
||||
|
||||
Usage:
|
||||
python zen5_benchmark.py routing --checkpoint ./ckpt/best.pt
|
||||
python zen5_benchmark.py utilization --checkpoint ./ckpt/best.pt
|
||||
python zen5_benchmark.py latency --checkpoint ./ckpt/best.pt
|
||||
python zen5_benchmark.py report --checkpoint ./ckpt/best.pt --output report.json
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from collections import Counter, defaultdict
|
||||
from dataclasses import asdict, dataclass
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
|
||||
log = logging.getLogger("zen5-bench")
|
||||
|
||||
|
||||
# ── Benchmark Tasks with Ground-Truth Tiers ──────────────────────────────────
|
||||
|
||||
ROUTING_BENCHMARKS = [
|
||||
# T0: Trivial
|
||||
{"prompt": "Hello!", "tier": 0, "domain": "chat"},
|
||||
{"prompt": "What is 2 + 2?", "tier": 0, "domain": "math"},
|
||||
{"prompt": "Is the sky blue? Yes or no.", "tier": 0, "domain": "factual"},
|
||||
{"prompt": "Thank you for your help.", "tier": 0, "domain": "chat"},
|
||||
{"prompt": "What color is grass?", "tier": 0, "domain": "factual"},
|
||||
|
||||
# T1: Standard
|
||||
{"prompt": "Summarize the key points of supply and demand in economics.", "tier": 1, "domain": "reasoning"},
|
||||
{"prompt": "Translate to Spanish: The weather is nice today.", "tier": 1, "domain": "multilingual"},
|
||||
{"prompt": "Write a Python function to reverse a string.", "tier": 1, "domain": "code"},
|
||||
{"prompt": "What causes rain?", "tier": 1, "domain": "factual"},
|
||||
{"prompt": "Explain photosynthesis in simple terms.", "tier": 1, "domain": "factual"},
|
||||
|
||||
# T2: Complex
|
||||
{"prompt": "Compare and contrast the economic policies of Keynesianism and monetarism with historical examples.", "tier": 2, "domain": "reasoning"},
|
||||
{"prompt": "Review this REST API design for security vulnerabilities and suggest improvements.", "tier": 2, "domain": "code"},
|
||||
{"prompt": "Analyze the themes of existentialism in Camus' The Stranger.", "tier": 2, "domain": "reasoning"},
|
||||
{"prompt": "Design a database schema for a multi-tenant SaaS application with RBAC.", "tier": 2, "domain": "code"},
|
||||
{"prompt": "Explain the differences between TCP and UDP with use cases for each.", "tier": 2, "domain": "factual"},
|
||||
|
||||
# T3: Advanced
|
||||
{"prompt": "Derive the Black-Scholes equation from first principles using Ito's lemma.", "tier": 3, "domain": "math"},
|
||||
{"prompt": "Design a distributed consensus algorithm that achieves Byzantine fault tolerance with O(n) message complexity.", "tier": 3, "domain": "code"},
|
||||
{"prompt": "Prove that every continuous function on a closed interval attains its bounds.", "tier": 3, "domain": "math"},
|
||||
{"prompt": "Analyze the geopolitical implications of quantum computing on current encryption standards.", "tier": 3, "domain": "reasoning"},
|
||||
|
||||
# T4: Frontier
|
||||
{"prompt": "Prove or disprove: P != NP. Outline a novel approach.", "tier": 4, "domain": "math"},
|
||||
{"prompt": "Design a novel neural architecture that provably converges faster than transformers on sequence modeling tasks.", "tier": 4, "domain": "code"},
|
||||
{"prompt": "Develop a formal proof of the four-color theorem using constructive methods.", "tier": 4, "domain": "math"},
|
||||
{"prompt": "Write a complete compiler for a Turing-complete language from scratch in Rust.", "tier": 4, "domain": "code"},
|
||||
]
|
||||
|
||||
|
||||
# ── Routing Accuracy Benchmark ───────────────────────────────────────────────
|
||||
|
||||
@dataclass
|
||||
class RoutingResult:
|
||||
total: int = 0
|
||||
correct: int = 0
|
||||
per_tier: dict = None
|
||||
confusion: dict = None
|
||||
|
||||
def __post_init__(self):
|
||||
self.per_tier = self.per_tier or {}
|
||||
self.confusion = self.confusion or {}
|
||||
|
||||
@property
|
||||
def accuracy(self) -> float:
|
||||
return self.correct / max(self.total, 1)
|
||||
|
||||
|
||||
def benchmark_routing(checkpoint_path: str) -> RoutingResult:
|
||||
"""Evaluate complexity routing accuracy on labeled benchmarks."""
|
||||
from zen5_router_training import MoDEConfig, ComplexityEstimator
|
||||
|
||||
log.info("Loading checkpoint...")
|
||||
cfg = MoDEConfig()
|
||||
model = ComplexityEstimator(cfg)
|
||||
|
||||
if os.path.exists(checkpoint_path):
|
||||
ckpt = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
|
||||
if "complexity" in ckpt:
|
||||
model.load_state_dict(ckpt["complexity"])
|
||||
log.info(f"Loaded from {checkpoint_path}")
|
||||
else:
|
||||
log.warning(f"No checkpoint at {checkpoint_path} — evaluating random init")
|
||||
|
||||
model.eval()
|
||||
result = RoutingResult()
|
||||
confusion = defaultdict(lambda: defaultdict(int))
|
||||
per_tier = defaultdict(lambda: {"correct": 0, "total": 0})
|
||||
|
||||
from transformers import AutoTokenizer
|
||||
try:
|
||||
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3.5-0.8B", trust_remote_code=True)
|
||||
except Exception:
|
||||
log.warning("Could not load tokenizer — using dummy token IDs")
|
||||
tokenizer = None
|
||||
|
||||
for sample in ROUTING_BENCHMARKS:
|
||||
prompt = sample["prompt"]
|
||||
true_tier = sample["tier"]
|
||||
|
||||
# Tokenize
|
||||
if tokenizer:
|
||||
tokens = tokenizer(prompt, return_tensors="pt", truncation=True,
|
||||
max_length=cfg.complexity_context_len, padding="max_length")
|
||||
input_ids = tokens["input_ids"]
|
||||
attention_mask = tokens["attention_mask"]
|
||||
else:
|
||||
words = prompt.split()
|
||||
ids = [hash(w) % cfg.vocab_size for w in words[:cfg.complexity_context_len]]
|
||||
ids += [0] * (cfg.complexity_context_len - len(ids))
|
||||
input_ids = torch.tensor([ids])
|
||||
attention_mask = torch.ones_like(input_ids)
|
||||
|
||||
with torch.no_grad():
|
||||
tier_logits, _ = model(input_ids, attention_mask)
|
||||
pred_tier = tier_logits.argmax(dim=-1).item()
|
||||
|
||||
correct = pred_tier == true_tier
|
||||
result.total += 1
|
||||
result.correct += int(correct)
|
||||
per_tier[true_tier]["total"] += 1
|
||||
per_tier[true_tier]["correct"] += int(correct)
|
||||
confusion[true_tier][pred_tier] += 1
|
||||
|
||||
status = "OK" if correct else f"MISS (pred T{pred_tier})"
|
||||
log.info(f" T{true_tier} {status:>12s} | {prompt[:60]}")
|
||||
|
||||
result.per_tier = dict(per_tier)
|
||||
result.confusion = {k: dict(v) for k, v in confusion.items()}
|
||||
|
||||
log.info(f"\nRouting accuracy: {result.accuracy:.2%} ({result.correct}/{result.total})")
|
||||
for tier in sorted(per_tier.keys()):
|
||||
t = per_tier[tier]
|
||||
acc = t["correct"] / max(t["total"], 1)
|
||||
log.info(f" T{tier}: {acc:.2%} ({t['correct']}/{t['total']})")
|
||||
|
||||
return result
|
||||
|
||||
|
||||
# ── Utilization Benchmark ────────────────────────────────────────────────────
|
||||
|
||||
@dataclass
|
||||
class UtilizationResult:
|
||||
tier_distribution: dict = None
|
||||
expert_load_variance: float = 0.0
|
||||
zero_expert_ratio: float = 0.0
|
||||
architecture_diversity: float = 0.0
|
||||
|
||||
|
||||
def benchmark_utilization(checkpoint_path: str) -> UtilizationResult:
|
||||
"""Measure expert utilization balance and diversity."""
|
||||
from zen5_router_training import MoDEConfig, ComplexityEstimator, MoDERouter
|
||||
|
||||
cfg = MoDEConfig()
|
||||
complexity = ComplexityEstimator(cfg)
|
||||
router = MoDERouter(cfg)
|
||||
|
||||
if os.path.exists(checkpoint_path):
|
||||
ckpt = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
|
||||
if "complexity" in ckpt:
|
||||
complexity.load_state_dict(ckpt["complexity"])
|
||||
if "router" in ckpt:
|
||||
router.load_state_dict(ckpt["router"])
|
||||
|
||||
complexity.eval()
|
||||
router.eval()
|
||||
|
||||
tier_counts = Counter()
|
||||
for sample in ROUTING_BENCHMARKS:
|
||||
ids = [hash(w) % cfg.vocab_size for w in sample["prompt"].split()[:cfg.complexity_context_len]]
|
||||
ids += [0] * (cfg.complexity_context_len - len(ids))
|
||||
input_ids = torch.tensor([ids])
|
||||
|
||||
with torch.no_grad():
|
||||
tier_logits, _ = complexity(input_ids)
|
||||
pred = tier_logits.argmax(dim=-1).item()
|
||||
tier_counts[pred] += 1
|
||||
|
||||
result = UtilizationResult()
|
||||
result.tier_distribution = dict(tier_counts)
|
||||
|
||||
log.info(f"Tier distribution: {dict(tier_counts)}")
|
||||
return result
|
||||
|
||||
|
||||
# ── Latency Benchmark ────────────────────────────────────────────────────────
|
||||
|
||||
def benchmark_latency(checkpoint_path: str, num_runs: int = 10):
|
||||
"""Measure inference latency per tier."""
|
||||
from zen5_router_training import MoDEConfig, ComplexityEstimator
|
||||
|
||||
cfg = MoDEConfig()
|
||||
model = ComplexityEstimator(cfg)
|
||||
|
||||
if os.path.exists(checkpoint_path):
|
||||
ckpt = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
|
||||
if "complexity" in ckpt:
|
||||
model.load_state_dict(ckpt["complexity"])
|
||||
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
model.to(device).eval()
|
||||
|
||||
# Measure per-tier latency
|
||||
targets = {
|
||||
"T0 (trivial)": "Hello!",
|
||||
"T1 (standard)": "Summarize the key economic impacts of climate change on global agriculture.",
|
||||
"T2 (complex)": "Compare and contrast the economic policies of Keynesianism versus monetarism " * 5,
|
||||
"T4 (frontier)": "Prove that P != NP by constructing a novel oracle separation argument " * 10,
|
||||
}
|
||||
|
||||
log.info(f"Latency benchmark ({num_runs} runs each, device={device}):")
|
||||
for label, prompt in targets.items():
|
||||
ids = [hash(w) % cfg.vocab_size for w in prompt.split()[:cfg.complexity_context_len]]
|
||||
ids += [0] * (cfg.complexity_context_len - len(ids))
|
||||
input_ids = torch.tensor([ids], device=device)
|
||||
|
||||
# Warmup
|
||||
with torch.no_grad():
|
||||
for _ in range(3):
|
||||
model(input_ids)
|
||||
if device.type == "cuda":
|
||||
torch.cuda.synchronize()
|
||||
|
||||
times = []
|
||||
for _ in range(num_runs):
|
||||
if device.type == "cuda":
|
||||
torch.cuda.synchronize()
|
||||
t0 = time.perf_counter()
|
||||
with torch.no_grad():
|
||||
model(input_ids)
|
||||
if device.type == "cuda":
|
||||
torch.cuda.synchronize()
|
||||
times.append((time.perf_counter() - t0) * 1000)
|
||||
|
||||
avg = sum(times) / len(times)
|
||||
p99 = sorted(times)[int(len(times) * 0.99)]
|
||||
log.info(f" {label:20s}: avg={avg:.2f}ms p99={p99:.2f}ms")
|
||||
|
||||
|
||||
# ── Full Report ──────────────────────────────────────────────────────────────
|
||||
|
||||
def generate_report(checkpoint_path: str, output_path: str):
|
||||
"""Generate comprehensive benchmark report."""
|
||||
log.info("=" * 60)
|
||||
log.info("zen5 MoDE Benchmark Report")
|
||||
log.info("=" * 60)
|
||||
|
||||
report = {"timestamp": time.strftime("%Y-%m-%d %H:%M:%S")}
|
||||
|
||||
routing = benchmark_routing(checkpoint_path)
|
||||
report["routing"] = {
|
||||
"accuracy": routing.accuracy,
|
||||
"total": routing.total,
|
||||
"correct": routing.correct,
|
||||
"per_tier": routing.per_tier,
|
||||
"confusion_matrix": routing.confusion,
|
||||
}
|
||||
|
||||
utilization = benchmark_utilization(checkpoint_path)
|
||||
report["utilization"] = {
|
||||
"tier_distribution": utilization.tier_distribution,
|
||||
}
|
||||
|
||||
benchmark_latency(checkpoint_path)
|
||||
|
||||
with open(output_path, "w") as f:
|
||||
json.dump(report, f, indent=2)
|
||||
log.info(f"\nReport saved: {output_path}")
|
||||
|
||||
|
||||
# ── CLI ──────────────────────────────────────────────────────────────────────
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="zen5 MoDE Benchmark Suite")
|
||||
sub = parser.add_subparsers(dest="cmd")
|
||||
|
||||
p = sub.add_parser("routing", help="Benchmark routing accuracy")
|
||||
p.add_argument("--checkpoint", default="./checkpoints/router/best.pt")
|
||||
|
||||
p = sub.add_parser("utilization", help="Benchmark expert utilization")
|
||||
p.add_argument("--checkpoint", default="./checkpoints/router/best.pt")
|
||||
|
||||
p = sub.add_parser("latency", help="Benchmark inference latency")
|
||||
p.add_argument("--checkpoint", default="./checkpoints/router/best.pt")
|
||||
p.add_argument("--runs", type=int, default=10)
|
||||
|
||||
p = sub.add_parser("report", help="Generate full report")
|
||||
p.add_argument("--checkpoint", default="./checkpoints/router/best.pt")
|
||||
p.add_argument("--output", default="./benchmark_report.json")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.cmd == "routing":
|
||||
benchmark_routing(args.checkpoint)
|
||||
elif args.cmd == "utilization":
|
||||
benchmark_utilization(args.checkpoint)
|
||||
elif args.cmd == "latency":
|
||||
benchmark_latency(args.checkpoint, args.runs)
|
||||
elif args.cmd == "report":
|
||||
generate_report(args.checkpoint, args.output)
|
||||
else:
|
||||
parser.print_help()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,679 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
zen5 MoDE — Expert Extraction Pipeline
|
||||
|
||||
Extracts expert FFN modules from source models, profiles their capabilities
|
||||
via activation analysis, and creates an indexed expert catalog for MoDE routing.
|
||||
|
||||
Supports both dense models (extract full FFN per layer) and MoE models
|
||||
(extract individual expert FFNs from MoE blocks).
|
||||
|
||||
Source models:
|
||||
T0: Qwen3.5-0.8B — dense, Gated DeltaNet
|
||||
T1: Qwen3.5-9B — dense, Gated DeltaNet
|
||||
T2: MiniMax-M2.5 — 230B, Lightning Attention MoE
|
||||
T3: GLM-5 — 744B, GLM MoE
|
||||
T4: Kimi K2.5 — 1.04T, DeepseekV3 MoE, 384 experts
|
||||
T4: Ling-1T — 1T, FP8 MoE
|
||||
|
||||
Usage:
|
||||
python zen5_expert_extraction.py list
|
||||
python zen5_expert_extraction.py extract --model qwen35-0.8b --output ./experts/
|
||||
python zen5_expert_extraction.py extract --model kimi-k2.5 --output ./experts/ --shard 0 --num-shards 4
|
||||
python zen5_expert_extraction.py fingerprint --model qwen35-0.8b --output ./experts/ --device cuda
|
||||
python zen5_expert_extraction.py extract-all --output ./experts/
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import gc
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from dataclasses import asdict, dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any, Optional
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
|
||||
log = logging.getLogger("zen5-extract")
|
||||
|
||||
|
||||
# ── Source Model Registry ────────────────────────────────────────────────────
|
||||
|
||||
MODELS: dict[str, dict[str, Any]] = {
|
||||
"qwen35-0.8b": {
|
||||
"hf_id": "Qwen/Qwen3.5-0.8B",
|
||||
"tier": "T0",
|
||||
"architecture": "gated_deltanet",
|
||||
"params": "0.8B",
|
||||
"expert_type": "dense",
|
||||
"num_layers": 28,
|
||||
"hidden_size": 1024,
|
||||
"intermediate_size": 3072,
|
||||
"moe": False,
|
||||
"description": "Trivial-tier dense expert for simple tasks",
|
||||
},
|
||||
"qwen35-9b": {
|
||||
"hf_id": "Qwen/Qwen3.5-9B",
|
||||
"tier": "T1",
|
||||
"architecture": "gated_deltanet",
|
||||
"params": "9B",
|
||||
"expert_type": "dense",
|
||||
"num_layers": 40,
|
||||
"hidden_size": 4096,
|
||||
"intermediate_size": 11008,
|
||||
"moe": False,
|
||||
"description": "Standard-tier dense expert for general tasks",
|
||||
},
|
||||
"minimax-m2.5": {
|
||||
"hf_id": "MiniMaxAI/MiniMax-M2.5",
|
||||
"tier": "T2",
|
||||
"architecture": "lightning_attention_moe",
|
||||
"params": "230B",
|
||||
"active_params": "10B",
|
||||
"expert_type": "hybrid_moe",
|
||||
"moe": True,
|
||||
"moe_block_pattern": "block_sparse_moe",
|
||||
"expert_attr": "experts",
|
||||
"description": "Complex-tier MoE with Lightning Attention for ultra-long context",
|
||||
},
|
||||
"glm-5": {
|
||||
"hf_id": "THUDM/GLM-5",
|
||||
"tier": "T3",
|
||||
"architecture": "glm_moe",
|
||||
"params": "744B",
|
||||
"active_params": "40B",
|
||||
"expert_type": "standard_moe",
|
||||
"num_experts": 256,
|
||||
"moe": True,
|
||||
"moe_block_pattern": "mlp", # GLM uses mlp.experts
|
||||
"expert_attr": "experts",
|
||||
"description": "Advanced-tier MoE for complex reasoning and multilingual",
|
||||
},
|
||||
"kimi-k2.5": {
|
||||
"hf_id": "moonshotai/Kimi-K2.5",
|
||||
"tier": "T4",
|
||||
"architecture": "deepseek_v3_moe",
|
||||
"params": "1.04T",
|
||||
"active_params": "32B",
|
||||
"expert_type": "deepseek_moe",
|
||||
"num_experts": 384,
|
||||
"top_k": 8,
|
||||
"moe": True,
|
||||
"moe_block_pattern": "mlp", # DeepseekV3: model.layers[i].mlp.experts[j]
|
||||
"expert_attr": "experts",
|
||||
"has_shared_experts": True,
|
||||
"description": "Frontier-tier MoE with 384 experts, agentic reasoning",
|
||||
},
|
||||
"ling-1t": {
|
||||
"hf_id": "inclusionAI/Ling-1T",
|
||||
"tier": "T4",
|
||||
"architecture": "fp8_moe",
|
||||
"params": "1T",
|
||||
"active_params": "50B",
|
||||
"expert_type": "fp8_moe",
|
||||
"moe": True,
|
||||
"moe_block_pattern": "mlp",
|
||||
"expert_attr": "experts",
|
||||
"description": "Frontier-tier FP8 MoE, largest open-source model",
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
# ── Expert Fingerprint ───────────────────────────────────────────────────────
|
||||
|
||||
@dataclass
|
||||
class ExpertFingerprint:
|
||||
"""Activation profile for a single expert."""
|
||||
expert_id: str
|
||||
model_name: str
|
||||
layer_idx: int
|
||||
expert_idx: int # -1 for dense models (entire layer FFN)
|
||||
architecture: str
|
||||
tier: str
|
||||
|
||||
# Activation scores per domain (0.0-1.0), computed via fingerprinting
|
||||
math_score: float = 0.0
|
||||
code_score: float = 0.0
|
||||
reasoning_score: float = 0.0
|
||||
multilingual_score: float = 0.0
|
||||
creative_score: float = 0.0
|
||||
factual_score: float = 0.0
|
||||
|
||||
# Shape info
|
||||
input_dim: int = 0
|
||||
output_dim: int = 0
|
||||
param_count: int = 0
|
||||
|
||||
# Storage
|
||||
weight_path: str = ""
|
||||
dtype: str = "bfloat16"
|
||||
requires_alignment: bool = True
|
||||
|
||||
|
||||
@dataclass
|
||||
class ExpertCatalog:
|
||||
"""Indexed collection of extracted expert fingerprints."""
|
||||
model_name: str
|
||||
total_experts: int = 0
|
||||
total_params: int = 0
|
||||
experts: list[dict] = field(default_factory=list)
|
||||
|
||||
def add(self, fp: ExpertFingerprint):
|
||||
self.experts.append(asdict(fp))
|
||||
self.total_experts += 1
|
||||
self.total_params += fp.param_count
|
||||
|
||||
def save(self, path: str):
|
||||
os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
|
||||
with open(path, "w") as f:
|
||||
json.dump(asdict(self), f, indent=2)
|
||||
log.info(f"Catalog saved: {path} ({self.total_experts} experts, {self.total_params:,} params)")
|
||||
|
||||
@classmethod
|
||||
def load(cls, path: str) -> "ExpertCatalog":
|
||||
with open(path) as f:
|
||||
data = json.load(f)
|
||||
cat = cls(model_name=data["model_name"])
|
||||
cat.total_experts = data["total_experts"]
|
||||
cat.total_params = data["total_params"]
|
||||
cat.experts = data["experts"]
|
||||
return cat
|
||||
|
||||
|
||||
# ── MoE Layer Navigator ─────────────────────────────────────────────────────
|
||||
|
||||
def find_moe_layers(model: nn.Module, config: dict) -> list[tuple[int, nn.Module, list[nn.Module]]]:
|
||||
"""
|
||||
Navigate model architecture to find MoE blocks and their expert modules.
|
||||
|
||||
Returns list of (layer_idx, moe_block, [expert_modules]).
|
||||
Handles: DeepseekV3 (Kimi), GLM, Qwen3-MoE, MiniMax, Ling.
|
||||
"""
|
||||
results = []
|
||||
layers = None
|
||||
|
||||
# Navigate to the decoder layers
|
||||
for attr in ("model.layers", "transformer.layers", "gpt_neox.layers"):
|
||||
obj = model
|
||||
try:
|
||||
for part in attr.split("."):
|
||||
obj = getattr(obj, part)
|
||||
layers = obj
|
||||
break
|
||||
except AttributeError:
|
||||
continue
|
||||
|
||||
if layers is None:
|
||||
log.warning("Could not find decoder layers — trying model.layers")
|
||||
layers = getattr(model, "model", model)
|
||||
layers = getattr(layers, "layers", [])
|
||||
|
||||
for layer_idx, layer in enumerate(layers):
|
||||
# Find the MoE block within this layer
|
||||
moe_block = None
|
||||
experts = []
|
||||
|
||||
# Strategy 1: layer.mlp.experts (DeepseekV3, GLM, Ling)
|
||||
mlp = getattr(layer, "mlp", None)
|
||||
if mlp is not None:
|
||||
exp = getattr(mlp, "experts", None)
|
||||
if exp is not None:
|
||||
moe_block = mlp
|
||||
experts = list(exp) if hasattr(exp, "__iter__") else [exp]
|
||||
|
||||
# Strategy 2: layer.block_sparse_moe.experts (Mixtral-style, MiniMax)
|
||||
if not experts:
|
||||
bsm = getattr(layer, "block_sparse_moe", None)
|
||||
if bsm is not None:
|
||||
exp = getattr(bsm, "experts", None)
|
||||
if exp is not None:
|
||||
moe_block = bsm
|
||||
experts = list(exp) if hasattr(exp, "__iter__") else [exp]
|
||||
|
||||
# Strategy 3: layer.ffn.experts (alternative naming)
|
||||
if not experts:
|
||||
ffn = getattr(layer, "ffn", None)
|
||||
if ffn is not None:
|
||||
exp = getattr(ffn, "experts", None)
|
||||
if exp is not None:
|
||||
moe_block = ffn
|
||||
experts = list(exp) if hasattr(exp, "__iter__") else [exp]
|
||||
|
||||
if experts:
|
||||
results.append((layer_idx, moe_block, experts))
|
||||
|
||||
return results
|
||||
|
||||
|
||||
def find_dense_ffn(model: nn.Module) -> list[tuple[int, nn.Module]]:
|
||||
"""Find FFN/MLP blocks in dense (non-MoE) models."""
|
||||
results = []
|
||||
layers = getattr(getattr(model, "model", model), "layers", [])
|
||||
|
||||
for layer_idx, layer in enumerate(layers):
|
||||
mlp = getattr(layer, "mlp", None)
|
||||
if mlp is not None:
|
||||
results.append((layer_idx, mlp))
|
||||
return results
|
||||
|
||||
|
||||
# ── Expert Extraction ────────────────────────────────────────────────────────
|
||||
|
||||
def _module_stats(module: nn.Module) -> tuple[int, int, int]:
|
||||
"""Return (param_count, input_dim, output_dim) for a module."""
|
||||
param_count = sum(p.numel() for p in module.parameters())
|
||||
input_dim, output_dim = 0, 0
|
||||
for m in module.modules():
|
||||
if isinstance(m, nn.Linear):
|
||||
input_dim = m.in_features
|
||||
output_dim = m.out_features
|
||||
break
|
||||
return param_count, input_dim, output_dim
|
||||
|
||||
|
||||
def extract_dense(model_name: str, cfg: dict, output_dir: str, device: str = "cpu") -> ExpertCatalog:
|
||||
"""Extract FFN blocks from dense models as single-expert-per-layer."""
|
||||
from transformers import AutoModelForCausalLM
|
||||
|
||||
log.info(f"Loading {model_name} ({cfg['params']}) on {device}...")
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
cfg["hf_id"], torch_dtype=torch.bfloat16, device_map=device, trust_remote_code=True,
|
||||
)
|
||||
model.eval()
|
||||
|
||||
catalog = ExpertCatalog(model_name=model_name)
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
ffn_layers = find_dense_ffn(model)
|
||||
log.info(f"Found {len(ffn_layers)} FFN layers")
|
||||
|
||||
for layer_idx, ffn in ffn_layers:
|
||||
weight_path = os.path.join(output_dir, f"layer{layer_idx:03d}_expert000.pt")
|
||||
torch.save(ffn.state_dict(), weight_path)
|
||||
|
||||
pc, in_d, out_d = _module_stats(ffn)
|
||||
fp = ExpertFingerprint(
|
||||
expert_id=f"{model_name}/L{layer_idx}/E0",
|
||||
model_name=model_name, layer_idx=layer_idx, expert_idx=0,
|
||||
architecture=cfg["architecture"], tier=cfg["tier"],
|
||||
input_dim=in_d, output_dim=out_d, param_count=pc,
|
||||
weight_path=weight_path,
|
||||
)
|
||||
catalog.add(fp)
|
||||
log.info(f" L{layer_idx:3d} E000: {pc:>12,} params → {os.path.basename(weight_path)}")
|
||||
|
||||
catalog.save(os.path.join(output_dir, "catalog.json"))
|
||||
del model; gc.collect(); torch.cuda.empty_cache() if torch.cuda.is_available() else None
|
||||
return catalog
|
||||
|
||||
|
||||
def extract_moe(
|
||||
model_name: str, cfg: dict, output_dir: str,
|
||||
device: str = "cpu",
|
||||
shard: int = 0, num_shards: int = 1,
|
||||
) -> ExpertCatalog:
|
||||
"""
|
||||
Extract individual expert FFN modules from MoE models.
|
||||
|
||||
For large models (>100B), use sharding: split layers across jobs.
|
||||
Each shard processes layers [shard::num_shards].
|
||||
|
||||
Example — 4-GPU parallel extraction of Kimi K2.5:
|
||||
GPU 0: python extract.py --model kimi-k2.5 --shard 0 --num-shards 4
|
||||
GPU 1: python extract.py --model kimi-k2.5 --shard 1 --num-shards 4
|
||||
GPU 2: python extract.py --model kimi-k2.5 --shard 2 --num-shards 4
|
||||
GPU 3: python extract.py --model kimi-k2.5 --shard 3 --num-shards 4
|
||||
"""
|
||||
from transformers import AutoModelForCausalLM, AutoConfig
|
||||
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
catalog = ExpertCatalog(model_name=model_name)
|
||||
|
||||
# Phase 1: Load config to determine architecture
|
||||
log.info(f"Loading config for {model_name} ({cfg['params']})...")
|
||||
hf_config = AutoConfig.from_pretrained(cfg["hf_id"], trust_remote_code=True)
|
||||
num_layers = getattr(hf_config, "num_hidden_layers", 0)
|
||||
num_experts_cfg = getattr(hf_config, "num_local_experts",
|
||||
getattr(hf_config, "n_routed_experts",
|
||||
cfg.get("num_experts", 0)))
|
||||
|
||||
log.info(f" Layers: {num_layers}, Experts/layer: {num_experts_cfg}")
|
||||
|
||||
# Compute shard assignment
|
||||
my_layers = list(range(shard, num_layers, num_shards))
|
||||
log.info(f" Shard {shard}/{num_shards}: processing layers {my_layers[:5]}{'...' if len(my_layers) > 5 else ''} ({len(my_layers)} total)")
|
||||
|
||||
# Phase 2: Load model — for frontier models, use quantized loading
|
||||
total_gb = _estimate_model_gb(cfg)
|
||||
if total_gb > 200:
|
||||
log.info(f" Model size ~{total_gb:.0f}GB — using INT4 quantized loading + CPU offload")
|
||||
load_kwargs = dict(
|
||||
torch_dtype=torch.bfloat16,
|
||||
device_map="auto",
|
||||
trust_remote_code=True,
|
||||
load_in_4bit=True,
|
||||
)
|
||||
elif total_gb > 40:
|
||||
log.info(f" Model size ~{total_gb:.0f}GB — using auto device map")
|
||||
load_kwargs = dict(
|
||||
torch_dtype=torch.bfloat16,
|
||||
device_map="auto",
|
||||
trust_remote_code=True,
|
||||
)
|
||||
else:
|
||||
load_kwargs = dict(
|
||||
torch_dtype=torch.bfloat16,
|
||||
device_map=device,
|
||||
trust_remote_code=True,
|
||||
)
|
||||
|
||||
log.info(f"Loading model weights...")
|
||||
t0 = time.time()
|
||||
model = AutoModelForCausalLM.from_pretrained(cfg["hf_id"], **load_kwargs)
|
||||
model.eval()
|
||||
log.info(f" Loaded in {time.time() - t0:.1f}s")
|
||||
|
||||
# Phase 3: Find and extract MoE experts
|
||||
moe_layers = find_moe_layers(model, cfg)
|
||||
log.info(f"Found {len(moe_layers)} MoE layers")
|
||||
|
||||
if not moe_layers:
|
||||
log.warning("No MoE layers found — falling back to dense extraction")
|
||||
return extract_dense(model_name, cfg, output_dir, device)
|
||||
|
||||
for layer_idx, moe_block, experts in moe_layers:
|
||||
if layer_idx not in my_layers:
|
||||
continue
|
||||
|
||||
# Also extract shared expert if present (DeepseekV3)
|
||||
shared_expert = getattr(moe_block, "shared_experts", None)
|
||||
if shared_expert is not None:
|
||||
weight_path = os.path.join(output_dir, f"layer{layer_idx:03d}_shared.pt")
|
||||
torch.save(shared_expert.state_dict(), weight_path)
|
||||
pc, in_d, out_d = _module_stats(shared_expert)
|
||||
fp = ExpertFingerprint(
|
||||
expert_id=f"{model_name}/L{layer_idx}/shared",
|
||||
model_name=model_name, layer_idx=layer_idx, expert_idx=-1,
|
||||
architecture=cfg["architecture"], tier=cfg["tier"],
|
||||
input_dim=in_d, output_dim=out_d, param_count=pc,
|
||||
weight_path=weight_path,
|
||||
)
|
||||
catalog.add(fp)
|
||||
log.info(f" L{layer_idx:3d} SHARED: {pc:>12,} params")
|
||||
|
||||
# Extract each expert FFN
|
||||
for expert_idx, expert in enumerate(experts):
|
||||
weight_path = os.path.join(output_dir, f"layer{layer_idx:03d}_expert{expert_idx:03d}.pt")
|
||||
state = {k: v.cpu() for k, v in expert.state_dict().items()}
|
||||
torch.save(state, weight_path)
|
||||
|
||||
pc, in_d, out_d = _module_stats(expert)
|
||||
fp = ExpertFingerprint(
|
||||
expert_id=f"{model_name}/L{layer_idx}/E{expert_idx}",
|
||||
model_name=model_name, layer_idx=layer_idx, expert_idx=expert_idx,
|
||||
architecture=cfg["architecture"], tier=cfg["tier"],
|
||||
input_dim=in_d, output_dim=out_d, param_count=pc,
|
||||
weight_path=weight_path,
|
||||
)
|
||||
catalog.add(fp)
|
||||
|
||||
log.info(f" L{layer_idx:3d}: {len(experts)} experts extracted")
|
||||
|
||||
# Also extract gate/router weights
|
||||
for layer_idx, moe_block, _ in moe_layers:
|
||||
if layer_idx not in my_layers:
|
||||
continue
|
||||
gate = getattr(moe_block, "gate", getattr(moe_block, "router", None))
|
||||
if gate is not None:
|
||||
gate_path = os.path.join(output_dir, f"layer{layer_idx:03d}_gate.pt")
|
||||
torch.save(gate.state_dict(), gate_path)
|
||||
log.info(f" L{layer_idx:3d} GATE saved")
|
||||
|
||||
shard_suffix = f"_shard{shard}" if num_shards > 1 else ""
|
||||
catalog.save(os.path.join(output_dir, f"catalog{shard_suffix}.json"))
|
||||
|
||||
del model; gc.collect()
|
||||
torch.cuda.empty_cache() if torch.cuda.is_available() else None
|
||||
return catalog
|
||||
|
||||
|
||||
def _estimate_model_gb(cfg: dict) -> float:
|
||||
"""Rough estimate of model size in GB (bf16)."""
|
||||
params_str = cfg.get("params", "0")
|
||||
multiplier = {"B": 1e9, "T": 1e12, "M": 1e6}
|
||||
for suffix, mult in multiplier.items():
|
||||
if params_str.upper().endswith(suffix):
|
||||
num = float(params_str[:-1].replace("+", ""))
|
||||
return num * mult * 2 / 1e9 # bf16 = 2 bytes/param
|
||||
return 0
|
||||
|
||||
|
||||
# ── Expert Fingerprinting ────────────────────────────────────────────────────
|
||||
|
||||
BENCHMARK_PROMPTS = {
|
||||
"math": [
|
||||
"Solve: What is the integral of x^2 dx?",
|
||||
"If f(x) = 3x^2 + 2x - 1, find f'(x).",
|
||||
"What is 47 * 83?",
|
||||
"Prove that the square root of 2 is irrational.",
|
||||
],
|
||||
"code": [
|
||||
"Write a Python function to find the longest common subsequence.",
|
||||
"Implement a binary search tree in Rust.",
|
||||
"Write a SQL query to find duplicate rows in a table.",
|
||||
"Explain the difference between async and threads in Python.",
|
||||
],
|
||||
"reasoning": [
|
||||
"If all roses are flowers and some flowers fade quickly, can we conclude all roses fade quickly?",
|
||||
"A is taller than B. C is shorter than B. D is taller than A. Who is the shortest?",
|
||||
"Three boxes: one has gold, two have rocks. You pick box A. Host opens box C (rocks). Switch?",
|
||||
],
|
||||
"multilingual": [
|
||||
"Translate to Chinese: The future of AI is open-source.",
|
||||
"Quelle est la capitale de la France?",
|
||||
"日本の人口は何人ですか？",
|
||||
],
|
||||
"creative": [
|
||||
"Write a haiku about quantum computing.",
|
||||
"Describe a color that doesn't exist.",
|
||||
],
|
||||
"factual": [
|
||||
"What is the speed of light in a vacuum?",
|
||||
"Who wrote 'The Republic'?",
|
||||
"What causes tides on Earth?",
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
def fingerprint_model(
|
||||
model_name: str, cfg: dict, catalog_dir: str,
|
||||
device: str = "cuda", max_samples: int = 4,
|
||||
):
|
||||
"""
|
||||
Run benchmark prompts through the model, recording per-layer expert activations.
|
||||
Uses forward hooks to measure which experts activate on which domains.
|
||||
"""
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
|
||||
catalog_path = os.path.join(catalog_dir, "catalog.json")
|
||||
if not os.path.exists(catalog_path):
|
||||
log.error(f"No catalog at {catalog_path} — run extract first")
|
||||
return
|
||||
|
||||
catalog = ExpertCatalog.load(catalog_path)
|
||||
log.info(f"Loaded catalog: {catalog.total_experts} experts")
|
||||
|
||||
log.info(f"Loading model + tokenizer on {device}...")
|
||||
tokenizer = AutoTokenizer.from_pretrained(cfg["hf_id"], trust_remote_code=True)
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
cfg["hf_id"], torch_dtype=torch.bfloat16, device_map=device, trust_remote_code=True,
|
||||
)
|
||||
model.eval()
|
||||
|
||||
# Hook storage: domain → layer → activation_magnitude
|
||||
activations: dict[str, dict[int, float]] = {}
|
||||
|
||||
def make_hook(domain: str, layer_idx: int):
|
||||
def hook_fn(module, input, output):
|
||||
if isinstance(output, tuple):
|
||||
out = output[0]
|
||||
else:
|
||||
out = output
|
||||
mag = out.float().abs().mean().item()
|
||||
activations.setdefault(domain, {})
|
||||
activations[domain].setdefault(layer_idx, 0.0)
|
||||
activations[domain][layer_idx] += mag
|
||||
return hook_fn
|
||||
|
||||
# Register hooks on FFN/MoE blocks
|
||||
layers = getattr(getattr(model, "model", model), "layers", [])
|
||||
hooks = []
|
||||
for domain in BENCHMARK_PROMPTS:
|
||||
for layer_idx, layer in enumerate(layers):
|
||||
mlp = getattr(layer, "mlp", None)
|
||||
if mlp is not None:
|
||||
h = mlp.register_forward_hook(make_hook(domain, layer_idx))
|
||||
hooks.append((h, domain))
|
||||
|
||||
# Run inference per domain
|
||||
for domain, prompts in BENCHMARK_PROMPTS.items():
|
||||
log.info(f" Profiling domain: {domain} ({len(prompts[:max_samples])} prompts)")
|
||||
for prompt in prompts[:max_samples]:
|
||||
inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=128)
|
||||
inputs = {k: v.to(model.device) for k, v in inputs.items()}
|
||||
|
||||
# Only register hooks for current domain
|
||||
with torch.no_grad():
|
||||
model(**inputs)
|
||||
|
||||
# Remove hooks
|
||||
for h, _ in hooks:
|
||||
h.remove()
|
||||
|
||||
# Normalize activations and update catalog
|
||||
for expert_data in catalog.experts:
|
||||
layer_idx = expert_data["layer_idx"]
|
||||
for domain in BENCHMARK_PROMPTS:
|
||||
acts = activations.get(domain, {})
|
||||
score = acts.get(layer_idx, 0.0)
|
||||
# Normalize to 0-1 range per domain
|
||||
all_scores = list(acts.values()) or [0]
|
||||
max_score = max(all_scores) or 1.0
|
||||
normalized = score / max_score
|
||||
expert_data[f"{domain}_score"] = round(normalized, 4)
|
||||
|
||||
# Save updated catalog
|
||||
catalog.save(os.path.join(catalog_dir, "catalog_fingerprinted.json"))
|
||||
log.info(f"Fingerprinted catalog saved with {len(BENCHMARK_PROMPTS)} domain scores")
|
||||
|
||||
del model; gc.collect()
|
||||
torch.cuda.empty_cache() if torch.cuda.is_available() else None
|
||||
|
||||
|
||||
# ── Catalog Merger ───────────────────────────────────────────────────────────
|
||||
|
||||
def merge_catalogs(catalog_dir: str, output_path: str):
|
||||
"""Merge shard catalogs and/or multiple model catalogs into unified index."""
|
||||
merged = ExpertCatalog(model_name="zen5-unified")
|
||||
catalog_files = sorted(Path(catalog_dir).rglob("catalog*.json"))
|
||||
|
||||
for cf in catalog_files:
|
||||
cat = ExpertCatalog.load(str(cf))
|
||||
for expert in cat.experts:
|
||||
merged.experts.append(expert)
|
||||
merged.total_experts += 1
|
||||
merged.total_params += expert.get("param_count", 0)
|
||||
log.info(f" Merged {cf.name}: {cat.total_experts} experts from {cat.model_name}")
|
||||
|
||||
merged.save(output_path)
|
||||
log.info(f"Unified catalog: {merged.total_experts} experts, {merged.total_params:,} params")
|
||||
return merged
|
||||
|
||||
|
||||
# ── CLI ──────────────────────────────────────────────────────────────────────
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="zen5 MoDE Expert Extraction Pipeline",
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
)
|
||||
sub = parser.add_subparsers(dest="cmd")
|
||||
|
||||
# list
|
||||
sub.add_parser("list", help="List source models")
|
||||
|
||||
# extract
|
||||
p = sub.add_parser("extract", help="Extract experts from a model")
|
||||
p.add_argument("--model", required=True, choices=list(MODELS.keys()))
|
||||
p.add_argument("--output", default="./experts")
|
||||
p.add_argument("--device", default="cpu", help="cpu, cuda, cuda:0, auto")
|
||||
p.add_argument("--shard", type=int, default=0, help="Shard index for distributed extraction")
|
||||
p.add_argument("--num-shards", type=int, default=1, help="Total shards")
|
||||
|
||||
# extract-all
|
||||
p = sub.add_parser("extract-all", help="Extract from all models")
|
||||
p.add_argument("--output", default="./experts")
|
||||
p.add_argument("--device", default="cpu")
|
||||
|
||||
# fingerprint
|
||||
p = sub.add_parser("fingerprint", help="Profile expert capabilities via activation analysis")
|
||||
p.add_argument("--model", required=True, choices=list(MODELS.keys()))
|
||||
p.add_argument("--output", default="./experts")
|
||||
p.add_argument("--device", default="cuda")
|
||||
p.add_argument("--max-samples", type=int, default=4)
|
||||
|
||||
# merge
|
||||
p = sub.add_parser("merge", help="Merge catalogs into unified index")
|
||||
p.add_argument("--input", default="./experts", help="Directory containing catalogs")
|
||||
p.add_argument("--output", default="./experts/unified_catalog.json")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.cmd == "list":
|
||||
print(f"\n{'Model':<20s} {'Tier':<4s} {'Params':<8s} {'Type':<12s} Architecture")
|
||||
print("=" * 80)
|
||||
for name, cfg in MODELS.items():
|
||||
moe_tag = "MoE" if cfg.get("moe") else "Dense"
|
||||
print(f" {name:<18s} {cfg['tier']:<4s} {cfg['params']:<8s} {moe_tag:<12s} {cfg['architecture']}")
|
||||
print(f"\nCombined expert pool: ~3.0T+ parameters")
|
||||
|
||||
elif args.cmd == "extract":
|
||||
cfg = MODELS[args.model]
|
||||
out = os.path.join(args.output, args.model)
|
||||
if cfg.get("moe"):
|
||||
extract_moe(args.model, cfg, out, args.device, args.shard, args.num_shards)
|
||||
else:
|
||||
extract_dense(args.model, cfg, out, args.device)
|
||||
|
||||
elif args.cmd == "extract-all":
|
||||
for name, cfg in MODELS.items():
|
||||
out = os.path.join(args.output, name)
|
||||
try:
|
||||
if cfg.get("moe"):
|
||||
extract_moe(name, cfg, out, args.device)
|
||||
else:
|
||||
extract_dense(name, cfg, out, args.device)
|
||||
except Exception as e:
|
||||
log.error(f"FAILED {name}: {e}")
|
||||
|
||||
elif args.cmd == "fingerprint":
|
||||
cfg = MODELS[args.model]
|
||||
fingerprint_model(args.model, cfg, os.path.join(args.output, args.model),
|
||||
args.device, args.max_samples)
|
||||
|
||||
elif args.cmd == "merge":
|
||||
merge_catalogs(args.input, args.output)
|
||||
|
||||
else:
|
||||
parser.print_help()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,218 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
zen5 MoDE — HuggingFace Repository Setup
|
||||
|
||||
Creates HuggingFace model repos for the zen5 lineup with proper model cards.
|
||||
|
||||
Models:
|
||||
zen5 — General (750B, 0.8-50B active)
|
||||
zen5-coder — Code-focused (1.8T, 0.8-80B active)
|
||||
zen5-omni — Omnimodal (2.5T, 0.8-100B active)
|
||||
zen5-max — Frontier (3.1T, 0.8-100B+ active)
|
||||
|
||||
Usage:
|
||||
python zen5_hf_repos.py # Create all repos
|
||||
python zen5_hf_repos.py zen5 # Create specific repo
|
||||
python zen5_hf_repos.py --dry-run # Preview READMEs without creating
|
||||
"""
|
||||
|
||||
import io
|
||||
import sys
|
||||
from huggingface_hub import HfApi
|
||||
|
||||
api = HfApi()
|
||||
|
||||
ZEN5_MODELS = [
|
||||
{
|
||||
"zen_name": "zen5",
|
||||
"display": "Zen5",
|
||||
"params": "750B",
|
||||
"active": "0.8-50B",
|
||||
"expert_pool": "Qwen3.5-0.8B + Qwen3.5-9B + GLM-5",
|
||||
"target": "General",
|
||||
"modalities": "Text",
|
||||
"context": "262K",
|
||||
"license": "apache-2.0",
|
||||
"desc": "General-purpose MoDE model with complexity-aware routing. Uses 0.8B for simple tasks, scales to 50B for complex reasoning.",
|
||||
},
|
||||
{
|
||||
"zen_name": "zen5-coder",
|
||||
"display": "Zen5 Coder",
|
||||
"params": "1.8T",
|
||||
"active": "0.8-80B",
|
||||
"expert_pool": "Qwen3.5 + GLM-5 + Kimi K2.5 code experts",
|
||||
"target": "Code",
|
||||
"modalities": "Text + Code",
|
||||
"context": "262K",
|
||||
"license": "apache-2.0",
|
||||
"desc": "Code-focused MoDE with 384 Kimi K2.5 experts for frontier coding capability. SWE-bench class performance with adaptive compute.",
|
||||
},
|
||||
{
|
||||
"zen_name": "zen5-omni",
|
||||
"display": "Zen5 Omni",
|
||||
"params": "2.5T",
|
||||
"active": "0.8-100B",
|
||||
"expert_pool": "All text + Vision + Video + 3D + Audio",
|
||||
"target": "Omnimodal",
|
||||
"modalities": "Text + Vision + Video + 3D + Audio",
|
||||
"context": "262K",
|
||||
"license": "apache-2.0",
|
||||
"desc": "True omnimodal AI: text, vision, video, 3D generation, and audio in a single unified model with Transfusion architecture.",
|
||||
},
|
||||
{
|
||||
"zen_name": "zen5-max",
|
||||
"display": "Zen5 Max",
|
||||
"params": "3.1T",
|
||||
"active": "0.8-100B+",
|
||||
"expert_pool": "All sources, all modalities",
|
||||
"target": "Frontier",
|
||||
"modalities": "Text + Vision + Video + 3D + Audio",
|
||||
"context": "262K",
|
||||
"license": "apache-2.0",
|
||||
"desc": "Frontier 3.1T MoDE — the largest open-source model. Full expert pool with complexity-aware routing across 6 model families and 5 modalities.",
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
def generate_readme(m):
|
||||
"""Generate zen5 model card."""
|
||||
|
||||
# Build family table
|
||||
family_rows = []
|
||||
for model in ZEN5_MODELS:
|
||||
bold = "**" if model["zen_name"] == m["zen_name"] else ""
|
||||
family_rows.append(
|
||||
f"| {bold}{model['display']}{bold} | {bold}{model['params']}{bold} | "
|
||||
f"{bold}{model['active']}{bold} | {bold}{model['modalities']}{bold} | "
|
||||
f"[zenlm/{model['zen_name']}](https://huggingface.co/zenlm/{model['zen_name']}) |"
|
||||
)
|
||||
family_table = "\n".join(family_rows)
|
||||
|
||||
return f"""---
|
||||
license: {m['license']}
|
||||
language:
|
||||
- en
|
||||
- zh
|
||||
- ja
|
||||
- ko
|
||||
- fr
|
||||
- de
|
||||
- es
|
||||
- pt
|
||||
- ru
|
||||
- ar
|
||||
tags:
|
||||
- zen5
|
||||
- zenlm
|
||||
- hanzo
|
||||
- mode
|
||||
- mixture-of-diverse-experts
|
||||
- frontier-ai
|
||||
- {'omnimodal' if 'Vision' in m['modalities'] else 'text-generation'}
|
||||
pipeline_tag: text-generation
|
||||
library_name: transformers
|
||||
---
|
||||
|
||||
# {m['display']}
|
||||
|
||||
**{m['display']}** is a {m['params']} parameter AI model from the [Zen5 family](https://zenlm.org) by [Zen LM](https://huggingface.co/zenlm) and [Hanzo AI](https://hanzo.ai).
|
||||
|
||||
{m['desc']}
|
||||
|
||||
## Architecture: MoDE (Mixture of Diverse Experts)
|
||||
|
||||
zen5 introduces **MoDE** — a novel architecture that routes across expert modules harvested from the largest open-source models with **complexity-aware hierarchical routing**.
|
||||
|
||||
| Feature | Value |
|
||||
|---------|-------|
|
||||
| **Total Parameters** | {m['params']} |
|
||||
| **Active Parameters** | {m['active']} (adaptive) |
|
||||
| **Architecture** | MoDE (Mixture of Diverse Experts) |
|
||||
| **Expert Pool** | {m['expert_pool']} |
|
||||
| **Modalities** | {m['modalities']} |
|
||||
| **Context** | {m['context']} tokens |
|
||||
| **Complexity Tiers** | T0 (0.8B) → T1 (9B) → T2 (10-40B) → T3 (40-80B) → T4 (80B+) |
|
||||
| **License** | {m['license'].upper()} |
|
||||
| **Family** | Zen5 |
|
||||
| **Target** | {m['target']} |
|
||||
| **Creator** | Zen LM / Hanzo AI |
|
||||
|
||||
## How It Works
|
||||
|
||||
Unlike standard models that always use the same compute, zen5 **adapts compute to task difficulty**:
|
||||
|
||||
- **"Hello, how are you?"** → T0 tier, 0.8B active params, <50ms
|
||||
- **"Summarize this article"** → T1 tier, 9B active params, <200ms
|
||||
- **"Prove the Riemann hypothesis approach"** → T4 tier, 100B+ active params, full reasoning
|
||||
|
||||
The complexity estimator runs on the first 64 tokens and routes to the appropriate expert tier.
|
||||
|
||||
## Key Innovations
|
||||
|
||||
1. **Complexity-Aware Routing**: Lightweight estimator predicts task difficulty, activates appropriate tier
|
||||
2. **Cross-Architecture Experts**: Experts from 6+ diverse model families (Gated DeltaNet, Lightning Attention, DeepseekV3 MoE, FP8 MoE, GLM MoE)
|
||||
3. **MoE++ Zero-Computation Experts**: Zero/Copy/Constant experts for trivial tokens (1.1-2.1x throughput)
|
||||
4. **ReMoE ReLU Routing**: Adaptive expert count via ReLU gating (no fixed Top-K)
|
||||
5. **Adaptive Escalation**: Mid-generation promotion to higher tiers if confidence drops
|
||||
{'6. **Transfusion**: Unified autoregressive (text) + diffusion (3D/video) in single model' if 'Vision' in m['modalities'] else ''}
|
||||
|
||||
## Zen5 Family
|
||||
|
||||
| Model | Total Params | Active Params | Modalities | HuggingFace |
|
||||
|-------|-------------|--------------|-----------|-------------|
|
||||
{family_table}
|
||||
|
||||
## Links
|
||||
|
||||
- [Zen LM](https://zenlm.org) | [Hanzo AI](https://hanzo.ai) | [Hanzo Chat](https://hanzo.chat)
|
||||
- [All Zen Models](https://huggingface.co/zenlm)
|
||||
- [Architecture Paper](https://zenlm.org/zen5-mode) (coming soon)
|
||||
|
||||
---
|
||||
|
||||
*Zen5: Mixture of Diverse Experts — Clarity Through Diversity*
|
||||
"""
|
||||
|
||||
|
||||
def main():
|
||||
dry_run = "--dry-run" in sys.argv
|
||||
args = [a for a in sys.argv[1:] if not a.startswith("--")]
|
||||
target = args[0] if args else "all"
|
||||
|
||||
if target == "all":
|
||||
models = ZEN5_MODELS
|
||||
else:
|
||||
models = [m for m in ZEN5_MODELS if m["zen_name"] == target]
|
||||
|
||||
if not models:
|
||||
print(f"No model matched: {target}")
|
||||
print(f"Available: {', '.join(m['zen_name'] for m in ZEN5_MODELS)}")
|
||||
sys.exit(1)
|
||||
|
||||
for m in models:
|
||||
repo_id = f"zenlm/{m['zen_name']}"
|
||||
readme = generate_readme(m)
|
||||
|
||||
if dry_run:
|
||||
print(f"\n{'='*60}")
|
||||
print(f"DRY RUN: {repo_id}")
|
||||
print(f"{'='*60}")
|
||||
print(readme[:500] + "...")
|
||||
continue
|
||||
|
||||
try:
|
||||
api.create_repo(repo_id, repo_type="model", exist_ok=True)
|
||||
api.upload_file(
|
||||
path_or_fileobj=io.BytesIO(readme.encode("utf-8")),
|
||||
path_in_repo="README.md",
|
||||
repo_id=repo_id,
|
||||
repo_type="model",
|
||||
commit_message=f"feat: {m['display']} MoDE model card",
|
||||
)
|
||||
print(f"OK: {repo_id} ({m['params']})")
|
||||
except Exception as e:
|
||||
print(f"FAIL: {repo_id} — {e}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,666 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
zen5 MoDE — Router & Alignment Training Pipeline
|
||||
|
||||
Trains the complexity-aware hierarchical router and cross-architecture alignment
|
||||
layers. This covers Phases 2-3 of the zen5 training roadmap.
|
||||
|
||||
Components trained:
|
||||
1. AlignmentLayer — projects each source architecture to shared latent space
|
||||
2. ComplexityEstimator — lightweight classifier predicting T0-T4 tier
|
||||
3. MoDERouter — per-tier expert selection with ReLU gating (ReMoE)
|
||||
4. CrossModalRouter — omnimodal expert chain selection
|
||||
|
||||
Training objectives:
|
||||
L = L_task + λ₁·L_complexity + λ₂·L_diversity + λ₃·L_efficiency + λ₄·L_anticollapse
|
||||
|
||||
Usage:
|
||||
python zen5_router_training.py info
|
||||
python zen5_router_training.py train --config configs/zen5_router.yaml
|
||||
python zen5_router_training.py train-complexity --data ./data/ --output ./ckpt/
|
||||
python zen5_router_training.py eval --checkpoint ./ckpt/router_best.pt
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import logging
|
||||
import math
|
||||
import os
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from torch.utils.data import DataLoader, Dataset
|
||||
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
|
||||
log = logging.getLogger("zen5-router")
|
||||
|
||||
|
||||
# ── Configuration ────────────────────────────────────────────────────────────
|
||||
|
||||
@dataclass
|
||||
class MoDEConfig:
|
||||
"""Full configuration for zen5 MoDE router architecture."""
|
||||
|
||||
# Shared representation
|
||||
hidden_size: int = 4096
|
||||
intermediate_size: int = 11008
|
||||
num_attention_heads: int = 32
|
||||
num_kv_heads: int = 8
|
||||
vocab_size: int = 152064 # Unified vocabulary
|
||||
max_position_embeddings: int = 262144
|
||||
|
||||
# Complexity tiers: T0 (trivial) → T4 (frontier)
|
||||
num_tiers: int = 5
|
||||
complexity_hidden: int = 1024
|
||||
complexity_layers: int = 4
|
||||
complexity_context_len: int = 64 # First N tokens for complexity estimation
|
||||
|
||||
# Router
|
||||
router_hidden: int = 2048
|
||||
num_experts_per_tier: tuple = (1, 1, 64, 256, 384)
|
||||
top_k_per_tier: tuple = (1, 1, 4, 8, 8)
|
||||
|
||||
# Modalities
|
||||
num_modalities: int = 5 # text, vision, video, 3d, audio
|
||||
modality_embed_dim: int = 256
|
||||
|
||||
# Loss coefficients
|
||||
lambda_complexity: float = 1.0
|
||||
lambda_diversity: float = 0.05
|
||||
lambda_efficiency: float = 0.02
|
||||
lambda_anticollapse: float = 0.01
|
||||
|
||||
# MoE++ zero-computation experts
|
||||
enable_zero_experts: bool = True
|
||||
num_zero_types: int = 3 # zero, copy, constant
|
||||
|
||||
# ReMoE
|
||||
use_relu_routing: bool = True
|
||||
|
||||
# Training
|
||||
lr: float = 3e-4
|
||||
weight_decay: float = 0.01
|
||||
warmup_steps: int = 1000
|
||||
max_steps: int = 100000
|
||||
batch_size: int = 64
|
||||
gradient_accumulation: int = 4
|
||||
fp16: bool = True
|
||||
save_every: int = 5000
|
||||
eval_every: int = 1000
|
||||
log_every: int = 100
|
||||
|
||||
|
||||
# ── Model Components ─────────────────────────────────────────────────────────
|
||||
|
||||
class ComplexityEstimator(nn.Module):
|
||||
"""
|
||||
Lightweight classifier (~100-200M params) that predicts task complexity
|
||||
from the first N tokens of input. Outputs T0-T4 tier probabilities.
|
||||
"""
|
||||
|
||||
def __init__(self, cfg: MoDEConfig):
|
||||
super().__init__()
|
||||
self.cfg = cfg
|
||||
self.embed = nn.Embedding(cfg.vocab_size, cfg.complexity_hidden)
|
||||
self.pos_embed = nn.Embedding(cfg.complexity_context_len, cfg.complexity_hidden)
|
||||
|
||||
encoder_layer = nn.TransformerEncoderLayer(
|
||||
d_model=cfg.complexity_hidden, nhead=8,
|
||||
dim_feedforward=cfg.complexity_hidden * 4,
|
||||
dropout=0.1, activation="gelu", batch_first=True,
|
||||
norm_first=True,
|
||||
)
|
||||
self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=cfg.complexity_layers)
|
||||
self.norm = nn.LayerNorm(cfg.complexity_hidden)
|
||||
|
||||
self.tier_head = nn.Sequential(
|
||||
nn.Linear(cfg.complexity_hidden, cfg.complexity_hidden),
|
||||
nn.GELU(),
|
||||
nn.Dropout(0.1),
|
||||
nn.Linear(cfg.complexity_hidden, cfg.num_tiers),
|
||||
)
|
||||
self.modality_head = nn.Linear(cfg.complexity_hidden, cfg.num_modalities)
|
||||
|
||||
def forward(self, input_ids: torch.Tensor, attention_mask: Optional[torch.Tensor] = None):
|
||||
B, S = input_ids.shape
|
||||
positions = torch.arange(S, device=input_ids.device).unsqueeze(0).expand(B, -1)
|
||||
x = self.embed(input_ids) + self.pos_embed(positions)
|
||||
|
||||
# Causal mask
|
||||
causal_mask = nn.Transformer.generate_square_subsequent_mask(S, device=x.device)
|
||||
if attention_mask is not None:
|
||||
key_padding_mask = ~attention_mask.bool()
|
||||
else:
|
||||
key_padding_mask = None
|
||||
|
||||
x = self.encoder(x, mask=causal_mask, src_key_padding_mask=key_padding_mask)
|
||||
x = self.norm(x)
|
||||
|
||||
# Pool: last non-padding token
|
||||
if attention_mask is not None:
|
||||
lengths = attention_mask.sum(dim=1, keepdim=True) - 1
|
||||
pooled = x.gather(1, lengths.unsqueeze(-1).expand(-1, -1, x.size(-1))).squeeze(1)
|
||||
else:
|
||||
pooled = x[:, -1]
|
||||
|
||||
return self.tier_head(pooled), self.modality_head(pooled)
|
||||
|
||||
|
||||
class AlignmentLayer(nn.Module):
|
||||
"""
|
||||
Per-architecture projection to shared latent space.
|
||||
Handles hidden-dimension mismatches and representation normalization.
|
||||
"""
|
||||
|
||||
ARCH_DIMS = {
|
||||
"gated_deltanet_08b": 1024,
|
||||
"gated_deltanet_9b": 4096,
|
||||
"lightning_attention": 4096,
|
||||
"glm_moe": 4096,
|
||||
"deepseek_v3": 7168,
|
||||
"fp8_moe": 4096,
|
||||
}
|
||||
|
||||
def __init__(self, cfg: MoDEConfig):
|
||||
super().__init__()
|
||||
self.projections = nn.ModuleDict()
|
||||
for arch, src_dim in self.ARCH_DIMS.items():
|
||||
self.projections[arch] = nn.Sequential(
|
||||
nn.Linear(src_dim, cfg.hidden_size),
|
||||
nn.LayerNorm(cfg.hidden_size),
|
||||
nn.GELU(),
|
||||
nn.Linear(cfg.hidden_size, cfg.hidden_size),
|
||||
nn.LayerNorm(cfg.hidden_size),
|
||||
) if src_dim != cfg.hidden_size else nn.Sequential(
|
||||
nn.LayerNorm(cfg.hidden_size),
|
||||
nn.Linear(cfg.hidden_size, cfg.hidden_size),
|
||||
nn.LayerNorm(cfg.hidden_size),
|
||||
)
|
||||
|
||||
def forward(self, hidden: torch.Tensor, arch: str) -> torch.Tensor:
|
||||
return self.projections[arch](hidden)
|
||||
|
||||
|
||||
class ZeroExpert(nn.Module):
|
||||
"""MoE++ zero expert: output zeros (discard token)."""
|
||||
def forward(self, x): return torch.zeros_like(x)
|
||||
|
||||
|
||||
class CopyExpert(nn.Module):
|
||||
"""MoE++ copy expert: identity (skip layer)."""
|
||||
def forward(self, x): return x
|
||||
|
||||
|
||||
class ConstantExpert(nn.Module):
|
||||
"""MoE++ constant expert: learned bias vector."""
|
||||
def __init__(self, dim: int):
|
||||
super().__init__()
|
||||
self.bias = nn.Parameter(torch.randn(dim) * 0.02)
|
||||
def forward(self, x): return self.bias.expand_as(x)
|
||||
|
||||
|
||||
class MoDERouter(nn.Module):
|
||||
"""
|
||||
Hierarchical router with complexity-aware tier selection.
|
||||
Per-tier expert gating with ReLU (ReMoE) for adaptive expert count.
|
||||
"""
|
||||
|
||||
def __init__(self, cfg: MoDEConfig):
|
||||
super().__init__()
|
||||
self.cfg = cfg
|
||||
|
||||
self.tier_routers = nn.ModuleList()
|
||||
for tier_idx in range(cfg.num_tiers):
|
||||
n_exp = cfg.num_experts_per_tier[tier_idx]
|
||||
n_total = n_exp + (cfg.num_zero_types if cfg.enable_zero_experts else 0)
|
||||
|
||||
layers = [
|
||||
nn.Linear(cfg.hidden_size, cfg.router_hidden),
|
||||
nn.GELU(),
|
||||
nn.Linear(cfg.router_hidden, n_total),
|
||||
]
|
||||
if cfg.use_relu_routing:
|
||||
layers.append(nn.ReLU()) # ReMoE: sparsity via ReLU
|
||||
|
||||
self.tier_routers.append(nn.Sequential(*layers))
|
||||
|
||||
# Cross-modal router
|
||||
self.cross_modal_gate = nn.Sequential(
|
||||
nn.Linear(cfg.hidden_size + cfg.modality_embed_dim, cfg.router_hidden),
|
||||
nn.GELU(),
|
||||
nn.Linear(cfg.router_hidden, cfg.num_modalities),
|
||||
nn.Sigmoid(),
|
||||
)
|
||||
self.modality_embed = nn.Embedding(cfg.num_modalities, cfg.modality_embed_dim)
|
||||
|
||||
def forward(self, hidden: torch.Tensor, tier: int,
|
||||
modality_ids: Optional[torch.Tensor] = None):
|
||||
router_logits = self.tier_routers[tier](hidden)
|
||||
|
||||
cross_modal = None
|
||||
if modality_ids is not None:
|
||||
mod_e = self.modality_embed(modality_ids)
|
||||
pooled = hidden.mean(dim=1) if hidden.dim() == 3 else hidden
|
||||
cross_modal = self.cross_modal_gate(torch.cat([pooled, mod_e], dim=-1))
|
||||
|
||||
return router_logits, cross_modal
|
||||
|
||||
|
||||
# ── Loss Functions ───────────────────────────────────────────────────────────
|
||||
|
||||
def complexity_loss(logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
|
||||
"""Cross-entropy for tier prediction."""
|
||||
return F.cross_entropy(logits, labels)
|
||||
|
||||
|
||||
def diversity_loss(router_logits: torch.Tensor, arch_ids: torch.Tensor) -> torch.Tensor:
|
||||
"""Encourage routing across different source architectures."""
|
||||
probs = F.softmax(router_logits.view(-1, router_logits.size(-1)), dim=-1)
|
||||
unique_archs = arch_ids.unique()
|
||||
arch_loads = []
|
||||
for a in unique_archs:
|
||||
mask = (arch_ids == a).float()
|
||||
load = (probs * mask.unsqueeze(0)).sum(dim=-1).mean()
|
||||
arch_loads.append(load)
|
||||
arch_loads = torch.stack(arch_loads)
|
||||
arch_probs = arch_loads / (arch_loads.sum() + 1e-8)
|
||||
return -(arch_probs * torch.log(arch_probs + 1e-8)).sum() # Negative entropy
|
||||
|
||||
|
||||
def efficiency_loss(tier_logits: torch.Tensor, tier_labels: torch.Tensor) -> torch.Tensor:
|
||||
"""Penalize over-estimation of complexity tier."""
|
||||
probs = F.softmax(tier_logits, dim=-1)
|
||||
tiers = torch.arange(probs.size(-1), device=probs.device, dtype=torch.float)
|
||||
expected = (probs * tiers).sum(dim=-1)
|
||||
return F.relu(expected - tier_labels.float()).mean()
|
||||
|
||||
|
||||
def anticollapse_loss(router_logits: torch.Tensor) -> torch.Tensor:
|
||||
"""Switch Transformer load-balancing: prevent expert starvation."""
|
||||
probs = F.softmax(router_logits.view(-1, router_logits.size(-1)), dim=-1)
|
||||
load = probs.mean(dim=0)
|
||||
ideal = 1.0 / load.size(0)
|
||||
return ((load - ideal) ** 2).sum()
|
||||
|
||||
|
||||
# ── Dataset ──────────────────────────────────────────────────────────────────
|
||||
|
||||
class ComplexityDataset(Dataset):
|
||||
"""
|
||||
Dataset for training the complexity estimator.
|
||||
Each sample: (input_ids, attention_mask, tier_label, modality_id)
|
||||
|
||||
Complexity labels can come from:
|
||||
- Synthetic: run each prompt on T0-T4, label with smallest tier achieving 95% of T4 quality
|
||||
- Manual: human-labeled 50K samples
|
||||
- Benchmark mapping: MMLU→T1, GPQA→T3, AIME→T4
|
||||
"""
|
||||
|
||||
def __init__(self, data_path: str, max_len: int = 64):
|
||||
self.max_len = max_len
|
||||
self.samples = []
|
||||
|
||||
if os.path.exists(data_path):
|
||||
with open(data_path) as f:
|
||||
self.samples = json.load(f)
|
||||
log.info(f"Loaded {len(self.samples)} samples from {data_path}")
|
||||
else:
|
||||
log.warning(f"Data path not found: {data_path} — generating synthetic data")
|
||||
self.samples = self._generate_synthetic()
|
||||
|
||||
def _generate_synthetic(self) -> list[dict]:
|
||||
"""Generate synthetic complexity-labeled data for prototyping."""
|
||||
samples = []
|
||||
|
||||
tier_examples = {
|
||||
0: ["Hello", "Hi there", "What's 2+2?", "Yes or no: is the sky blue?", "Thanks!"],
|
||||
1: ["Summarize this paragraph about climate change.", "Translate to French: Good morning.",
|
||||
"What is photosynthesis?", "Write a short email declining a meeting."],
|
||||
2: ["Analyze the causes of World War I and their modern parallels.",
|
||||
"Review this Python code for security vulnerabilities.",
|
||||
"Compare and contrast REST and GraphQL APIs."],
|
||||
3: ["Derive the Euler-Lagrange equation from first principles.",
|
||||
"Design a distributed system for real-time fraud detection.",
|
||||
"Prove that every finite group of order p^2 is abelian."],
|
||||
4: ["Prove the Riemann hypothesis implies the prime number theorem.",
|
||||
"Design a novel algorithm for protein folding that improves on AlphaFold.",
|
||||
"Write a formal proof of the four-color theorem."],
|
||||
}
|
||||
|
||||
for tier, prompts in tier_examples.items():
|
||||
for prompt in prompts:
|
||||
tokens = list(range(min(len(prompt.split()), self.max_len)))
|
||||
samples.append({
|
||||
"input_ids": tokens + [0] * (self.max_len - len(tokens)),
|
||||
"attention_mask": [1] * len(tokens) + [0] * (self.max_len - len(tokens)),
|
||||
"tier": tier,
|
||||
"modality": 0, # text
|
||||
})
|
||||
return samples
|
||||
|
||||
def __len__(self):
|
||||
return len(self.samples)
|
||||
|
||||
def __getitem__(self, idx):
|
||||
s = self.samples[idx]
|
||||
return {
|
||||
"input_ids": torch.tensor(s["input_ids"][:self.max_len], dtype=torch.long),
|
||||
"attention_mask": torch.tensor(s["attention_mask"][:self.max_len], dtype=torch.long),
|
||||
"tier": torch.tensor(s["tier"], dtype=torch.long),
|
||||
"modality": torch.tensor(s.get("modality", 0), dtype=torch.long),
|
||||
}
|
||||
|
||||
|
||||
# ── Training Loop ────────────────────────────────────────────────────────────
|
||||
|
||||
class MoDETrainer:
|
||||
"""Trains the complexity estimator and MoDE router."""
|
||||
|
||||
def __init__(self, cfg: MoDEConfig, output_dir: str):
|
||||
self.cfg = cfg
|
||||
self.output_dir = output_dir
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
# Initialize models
|
||||
self.complexity = ComplexityEstimator(cfg)
|
||||
self.alignment = AlignmentLayer(cfg)
|
||||
self.router = MoDERouter(cfg)
|
||||
|
||||
# Move to device
|
||||
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
self.complexity.to(self.device)
|
||||
self.alignment.to(self.device)
|
||||
self.router.to(self.device)
|
||||
|
||||
# Optimizer
|
||||
all_params = (
|
||||
list(self.complexity.parameters()) +
|
||||
list(self.alignment.parameters()) +
|
||||
list(self.router.parameters())
|
||||
)
|
||||
self.optimizer = torch.optim.AdamW(
|
||||
all_params, lr=cfg.lr, weight_decay=cfg.weight_decay,
|
||||
betas=(0.9, 0.95),
|
||||
)
|
||||
|
||||
# LR scheduler: linear warmup + cosine decay
|
||||
self.scheduler = torch.optim.lr_scheduler.OneCycleLR(
|
||||
self.optimizer,
|
||||
max_lr=cfg.lr,
|
||||
total_steps=cfg.max_steps,
|
||||
pct_start=cfg.warmup_steps / cfg.max_steps,
|
||||
anneal_strategy="cos",
|
||||
)
|
||||
|
||||
# Mixed precision
|
||||
self.scaler = torch.amp.GradScaler("cuda", enabled=cfg.fp16 and self.device.type == "cuda")
|
||||
|
||||
self._log_param_count()
|
||||
|
||||
def _log_param_count(self):
|
||||
c = sum(p.numel() for p in self.complexity.parameters())
|
||||
a = sum(p.numel() for p in self.alignment.parameters())
|
||||
r = sum(p.numel() for p in self.router.parameters())
|
||||
total = c + a + r
|
||||
log.info(f"Trainable parameters:")
|
||||
log.info(f" ComplexityEstimator: {c:>12,}")
|
||||
log.info(f" AlignmentLayer: {a:>12,}")
|
||||
log.info(f" MoDERouter: {r:>12,}")
|
||||
log.info(f" Total: {total:>12,} ({total/1e6:.1f}M)")
|
||||
log.info(f" vs 3T+ frozen experts = {total/3e12*100:.4f}% trainable")
|
||||
|
||||
def train_step(self, batch: dict) -> dict[str, float]:
|
||||
"""Single training step. Returns loss dict."""
|
||||
input_ids = batch["input_ids"].to(self.device)
|
||||
attention_mask = batch["attention_mask"].to(self.device)
|
||||
tier_labels = batch["tier"].to(self.device)
|
||||
modality_ids = batch["modality"].to(self.device)
|
||||
|
||||
with torch.amp.autocast("cuda", enabled=self.cfg.fp16 and self.device.type == "cuda"):
|
||||
# Complexity estimation
|
||||
tier_logits, mod_logits = self.complexity(input_ids, attention_mask)
|
||||
|
||||
# Losses
|
||||
l_complexity = complexity_loss(tier_logits, tier_labels)
|
||||
l_efficiency = efficiency_loss(tier_logits, tier_labels)
|
||||
|
||||
# For router loss, we need hidden states — use complexity encoder output
|
||||
# In production, this would use actual expert outputs
|
||||
hidden = self.complexity.embed(input_ids)
|
||||
|
||||
# Pick predicted tier and route
|
||||
pred_tier = tier_logits.argmax(dim=-1)
|
||||
# Use most common tier in batch for router forward
|
||||
most_common = pred_tier.mode().values.item()
|
||||
router_logits, cross_modal = self.router(hidden, most_common, modality_ids)
|
||||
|
||||
l_anticollapse = anticollapse_loss(router_logits)
|
||||
|
||||
# Total loss
|
||||
total = (
|
||||
self.cfg.lambda_complexity * l_complexity
|
||||
+ self.cfg.lambda_efficiency * l_efficiency
|
||||
+ self.cfg.lambda_anticollapse * l_anticollapse
|
||||
)
|
||||
|
||||
return {
|
||||
"total": total,
|
||||
"complexity": l_complexity.item(),
|
||||
"efficiency": l_efficiency.item(),
|
||||
"anticollapse": l_anticollapse.item(),
|
||||
"tier_accuracy": (pred_tier == tier_labels).float().mean().item(),
|
||||
}
|
||||
|
||||
def train(self, train_data: str, eval_data: Optional[str] = None):
|
||||
"""Full training loop with logging, eval, and checkpointing."""
|
||||
train_ds = ComplexityDataset(train_data, self.cfg.complexity_context_len)
|
||||
train_loader = DataLoader(
|
||||
train_ds, batch_size=self.cfg.batch_size, shuffle=True,
|
||||
num_workers=4, pin_memory=True, drop_last=True,
|
||||
)
|
||||
|
||||
eval_loader = None
|
||||
if eval_data and os.path.exists(eval_data):
|
||||
eval_ds = ComplexityDataset(eval_data, self.cfg.complexity_context_len)
|
||||
eval_loader = DataLoader(eval_ds, batch_size=self.cfg.batch_size * 2)
|
||||
|
||||
log.info(f"Training: {len(train_ds)} samples, {self.cfg.max_steps} steps")
|
||||
log.info(f"Device: {self.device}, FP16: {self.cfg.fp16}")
|
||||
|
||||
step = 0
|
||||
best_acc = 0.0
|
||||
running_loss = 0.0
|
||||
running_acc = 0.0
|
||||
t0 = time.time()
|
||||
|
||||
self.complexity.train()
|
||||
self.router.train()
|
||||
|
||||
while step < self.cfg.max_steps:
|
||||
for batch in train_loader:
|
||||
if step >= self.cfg.max_steps:
|
||||
break
|
||||
|
||||
losses = self.train_step(batch)
|
||||
loss = losses["total"]
|
||||
|
||||
# Backward
|
||||
self.scaler.scale(loss / self.cfg.gradient_accumulation).backward()
|
||||
|
||||
if (step + 1) % self.cfg.gradient_accumulation == 0:
|
||||
self.scaler.unscale_(self.optimizer)
|
||||
torch.nn.utils.clip_grad_norm_(
|
||||
list(self.complexity.parameters()) + list(self.router.parameters()),
|
||||
max_norm=1.0,
|
||||
)
|
||||
self.scaler.step(self.optimizer)
|
||||
self.scaler.update()
|
||||
self.optimizer.zero_grad()
|
||||
self.scheduler.step()
|
||||
|
||||
running_loss += losses["complexity"]
|
||||
running_acc += losses["tier_accuracy"]
|
||||
step += 1
|
||||
|
||||
# Logging
|
||||
if step % self.cfg.log_every == 0:
|
||||
avg_loss = running_loss / self.cfg.log_every
|
||||
avg_acc = running_acc / self.cfg.log_every
|
||||
lr = self.scheduler.get_last_lr()[0]
|
||||
elapsed = time.time() - t0
|
||||
log.info(
|
||||
f"Step {step:>6d}/{self.cfg.max_steps} | "
|
||||
f"loss={avg_loss:.4f} | acc={avg_acc:.4f} | "
|
||||
f"lr={lr:.2e} | {elapsed:.0f}s"
|
||||
)
|
||||
running_loss = 0.0
|
||||
running_acc = 0.0
|
||||
|
||||
# Eval
|
||||
if eval_loader and step % self.cfg.eval_every == 0:
|
||||
eval_acc = self.evaluate(eval_loader)
|
||||
if eval_acc > best_acc:
|
||||
best_acc = eval_acc
|
||||
self.save_checkpoint(os.path.join(self.output_dir, "best.pt"), step)
|
||||
log.info(f" New best: {best_acc:.4f}")
|
||||
|
||||
# Checkpoint
|
||||
if step % self.cfg.save_every == 0:
|
||||
self.save_checkpoint(
|
||||
os.path.join(self.output_dir, f"step_{step:06d}.pt"), step
|
||||
)
|
||||
|
||||
# Final save
|
||||
self.save_checkpoint(os.path.join(self.output_dir, "final.pt"), step)
|
||||
log.info(f"Training complete. {step} steps, best accuracy: {best_acc:.4f}")
|
||||
|
||||
@torch.no_grad()
|
||||
def evaluate(self, eval_loader: DataLoader) -> float:
|
||||
"""Evaluate tier prediction accuracy."""
|
||||
self.complexity.eval()
|
||||
correct, total = 0, 0
|
||||
for batch in eval_loader:
|
||||
input_ids = batch["input_ids"].to(self.device)
|
||||
attention_mask = batch["attention_mask"].to(self.device)
|
||||
tier_labels = batch["tier"].to(self.device)
|
||||
|
||||
tier_logits, _ = self.complexity(input_ids, attention_mask)
|
||||
pred = tier_logits.argmax(dim=-1)
|
||||
correct += (pred == tier_labels).sum().item()
|
||||
total += tier_labels.size(0)
|
||||
|
||||
acc = correct / max(total, 1)
|
||||
log.info(f" Eval accuracy: {acc:.4f} ({correct}/{total})")
|
||||
self.complexity.train()
|
||||
return acc
|
||||
|
||||
def save_checkpoint(self, path: str, step: int):
|
||||
torch.save({
|
||||
"step": step,
|
||||
"complexity": self.complexity.state_dict(),
|
||||
"alignment": self.alignment.state_dict(),
|
||||
"router": self.router.state_dict(),
|
||||
"optimizer": self.optimizer.state_dict(),
|
||||
"scheduler": self.scheduler.state_dict(),
|
||||
"config": self.cfg.__dict__,
|
||||
}, path)
|
||||
log.info(f" Checkpoint saved: {path}")
|
||||
|
||||
def load_checkpoint(self, path: str):
|
||||
ckpt = torch.load(path, map_location=self.device, weights_only=False)
|
||||
self.complexity.load_state_dict(ckpt["complexity"])
|
||||
self.alignment.load_state_dict(ckpt["alignment"])
|
||||
self.router.load_state_dict(ckpt["router"])
|
||||
if "optimizer" in ckpt:
|
||||
self.optimizer.load_state_dict(ckpt["optimizer"])
|
||||
log.info(f"Loaded checkpoint: {path} (step {ckpt.get('step', '?')})")
|
||||
return ckpt.get("step", 0)
|
||||
|
||||
|
||||
# ── CLI ──────────────────────────────────────────────────────────────────────
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="zen5 MoDE Router Training")
|
||||
sub = parser.add_subparsers(dest="cmd")
|
||||
|
||||
# info
|
||||
sub.add_parser("info", help="Show architecture info and param counts")
|
||||
|
||||
# train
|
||||
p = sub.add_parser("train", help="Train full pipeline (alignment + router)")
|
||||
p.add_argument("--data", default="./data/complexity_train.json")
|
||||
p.add_argument("--eval-data", default="./data/complexity_eval.json")
|
||||
p.add_argument("--output", default="./checkpoints/router/")
|
||||
p.add_argument("--max-steps", type=int, default=100000)
|
||||
p.add_argument("--batch-size", type=int, default=64)
|
||||
p.add_argument("--lr", type=float, default=3e-4)
|
||||
p.add_argument("--resume", default=None, help="Resume from checkpoint")
|
||||
|
||||
# eval
|
||||
p = sub.add_parser("eval", help="Evaluate checkpoint")
|
||||
p.add_argument("--checkpoint", required=True)
|
||||
p.add_argument("--data", default="./data/complexity_eval.json")
|
||||
|
||||
args = parser.parse_args()
|
||||
cfg = MoDEConfig()
|
||||
|
||||
if args.cmd == "info":
|
||||
ce = ComplexityEstimator(cfg)
|
||||
al = AlignmentLayer(cfg)
|
||||
rt = MoDERouter(cfg)
|
||||
|
||||
c = sum(p.numel() for p in ce.parameters())
|
||||
a = sum(p.numel() for p in al.parameters())
|
||||
r = sum(p.numel() for p in rt.parameters())
|
||||
total = c + a + r
|
||||
|
||||
print(f"\nzen5 MoDE Router Architecture")
|
||||
print("=" * 60)
|
||||
print(f"Hidden size: {cfg.hidden_size}")
|
||||
print(f"Vocab size: {cfg.vocab_size:,}")
|
||||
print(f"Max context: {cfg.max_position_embeddings:,}")
|
||||
print(f"Complexity tiers: {cfg.num_tiers} (T0-T4)")
|
||||
print(f"Experts per tier: {cfg.num_experts_per_tier}")
|
||||
print(f"Top-K per tier: {cfg.top_k_per_tier}")
|
||||
print(f"ReLU routing: {cfg.use_relu_routing}")
|
||||
print(f"Zero experts: {cfg.enable_zero_experts}")
|
||||
print(f"Modalities: {cfg.num_modalities}")
|
||||
print()
|
||||
print(f"Trainable Parameters:")
|
||||
print(f" ComplexityEstimator: {c:>12,}")
|
||||
print(f" AlignmentLayer: {a:>12,}")
|
||||
print(f" MoDERouter: {r:>12,}")
|
||||
print(f" Total: {total:>12,} ({total/1e6:.1f}M)")
|
||||
print(f"\nFrozen expert pool: ~3.0T+ params")
|
||||
print(f"Training efficiency: {total/3e12*100:.4f}% of total model")
|
||||
|
||||
elif args.cmd == "train":
|
||||
cfg.max_steps = args.max_steps
|
||||
cfg.batch_size = args.batch_size
|
||||
cfg.lr = args.lr
|
||||
|
||||
trainer = MoDETrainer(cfg, args.output)
|
||||
if args.resume:
|
||||
trainer.load_checkpoint(args.resume)
|
||||
trainer.train(args.data, args.eval_data)
|
||||
|
||||
elif args.cmd == "eval":
|
||||
trainer = MoDETrainer(cfg, ".")
|
||||
trainer.load_checkpoint(args.checkpoint)
|
||||
|
||||
eval_ds = ComplexityDataset(args.data, cfg.complexity_context_len)
|
||||
eval_loader = DataLoader(eval_ds, batch_size=cfg.batch_size * 2)
|
||||
trainer.evaluate(eval_loader)
|
||||
|
||||
else:
|
||||
parser.print_help()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user