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.11.tar.gz (24.0 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.11-py3-none-any.whl (26.7 kB view details)

Uploaded Python 3

File details

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

File metadata

  • Download URL: slm388_torchdiff-0.1.11.tar.gz
  • Upload date:
  • Size: 24.0 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.11.tar.gz
Algorithm Hash digest
SHA256 322fa90103bb7b591e9f2dc4053c852c6a10d6ef2da4ec3b63204c9fbebebbdc
MD5 f8413dd64473c2b92fa315f42d0bd7e7
BLAKE2b-256 2d7938745500f1ecbf60b6be285d067aa2b7913d6e07fdec7d73be70402bb126

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for slm388_torchdiff-0.1.11-py3-none-any.whl
Algorithm Hash digest
SHA256 12cf51409e151524f44d5a6c3c77c4956079eec3f94d23deff697aa6f88e462c
MD5 e0ba301ca95e31109538700a2cb5ee09
BLAKE2b-256 852f5e9695550e784d9547cfc23cd1bcd14d7cfb1cc88660645c260b84a142b3

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

This release

0.1.11 This release

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