MoEWatch
A diagnostic and audit library for Mixture-of-Experts models in HuggingFace Transformers that instruments router modules to catch expert collapse, routing entropy collapse and load imbalance during training.
MoEWatch attaches zero-weight-modification PyTorch hooks to router modules in a Mixture-of-Experts model, aggregates routing statistics, and surfaces structured alerts as problems appear, rather than after a training run has silently degraded. It tracks dead and cold experts per layer, detects routing-distribution collapse relative to theoretical maximum entropy, and fires an alert when a single expert dominates token dispatch.
It auto-detects common MoE architectures including Mixtral, OLMoE, DeepSeek-MoE, Qwen-MoE, Phi-MoE, Switch Transformer, NLLB-MoE, Arctic and Jamba, falling back to a heuristic scan for unrecognized architectures, and any custom architecture can be targeted explicitly. It supports a one-shot offline audit call or live training-time monitoring, and attaches to a HuggingFace Trainer as a callback in one line.
Output can be printed as colored console text, structured JSON for log pipelines, or collected silently into a report object. Sampling frequency is configurable to trade monitoring overhead against fidelity, and a fixed-size ring buffer keeps memory use bounded over long runs.