Metadata-Version: 2.4
Name: d9d
Version: 0.22.0
Summary: d9d - d[istribute]d - distributed training framework based on PyTorch that tries to be efficient yet hackable
Author: Maksim Afanasyev
Author-email: Maksim Afanasyev <mr.applexz@gmail.com>
License-Expression: Apache-2.0
License-File: LICENSE
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Education
Classifier: Intended Audience :: Science/Research
Classifier: Topic :: Scientific/Engineering
Classifier: Topic :: Scientific/Engineering :: Mathematics
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Software Development
Classifier: Topic :: Software Development :: Libraries
Classifier: Topic :: Software Development :: Libraries :: Python Modules
Requires-Dist: pydantic>=2.0.0
Requires-Dist: torchdata>=0.11.0
Requires-Dist: safetensors>=0.7.0
Requires-Dist: tqdm>=4.0.0
Requires-Dist: triton>=3.6.0
Requires-Dist: torch>=2.14.0
Requires-Dist: optree>=0.19.0
Requires-Dist: aim>=3.0.0,<4.0.0 ; extra == 'aim'
Requires-Dist: setuptools>=70.0.0 ; extra == 'aim'
Requires-Dist: flash-attn>=2.0.0 ; extra == 'backend-sdpa-flash-attention-2'
Requires-Dist: flash-attn-4>=4.0.0b32 ; extra == 'backend-sdpa-flash-attention-4'
Requires-Dist: cut-cross-entropy>=25.9.3 ; extra == 'cce'
Requires-Dist: fla-core>=0.4.2 ; extra == 'linear-attention'
Requires-Dist: deep-ep>=1.2.1 ; extra == 'moe'
Requires-Dist: nv-grouped-gemm>=1.1.4 ; extra == 'moe'
Requires-Dist: plotly>=6.0.0 ; extra == 'visualization'
Requires-Python: >=3.11, <3.15
Project-URL: Repository, https://github.com/d9d-project/d9d
Project-URL: Documentation, https://d9d-project.github.io/d9d
Project-URL: Homepage, https://d9d-project.github.io/d9d
Project-URL: Issues, https://github.com/d9d-project/d9d/issues
Project-URL: Discord Server, https://discord.gg/sNRjDbxVrg
Provides-Extra: aim
Provides-Extra: backend-sdpa-flash-attention-2
Provides-Extra: backend-sdpa-flash-attention-4
Provides-Extra: cce
Provides-Extra: linear-attention
Provides-Extra: moe
Provides-Extra: visualization
Description-Content-Type: text/markdown

# The d9d Project

