slm388-torchdiff 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- slm388_torchdiff-0.1.0/.gitignore +27 -0
- slm388_torchdiff-0.1.0/LICENSE +21 -0
- slm388_torchdiff-0.1.0/PKG-INFO +234 -0
- slm388_torchdiff-0.1.0/README.md +201 -0
- slm388_torchdiff-0.1.0/pyproject.toml +71 -0
- slm388_torchdiff-0.1.0/src/torchdiff/__init__.py +54 -0
- slm388_torchdiff-0.1.0/src/torchdiff/data.py +257 -0
- slm388_torchdiff-0.1.0/src/torchdiff/model.py +388 -0
- slm388_torchdiff-0.1.0/src/torchdiff/optim.py +89 -0
- slm388_torchdiff-0.1.0/src/torchdiff/train.py +821 -0
|
@@ -0,0 +1,27 @@
|
|
|
1
|
+
# Byte-compiled / optimized / DLL files
|
|
2
|
+
__pycache__/
|
|
3
|
+
*.py[cod]
|
|
4
|
+
*$py.class
|
|
5
|
+
|
|
6
|
+
# Distribution / packaging
|
|
7
|
+
dist/
|
|
8
|
+
build/
|
|
9
|
+
*.egg-info/
|
|
10
|
+
*.egg
|
|
11
|
+
|
|
12
|
+
# Virtual environments
|
|
13
|
+
.venv/
|
|
14
|
+
venv/
|
|
15
|
+
ENV/
|
|
16
|
+
|
|
17
|
+
# Data & Checkpoints
|
|
18
|
+
checkpoints/
|
|
19
|
+
data/
|
|
20
|
+
*.bin
|
|
21
|
+
*.safetensors
|
|
22
|
+
|
|
23
|
+
# IDE & OS
|
|
24
|
+
.vscode/
|
|
25
|
+
.idea/
|
|
26
|
+
.DS_Store
|
|
27
|
+
Thumbs.db
|
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 vovaRL
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
|
@@ -0,0 +1,234 @@
|
|
|
1
|
+
Metadata-Version: 2.5
|
|
2
|
+
Name: slm388-torchdiff
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Masked Diffusion Language Model (MDLM) pretraining in PyTorch with FSDP2, Muon optimizer, and Hugging Face Jobs support.
|
|
5
|
+
Project-URL: Homepage, https://github.com/vovaRL/torchdiff
|
|
6
|
+
Project-URL: Repository, https://github.com/vovaRL/torchdiff
|
|
7
|
+
Project-URL: Issues, https://github.com/vovaRL/torchdiff/issues
|
|
8
|
+
Project-URL: Dataset, https://huggingface.co/datasets/vovaRL/slm388-corpus
|
|
9
|
+
Author: vovaRL
|
|
10
|
+
License: MIT
|
|
11
|
+
License-File: LICENSE
|
|
12
|
+
Keywords: deep-learning,diffusion-models,discrete-diffusion,fsdp2,huggingface-jobs,language-model,mdlm,muon,pytorch
|
|
13
|
+
Classifier: Development Status :: 4 - Beta
|
|
14
|
+
Classifier: Intended Audience :: Developers
|
|
15
|
+
Classifier: Intended Audience :: Science/Research
|
|
16
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
17
|
+
Classifier: Programming Language :: Python :: 3
|
|
18
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
19
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
20
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
21
|
+
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
22
|
+
Requires-Python: >=3.10
|
|
23
|
+
Requires-Dist: cut-cross-entropy>=25.1.1
|
|
24
|
+
Requires-Dist: huggingface-hub>=0.28.0
|
|
25
|
+
Requires-Dist: numpy>=1.24.0
|
|
26
|
+
Requires-Dist: safetensors>=0.4.0
|
|
27
|
+
Requires-Dist: torch>=2.4
|
|
28
|
+
Provides-Extra: dev
|
|
29
|
+
Requires-Dist: build>=1.0.0; extra == 'dev'
|
|
30
|
+
Requires-Dist: pytest>=8.0.0; extra == 'dev'
|
|
31
|
+
Requires-Dist: twine>=5.0.0; extra == 'dev'
|
|
32
|
+
Description-Content-Type: text/markdown
|
|
33
|
+
|
|
34
|
+
# TorchDiff (`slm388-torchdiff`)
|
|
35
|
+
|
|
36
|
+
[](https://pypi.org/project/slm388-torchdiff/)
|
|
37
|
+
[](https://pypi.org/project/slm388-torchdiff/)
|
|
38
|
+
[](https://opensource.org/licenses/MIT)
|
|
39
|
+
[](https://huggingface.co/datasets/vovaRL/slm388-corpus)
|
|
40
|
+
|
|
41
|
+
**TorchDiff** (`slm388-torchdiff`) is a high-performance PyTorch implementation of a **Masked Diffusion Language Model (MDLM)** with:
|
|
42
|
+
- **Distributed Training:** Native PyTorch FSDP2 (Fully Sharded Data Parallel 2) multi-GPU sharding.
|
|
43
|
+
- **Hybrid Optimizer:** Muon (Newton-Schulz orthogonalization) for 2D hidden matrices + AdamW for embeddings, norms, and heads.
|
|
44
|
+
- **Curriculum Learning:** Automatic pretraining vs. SFT dynamic mixing and context extension (2K $\to$ 4K).
|
|
45
|
+
- **Fast Loss:** Memory-efficient token loss via Apple's [Cut Cross-Entropy](https://github.com/apple/ml-cross-entropy) with automatic PyTorch fallback.
|
|
46
|
+
- **Hugging Face Jobs Ready:** Seamless zero-setup execution with `hf jobs uv run` and dataset volume mounting.
|
|
47
|
+
|
|
48
|
+
---
|
|
49
|
+
|
|
50
|
+
## 1. Installation
|
|
51
|
+
|
|
52
|
+
### From PyPI
|
|
53
|
+
```bash
|
|
54
|
+
pip install slm388-torchdiff
|
|
55
|
+
```
|
|
56
|
+
|
|
57
|
+
### From GitHub Repository
|
|
58
|
+
```bash
|
|
59
|
+
pip install "git+https://github.com/vovaRL/torchdiff.git"
|
|
60
|
+
```
|
|
61
|
+
|
|
62
|
+
### Editable Development Installation
|
|
63
|
+
```bash
|
|
64
|
+
git clone https://github.com/vovaRL/torchdiff.git
|
|
65
|
+
cd torchdiff
|
|
66
|
+
pip install -e ".[dev]"
|
|
67
|
+
```
|
|
68
|
+
|
|
69
|
+
---
|
|
70
|
+
|
|
71
|
+
## 2. Dataset: `vovaRL/slm388-corpus`
|
|
72
|
+
|
|
73
|
+
TorchDiff is pre-configured to train on the [vovaRL/slm388-corpus](https://huggingface.co/datasets/vovaRL/slm388-corpus) dataset.
|
|
74
|
+
|
|
75
|
+
### Dataset Overview
|
|
76
|
+
- **Vocabulary:** SmolLM2-compatible (49,152 vocab size, token ID 49,151 as mask token).
|
|
77
|
+
- **Pretraining Shards:** Cosmopedia, DCLM, FineMath, FineWeb-Edu, and Python-Edu `.bin` binary token shards.
|
|
78
|
+
- **SFT Shards:** SmolTalk instruction tuning shards (`*smoltalk*.bin`), dynamically phased in via curriculum scheduling.
|
|
79
|
+
- **Validation Shards:** Dedicated held-out validation shards (`*_val_*.bin`).
|
|
80
|
+
|
|
81
|
+
### Dynamic Shard Discovery
|
|
82
|
+
TorchDiff's `CurriculumBinLoader` automatically discovers shards and partitions them:
|
|
83
|
+
```
|
|
84
|
+
[Dataset Discovery]
|
|
85
|
+
Directory: /dataset
|
|
86
|
+
Total shards found: 64
|
|
87
|
+
├── Pretraining shards: 56
|
|
88
|
+
├── SFT shards: 1
|
|
89
|
+
└── Validation shards: 7 (held out)
|
|
90
|
+
```
|
|
91
|
+
|
|
92
|
+
---
|
|
93
|
+
|
|
94
|
+
## 3. Training on Hugging Face Jobs (`hf jobs uv run`)
|
|
95
|
+
|
|
96
|
+
Hugging Face Jobs allows you to train on high-end cloud GPUs (A100, H200, A10G) without managing servers or Docker images.
|
|
97
|
+
|
|
98
|
+
### Prerequisites
|
|
99
|
+
1. Install the Hugging Face CLI:
|
|
100
|
+
```bash
|
|
101
|
+
pip install -U huggingface_hub
|
|
102
|
+
```
|
|
103
|
+
2. Log in with your write token:
|
|
104
|
+
```bash
|
|
105
|
+
hf auth login
|
|
106
|
+
```
|
|
107
|
+
|
|
108
|
+
### Method A: Run with Published PyPI Package (Recommended)
|
|
109
|
+
Once `torchdiff` is on PyPI, mount the dataset volume at `/dataset` and launch directly:
|
|
110
|
+
```bash
|
|
111
|
+
hf jobs uv run \
|
|
112
|
+
--flavor a100-large \
|
|
113
|
+
--secrets HF_TOKEN \
|
|
114
|
+
--with slm388-torchdiff \
|
|
115
|
+
-v hf://datasets/vovaRL/slm388-corpus:/dataset \
|
|
116
|
+
torchdiff-train \
|
|
117
|
+
--data_dir /dataset \
|
|
118
|
+
--batch_size 128 \
|
|
119
|
+
--repo_id vovaRL/DiffLM
|
|
120
|
+
```
|
|
121
|
+
|
|
122
|
+
### Method B: Run from Git Repository
|
|
123
|
+
You can run directly from GitHub without publishing to PyPI:
|
|
124
|
+
```bash
|
|
125
|
+
hf jobs uv run \
|
|
126
|
+
--flavor a100-large \
|
|
127
|
+
--secrets HF_TOKEN \
|
|
128
|
+
--with "torchdiff @ git+https://github.com/vovaRL/torchdiff.git" \
|
|
129
|
+
-v hf://datasets/vovaRL/slm388-corpus:/dataset \
|
|
130
|
+
torchdiff-train \
|
|
131
|
+
--data_dir /dataset \
|
|
132
|
+
--batch_size 128 \
|
|
133
|
+
--repo_id vovaRL/DiffLM
|
|
134
|
+
```
|
|
135
|
+
|
|
136
|
+
### Method C: Run Standalone Script
|
|
137
|
+
Run the self-contained script `train_standalone.py` locally from this repository:
|
|
138
|
+
```bash
|
|
139
|
+
hf jobs uv run \
|
|
140
|
+
--flavor a100-large \
|
|
141
|
+
--secrets HF_TOKEN \
|
|
142
|
+
-v hf://datasets/vovaRL/slm388-corpus:/dataset \
|
|
143
|
+
train_standalone.py \
|
|
144
|
+
--data_dir /dataset \
|
|
145
|
+
--batch_size 128 \
|
|
146
|
+
--repo_id vovaRL/DiffLM
|
|
147
|
+
```
|
|
148
|
+
|
|
149
|
+
### Multi-GPU Training (FSDP2)
|
|
150
|
+
Scale to multi-GPU flavors (e.g. 4x or 8x A100):
|
|
151
|
+
```bash
|
|
152
|
+
hf jobs uv run \
|
|
153
|
+
--flavor a100x4 \
|
|
154
|
+
--secrets HF_TOKEN \
|
|
155
|
+
--with slm388-torchdiff \
|
|
156
|
+
-v hf://datasets/vovaRL/slm388-corpus:/dataset \
|
|
157
|
+
torchdiff-train \
|
|
158
|
+
--data_dir /dataset \
|
|
159
|
+
--batch_size 256 \
|
|
160
|
+
--repo_id vovaRL/DiffLM
|
|
161
|
+
```
|
|
162
|
+
|
|
163
|
+
### Persistent Bucket Checkpointing
|
|
164
|
+
Mount a Hugging Face Bucket as a read-write volume to save checkpoints continuously:
|
|
165
|
+
```bash
|
|
166
|
+
hf jobs uv run \
|
|
167
|
+
--flavor a100-large \
|
|
168
|
+
--secrets HF_TOKEN \
|
|
169
|
+
--with slm388-torchdiff \
|
|
170
|
+
-v hf://datasets/vovaRL/slm388-corpus:/dataset \
|
|
171
|
+
-v hf://buckets/vovaRL/diff-checkpoints:/checkpoints \
|
|
172
|
+
torchdiff-train \
|
|
173
|
+
--data_dir /dataset \
|
|
174
|
+
--checkpoint_dir /checkpoints \
|
|
175
|
+
--checkpoint_interval 1000
|
|
176
|
+
```
|
|
177
|
+
|
|
178
|
+
---
|
|
179
|
+
|
|
180
|
+
## 4. Local Execution
|
|
181
|
+
|
|
182
|
+
### Run Training Locally
|
|
183
|
+
```bash
|
|
184
|
+
# Using installed CLI
|
|
185
|
+
torchdiff-train --data_dir ./data --batch_size 32
|
|
186
|
+
|
|
187
|
+
# Or using uv run
|
|
188
|
+
uv run train.py --data_dir ./data --batch_size 32
|
|
189
|
+
```
|
|
190
|
+
|
|
191
|
+
### Benchmark Mode
|
|
192
|
+
Benchmark raw forward/backward throughput without downloading the full dataset:
|
|
193
|
+
```bash
|
|
194
|
+
torchdiff-train --mode bench --synthetic_bench --batch_size 64 --seq_len 2048
|
|
195
|
+
```
|
|
196
|
+
|
|
197
|
+
---
|
|
198
|
+
|
|
199
|
+
## 5. CLI Arguments
|
|
200
|
+
|
|
201
|
+
| Flag | Default | Description |
|
|
202
|
+
|---|---|---|
|
|
203
|
+
| `--mode` | `train` | Mode: `'train'` or `'bench'`. |
|
|
204
|
+
| `--data_dir` | `None` | Directory containing `.bin` shards. Auto-detects `/dataset`, `/data`, `./data`. |
|
|
205
|
+
| `--dataset` | `vovaRL/slm388-corpus` | Hugging Face dataset repo ID to download if local shards are not found. |
|
|
206
|
+
| `--download_pattern` | `None` | Glob pattern for partial dataset download (e.g. `'*.bin'`). |
|
|
207
|
+
| `--batch_size` | `512` | Batch size across GPUs. |
|
|
208
|
+
| `--seq_len` | `2048` | Initial token context length (extends to 4096 in Phase 4). |
|
|
209
|
+
| `--total_steps` | `85000` | Total pretraining steps. |
|
|
210
|
+
| `--checkpoint` | `None` | Path to `.safetensors` checkpoint to resume training. |
|
|
211
|
+
| `--checkpoint_dir` | `./checkpoints` | Local/volume checkpoint directory. |
|
|
212
|
+
| `--checkpoint_interval` | `2500` | Steps between saving checkpoints. |
|
|
213
|
+
| `--repo_id` | `$HF_REPO_ID` | Hugging Face model repository to upload saved checkpoints. |
|
|
214
|
+
| `--no_compile` | `False` | Disable `torch.compile`. |
|
|
215
|
+
|
|
216
|
+
---
|
|
217
|
+
|
|
218
|
+
## 6. Building and Publishing to PyPI
|
|
219
|
+
|
|
220
|
+
Build sdist and wheel:
|
|
221
|
+
```bash
|
|
222
|
+
uv build
|
|
223
|
+
# or: python -m build
|
|
224
|
+
```
|
|
225
|
+
|
|
226
|
+
Upload to PyPI using Twine:
|
|
227
|
+
```bash
|
|
228
|
+
twine upload dist/*
|
|
229
|
+
```
|
|
230
|
+
|
|
231
|
+
---
|
|
232
|
+
|
|
233
|
+
## License
|
|
234
|
+
MIT License. See [LICENSE](LICENSE) for details.
|
|
@@ -0,0 +1,201 @@
|
|
|
1
|
+
# TorchDiff (`slm388-torchdiff`)
|
|
2
|
+
|
|
3
|
+
[](https://pypi.org/project/slm388-torchdiff/)
|
|
4
|
+
[](https://pypi.org/project/slm388-torchdiff/)
|
|
5
|
+
[](https://opensource.org/licenses/MIT)
|
|
6
|
+
[](https://huggingface.co/datasets/vovaRL/slm388-corpus)
|
|
7
|
+
|
|
8
|
+
**TorchDiff** (`slm388-torchdiff`) is a high-performance PyTorch implementation of a **Masked Diffusion Language Model (MDLM)** with:
|
|
9
|
+
- **Distributed Training:** Native PyTorch FSDP2 (Fully Sharded Data Parallel 2) multi-GPU sharding.
|
|
10
|
+
- **Hybrid Optimizer:** Muon (Newton-Schulz orthogonalization) for 2D hidden matrices + AdamW for embeddings, norms, and heads.
|
|
11
|
+
- **Curriculum Learning:** Automatic pretraining vs. SFT dynamic mixing and context extension (2K $\to$ 4K).
|
|
12
|
+
- **Fast Loss:** Memory-efficient token loss via Apple's [Cut Cross-Entropy](https://github.com/apple/ml-cross-entropy) with automatic PyTorch fallback.
|
|
13
|
+
- **Hugging Face Jobs Ready:** Seamless zero-setup execution with `hf jobs uv run` and dataset volume mounting.
|
|
14
|
+
|
|
15
|
+
---
|
|
16
|
+
|
|
17
|
+
## 1. Installation
|
|
18
|
+
|
|
19
|
+
### From PyPI
|
|
20
|
+
```bash
|
|
21
|
+
pip install slm388-torchdiff
|
|
22
|
+
```
|
|
23
|
+
|
|
24
|
+
### From GitHub Repository
|
|
25
|
+
```bash
|
|
26
|
+
pip install "git+https://github.com/vovaRL/torchdiff.git"
|
|
27
|
+
```
|
|
28
|
+
|
|
29
|
+
### Editable Development Installation
|
|
30
|
+
```bash
|
|
31
|
+
git clone https://github.com/vovaRL/torchdiff.git
|
|
32
|
+
cd torchdiff
|
|
33
|
+
pip install -e ".[dev]"
|
|
34
|
+
```
|
|
35
|
+
|
|
36
|
+
---
|
|
37
|
+
|
|
38
|
+
## 2. Dataset: `vovaRL/slm388-corpus`
|
|
39
|
+
|
|
40
|
+
TorchDiff is pre-configured to train on the [vovaRL/slm388-corpus](https://huggingface.co/datasets/vovaRL/slm388-corpus) dataset.
|
|
41
|
+
|
|
42
|
+
### Dataset Overview
|
|
43
|
+
- **Vocabulary:** SmolLM2-compatible (49,152 vocab size, token ID 49,151 as mask token).
|
|
44
|
+
- **Pretraining Shards:** Cosmopedia, DCLM, FineMath, FineWeb-Edu, and Python-Edu `.bin` binary token shards.
|
|
45
|
+
- **SFT Shards:** SmolTalk instruction tuning shards (`*smoltalk*.bin`), dynamically phased in via curriculum scheduling.
|
|
46
|
+
- **Validation Shards:** Dedicated held-out validation shards (`*_val_*.bin`).
|
|
47
|
+
|
|
48
|
+
### Dynamic Shard Discovery
|
|
49
|
+
TorchDiff's `CurriculumBinLoader` automatically discovers shards and partitions them:
|
|
50
|
+
```
|
|
51
|
+
[Dataset Discovery]
|
|
52
|
+
Directory: /dataset
|
|
53
|
+
Total shards found: 64
|
|
54
|
+
├── Pretraining shards: 56
|
|
55
|
+
├── SFT shards: 1
|
|
56
|
+
└── Validation shards: 7 (held out)
|
|
57
|
+
```
|
|
58
|
+
|
|
59
|
+
---
|
|
60
|
+
|
|
61
|
+
## 3. Training on Hugging Face Jobs (`hf jobs uv run`)
|
|
62
|
+
|
|
63
|
+
Hugging Face Jobs allows you to train on high-end cloud GPUs (A100, H200, A10G) without managing servers or Docker images.
|
|
64
|
+
|
|
65
|
+
### Prerequisites
|
|
66
|
+
1. Install the Hugging Face CLI:
|
|
67
|
+
```bash
|
|
68
|
+
pip install -U huggingface_hub
|
|
69
|
+
```
|
|
70
|
+
2. Log in with your write token:
|
|
71
|
+
```bash
|
|
72
|
+
hf auth login
|
|
73
|
+
```
|
|
74
|
+
|
|
75
|
+
### Method A: Run with Published PyPI Package (Recommended)
|
|
76
|
+
Once `torchdiff` is on PyPI, mount the dataset volume at `/dataset` and launch directly:
|
|
77
|
+
```bash
|
|
78
|
+
hf jobs uv run \
|
|
79
|
+
--flavor a100-large \
|
|
80
|
+
--secrets HF_TOKEN \
|
|
81
|
+
--with slm388-torchdiff \
|
|
82
|
+
-v hf://datasets/vovaRL/slm388-corpus:/dataset \
|
|
83
|
+
torchdiff-train \
|
|
84
|
+
--data_dir /dataset \
|
|
85
|
+
--batch_size 128 \
|
|
86
|
+
--repo_id vovaRL/DiffLM
|
|
87
|
+
```
|
|
88
|
+
|
|
89
|
+
### Method B: Run from Git Repository
|
|
90
|
+
You can run directly from GitHub without publishing to PyPI:
|
|
91
|
+
```bash
|
|
92
|
+
hf jobs uv run \
|
|
93
|
+
--flavor a100-large \
|
|
94
|
+
--secrets HF_TOKEN \
|
|
95
|
+
--with "torchdiff @ git+https://github.com/vovaRL/torchdiff.git" \
|
|
96
|
+
-v hf://datasets/vovaRL/slm388-corpus:/dataset \
|
|
97
|
+
torchdiff-train \
|
|
98
|
+
--data_dir /dataset \
|
|
99
|
+
--batch_size 128 \
|
|
100
|
+
--repo_id vovaRL/DiffLM
|
|
101
|
+
```
|
|
102
|
+
|
|
103
|
+
### Method C: Run Standalone Script
|
|
104
|
+
Run the self-contained script `train_standalone.py` locally from this repository:
|
|
105
|
+
```bash
|
|
106
|
+
hf jobs uv run \
|
|
107
|
+
--flavor a100-large \
|
|
108
|
+
--secrets HF_TOKEN \
|
|
109
|
+
-v hf://datasets/vovaRL/slm388-corpus:/dataset \
|
|
110
|
+
train_standalone.py \
|
|
111
|
+
--data_dir /dataset \
|
|
112
|
+
--batch_size 128 \
|
|
113
|
+
--repo_id vovaRL/DiffLM
|
|
114
|
+
```
|
|
115
|
+
|
|
116
|
+
### Multi-GPU Training (FSDP2)
|
|
117
|
+
Scale to multi-GPU flavors (e.g. 4x or 8x A100):
|
|
118
|
+
```bash
|
|
119
|
+
hf jobs uv run \
|
|
120
|
+
--flavor a100x4 \
|
|
121
|
+
--secrets HF_TOKEN \
|
|
122
|
+
--with slm388-torchdiff \
|
|
123
|
+
-v hf://datasets/vovaRL/slm388-corpus:/dataset \
|
|
124
|
+
torchdiff-train \
|
|
125
|
+
--data_dir /dataset \
|
|
126
|
+
--batch_size 256 \
|
|
127
|
+
--repo_id vovaRL/DiffLM
|
|
128
|
+
```
|
|
129
|
+
|
|
130
|
+
### Persistent Bucket Checkpointing
|
|
131
|
+
Mount a Hugging Face Bucket as a read-write volume to save checkpoints continuously:
|
|
132
|
+
```bash
|
|
133
|
+
hf jobs uv run \
|
|
134
|
+
--flavor a100-large \
|
|
135
|
+
--secrets HF_TOKEN \
|
|
136
|
+
--with slm388-torchdiff \
|
|
137
|
+
-v hf://datasets/vovaRL/slm388-corpus:/dataset \
|
|
138
|
+
-v hf://buckets/vovaRL/diff-checkpoints:/checkpoints \
|
|
139
|
+
torchdiff-train \
|
|
140
|
+
--data_dir /dataset \
|
|
141
|
+
--checkpoint_dir /checkpoints \
|
|
142
|
+
--checkpoint_interval 1000
|
|
143
|
+
```
|
|
144
|
+
|
|
145
|
+
---
|
|
146
|
+
|
|
147
|
+
## 4. Local Execution
|
|
148
|
+
|
|
149
|
+
### Run Training Locally
|
|
150
|
+
```bash
|
|
151
|
+
# Using installed CLI
|
|
152
|
+
torchdiff-train --data_dir ./data --batch_size 32
|
|
153
|
+
|
|
154
|
+
# Or using uv run
|
|
155
|
+
uv run train.py --data_dir ./data --batch_size 32
|
|
156
|
+
```
|
|
157
|
+
|
|
158
|
+
### Benchmark Mode
|
|
159
|
+
Benchmark raw forward/backward throughput without downloading the full dataset:
|
|
160
|
+
```bash
|
|
161
|
+
torchdiff-train --mode bench --synthetic_bench --batch_size 64 --seq_len 2048
|
|
162
|
+
```
|
|
163
|
+
|
|
164
|
+
---
|
|
165
|
+
|
|
166
|
+
## 5. CLI Arguments
|
|
167
|
+
|
|
168
|
+
| Flag | Default | Description |
|
|
169
|
+
|---|---|---|
|
|
170
|
+
| `--mode` | `train` | Mode: `'train'` or `'bench'`. |
|
|
171
|
+
| `--data_dir` | `None` | Directory containing `.bin` shards. Auto-detects `/dataset`, `/data`, `./data`. |
|
|
172
|
+
| `--dataset` | `vovaRL/slm388-corpus` | Hugging Face dataset repo ID to download if local shards are not found. |
|
|
173
|
+
| `--download_pattern` | `None` | Glob pattern for partial dataset download (e.g. `'*.bin'`). |
|
|
174
|
+
| `--batch_size` | `512` | Batch size across GPUs. |
|
|
175
|
+
| `--seq_len` | `2048` | Initial token context length (extends to 4096 in Phase 4). |
|
|
176
|
+
| `--total_steps` | `85000` | Total pretraining steps. |
|
|
177
|
+
| `--checkpoint` | `None` | Path to `.safetensors` checkpoint to resume training. |
|
|
178
|
+
| `--checkpoint_dir` | `./checkpoints` | Local/volume checkpoint directory. |
|
|
179
|
+
| `--checkpoint_interval` | `2500` | Steps between saving checkpoints. |
|
|
180
|
+
| `--repo_id` | `$HF_REPO_ID` | Hugging Face model repository to upload saved checkpoints. |
|
|
181
|
+
| `--no_compile` | `False` | Disable `torch.compile`. |
|
|
182
|
+
|
|
183
|
+
---
|
|
184
|
+
|
|
185
|
+
## 6. Building and Publishing to PyPI
|
|
186
|
+
|
|
187
|
+
Build sdist and wheel:
|
|
188
|
+
```bash
|
|
189
|
+
uv build
|
|
190
|
+
# or: python -m build
|
|
191
|
+
```
|
|
192
|
+
|
|
193
|
+
Upload to PyPI using Twine:
|
|
194
|
+
```bash
|
|
195
|
+
twine upload dist/*
|
|
196
|
+
```
|
|
197
|
+
|
|
198
|
+
---
|
|
199
|
+
|
|
200
|
+
## License
|
|
201
|
+
MIT License. See [LICENSE](LICENSE) for details.
|
|
@@ -0,0 +1,71 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["hatchling"]
|
|
3
|
+
build-backend = "hatchling.build"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "slm388-torchdiff"
|
|
7
|
+
version = "0.1.0"
|
|
8
|
+
description = "Masked Diffusion Language Model (MDLM) pretraining in PyTorch with FSDP2, Muon optimizer, and Hugging Face Jobs support."
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
requires-python = ">=3.10"
|
|
11
|
+
license = { text = "MIT" }
|
|
12
|
+
authors = [
|
|
13
|
+
{ name = "vovaRL" }
|
|
14
|
+
]
|
|
15
|
+
keywords = [
|
|
16
|
+
"deep-learning",
|
|
17
|
+
"diffusion-models",
|
|
18
|
+
"discrete-diffusion",
|
|
19
|
+
"language-model",
|
|
20
|
+
"mdlm",
|
|
21
|
+
"pytorch",
|
|
22
|
+
"huggingface-jobs",
|
|
23
|
+
"fsdp2",
|
|
24
|
+
"muon"
|
|
25
|
+
]
|
|
26
|
+
classifiers = [
|
|
27
|
+
"Development Status :: 4 - Beta",
|
|
28
|
+
"Intended Audience :: Developers",
|
|
29
|
+
"Intended Audience :: Science/Research",
|
|
30
|
+
"License :: OSI Approved :: MIT License",
|
|
31
|
+
"Programming Language :: Python :: 3",
|
|
32
|
+
"Programming Language :: Python :: 3.10",
|
|
33
|
+
"Programming Language :: Python :: 3.11",
|
|
34
|
+
"Programming Language :: Python :: 3.12",
|
|
35
|
+
"Topic :: Scientific/Engineering :: Artificial Intelligence",
|
|
36
|
+
]
|
|
37
|
+
dependencies = [
|
|
38
|
+
"torch>=2.4",
|
|
39
|
+
"cut-cross-entropy>=25.1.1",
|
|
40
|
+
"safetensors>=0.4.0",
|
|
41
|
+
"huggingface-hub>=0.28.0",
|
|
42
|
+
"numpy>=1.24.0",
|
|
43
|
+
]
|
|
44
|
+
|
|
45
|
+
[project.optional-dependencies]
|
|
46
|
+
dev = [
|
|
47
|
+
"pytest>=8.0.0",
|
|
48
|
+
"build>=1.0.0",
|
|
49
|
+
"twine>=5.0.0",
|
|
50
|
+
]
|
|
51
|
+
|
|
52
|
+
[project.scripts]
|
|
53
|
+
torchdiff = "torchdiff.train:main"
|
|
54
|
+
torchdiff-train = "torchdiff.train:main"
|
|
55
|
+
|
|
56
|
+
[project.urls]
|
|
57
|
+
Homepage = "https://github.com/vovaRL/torchdiff"
|
|
58
|
+
Repository = "https://github.com/vovaRL/torchdiff"
|
|
59
|
+
Issues = "https://github.com/vovaRL/torchdiff/issues"
|
|
60
|
+
Dataset = "https://huggingface.co/datasets/vovaRL/slm388-corpus"
|
|
61
|
+
|
|
62
|
+
[tool.hatch.build.targets.wheel]
|
|
63
|
+
packages = ["src/torchdiff"]
|
|
64
|
+
|
|
65
|
+
[tool.hatch.build.targets.sdist]
|
|
66
|
+
include = [
|
|
67
|
+
"/src",
|
|
68
|
+
"/README.md",
|
|
69
|
+
"/LICENSE",
|
|
70
|
+
"/pyproject.toml",
|
|
71
|
+
]
|
|
@@ -0,0 +1,54 @@
|
|
|
1
|
+
"""TorchDiff: Masked Diffusion Language Model (MDLM) pretraining in PyTorch."""
|
|
2
|
+
|
|
3
|
+
from torchdiff.data import CurriculumBinLoader, ShardedStream, resolve_dataset_dir
|
|
4
|
+
from torchdiff.model import (
|
|
5
|
+
AdaLN,
|
|
6
|
+
MDLMAttention,
|
|
7
|
+
MDLMBlock,
|
|
8
|
+
MDLMDenseMLP,
|
|
9
|
+
MDLMModel,
|
|
10
|
+
MDLModelConfig,
|
|
11
|
+
cosine_auxiliary_weight_schedule,
|
|
12
|
+
get_block_noise_bounds,
|
|
13
|
+
mask_batch_for_block,
|
|
14
|
+
modulate,
|
|
15
|
+
timestep_embedding,
|
|
16
|
+
)
|
|
17
|
+
from torchdiff.optim import Muon
|
|
18
|
+
from torchdiff.train import (
|
|
19
|
+
build_optimizer,
|
|
20
|
+
compute_loss,
|
|
21
|
+
main,
|
|
22
|
+
run_benchmark,
|
|
23
|
+
run_training,
|
|
24
|
+
save_checkpoint,
|
|
25
|
+
train_step,
|
|
26
|
+
)
|
|
27
|
+
|
|
28
|
+
__version__ = "0.1.0"
|
|
29
|
+
|
|
30
|
+
__all__ = [
|
|
31
|
+
"MDLMModel",
|
|
32
|
+
"MDLModelConfig",
|
|
33
|
+
"MDLMBlock",
|
|
34
|
+
"MDLMAttention",
|
|
35
|
+
"MDLMDenseMLP",
|
|
36
|
+
"AdaLN",
|
|
37
|
+
"CurriculumBinLoader",
|
|
38
|
+
"ShardedStream",
|
|
39
|
+
"resolve_dataset_dir",
|
|
40
|
+
"Muon",
|
|
41
|
+
"build_optimizer",
|
|
42
|
+
"compute_loss",
|
|
43
|
+
"run_training",
|
|
44
|
+
"run_benchmark",
|
|
45
|
+
"save_checkpoint",
|
|
46
|
+
"train_step",
|
|
47
|
+
"main",
|
|
48
|
+
"timestep_embedding",
|
|
49
|
+
"modulate",
|
|
50
|
+
"get_block_noise_bounds",
|
|
51
|
+
"mask_batch_for_block",
|
|
52
|
+
"cosine_auxiliary_weight_schedule",
|
|
53
|
+
"__version__",
|
|
54
|
+
]
|