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 (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.

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 flash4

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.49.tar.gz (45.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.49-py3-none-any.whl (49.9 kB view details)

Uploaded Python 3

File details

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

File metadata

  • Download URL: slm388_torchdiff-0.1.49.tar.gz
  • Upload date:
  • Size: 45.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.49.tar.gz
Algorithm Hash digest
SHA256 456583cdc07b9888897ef85fb8fadb0b255de2bc8d2b1cbe0fbc0638ba5774fd
MD5 e2a260e644a19026fff5bfbb8685cb6d
BLAKE2b-256 12143f01055a229863c9936e2a724010e25e97a23fae2be56767c6f307ae412e

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for slm388_torchdiff-0.1.49-py3-none-any.whl
Algorithm Hash digest
SHA256 ac927131cc0172e163e64111968d82b89beceea6b7b1aea77a17fa80c00aff5c
MD5 655b05a317671b5229461238db92aa52
BLAKE2b-256 42301fdad28b104adb29b75d24e091dbe32837263927c2190da61fc5dc060e05

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

This release

0.1.49 This release

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

0.1.7

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