Metadata-Version: 2.4
Name: cupydiff
Version: 0.2.10
Summary: High-Performance Discrete Diffusion Language Model in Triton + CuPy (Functional, No PyTorch)
Author: vovaRL
License-Expression: Apache-2.0
Classifier: Programming Language :: Python :: 3
Classifier: Operating System :: OS Independent
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.9
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: numpy>=1.22.0
Requires-Dist: safetensors>=0.4.0
Requires-Dist: huggingface_hub>=0.20.0
Provides-Extra: cuda
Requires-Dist: cupy-cuda12x>=13.0.0; extra == "cuda"
Requires-Dist: triton>=3.0.0; extra == "cuda"
Requires-Dist: nvidia-nccl-cu12>=2.18.1; platform_system == "Linux" and extra == "cuda"
Requires-Dist: nvidia-cublas-cu12>=12.1.0; platform_system == "Linux" and extra == "cuda"
Requires-Dist: nvidia-cuda-runtime-cu12>=12.1.0; platform_system == "Linux" and extra == "cuda"
Requires-Dist: nvidia-cuda-nvrtc-cu12>=12.1.0; platform_system == "Linux" and extra == "cuda"
Provides-Extra: test
Requires-Dist: pytest>=7.0.0; extra == "test"
Dynamic: license-file

# cupydiff

> **High-Performance Discrete Diffusion Language Models in Triton + CuPy**  
> *Strictly NO PyTorch dependency. Pure functional architecture. Optimized for 8x NVIDIA RTX PRO 6000 96GB (Blackwell).*

---

## Highlights

- **Strictly No PyTorch**: Built entirely on **Triton** kernels and **CuPy** array operations with pure NumPy fallback for CPU testing.
- **Pure Functional Design**: Stateless mathematical functions, explicit weight dictionaries, and deterministic analytical forward and backward passes.
- **Triton GPU Kernels**:
  - **Fused RMSNorm & AdaLN**: High-bandwidth elementwise normalization and adaptive modulation.
  - **Fused RoPE**: Analytical rotary position embedding rotation.
  - **Fused SwiGLU**: `silu(gate) * up` with analytical backward gradients.
  - **Bidirectional FlashAttention**: Non-causal scaled dot-product attention executing in $O(L)$ SRAM instead of $O(L^2)$ global memory.
  - **Cut Cross-Entropy (CCE)**: Chunked vocabulary projection that avoids materializing $(N, 49152)$ float32 logit tensors, preventing out-of-memory errors on massive context lengths.
- **Hybrid Muon + AdamW Optimizer**:
  - **Muon**: Quintic Newton-Schulz iteration (polar orthogonalization) for internal 2D weight matrices with Nesterov momentum.
  - **AdamW**: First and second moment buffers with decoupled weight decay for 1D vectors, biases, and token embeddings.
  - **Cosine Learning Rate Schedule**: Warmup + cosine decay.
- **Multi-GPU Scalability (8x RTX PRO 6000 96GB Blackwell)**:
  - Native **CuPy NCCL** communicator with in-place gradient all-reduce.
  - Multi-process launcher (`cupydiff-launch` or `python -m cupydiff.cli`).
- **Curriculum Shard Pipeline**:
  - Compatible with binary token shards (`.bin`, uint16) and direct streaming from Hugging Face: [`vovaRL/slm388-corpus`](https://huggingface.co/datasets/vovaRL/slm388-corpus).

---

## Installation

### From Source / Local Development
```bash
git clone https://github.com/vovaRL/cupydiff.git
cd cupydiff
pip install -e .
```

### With CUDA / Triton Acceleration (GPU)
```bash
pip install cupydiff[cuda]
```
*(Installs `cupy-cuda12x` and `triton`)*

---

## Testing

`cupydiff` includes a complete standalone test suite that runs on **CPU** (using NumPy reference kernels) when no GPU is present, and automatically verifies **Triton + CuPy** GPU kernels when CUDA is available:

```bash
python tests/run_all.py
```
or with pytest:
```bash
pytest -v tests
```

---

## Usage

### 1. Profiling & Benchmark Mode
Measure exact kernel latencies and throughput (tokens/sec):

```bash
# Single GPU
cupydiff --mode bench --bench_steps 20 --batch_size 64 --seq_len 2048

# 8x RTX PRO 6000 96GB
cupydiff --mode bench --nproc_per_node 8 --batch_size 512 --seq_len 2048
```

### 2. Pretraining on `vovaRL/slm388-corpus`
Trains the MDLM model using the curriculum binary shard loader:

```bash
cupydiff --mode train \
    --hf_dataset vovaRL/slm388-corpus \
    --data_dir /data/slm388-corpus \
    --batch_size 512 \
    --seq_len 2048 \
    --nproc_per_node 8 \
    --total_steps 85000 \
    --repo_id vovaRL/DiffLM-388M \
    --checkpoint_interval 2500
```

---

## Hugging Face Jobs (8x RTX PRO 6000 96GB Blackwell)

To run distributed training on Hugging Face Jobs with 8x RTX PRO 6000:

```bash
hf jobs run \
    --flavor 8x-rtx-pro-6000 \
    --image nvidia/cuda:12.4.1-devel-ubuntu22.04 \
    --env HF_TOKEN="hf_your_token_here" \
    --command "pip install cupydiff[cuda] && cupydiff --mode train --hf_dataset vovaRL/slm388-corpus --batch_size 512 --seq_len 2048 --nproc_per_node 8"
```

---

## Building and Pushing to PyPI

```bash
pip install build twine
python -m build
twine check dist/*
twine upload dist/*
```

---

## Architecture Specification

| Hyperparameter | Value | Description |
|---|---|---|
| Model Dimension ($d$) | 512 | Hidden dimension |
| Diffusion Blocks | 3 | Block specializations ($t \in [0, 1]$) |
| Layers per Block | 6 | 18 total transformer layers |
| Attention Heads | 8 | Multi-head attention |
| Head Dimension | 64 | Query/Key/Value dimension ($8 \times 64 = 512$) |
| RoPE Theta ($\theta$) | 50,000.0 | Rotary frequency base |
| SwiGLU Hidden Dim | 1408 | MLP intermediate expansion |
| Vocabulary Size | 49,152 | Token vocabulary |
| Mask Token ID | 49,151 | `[MASK]` token for diffusion corruption |
| CCE Chunk Size | 8,192 | Cut Cross-Entropy projection chunk |
| Optimizer | Hybrid Muon + AdamW | Fused quintic Newton-Schulz for 2D weights, AdamW for 1D/heads |

---

## License

Apache License 2.0.
