PithTrain is a compact, agent-native Mixture-of-Experts (MoE) training system designed for production-grade performance and enhanced human/AI navigability. It addresses the high cost of evolving complex MoE frameworks by offering a ~11K line Python codebase, minimizing implicit indirection. This allows AI coding agents and developers to understand and extend the system more easily while still delivering advanced features like 4D parallelism and FP8 training.

How It Works

PithTrain employs a layered architecture comprising Application, Engine, and Operator layers. The Engine is central, featuring modules for Model protocols (supporting Qwen, DeepSeek, GPT-OSS), Building Blocks (FP8, quantization, ring attention), a Pipeline Engine (DualPipeV scheduler for 5-stage overlapped execution), Distributed Training (4D parallelism: PP x FSDP x CP x EP), and Training Infrastructure. This design prioritizes agent-native patterns and explicit control, enabling efficient compute-communication overlap and deep customization within a readable codebase.

Quick Start & Requirements

• Installation:

  git clone https://github.com/mlc-ai/pith-train.git && cd pith-train
  uv venv
  uv pip install .
  

  For developers: uv sync.
• Prerequisites: NVIDIA Hopper (SM90) or Blackwell (SM100) GPUs, CUDA >= 13.0, Python >= 3.12. uv is used for dependency management.
• Getting Started: Involves dataset tokenization (examples/tokenize_corpus/launch.sh), training configuration (examples/pretrain_lm/script.py), launching training (examples/pretrain_lm/launch.sh), and checkpoint conversion (examples/convert_checkpoint/launch.sh).
• Links: Project repository: https://github.com/mlc-ai/pith-train.git. Detailed architecture: docs/architecture.md. Contribution guidelines: CONTRIBUTING.md.

Highlighted Details

• Production-grade performance with 4D parallelism (PP x FSDP x CP x EP).
• Supports FP8 training, compute-communication overlap, and ring attention.
• Compact (~11K Python lines) and agent-native design for improved navigability.
• DualPipeV scheduler enables 5-stage overlapped forward-backward execution.
• Automatic training resumption and reshardable checkpoints.
• Tools for converting checkpoints to/from Hugging Face format.

Maintenance & Community

The project is developed by contributors from CMU. Specific community channels (e.g., Discord, Slack) or a public roadmap are not detailed in the provided README.

Licensing & Compatibility

PithTrain is released under the Apache 2.0 License, which generally permits commercial use and integration into closed-source projects.

Limitations & Caveats

Strict hardware requirements mandate NVIDIA Hopper (SM90) or Blackwell (SM100) GPUs. Advanced configuration and troubleshooting details are primarily directed to an external "User Guide," suggesting a learning curve for complex setups. The project appears relatively new, with an associated arXiv paper dated 2026.