TorchDiff (slm388-torchdiff)

PyPI Version Python Versions License: MIT Hugging Face Dataset

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 run and 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 .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: 
    pip install -U huggingface_hub
  2. 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

# 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:

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:

uv build
# or: python -m build

Upload to PyPI using Twine:

twine upload dist/*

License

MIT License. See LICENSE for details.

Project links

Key dates

PyPI data
Data sourced directly from PyPI's database.
  • Released:
Newer release available (0.1.60)

1 maintainer

PyPI data
Data sourced directly from PyPI's database.
Avatar for Berrny from gravatar.com Berrny

Credits

Author: vovaRL

License

MIT License (MIT)

Requires

Python >=3.10

Provides Extra

dev

Tags

deep-learning diffusion-models discrete-diffusion fsdp2 huggingface-jobs language-model mdlm muon pytorch

Classifiers

Report project as malware

Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

slm388_torchdiff-0.1.7.tar.gz (21.9 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

slm388_torchdiff-0.1.7-py3-none-any.whl (24.4 kB view details)

Uploaded Python 3

File details

Details for the file slm388_torchdiff-0.1.7.tar.gz.

File metadata

  • Download URL: slm388_torchdiff-0.1.7.tar.gz
  • Upload date:
  • Size: 21.9 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/7.0.0 CPython/3.14.4

File hashes

Hashes for slm388_torchdiff-0.1.7.tar.gz
Algorithm Hash digest
SHA256 d7f553ca22f9d3620ba12a60ed29909897c9a9d67109112146cf0414fe72f6a3
MD5 6c9b466ec435e293de92e54e9aa3bc14
BLAKE2b-256 8fcf605fd8563bcfcef602109875c75d79db7281f9e6ee5902b3fd33c574cb03

See more details on using hashes here.

File details

Details for the file slm388_torchdiff-0.1.7-py3-none-any.whl.

File metadata

File hashes

Hashes for slm388_torchdiff-0.1.7-py3-none-any.whl
Algorithm Hash digest
SHA256 832c1c30ce10dbd0bf679efad36df630f1bb2f97bdeea8ea8ae9a39ec8a3de67
MD5 c1444201faf2403ed3c7ee5d596a6c33
BLAKE2b-256 f5ba23749676a56c76b35afc37b963ec767e670723cf18f36195f15f52d16e25

See more details on using hashes here.

Release history Release notifications | RSS feed

0.1.60

2 files

0.1.59

2 files

0.1.58

2 files

0.1.57

2 files

0.1.56

2 files

0.1.55

2 files

0.1.54

2 files

0.1.53

2 files

0.1.52

2 files

0.1.51

2 files

0.1.50

2 files

0.1.49

2 files

0.1.47

2 files

0.1.46

2 files

0.1.45

2 files

0.1.44

2 files

0.1.43

2 files

0.1.42

2 files

0.1.41

2 files

0.1.39

2 files

0.1.38

2 files

0.1.37

2 files

0.1.36

2 files

0.1.35

2 files

0.1.33

2 files

0.1.32

2 files

0.1.31

2 files

0.1.30

2 files

0.1.28

2 files

0.1.27

2 files

0.1.26

2 files

0.1.25

2 files

0.1.24

2 files

0.1.23

2 files

0.1.22

2 files

0.1.21

2 files

0.1.20

2 files

0.1.19

2 files

0.1.18

2 files

0.1.17

2 files

0.1.16

2 files

0.1.15

2 files

0.1.14

2 files

0.1.13

2 files

0.1.12

2 files

0.1.11

2 files

0.1.10

2 files

0.1.9

2 files

0.1.8

2 files

This release

0.1.7 This release

2 files

0.1.6

2 files

0.1.5

2 files

0.1.4

2 files

0.1.3

2 files

0.1.2

2 files

0.1.1

2 files

0.1.0

2 files

Anthropic, PBC Visionary sponsor Bloomberg Visionary sponsor Hudson River Trading Visionary sponsor Meta Visionary sponsor NVIDIA Visionary sponsor Microsoft Sustainability sponsor Depot Continuous Integration AWS Cloud computing and Security Sponsor Datadog Monitoring Fastly CDN Google Download Analytics Sentry Error logging StatusPage Status page