Metadata-Version: 2.3
Name: slm388-torchdiff
Version: 0.1.4
Summary: Masked Diffusion Language Model (MDLM) pretraining in PyTorch with FSDP2, Muon optimizer, and Hugging Face Jobs support.
Project-URL: Homepage, https://github.com/vovaRL/torchdiff
Project-URL: Repository, https://github.com/vovaRL/torchdiff
Project-URL: Issues, https://github.com/vovaRL/torchdiff/issues
Project-URL: Dataset, https://huggingface.co/datasets/vovaRL/slm388-corpus
Author: vovaRL
License: MIT
Keywords: deep-learning,diffusion-models,discrete-diffusion,fsdp2,huggingface-jobs,language-model,mdlm,muon,pytorch
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.10
Requires-Dist: cut-cross-entropy>=25.1.1
Requires-Dist: filelock>=3.12.0
Requires-Dist: huggingface-hub>=0.28.0
Requires-Dist: numpy>=1.24.0
Requires-Dist: safetensors>=0.4.0
Requires-Dist: torch>=2.4
Provides-Extra: dev
Requires-Dist: build>=1.0.0; extra == 'dev'
Requires-Dist: pytest>=8.0.0; extra == 'dev'
Requires-Dist: twine>=5.0.0; extra == 'dev'
Description-Content-Type: text/markdown

# TorchDiff (`slm388-torchdiff`)