![PyPI - Version](https://img.shields.io/pypi/v/d9d)
![PyPI - License](https://img.shields.io/pypi/l/d9d)
[![Discord](https://img.shields.io/badge/Discord-%235865F2.svg?logo=discord&logoColor=white)](https://discord.gg/sNRjDbxVrg) [![Docs](https://img.shields.io/badge/docs-latest-blue.svg)](https://d9d-project.github.io/d9d/)

**d9d** is a distributed training framework built on top of modern PyTorch. It aims to be hackable, modular, and efficient, designed to scale from single-GPU debugging to massive clusters running 6D parallelism.

[LET'S START TRAINING 🚀](https://d9d-project.github.io/d9d/)

## Installation

Install it with your package manager:
```bash
pip install d9d
poetry add d9d
uv add d9d
```

### Extras

*   `d9d[aim]`: [Aim](https://aimstack.io/) experiment tracker integration.
*   `d9d[visualization]`: [Plotly](https://plotly.com/python/), for plotting learning rate schedules.
*   `d9d[linear-attention]`: [Flash Linear Attention](https://github.com/fla-org/flash-linear-attention) kernels for linear attention.
*   `d9d[backend-sdpa-flash-attention-2]`: [FlashAttention 2](https://github.com/Dao-AILab/flash-attention) SDPA backend kernels. You may need to build FlashAttention 2 from source if no prebuilt wheel matches your setup.
*   `d9d[backend-sdpa-flash-attention-4]`: [FlashAttention 4](https://github.com/Dao-AILab/flash-attention) SDPA backend kernels.
*   `d9d[moe]`: Mixture-of-Experts (MoE) GPU kernels. You must build and install [DeepEP](https://github.com/deepseek-ai/DeepEP) and [grouped-gemm](https://github.com/fanshiqing/grouped_gemm/) manually first.
*   `d9d[cce]`: Fused cross-entropy kernels. You must install [Cut Cross Entropy](https://github.com/apple/ml-cross-entropy) from its repository first.

## Examples

*   **[Qwen3-MoE Pretraining](https://github.com/d9d-project/d9d/blob/main/example/qwen3_moe/pretrain.py)**: Causal LM pretraining of a Qwen3-MoE model.

---

## About

### Why Another Framework?

Distributed training frameworks such as Megatron-LM are monolithic: you run a script from the command line to train one of a set of *predefined* models, using *predefined* regimes. These systems are hard to hack and to integrate into new research workflows. They aim to be a complete end-to-end solution, which limits flexibility for experiment-driven research.

Writing your own distributed training solution from scratch is also hard. You must implement many low-level components that are the same across setups, such as distributed checkpoints and synchronization. You must also fix common performance bottlenecks yourself.

d9d fills the gap between monolithic frameworks and homebrew setups. It gives you modular building blocks for distributed training.

### What d9d Is and Isn't

In terms of core concept:

*   **IS** a pluggable framework for implementing distributed training regimes for your deep learning models.
*   **IS** built on clear interfaces and building blocks that may be composed and implemented in your own way.
*   **IS NOT** an all-in-one CLI platform for setting up pre-training and post-training like torchtitan, Megatron-LM, or torchforge.

In terms of codebase and engineering:

*   **IS** built on a strong engineering foundation: we enforce strict type checking and linting to catch errors before execution.
*   **IS** reliable: a suite of over 900 tests covers unit logic, integration flows, and end-to-end distributed scenarios.
*   **IS** eager to use performance hacks (like DeepEP or custom kernels) if they improve MFU, even if they aren't PyTorch-native.
*   **IS NOT** for legacy setups: we do not maintain backward compatibility with older PyTorch versions or hardware. We prefer simplicity and modern APIs (like `DTensor`).

### Key Philosophies

To balance hackability and performance, d9d follows these design principles:

*   **Composition over Monoliths**: We avoid "God Classes" like `DistributedDataParallel` or `ParallelDims` that own the entire execution loop. Instead, we provide composable and extendable APIs, such as per-layer horizontal parallelism strategies (`parallelize_replicate`, `parallelize_expert_parallel`, ...).
*   **White-Box Modeling**: We encourage standard PyTorch code. Models are not wrapped in metadata specifications. They are standard `nn.Module`s that implement small protocols.
*   **Pragmatic Efficiency**: We prefer native PyTorch, but we integrate non-native solutions if they improve MFU. For example, our MoE uses DeepEP communications, reindexing kernels from Megatron-LM, and grouped GEMM kernels.
*   **Graph-Based State Management**: Our I/O system treats model checkpoints as directed acyclic graphs. You can transform architectures (e.g. merge `q`, `k`, `v` into `qkv`) on the fly while streaming from disk, without loading the whole checkpoint into memory.
*   **DTensors**: Distributed parameters must be `torch.distributed.tensor.DTensor`s. DTensors know their topology, which makes checkpointing simpler. We use modern PyTorch APIs (`DeviceMesh`) wherever we can.

## Community and Support

Join our community if you need help, want to discuss large-scale training strategies, or want to contribute to d9d.

*   💬 **Discord**: [Join the d9d server](https://discord.gg/sNRjDbxVrg) for real-time chat, support, and framework hacking.
*   🐛 **Issues**: For bug reports and feature requests, please use the [GitHub Issue Tracker](https://github.com/d9d-project/d9d/issues).
