TorchDiff (slm388-torchdiff)
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 with automatic PyTorch fallback.
- Hugging Face Jobs Ready: Seamless zero-setup execution with
hf jobs uv runand dataset volume mounting.
1. Installation
From PyPI
pip install slm388-torchdiff
From GitHub Repository
pip install "git+https://github.com/vovaRL/torchdiff.git"
Editable Development Installation
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 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
.binbinary 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
- Install the Hugging Face CLI:
pip install -U huggingface_hub
- Log in with your write token:
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:
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:
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:
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):
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:
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 (Multi-GPU or Single-GPU)
# Single GPU (e.g. RTX 6000 96GB Blackwell Server or 48GB Ada):
# Using gradient accumulation to preserve the full 512 effective global batch size:
python train.py --data_dir ./data --batch_size 512 --micro_batch_size 128
# Or with activation checkpointing to further reduce activation memory:
python train.py --data_dir ./data --batch_size 512 --micro_batch_size 128 --gradient_checkpointing
# Or using uv run:
uv run train.py --data_dir ./data --batch_size 512 --micro_batch_size 128
Benchmark Mode
Benchmark raw forward/backward throughput without downloading the full dataset:
python train.py --mode bench --synthetic_bench --batch_size 512 --micro_batch_size 128 --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 |
Effective global batch size across all GPUs and accumulation steps. |
--micro_batch_size |
None |
Micro-batch size per forward pass. When set (e.g. 128 for 96GB GPU, 64 for 48GB GPU), gradient accumulation is automatically computed. |
--grad_accum_steps |
None |
Explicit gradient accumulation steps per optimizer update. Overrides micro_batch_size calculation if provided. |
--gradient_checkpointing |
False |
Enable activation checkpointing across diffusion blocks to drastically reduce VRAM usage. |
--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:
uv build
# or: python -m build
Upload to PyPI using Twine:
twine upload dist/*
License
MIT License. See LICENSE for details.
1 maintainer
Metadata
1 maintainer
Release files for slm388-torchdiff 0.1.72
For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.
Source distribution (sdist)
| File | Size | Uploaded | |
|---|---|---|---|
| slm388_torchdiff-0.1.72.tar.gz | 72.6 kB | Details |
Built distribution (wheel)
| File | Interpreter | ABI | Platform | Reset |
|---|---|---|---|---|
| slm388_torchdiff-0.1.72-py3-none-any.whl | Python 3 | none | any | Details |
Total release size: 150.6 kB
Release files / slm388_torchdiff-0.1.72.tar.gz
| Download URL | slm388_torchdiff-0.1.72.tar.gz |
|---|---|
| Size | 72.6 kB |
| Tags | Source |
|
SHA-256 checksum How to use checksums |
8b92205e74a7609b250c385aa25ef500cb3aaf816fd13e8874bfd5327f291921
|
|
BLAKE2b-256 checksum How to use checksums |
2aa43a42dc1145936dbfec493d047749ed90340803d7e22aaf3f72099c10a96a
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/7.0.0 CPython/3.14.4
|
Release files / slm388_torchdiff-0.1.72-py3-none-any.whl
| Download URL | slm388_torchdiff-0.1.72-py3-none-any.whl |
|---|---|
| Size | 78.1 kB |
| Tags | Python 3 |
|
SHA-256 checksum How to use checksums |
f5df9c62ace0a14632c00e1e63093e998fce25fc35b96b7ba79cfbd50aa6d9d1
|
|
BLAKE2b-256 checksum How to use checksums |
dc0ab01ad4ec47e06b6a29f79a45234b527dd382f00d1988e98b3ad650d9f8cf
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/7.0.0 CPython/3.14.4
|