[![PyPI Version](https://img.shields.io/pypi/v/slm388-torchdiff.svg)](https://pypi.org/project/slm388-torchdiff/)
[![Python Versions](https://img.shields.io/pypi/pyversions/slm388-torchdiff.svg)](https://pypi.org/project/slm388-torchdiff/)
[![License: MIT](https://img.shields.io/badge/License-MIT-blue.svg)](https://opensource.org/licenses/MIT)
[![Hugging Face Dataset](https://img.shields.io/badge/HF%20Dataset-vovaRL%2Fslm388--corpus-yellow)](https://huggingface.co/datasets/vovaRL/slm388-corpus)

**TorchDiff** (`slm388-torchdiff`) is a high-performance PyTorch implementation of a **Masked Diffusion Language Model (MDLM)** with:
- **Distributed Training:** Native PyTorch FSDP2 (Fully Sharded Data Parallel 2) multi-GPU sharding.
- **Hybrid Optimizer:** Muon (Newton-Schulz orthogonalization) for 2D hidden matrices + AdamW for embeddings, norms, and heads.
- **Curriculum Learning:** Automatic pretraining vs. SFT dynamic mixing and context extension (2K $\to$ 4K).
- **Fast Loss:** Memory-efficient token loss via Apple's [Cut Cross-Entropy](https://github.com/apple/ml-cross-entropy) with automatic PyTorch fallback.
- **Hugging Face Jobs Ready:** Seamless zero-setup execution with `hf jobs uv run` and dataset volume mounting.

---

## 1. Installation

### From PyPI
```bash
pip install slm388-torchdiff
```

### From GitHub Repository
```bash
pip install "git+https://github.com/vovaRL/torchdiff.git"
```

### Editable Development Installation
```bash
git clone https://github.com/vovaRL/torchdiff.git
cd torchdiff
pip install -e ".[dev]"
```

---

## 2. Dataset: `vovaRL/slm388-corpus`

TorchDiff is pre-configured to train on the [vovaRL/slm388-corpus](https://huggingface.co/datasets/vovaRL/slm388-corpus) dataset.

### Dataset Overview
- **Vocabulary:** SmolLM2-compatible (49,152 vocab size, token ID 49,151 as mask token).
- **Pretraining Shards:** Cosmopedia, DCLM, FineMath, FineWeb-Edu, and Python-Edu `.bin` binary token shards.
- **SFT Shards:** SmolTalk instruction tuning shards (`*smoltalk*.bin`), dynamically phased in via curriculum scheduling.
- **Validation Shards:** Dedicated held-out validation shards (`*_val_*.bin`).

### Dynamic Shard Discovery
TorchDiff's `CurriculumBinLoader` automatically discovers shards and partitions them:
```
[Dataset Discovery]
Directory:             /dataset
Total shards found:    64
├── Pretraining shards: 56
├── SFT shards:         1
└── Validation shards:  7 (held out)
```

---

## 3. Training on Hugging Face Jobs (`hf jobs uv run`)

Hugging Face Jobs allows you to train on high-end cloud GPUs (A100, H200, A10G) without managing servers or Docker images.

### Prerequisites
1. Install the Hugging Face CLI:
   ```bash
   pip install -U huggingface_hub
   ```
2. Log in with your write token:
   ```bash
   hf auth login
   ```

### Method A: Run with Published PyPI Package (Recommended)
Mount the dataset volume and pass `--clone_to_ssd` to copy shards to local NVMe SSD (`/tmp/slm388-corpus`) for maximum throughput:
```bash
hf jobs uv run \
  --flavor a100-large \
  --secrets HF_TOKEN \
  --with slm388-torchdiff \
  -v hf://datasets/vovaRL/slm388-corpus:/dataset \
  torchdiff-train \
  --data_dir /dataset \
  --clone_to_ssd \
  --batch_size 128 \
  --repo_id vovaRL/DiffLM
```

### Method B: Clone directly from Hugging Face Hub to Local SSD (Zero Mounts)
HF Jobs nodes (e.g. `a100-large` with 1,000 GB NVMe SSD) download cloud-to-cloud directly to local SSD:
```bash
hf jobs uv run \
  --flavor a100-large \
  --secrets HF_TOKEN \
  --with slm388-torchdiff \
  torchdiff-train \
  --data_dir /tmp/slm388-corpus \
  --batch_size 128 \
  --repo_id vovaRL/DiffLM
```

### Method C: Run Standalone Script with Local SSD
Run the self-contained script `train_standalone.py` locally from this repository:
```bash
hf jobs uv run \
  --flavor a100-large \
  --secrets HF_TOKEN \
  -v hf://datasets/vovaRL/slm388-corpus:/dataset \
  train_standalone.py \
  --data_dir /dataset \
  --clone_to_ssd \
  --batch_size 128 \
  --repo_id vovaRL/DiffLM
```

### Multi-GPU Training (FSDP2)
Scale to multi-GPU flavors (e.g. 4x or 8x A100):
```bash
hf jobs uv run \
  --flavor a100x4 \
  --secrets HF_TOKEN \
  --with slm388-torchdiff \
  -v hf://datasets/vovaRL/slm388-corpus:/dataset \
  torchdiff-train \
  --data_dir /dataset \
  --batch_size 256 \
  --repo_id vovaRL/DiffLM
```

### Persistent Bucket Checkpointing
Mount a Hugging Face Bucket as a read-write volume to save checkpoints continuously:
```bash
hf jobs uv run \
  --flavor a100-large \
  --secrets HF_TOKEN \
  --with slm388-torchdiff \
  -v hf://datasets/vovaRL/slm388-corpus:/dataset \
  -v hf://buckets/vovaRL/diff-checkpoints:/checkpoints \
  torchdiff-train \
  --data_dir /dataset \
  --checkpoint_dir /checkpoints \
  --checkpoint_interval 1000
```

---

## 4. Local Execution

### Run Training Locally
```bash
# Using installed CLI
torchdiff-train --data_dir ./data --batch_size 32

# Or using uv run
uv run train.py --data_dir ./data --batch_size 32
```

### Benchmark Mode
Benchmark raw forward/backward throughput without downloading the full dataset:
```bash
torchdiff-train --mode bench --synthetic_bench --batch_size 64 --seq_len 2048
```

---

## 5. CLI Arguments

| Flag | Default | Description |
|---|---|---|
| `--mode` | `train` | Mode: `'train'` or `'bench'`. |
| `--data_dir` | `None` | Directory containing `.bin` shards. Auto-detects `/dataset`, `/data`, `./data`. |
| `--dataset` | `vovaRL/slm388-corpus` | Hugging Face dataset repo ID to download if local shards are not found. |
| `--download_pattern` | `None` | Glob pattern for partial dataset download (e.g. `'*.bin'`). |
| `--batch_size` | `512` | Batch size across GPUs. |
| `--seq_len` | `2048` | Initial token context length (extends to 4096 in Phase 4). |
| `--total_steps` | `85000` | Total pretraining steps. |
| `--checkpoint` | `None` | Path to `.safetensors` checkpoint to resume training. |
| `--checkpoint_dir` | `./checkpoints` | Local/volume checkpoint directory. |
| `--checkpoint_interval` | `2500` | Steps between saving checkpoints. |
| `--repo_id` | `$HF_REPO_ID` | Hugging Face model repository to upload saved checkpoints. |
| `--no_compile` | `False` | Disable `torch.compile`. |

---

## 6. Building and Publishing to PyPI

Build sdist and wheel:
```bash
uv build
# or: python -m build
```

Upload to PyPI using Twine:
```bash
twine upload dist/*
```

---

## License
MIT License. See [LICENSE](LICENSE) for details.
