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.51.tar.gz (47.2 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.51-py3-none-any.whl (51.1 kB view details)

Uploaded Python 3

File details

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

File metadata

  • Download URL: slm388_torchdiff-0.1.51.tar.gz
  • Upload date:
  • Size: 47.2 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.51.tar.gz
Algorithm Hash digest
SHA256 250c56e2ca64d48e54d9cfedcab656b07a854a1065549e149aa7cd39c7456323
MD5 d9cf055eb5709ced2e64f47c38e5ca94
BLAKE2b-256 ab31715b1c5583014171eed15bff03c27dfc4519ae13dfb228e95be32891336d

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for slm388_torchdiff-0.1.51-py3-none-any.whl
Algorithm Hash digest
SHA256 4570d9a4a6b934ee261c0094a90b238225931c458d8f739b51783d455b8cca40
MD5 d0ba7bcf4f46b69f1fc5ee911f519c84
BLAKE2b-256 3b3fd570f55b3d2f9cee656797a044907ac40a88b2ff2a41b5669355e8e17056

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

This release

0.1.51 This release

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

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