MoEWatch v0.1.0 — Initial Release 🎉

🎉 MoEWatch v0.1.0 — Initial Release

The pytest for Mixture-of-Experts models. Catch expert collapse, routing entropy collapse, and load imbalance — before they silently wreck your training run.

pip install moewatch

---

🧠 What is MoEWatch?

MoEWatch is a lightweight diagnostic and audit library for MoE models in HuggingFace Transformers. Drop it into any training loop — it instruments router modules with zero-weight-modification PyTorch hooks, aggregates routing statistics, and surfaces structured alerts the moment something goes wrong.

It never stops your training. It diagnoses. You decide.

---

✨ What's New in v0.1.0

🔍 Core Diagnostics Engine

Feature	Description
💀 Expert Collapse Detection	Tracks dead and cold experts per layer across the full training run
📊 Routing Entropy Analysis	Catches distribution collapse relative to theoretical maximum entropy
⚖️ Load Imbalance Alerts	Fires when any single expert dominates token dispatch (max/mean ratio)

🧩 Auto-Detection Registry

MoEWatch ships with a curated architecture registry — no configuration needed for 9 major model families:

🏗️ Family	🤖 Models
Mixtral	mistralai/Mixtral-*
OLMoE	allenai/OLMoE-*
DeepSeek-MoE	deepseek-ai/DeepSeek-V2, DeepSeek-V3
Qwen-MoE	Qwen/Qwen2-MoE-*, Qwen3-MoE-*
Phi-MoE	microsoft/Phi-*-MoE
Switch Transformer	Google's HuggingFace port
NLLB-MoE	facebook/nllb-moe-*
Arctic	Snowflake/snowflake-arctic-*
Jamba	ai21labs/Jamba-*

Unknown architecture? MoEWatch falls back to a heuristic scan using class name analysis. Or just specify WatchConfig(router_modules=[...]) for full manual control.

---

⚡ Two Integration Modes

🔬 Offline Audit — One-shot diagnostic

import moewatch
 
report = moewatch.audit(model, dataloader, steps=200)
print(report.summary())

Runs N forward passes, collects routing stats, and returns a structured AuditReport — no training loop changes needed.

🔴 Live Monitor — HuggingFace Trainer

from moewatch import MoEWatch, WatchConfig
 
watcher = MoEWatch(model, config=WatchConfig())
watcher.attach(trainer)   # injects as a TrainerCallback
trainer.train()
watcher.detach()

🔄 Live Monitor — Custom Loop

from moewatch import MoEWatch
 
watcher = MoEWatch(model)
watcher.start()
 
for step, batch in enumerate(dataloader):
    loss = model(**batch).loss
    loss.backward()
    optimizer.step()
    alerts = watcher.step(step)   # List[Alert]; empty when healthy
 
watcher.stop()

Also supports context manager syntax: with MoEWatch(model) as watcher:

---

⚙️ Configuration Presets

Three built-in presets cover most use cases:

from moewatch import WatchConfig
 
WatchConfig.default()      # ✅ Balanced — recommended starting point
WatchConfig.aggressive()   # 🔬 Tighter thresholds, every-step sampling — for debugging
WatchConfig.lightweight()  # 🪶 Minimal overhead — for large-scale production runs
WatchConfig.silent()       # 🔇 No output — results only via AuditReport

Or configure everything manually:

config = WatchConfig(
    dead_threshold=0.001,        # < 0.1% token share → expert is DEAD
    entropy_warn=0.60,           # < 60% of H_max → WARN
    entropy_critical=0.40,       # < 40% of H_max → ERROR
    load_imbalance_error=5.0,    # max/mean > 5× → ERROR
    sample_every=10,             # instrument every 10th forward pass
    output="json",               # "console" | "json" | "silent"
)

---

📢 Alert System

MoEWatch uses a 3-level alert ladder:

Level	Icon	Meaning
INFO	✅	Routine routing statistics — everything healthy
WARN	⚠️	Degraded routing — investigate soon
ERROR	❌	Severe collapse or imbalance — likely harming training

Every ERROR alert ships with a 💡 Suggestion — actionable advice for fixing the specific problem (e.g. raise aux_loss_coef, check router_jitter_noise, etc.).

---

📤 Output Modes

WatchConfig(output="console")   # 🖥️ Coloured ASCII — human-readable
WatchConfig(output="json")      # 📦 Newline-delimited JSON — pipe to Grafana, Splunk, etc.
WatchConfig(output="silent")    # 🔇 Silent — results only via AuditReport

Respects NO_COLOR environment variable automatically. 🎨

---

🛡️ Design Principles

• ⚡ Zero weight modifications — hooks never touch model parameters
• 🧱 Fixed memory footprint — ring buffer with configurable capacity; no unbounded growth
• 🔒 Always detach on exception — no leaked PyTorch hooks, ever
• 📉 Configurable overhead — sample_every=10 keeps instrumentation below 2% in production
• 🐍 Python 3.8 – 3.12 compatible
• 🔥 PyTorch ≥ 2.0, Transformers ≥ 4.36.0

---

📦 Installation

pip install moewatch

Optional extras:

pip install "moewatch[wandb]"        # WandB integration
pip install "moewatch[tensorboard]"  # TensorBoard integration
pip install "moewatch[dev]"          # Development tools

---

📚 Documentation

• 📖 Quick Start Guide
• ⚙️ Configuration Reference
• 🔌 API Reference
• 🧩 Custom Architecture Guide
• 🤝 Contributing

---

🤝 Contributing

Issues and PRs are welcome! To add a new architecture:

1. Open an issue or
2. Add your router class name(s) to _ARCHITECTURE_REGISTRY in hooks/detection.py and submit a PR

---

📜 License

Apache 2.0 — see LICENSE

---

Built with ❤️ by Abinesh · PyPI · GitHub