ravex 0.0.2__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.
ravex-0.0.2/LICENSE ADDED
@@ -0,0 +1,201 @@
1
+ Apache License
2
+ Version 2.0, January 2004
3
+ http://www.apache.org/licenses/
4
+
5
+ TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
6
+
7
+ 1. Definitions.
8
+
9
+ "License" shall mean the terms and conditions for use, reproduction,
10
+ and distribution as defined by Sections 1 through 9 of this document.
11
+
12
+ "Licensor" shall mean the copyright owner or entity authorized by
13
+ the copyright owner that is granting the License.
14
+
15
+ "Legal Entity" shall mean the union of the acting entity and all
16
+ other entities that control, are controlled by, or are under common
17
+ control with that entity. For the purposes of this definition,
18
+ "control" means (i) the power, direct or indirect, to cause the
19
+ direction or management of such entity, whether by contract or
20
+ otherwise, or (ii) ownership of fifty percent (50%) or more of the
21
+ outstanding shares, or (iii) beneficial ownership of such entity.
22
+
23
+ "You" (or "Your") shall mean an individual or Legal Entity
24
+ exercising permissions granted by this License.
25
+
26
+ "Source" form shall mean the preferred form for making modifications,
27
+ including but not limited to software source code, documentation
28
+ source, and configuration files.
29
+
30
+ "Object" form shall mean any form resulting from mechanical
31
+ transformation or translation of a Source form, including but
32
+ not limited to compiled object code, generated documentation,
33
+ and conversions to other media types.
34
+
35
+ "Work" shall mean the work of authorship, whether in Source or
36
+ Object form, made available under the License, as indicated by a
37
+ copyright notice that is included in or attached to the work
38
+ (an example is provided in the Appendix below).
39
+
40
+ "Derivative Works" shall mean any work, whether in Source or Object
41
+ form, that is based on (or derived from) the Work and for which the
42
+ editorial revisions, annotations, elaborations, or other modifications
43
+ represent, as a whole, an original work of authorship. For the purposes
44
+ of this License, Derivative Works shall not include works that remain
45
+ separable from, or merely link (or bind by name) to the interfaces of,
46
+ the Work and Derivative Works thereof.
47
+
48
+ "Contribution" shall mean any work of authorship, including
49
+ the original version of the Work and any modifications or additions
50
+ to that Work or Derivative Works thereof, that is intentionally
51
+ submitted to Licensor for inclusion in the Work by the copyright owner
52
+ or by an individual or Legal Entity authorized to submit on behalf of
53
+ the copyright owner. For the purposes of this definition, "submitted"
54
+ means any form of electronic, verbal, or written communication sent
55
+ to the Licensor or its representatives, including but not limited to
56
+ communication on electronic mailing lists, source code control systems,
57
+ and issue tracking systems that are managed by, or on behalf of, the
58
+ Licensor for the purpose of discussing and improving the Work, but
59
+ excluding communication that is conspicuously marked or otherwise
60
+ designated in writing by the copyright owner as "Not a Contribution."
61
+
62
+ "Contributor" shall mean Licensor and any individual or Legal Entity
63
+ on behalf of whom a Contribution has been received by Licensor and
64
+ subsequently incorporated within the Work.
65
+
66
+ 2. Grant of Copyright License. Subject to the terms and conditions of
67
+ this License, each Contributor hereby grants to You a perpetual,
68
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
69
+ copyright license to reproduce, prepare Derivative Works of,
70
+ publicly display, publicly perform, sublicense, and distribute the
71
+ Work and such Derivative Works in Source or Object form.
72
+
73
+ 3. Grant of Patent License. Subject to the terms and conditions of
74
+ this License, each Contributor hereby grants to You a perpetual,
75
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
76
+ (except as stated in this section) patent license to make, have made,
77
+ use, offer to sell, sell, import, and otherwise transfer the Work,
78
+ where such license applies only to those patent claims licensable
79
+ by such Contributor that are necessarily infringed by their
80
+ Contribution(s) alone or by combination of their Contribution(s)
81
+ with the Work to which such Contribution(s) was submitted. If You
82
+ institute patent litigation against any entity (including a
83
+ cross-claim or counterclaim in a lawsuit) alleging that the Work
84
+ or a Contribution incorporated within the Work constitutes direct
85
+ or contributory patent infringement, then any patent licenses
86
+ granted to You under this License for that Work shall terminate
87
+ as of the date such litigation is filed.
88
+
89
+ 4. Redistribution. You may reproduce and distribute copies of the
90
+ Work or Derivative Works thereof in any medium, with or without
91
+ modifications, and in Source or Object form, provided that You
92
+ meet the following conditions:
93
+
94
+ (a) You must give any other recipients of the Work or
95
+ Derivative Works a copy of this License; and
96
+
97
+ (b) You must cause any modified files to carry prominent notices
98
+ stating that You changed the files; and
99
+
100
+ (c) You must retain, in the Source form of any Derivative Works
101
+ that You distribute, all copyright, patent, trademark, and
102
+ attribution notices from the Source form of the Work,
103
+ excluding those notices that do not pertain to any part of
104
+ the Derivative Works; and
105
+
106
+ (d) If the Work includes a "NOTICE" text file as part of its
107
+ distribution, then any Derivative Works that You distribute must
108
+ include a readable copy of the attribution notices contained
109
+ within such NOTICE file, excluding those notices that do not
110
+ pertain to any part of the Derivative Works, in at least one
111
+ of the following places: within a NOTICE text file distributed
112
+ as part of the Derivative Works; within the Source form or
113
+ documentation, if provided along with the Derivative Works; or,
114
+ within a display generated by the Derivative Works, if and
115
+ wherever such third-party notices normally appear. The contents
116
+ of the NOTICE file are for informational purposes only and
117
+ do not modify the License. You may add Your own attribution
118
+ notices within Derivative Works that You distribute, alongside
119
+ or as an addendum to the NOTICE text from the Work, provided
120
+ that such additional attribution notices cannot be construed
121
+ as modifying the License.
122
+
123
+ You may add Your own copyright statement to Your modifications and
124
+ may provide additional or different license terms and conditions
125
+ for use, reproduction, or distribution of Your modifications, or
126
+ for any such Derivative Works as a whole, provided Your use,
127
+ reproduction, and distribution of the Work otherwise complies with
128
+ the conditions stated in this License.
129
+
130
+ 5. Submission of Contributions. Unless You explicitly state otherwise,
131
+ any Contribution intentionally submitted for inclusion in the Work
132
+ by You to the Licensor shall be under the terms and conditions of
133
+ this License, without any additional terms or conditions.
134
+ Notwithstanding the above, nothing herein shall supersede or modify
135
+ the terms of any separate license agreement you may have executed
136
+ with Licensor regarding such Contributions.
137
+
138
+ 6. Trademarks. This License does not grant permission to use the trade
139
+ names, trademarks, service marks, or product names of the Licensor,
140
+ except as required for reasonable and customary use in describing the
141
+ origin of the Work and reproducing the content of the NOTICE file.
142
+
143
+ 7. Disclaimer of Warranty. Unless required by applicable law or
144
+ agreed to in writing, Licensor provides the Work (and each
145
+ Contributor provides its Contributions) on an "AS IS" BASIS,
146
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
147
+ implied, including, without limitation, any warranties or conditions
148
+ of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
149
+ PARTICULAR PURPOSE. You are solely responsible for determining the
150
+ appropriateness of using or redistributing the Work and assume any
151
+ risks associated with Your exercise of permissions under this License.
152
+
153
+ 8. Limitation of Liability. In no event and under no legal theory,
154
+ whether in tort (including negligence), contract, or otherwise,
155
+ unless required by applicable law (such as deliberate and grossly
156
+ negligent acts) or agreed to in writing, shall any Contributor be
157
+ liable to You for damages, including any direct, indirect, special,
158
+ incidental, or consequential damages of any character arising as a
159
+ result of this License or out of the use or inability to use the
160
+ Work (including but not limited to damages for loss of goodwill,
161
+ work stoppage, computer failure or malfunction, or any and all
162
+ other commercial damages or losses), even if such Contributor
163
+ has been advised of the possibility of such damages.
164
+
165
+ 9. Accepting Warranty or Additional Liability. While redistributing
166
+ the Work or Derivative Works thereof, You may choose to offer,
167
+ and charge a fee for, acceptance of support, warranty, indemnity,
168
+ or other liability obligations and/or rights consistent with this
169
+ License. However, in accepting such obligations, You may act only
170
+ on Your own behalf and on Your sole responsibility, not on behalf
171
+ of any other Contributor, and only if You agree to indemnify,
172
+ defend, and hold each Contributor harmless for any liability
173
+ incurred by, or claims asserted against, such Contributor by reason
174
+ of your accepting any such warranty or additional liability.
175
+
176
+ END OF TERMS AND CONDITIONS
177
+
178
+ APPENDIX: How to apply the Apache License to your work.
179
+
180
+ To apply the Apache License to your work, attach the following
181
+ boilerplate notice, with the fields enclosed by brackets "[]"
182
+ replaced with your own identifying information. (Don't include
183
+ the brackets!) The text should be enclosed in the appropriate
184
+ comment syntax for the file format. We also recommend that a
185
+ file or class name and description of purpose be included on the
186
+ same "printed page" as the copyright notice for easier
187
+ identification within third-party archives.
188
+
189
+ Copyright [yyyy] [name of copyright owner]
190
+
191
+ Licensed under the Apache License, Version 2.0 (the "License");
192
+ you may not use this file except in compliance with the License.
193
+ You may obtain a copy of the License at
194
+
195
+ http://www.apache.org/licenses/LICENSE-2.0
196
+
197
+ Unless required by applicable law or agreed to in writing, software
198
+ distributed under the License is distributed on an "AS IS" BASIS,
199
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
200
+ See the License for the specific language governing permissions and
201
+ limitations under the License.
ravex-0.0.2/NOTICE ADDED
@@ -0,0 +1,7 @@
1
+ Ravex
2
+ Copyright 2026 Jonathan Vecchione
3
+
4
+ This product includes software developed as part of the GPU Zero project
5
+ (https://codeberg.org/JHNMACHINE).
6
+
7
+ Licensed under the Apache License, Version 2.0. See LICENSE for the full text.
ravex-0.0.2/PKG-INFO ADDED
@@ -0,0 +1,253 @@
1
+ Metadata-Version: 2.4
2
+ Name: ravex
3
+ Version: 0.0.2
4
+ Summary: Transparent checkpoint and resume for PyTorch training, with no changes to your training code
5
+ Author: JHNMACHINE
6
+ License: Apache-2.0
7
+ Project-URL: Homepage, https://codeberg.org/JHNMACHINE/ravex
8
+ Project-URL: Source, https://codeberg.org/JHNMACHINE/ravex
9
+ Project-URL: Issues, https://codeberg.org/JHNMACHINE/ravex/issues
10
+ Project-URL: Changelog, https://codeberg.org/JHNMACHINE/ravex/src/branch/main/CHANGELOG.md
11
+ Keywords: pytorch,checkpoint,training,resume,mlops,preemption
12
+ Classifier: Development Status :: 3 - Alpha
13
+ Classifier: Intended Audience :: Developers
14
+ Classifier: Intended Audience :: Science/Research
15
+ Classifier: License :: OSI Approved :: Apache Software License
16
+ Classifier: Programming Language :: Python :: 3
17
+ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
18
+ Requires-Python: >=3.9
19
+ Description-Content-Type: text/markdown
20
+ License-File: LICENSE
21
+ License-File: NOTICE
22
+ Requires-Dist: PyYAML>=5.1
23
+ Provides-Extra: moonclip
24
+ Requires-Dist: moonclip>=0.0.4; extra == "moonclip"
25
+ Provides-Extra: dev
26
+ Requires-Dist: pytest>=7; extra == "dev"
27
+ Requires-Dist: numpy; extra == "dev"
28
+ Dynamic: license-file
29
+
30
+ # Ravex
31
+
32
+ Transparent checkpoint and resume for PyTorch training. Your training script
33
+ does not change — not one line, not one import.
34
+
35
+ ```bash
36
+ pip install ravex
37
+ ravex enable
38
+ ```
39
+
40
+ Ravex itself is pure Python and installs anywhere. Its default checkpoint engine,
41
+ [Moonclip](https://codeberg.org/JHNMACHINE/moonclip), ships wheels for Linux
42
+ x86_64 only — so `pip install "ravex[moonclip]"` is a Linux thing, and on any
43
+ other platform Ravex falls back to `torch_save` on its own. What changed between
44
+ versions is in
45
+ [CHANGELOG.md](https://codeberg.org/JHNMACHINE/ravex/src/branch/main/CHANGELOG.md).
46
+
47
+ Drop a `ravex.yaml` next to your code and run what you always ran:
48
+
49
+ ```bash
50
+ python train.py
51
+ ```
52
+
53
+ The run now checkpoints itself. If the process dies — spot instance reclaimed,
54
+ node rebooted, OOM killer, power cut — running the same command again picks up
55
+ where it stopped: same weights, same optimizer moments, same LR schedule, same
56
+ RNG state, same position in the dataset.
57
+
58
+ ## Why this is not just `torch.save`
59
+
60
+ Saving the weights is the easy part. What makes a resumed run *continue* rather
61
+ than merely *restart from a good place* is everything around them:
62
+
63
+ | Restored | Why it matters |
64
+ |---|---|
65
+ | Model weights and buffers | the obvious part |
66
+ | Optimizer state | Adam moments; without them the first steps after resume are wrong |
67
+ | LR scheduler | resume at the wrong LR and the loss curve visibly kinks |
68
+ | AMP `GradScaler` | its loss scale is tuned state, not a constant |
69
+ | RNG: torch, CUDA, Python, NumPy | dropout masks and augmentations replay identically |
70
+ | Dataset position | you continue with batch 4001, not batch 1 |
71
+
72
+ Ravex's test suite asserts the strong form of this: a run killed at step 20 and
73
+ resumed produces losses that are **bit-identical**, step by step, to the run
74
+ that was never interrupted, and the same final weights.
75
+
76
+ ## Two ways to use it
77
+
78
+ **Zero code changes.** `ravex enable` installs a one-line `.pth` file in
79
+ site-packages, which Python executes at interpreter startup. From then on Ravex
80
+ attaches itself to any training process that has a `ravex.yaml`.
81
+
82
+ **One line**, when you would rather be explicit:
83
+
84
+ ```python
85
+ import ravex
86
+ ravex.activate()
87
+ ```
88
+
89
+ Both do the same thing. The `.pth` route exists so that a platform can enable
90
+ checkpointing for code it does not own.
91
+
92
+ ## Configuration
93
+
94
+ `ravex.yaml`, anywhere at or above the working directory:
95
+
96
+ ```yaml
97
+ checkpoint_every: 500 # optimizer steps between checkpoints
98
+ backend: moonclip # moonclip | torch_save
99
+ storage:
100
+ type: local # local | s3 | r2
101
+ path: ./checkpoints
102
+ keep_last: 5
103
+ max_steps: null # optional hard stop, see below
104
+ sharded_checkpoints: gather # gather | per_rank, for FSDP — see below
105
+ ```
106
+
107
+ Every option also reads from `RAVEX_*` environment variables, which win over
108
+ the file — so a scheduler can override a config committed to the repository:
109
+
110
+ ```bash
111
+ RAVEX_CHECKPOINT_EVERY=100 RAVEX_STORAGE_TYPE=r2 RAVEX_STORAGE_BUCKET=runs python train.py
112
+ ```
113
+
114
+ Credentials are never read from the config file. Set `RAVEX_S3_ACCESS_KEY` /
115
+ `RAVEX_S3_SECRET_KEY`, or the usual `AWS_*` pair.
116
+
117
+ Full reference: [docs/configuration.md](docs/configuration.md).
118
+
119
+ ## Backends
120
+
121
+ **`moonclip`** (default) — the [Moonclip](https://codeberg.org/JHNMACHINE/moonclip)
122
+ engine: per-tensor delta tracking, so unchanged weights cost zero I/O; zstd
123
+ compression; background writes; direct S3/R2 sync.
124
+
125
+ **`torch_save`** — one `.pt` file per checkpoint, written on a background
126
+ thread. Used automatically when Moonclip is not installed. Correct, just larger
127
+ and slower.
128
+
129
+ ## How it works
130
+
131
+ Ravex patches five things in PyTorch and nothing in your code:
132
+
133
+ - `nn.Module.__init__` and `.train()` — to notice your models
134
+ - `Optimizer.__init__` — to attach a step hook to every optimizer
135
+ - `DataLoader.__init__` and `.__iter__` — to track the dataset position and to
136
+ find the one moment where a resume can be applied
137
+
138
+ The step counter advances once per `optimizer.step()`, so gradient accumulation
139
+ needs no special handling. Checkpoints are collected at the *top of an
140
+ iteration*, never inside one: mid-iteration the LR scheduler has not stepped
141
+ yet, and a checkpoint taken there resumes with a stale learning rate.
142
+
143
+ Collection runs on the training thread — it has to, to be consistent with the
144
+ step that just finished — and copies the state; the write itself happens in the
145
+ background. What the loop pays for is the copy, not the I/O.
146
+
147
+ More detail: [docs/how-it-works.md](docs/how-it-works.md).
148
+
149
+ ## Safety
150
+
151
+ Ravex is designed to be un-noticeable when it works and harmless when it does
152
+ not:
153
+
154
+ - every hook is wrapped; if one raises, your call still returns normally
155
+ - if a checkpoint fails, Ravex disables itself and logs it — training continues
156
+ - nothing is ever written to stdout; logs go to `log_file`, or to stderr at
157
+ WARNING and above
158
+ - installing the package changes nothing on its own. Without `ravex enable`
159
+ there is no `.pth`; with it, Ravex still only wakes up for projects that have
160
+ a `ravex.yaml` or set `RAVEX_ENABLED=1`
161
+ - `ravex disable` removes the autoloader; `RAVEX_ENABLED=0` turns it off for a
162
+ single run
163
+
164
+ ## Status and limits
165
+
166
+ Alpha. Works with plain PyTorch loops, and with anything built on them, since
167
+ the hooks are on PyTorch itself.
168
+
169
+ ### With a framework driving the loop
170
+
171
+ HuggingFace `Trainer` and Lightning are covered by their own tests, and the
172
+ result deserves to be stated precisely rather than as "it works":
173
+
174
+ - **State restoration is exact.** Model, optimizer, LR scheduler and step count
175
+ all come back. With the per-step randomness removed, a killed run resumes
176
+ into a loss sequence identical to the uninterrupted one.
177
+ - **Replay is not.** With shuffling and dropout on, the resumed run continues
178
+ correctly from the checkpointed state but sees a different draw. Both
179
+ frameworks iterate the dataloader on their own schedule and consume the
180
+ global RNG around the loop, so the epoch-start snapshot no longer lines up.
181
+
182
+ Plain loops, DDP and FSDP *are* bit-exact with randomness on. This is a
183
+ framework-interaction limit, not a general one, and it costs you a different
184
+ shuffle from the resume point onwards — not a wrong model.
185
+
186
+ Verified: plain loops, gradient accumulation, LR schedulers, AMP loss-scale
187
+ state, `num_workers > 0`, DDP, and FSDP. A killed `torchrun` job resumes on
188
+ *every* rank with bit-identical losses, sharded or not, and the checkpoint it
189
+ leaves behind loads into a plain single-process model afterwards.
190
+
191
+ ### Sharded models
192
+
193
+ With FSDP each rank holds a slice of every parameter, so `state_dict()` returns
194
+ a fragment. Two ways to turn that into a checkpoint, picked with
195
+ `sharded_checkpoints`:
196
+
197
+ **`gather`** (default) rebuilds the whole state on rank 0, which writes it. The
198
+ checkpoint is then independent of the topology that produced it — eight GPUs in,
199
+ one out — and it does not scale: rank 0 has to hold the entire model and
200
+ optimizer in host memory, and it is the rank that then does the writing.
201
+
202
+ **`per_rank`** has every rank write its own shard into its own store,
203
+ `<storage.path>/rank_<n>`. Nothing is gathered, so nothing is bounded by one
204
+ rank's memory, and on a 1.48B model collecting the state went from 15.6 s to
205
+ 1.5 s. What you give up is the resharding: those shards are cut for one topology,
206
+ so the checkpoint resumes at the same world size and starts clean at any other.
207
+ Needs FSDP2 — under FSDP1 Ravex degrades to `gather` and says so.
208
+
209
+ Either way, collecting is a **collective**: every rank participates, and there is
210
+ **no final checkpoint at exit** for a sharded model. Shutdown is where ranks stop
211
+ being in lockstep, and a collective nobody else joins hangs. Losing the last few
212
+ steps is bounded; a hang is not. Set `checkpoint_every` accordingly.
213
+
214
+ Numbers and the FSDP1 details: [docs/configuration.md](docs/configuration.md).
215
+
216
+ Known limits today:
217
+
218
+ - **`IterableDataset`**: no index sampler exists, so the stream position cannot
219
+ be replayed. Everything else is still restored.
220
+ - **Your loop's bounds**: a resumed script runs its own `for epoch in
221
+ range(N)` again from the top; it has no idea 3000 steps already happened. Set
222
+ `max_steps` and Ravex ends the run at the right step regardless of how many
223
+ times the process restarted.
224
+
225
+ Under AMP, note that an overflowing gradient makes `scaler.step()` skip the
226
+ optimizer. Ravex counts optimizer steps, not loop iterations, so a skipped
227
+ iteration does not advance the counter — which is the right unit, since nothing
228
+ about the model changed, but it does mean the step count and the number of
229
+ batches you fed differ.
230
+
231
+ The GPU paths — AMP with real fp16 overflow, the CUDA RNG, FSDP1, NCCL — are
232
+ covered by `integration/test_cuda.py`, which skips without a GPU. They were
233
+ last verified on 8× RTX 5060 Ti with torch 2.12/cu130.
234
+
235
+ ## Development
236
+
237
+ ```bash
238
+ pip install -e ".[dev]"
239
+ pytest
240
+ ```
241
+
242
+ The unit suite runs in-process. The parts that only exist across a real process
243
+ boundary — the `.pth` autoloader, a resume starting from an empty interpreter,
244
+ `torchrun` — live in `integration/` and need Linux:
245
+
246
+ ```bash
247
+ docker build -f integration/Dockerfile -t ravex-integration .
248
+ docker run --rm ravex-integration
249
+ ```
250
+
251
+ ## Licence
252
+
253
+ Apache 2.0. See [LICENSE](LICENSE).
ravex-0.0.2/README.md ADDED
@@ -0,0 +1,224 @@
1
+ # Ravex
2
+
3
+ Transparent checkpoint and resume for PyTorch training. Your training script
4
+ does not change — not one line, not one import.
5
+
6
+ ```bash
7
+ pip install ravex
8
+ ravex enable
9
+ ```
10
+
11
+ Ravex itself is pure Python and installs anywhere. Its default checkpoint engine,
12
+ [Moonclip](https://codeberg.org/JHNMACHINE/moonclip), ships wheels for Linux
13
+ x86_64 only — so `pip install "ravex[moonclip]"` is a Linux thing, and on any
14
+ other platform Ravex falls back to `torch_save` on its own. What changed between
15
+ versions is in
16
+ [CHANGELOG.md](https://codeberg.org/JHNMACHINE/ravex/src/branch/main/CHANGELOG.md).
17
+
18
+ Drop a `ravex.yaml` next to your code and run what you always ran:
19
+
20
+ ```bash
21
+ python train.py
22
+ ```
23
+
24
+ The run now checkpoints itself. If the process dies — spot instance reclaimed,
25
+ node rebooted, OOM killer, power cut — running the same command again picks up
26
+ where it stopped: same weights, same optimizer moments, same LR schedule, same
27
+ RNG state, same position in the dataset.
28
+
29
+ ## Why this is not just `torch.save`
30
+
31
+ Saving the weights is the easy part. What makes a resumed run *continue* rather
32
+ than merely *restart from a good place* is everything around them:
33
+
34
+ | Restored | Why it matters |
35
+ |---|---|
36
+ | Model weights and buffers | the obvious part |
37
+ | Optimizer state | Adam moments; without them the first steps after resume are wrong |
38
+ | LR scheduler | resume at the wrong LR and the loss curve visibly kinks |
39
+ | AMP `GradScaler` | its loss scale is tuned state, not a constant |
40
+ | RNG: torch, CUDA, Python, NumPy | dropout masks and augmentations replay identically |
41
+ | Dataset position | you continue with batch 4001, not batch 1 |
42
+
43
+ Ravex's test suite asserts the strong form of this: a run killed at step 20 and
44
+ resumed produces losses that are **bit-identical**, step by step, to the run
45
+ that was never interrupted, and the same final weights.
46
+
47
+ ## Two ways to use it
48
+
49
+ **Zero code changes.** `ravex enable` installs a one-line `.pth` file in
50
+ site-packages, which Python executes at interpreter startup. From then on Ravex
51
+ attaches itself to any training process that has a `ravex.yaml`.
52
+
53
+ **One line**, when you would rather be explicit:
54
+
55
+ ```python
56
+ import ravex
57
+ ravex.activate()
58
+ ```
59
+
60
+ Both do the same thing. The `.pth` route exists so that a platform can enable
61
+ checkpointing for code it does not own.
62
+
63
+ ## Configuration
64
+
65
+ `ravex.yaml`, anywhere at or above the working directory:
66
+
67
+ ```yaml
68
+ checkpoint_every: 500 # optimizer steps between checkpoints
69
+ backend: moonclip # moonclip | torch_save
70
+ storage:
71
+ type: local # local | s3 | r2
72
+ path: ./checkpoints
73
+ keep_last: 5
74
+ max_steps: null # optional hard stop, see below
75
+ sharded_checkpoints: gather # gather | per_rank, for FSDP — see below
76
+ ```
77
+
78
+ Every option also reads from `RAVEX_*` environment variables, which win over
79
+ the file — so a scheduler can override a config committed to the repository:
80
+
81
+ ```bash
82
+ RAVEX_CHECKPOINT_EVERY=100 RAVEX_STORAGE_TYPE=r2 RAVEX_STORAGE_BUCKET=runs python train.py
83
+ ```
84
+
85
+ Credentials are never read from the config file. Set `RAVEX_S3_ACCESS_KEY` /
86
+ `RAVEX_S3_SECRET_KEY`, or the usual `AWS_*` pair.
87
+
88
+ Full reference: [docs/configuration.md](docs/configuration.md).
89
+
90
+ ## Backends
91
+
92
+ **`moonclip`** (default) — the [Moonclip](https://codeberg.org/JHNMACHINE/moonclip)
93
+ engine: per-tensor delta tracking, so unchanged weights cost zero I/O; zstd
94
+ compression; background writes; direct S3/R2 sync.
95
+
96
+ **`torch_save`** — one `.pt` file per checkpoint, written on a background
97
+ thread. Used automatically when Moonclip is not installed. Correct, just larger
98
+ and slower.
99
+
100
+ ## How it works
101
+
102
+ Ravex patches five things in PyTorch and nothing in your code:
103
+
104
+ - `nn.Module.__init__` and `.train()` — to notice your models
105
+ - `Optimizer.__init__` — to attach a step hook to every optimizer
106
+ - `DataLoader.__init__` and `.__iter__` — to track the dataset position and to
107
+ find the one moment where a resume can be applied
108
+
109
+ The step counter advances once per `optimizer.step()`, so gradient accumulation
110
+ needs no special handling. Checkpoints are collected at the *top of an
111
+ iteration*, never inside one: mid-iteration the LR scheduler has not stepped
112
+ yet, and a checkpoint taken there resumes with a stale learning rate.
113
+
114
+ Collection runs on the training thread — it has to, to be consistent with the
115
+ step that just finished — and copies the state; the write itself happens in the
116
+ background. What the loop pays for is the copy, not the I/O.
117
+
118
+ More detail: [docs/how-it-works.md](docs/how-it-works.md).
119
+
120
+ ## Safety
121
+
122
+ Ravex is designed to be un-noticeable when it works and harmless when it does
123
+ not:
124
+
125
+ - every hook is wrapped; if one raises, your call still returns normally
126
+ - if a checkpoint fails, Ravex disables itself and logs it — training continues
127
+ - nothing is ever written to stdout; logs go to `log_file`, or to stderr at
128
+ WARNING and above
129
+ - installing the package changes nothing on its own. Without `ravex enable`
130
+ there is no `.pth`; with it, Ravex still only wakes up for projects that have
131
+ a `ravex.yaml` or set `RAVEX_ENABLED=1`
132
+ - `ravex disable` removes the autoloader; `RAVEX_ENABLED=0` turns it off for a
133
+ single run
134
+
135
+ ## Status and limits
136
+
137
+ Alpha. Works with plain PyTorch loops, and with anything built on them, since
138
+ the hooks are on PyTorch itself.
139
+
140
+ ### With a framework driving the loop
141
+
142
+ HuggingFace `Trainer` and Lightning are covered by their own tests, and the
143
+ result deserves to be stated precisely rather than as "it works":
144
+
145
+ - **State restoration is exact.** Model, optimizer, LR scheduler and step count
146
+ all come back. With the per-step randomness removed, a killed run resumes
147
+ into a loss sequence identical to the uninterrupted one.
148
+ - **Replay is not.** With shuffling and dropout on, the resumed run continues
149
+ correctly from the checkpointed state but sees a different draw. Both
150
+ frameworks iterate the dataloader on their own schedule and consume the
151
+ global RNG around the loop, so the epoch-start snapshot no longer lines up.
152
+
153
+ Plain loops, DDP and FSDP *are* bit-exact with randomness on. This is a
154
+ framework-interaction limit, not a general one, and it costs you a different
155
+ shuffle from the resume point onwards — not a wrong model.
156
+
157
+ Verified: plain loops, gradient accumulation, LR schedulers, AMP loss-scale
158
+ state, `num_workers > 0`, DDP, and FSDP. A killed `torchrun` job resumes on
159
+ *every* rank with bit-identical losses, sharded or not, and the checkpoint it
160
+ leaves behind loads into a plain single-process model afterwards.
161
+
162
+ ### Sharded models
163
+
164
+ With FSDP each rank holds a slice of every parameter, so `state_dict()` returns
165
+ a fragment. Two ways to turn that into a checkpoint, picked with
166
+ `sharded_checkpoints`:
167
+
168
+ **`gather`** (default) rebuilds the whole state on rank 0, which writes it. The
169
+ checkpoint is then independent of the topology that produced it — eight GPUs in,
170
+ one out — and it does not scale: rank 0 has to hold the entire model and
171
+ optimizer in host memory, and it is the rank that then does the writing.
172
+
173
+ **`per_rank`** has every rank write its own shard into its own store,
174
+ `<storage.path>/rank_<n>`. Nothing is gathered, so nothing is bounded by one
175
+ rank's memory, and on a 1.48B model collecting the state went from 15.6 s to
176
+ 1.5 s. What you give up is the resharding: those shards are cut for one topology,
177
+ so the checkpoint resumes at the same world size and starts clean at any other.
178
+ Needs FSDP2 — under FSDP1 Ravex degrades to `gather` and says so.
179
+
180
+ Either way, collecting is a **collective**: every rank participates, and there is
181
+ **no final checkpoint at exit** for a sharded model. Shutdown is where ranks stop
182
+ being in lockstep, and a collective nobody else joins hangs. Losing the last few
183
+ steps is bounded; a hang is not. Set `checkpoint_every` accordingly.
184
+
185
+ Numbers and the FSDP1 details: [docs/configuration.md](docs/configuration.md).
186
+
187
+ Known limits today:
188
+
189
+ - **`IterableDataset`**: no index sampler exists, so the stream position cannot
190
+ be replayed. Everything else is still restored.
191
+ - **Your loop's bounds**: a resumed script runs its own `for epoch in
192
+ range(N)` again from the top; it has no idea 3000 steps already happened. Set
193
+ `max_steps` and Ravex ends the run at the right step regardless of how many
194
+ times the process restarted.
195
+
196
+ Under AMP, note that an overflowing gradient makes `scaler.step()` skip the
197
+ optimizer. Ravex counts optimizer steps, not loop iterations, so a skipped
198
+ iteration does not advance the counter — which is the right unit, since nothing
199
+ about the model changed, but it does mean the step count and the number of
200
+ batches you fed differ.
201
+
202
+ The GPU paths — AMP with real fp16 overflow, the CUDA RNG, FSDP1, NCCL — are
203
+ covered by `integration/test_cuda.py`, which skips without a GPU. They were
204
+ last verified on 8× RTX 5060 Ti with torch 2.12/cu130.
205
+
206
+ ## Development
207
+
208
+ ```bash
209
+ pip install -e ".[dev]"
210
+ pytest
211
+ ```
212
+
213
+ The unit suite runs in-process. The parts that only exist across a real process
214
+ boundary — the `.pth` autoloader, a resume starting from an empty interpreter,
215
+ `torchrun` — live in `integration/` and need Linux:
216
+
217
+ ```bash
218
+ docker build -f integration/Dockerfile -t ravex-integration .
219
+ docker run --rm ravex-integration
220
+ ```
221
+
222
+ ## Licence
223
+
224
+ Apache 2.0. See [LICENSE](LICENSE).