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

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.12.tar.gz (24.3 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.12-py3-none-any.whl (27.1 kB view details)

Uploaded Python 3

File details

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

File metadata

  • Download URL: slm388_torchdiff-0.1.12.tar.gz
  • Upload date:
  • Size: 24.3 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.12.tar.gz
Algorithm Hash digest
SHA256 089f7033cb23a76aaa436b955b194f08c7513cd88771d1cfa265174da91f4ed0
MD5 2d2ef771a7d8f6cb0ddcd72affb972ef
BLAKE2b-256 11dd703f00c6983a2e792cae67485d49985b62c19405ae83c5d047c7fe497aeb

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for slm388_torchdiff-0.1.12-py3-none-any.whl
Algorithm Hash digest
SHA256 e494a44a09e482f79b4e370fb85f206a11e665fc29b3e91e839f25432a2799cb
MD5 132441c3a0eee2e32a7e1a23a3a4bd65
BLAKE2b-256 85c4cfa80e6e8e5ca3ece07b267a5b4e22eb8e7f4fead27f074a3d5d735844fa

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

This release

0.1.12 This release

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