1module Checkpoint.Envelope
2
3import Data.Bytes
4import Data.SHA256Digest
5import Std.Foundation
6import Std.Natural
7
8-- THE CHECKPOINT ENVELOPE: what a native host writes around a system's raw
9-- training state, and refuses to resume from unless it holds.
10--
11-- A checkpoint file is a header of `headerBytes` (a whole number of pages)
12-- and then the payload, unmodified (for Bob and Coppelius: the P/M/V
13-- planes). The header, little-endian:
14-- 0 magic "ALPHACKP"
15-- 8 version (1)
16-- 16 header bytes
17-- 24 payload bytes
18-- 32 chunk bytes
19-- 40 full chunks (payload bytes / chunk bytes)
20-- 48 tail bytes (payload bytes mod chunk bytes)
21-- 56 completed updates (the learner's step count)
22-- 64 completed invocations (the sampler's position: batches consumed)
23-- 72 SHA-256 of the state schema's identity (model, state layout)
24-- 104 SHA-256 of the learner's identity (optimizer, hyperparameters)
25-- 136 SHA-256 of the data contract's identity (tokenizer, batch format)
26-- 168 SHA-256 of the initialization's identity (the RNG seed)
27-- 200 the SHA-256 of each chunk of the payload, in order (the tail last)
28-- then zeros to `headerBytes`. The file is complete when its size is
29-- exactly header + payload bytes.
30--
31-- A system states its contract (below); the first 56 bytes and the four
32-- identity digests of any checkpoint it resumes from must be the ones the
33-- contract gives, and every chunk must hash to its digest. The first
34-- failure is named (`CheckpointEnvelopeVerdict`): a short or long file is
35-- truncated, another version or layout is refused, another system's,
36-- learner's, data's or seed's state is foreign, and a chunk whose digest
37-- differs -- a swapped or corrupted file -- is refused by its index.
38
39family CheckpointEnvelopeContract : Type 0
40constructor CheckpointEnvelopeContractValue
41field unrestricted checkpointEnvelopeSchemaIdentity : Bytes
42field unrestricted checkpointEnvelopeLearnerIdentity : Bytes
43field unrestricted checkpointEnvelopeDataIdentity : Bytes
44field unrestricted checkpointEnvelopeSeedIdentity : Bytes
45field unrestricted checkpointEnvelopePayloadBytes : Nat
46field unrestricted checkpointEnvelopeChunkBytes : Nat
47field unrestricted checkpointEnvelopeUpdatesPerInvocation : Nat
48end-family
49
50family CheckpointEnvelopeVerdict : Type 0
51constructor CheckpointEnvelopeAccepted
52constructor CheckpointEnvelopeTruncated
53constructor CheckpointEnvelopeLayout
54constructor CheckpointEnvelopeForeignSchema
55constructor CheckpointEnvelopeForeignLearner
56constructor CheckpointEnvelopeForeignData
57constructor CheckpointEnvelopeForeignSeed
58constructor CheckpointEnvelopeChunkMismatch
59field unrestricted checkpointEnvelopeMismatchedChunk : Nat
60end-family
61
62def checkpointEnvelopeMagic : Bytes = b"ALPHACKP"
63def checkpointEnvelopeVersion : Nat = 1
64def checkpointEnvelopePageBytes : Nat = 4096
65def checkpointEnvelopeDigestBytes : Nat = 32
66
67-- the header's fields, as byte offsets
68def checkpointEnvelopeVersionAt : Nat = 8
69def checkpointEnvelopeHeaderBytesAt : Nat = 16
70def checkpointEnvelopePayloadBytesAt : Nat = 24
71def checkpointEnvelopeChunkBytesAt : Nat = 32
72def checkpointEnvelopeFullChunksAt : Nat = 40
73def checkpointEnvelopeTailBytesAt : Nat = 48
74def checkpointEnvelopeUpdatesAt : Nat = 56
75def checkpointEnvelopeInvocationsAt : Nat = 64
76def checkpointEnvelopeSchemaAt : Nat = 72
77def checkpointEnvelopeLearnerAt : Nat = 104
78def checkpointEnvelopeDataAt : Nat = 136
79def checkpointEnvelopeSeedAt : Nat = 168
80def checkpointEnvelopeDigestsAt : Nat = 200
81-- the bytes every checkpoint of a contract begins with: magic to tail
82def checkpointEnvelopeFixedBytes : Nat = 56
83
84def checkpointEnvelopeContractPayload =
85 (lambda unrestricted contract : (family CheckpointEnvelopeContract) .
86 (eliminate CheckpointEnvelopeContract (lambda unrestricted current : (family CheckpointEnvelopeContract) . Nat) contract
87 (branch CheckpointEnvelopeContractValue schema learner data seed payload chunk updates . payload)))
88
89def checkpointEnvelopeContractChunk =
90 (lambda unrestricted contract : (family CheckpointEnvelopeContract) .
91 (eliminate CheckpointEnvelopeContract (lambda unrestricted current : (family CheckpointEnvelopeContract) . Nat) contract
92 (branch CheckpointEnvelopeContractValue schema learner data seed payload chunk updates . chunk)))
93
94def checkpointEnvelopeContractUpdates =
95 (lambda unrestricted contract : (family CheckpointEnvelopeContract) .
96 (eliminate CheckpointEnvelopeContract (lambda unrestricted current : (family CheckpointEnvelopeContract) . Nat) contract
97 (branch CheckpointEnvelopeContractValue schema learner data seed payload chunk updates . updates)))
98
99-- the four identities' digests, in header order
100def checkpointEnvelopeIdentityDigests =
101 (lambda unrestricted contract : (family CheckpointEnvelopeContract) .
102 (eliminate CheckpointEnvelopeContract (lambda unrestricted current : (family CheckpointEnvelopeContract) . Bytes) contract
103 (branch CheckpointEnvelopeContractValue schema learner data seed payload chunk updates .
104 (bytes-append (sha256RawDigestOrEmpty schema)
105 (bytes-append (sha256RawDigestOrEmpty learner)
106 (bytes-append (sha256RawDigestOrEmpty data) (sha256RawDigestOrEmpty seed)))))))
107
108def checkpointEnvelopeFullChunks =
109 (lambda unrestricted contract : (family CheckpointEnvelopeContract) .
110 (naturalDivideUnchecked (checkpointEnvelopeContractPayload contract) (checkpointEnvelopeContractChunk contract)))
111
112def checkpointEnvelopeTailBytes =
113 (lambda unrestricted contract : (family CheckpointEnvelopeContract) .
114 (naturalModuloUnchecked (checkpointEnvelopeContractPayload contract) (checkpointEnvelopeContractChunk contract)))
115
116-- the chunks a payload has: the full ones and, when there is one, the tail
117def checkpointEnvelopeChunkCount =
118 (lambda unrestricted contract : (family CheckpointEnvelopeContract) .
119 (naturalAdd (checkpointEnvelopeFullChunks contract) (naturalNonzero (checkpointEnvelopeTailBytes contract))))
120
121-- The final digest covers only bytes present in the payload. The generic
122-- byte-prefix operation zero-pads, so passing the full chunk extent for the
123-- tail would disagree with native file hashing while still passing a codec
124-- round trip that made the same mistake on both sides.
125def checkpointEnvelopeChunkExtent =
126 (lambda unrestricted contract : (family CheckpointEnvelopeContract) .
127 (lambda unrestricted index : Nat .
128 (naturalSelect (naturalLess index (checkpointEnvelopeFullChunks contract))
129 (checkpointEnvelopeContractChunk contract)
130 (naturalSelect (naturalEqual index (checkpointEnvelopeFullChunks contract))
131 (checkpointEnvelopeTailBytes contract) 0))))
132
133def checkpointEnvelopeHeaderBytes =
134 (lambda unrestricted contract : (family CheckpointEnvelopeContract) .
135 (naturalMultiply
136 (naturalDivideUnchecked
137 (naturalAdd
138 (naturalAdd checkpointEnvelopeDigestsAt (naturalMultiply checkpointEnvelopeDigestBytes (checkpointEnvelopeChunkCount contract)))
139 (naturalSaturatingSubtract checkpointEnvelopePageBytes 1))
140 checkpointEnvelopePageBytes)
141 checkpointEnvelopePageBytes))
142
143def checkpointEnvelopeFileBytes =
144 (lambda unrestricted contract : (family CheckpointEnvelopeContract) .
145 (naturalAdd (checkpointEnvelopeHeaderBytes contract) (checkpointEnvelopeContractPayload contract)))
146
147-- a natural as eight little-endian bytes (it is below 2^64)
148def checkpointEnvelopeWord =
149 (lambda unrestricted value : Nat .
150 (bytes-builder-build
151 (app (nat-eliminate
152 (lambda unrestricted current : Nat . (pi unrestricted rest : Nat . BytesBuilder))
153 (lambda unrestricted rest : Nat . (bytes-builder-chunk b""))
154 (lambda unrestricted p : Nat .
155 (lambda unrestricted induction : (pi unrestricted rest : Nat . BytesBuilder) .
156 (lambda unrestricted rest : Nat .
157 (bytes-builder-append
158 (bytes-builder-chunk (bytes (nat-to-byte (naturalModuloUnchecked rest 256))))
159 (induction (naturalDivideUnchecked rest 256))))))
160 8)
161 value)))
162
163-- the fixed bytes: magic, version, header, payload, chunk, full, tail
164def checkpointEnvelopeFixed =
165 (lambda unrestricted contract : (family CheckpointEnvelopeContract) .
166 (bytes-append checkpointEnvelopeMagic
167 (bytes-append (checkpointEnvelopeWord checkpointEnvelopeVersion)
168 (bytes-append (checkpointEnvelopeWord (checkpointEnvelopeHeaderBytes contract))
169 (bytes-append (checkpointEnvelopeWord (checkpointEnvelopeContractPayload contract))
170 (bytes-append (checkpointEnvelopeWord (checkpointEnvelopeContractChunk contract))
171 (bytes-append (checkpointEnvelopeWord (checkpointEnvelopeFullChunks contract))
172 (checkpointEnvelopeWord (checkpointEnvelopeTailBytes contract)))))))))
173
174-- a choice between verdicts on a 0/1 natural
175def checkpointEnvelopeIf =
176 (lambda unrestricted condition : Nat .
177 (lambda unrestricted whenTrue : (family CheckpointEnvelopeVerdict) .
178 (lambda unrestricted whenFalse : (family CheckpointEnvelopeVerdict) .
179 (nat-eliminate
180 (lambda unrestricted current : Nat . (family CheckpointEnvelopeVerdict))
181 whenFalse
182 (lambda unrestricted p : Nat . (lambda unrestricted ignored : (family CheckpointEnvelopeVerdict) . whenTrue))
183 condition))))
184
185def checkpointEnvelopeZeros =
186 (lambda unrestricted count : Nat .
187 (bytes-builder-build
188 (nat-eliminate
189 (lambda unrestricted current : Nat . BytesBuilder)
190 (bytes-builder-chunk b"")
191 (lambda unrestricted p : Nat . (lambda unrestricted induction : BytesBuilder . (bytes-builder-append (bytes-builder-chunk (bytes 0)) induction)))
192 count)))
193
194-- ---- the model's reading of a checkpoint file ----
195def checkpointEnvelopeSlice =
196 (lambda unrestricted file : Bytes .
197 (lambda unrestricted offset : Nat .
198 (lambda unrestricted length : Nat .
199 (dataBytesTakeValidated length (dataBytesDropValidated offset file)))))
200
201-- the index of the first chunk whose digest differs, or the chunk count
202def checkpointEnvelopeFirstBadChunk =
203 (lambda unrestricted contract : (family CheckpointEnvelopeContract) .
204 (lambda unrestricted file : Bytes .
205 (let unrestricted count = (checkpointEnvelopeChunkCount contract) in
206 (let unrestricted header = (checkpointEnvelopeHeaderBytes contract) in
207 (let unrestricted chunk = (checkpointEnvelopeContractChunk contract) in
208 (app
209 (nat-eliminate
210 (lambda unrestricted remaining : Nat . (pi unrestricted index : Nat . Nat))
211 (lambda unrestricted index : Nat . index)
212 (lambda unrestricted p : Nat .
213 (lambda unrestricted induction : (pi unrestricted index : Nat . Nat) .
214 (lambda unrestricted index : Nat .
215 (nat-eliminate
216 (lambda unrestricted same : Nat . Nat)
217 index
218 (lambda unrestricted q : Nat . (lambda unrestricted ignored : Nat . (induction (succ index))))
219 (bytes-equal
220 (checkpointEnvelopeSlice file (naturalAdd checkpointEnvelopeDigestsAt (naturalMultiply checkpointEnvelopeDigestBytes index)) checkpointEnvelopeDigestBytes)
221 (sha256RawDigestOrEmpty (checkpointEnvelopeSlice file (naturalAdd header (naturalMultiply chunk index)) (checkpointEnvelopeChunkExtent contract index))))))))
222 count)
223 0))))))
224
225def checkpointEnvelopeVerdictOf =
226 (lambda unrestricted contract : (family CheckpointEnvelopeContract) .
227 (lambda unrestricted file : Bytes .
228 (let unrestricted digests = (checkpointEnvelopeIdentityDigests contract) in
229 (let unrestricted identity = (lambda unrestricted index : Nat . (lambda unrestricted at : Nat .
230 (bytes-equal (checkpointEnvelopeSlice file at checkpointEnvelopeDigestBytes)
231 (checkpointEnvelopeSlice digests (naturalMultiply index checkpointEnvelopeDigestBytes) checkpointEnvelopeDigestBytes)))) in
232 (let unrestricted bad = (checkpointEnvelopeFirstBadChunk contract file) in
233 (checkpointEnvelopeIf
234 (naturalEqual (bytes-length file) (checkpointEnvelopeFileBytes contract))
235 (checkpointEnvelopeIf
236 (bytes-equal (checkpointEnvelopeSlice file 0 checkpointEnvelopeFixedBytes) (checkpointEnvelopeFixed contract))
237 (checkpointEnvelopeIf (identity 0 checkpointEnvelopeSchemaAt)
238 (checkpointEnvelopeIf (identity 1 checkpointEnvelopeLearnerAt)
239 (checkpointEnvelopeIf (identity 2 checkpointEnvelopeDataAt)
240 (checkpointEnvelopeIf (identity 3 checkpointEnvelopeSeedAt)
241 (checkpointEnvelopeIf (naturalEqual bad (checkpointEnvelopeChunkCount contract))
242 (constructor CheckpointEnvelopeVerdict CheckpointEnvelopeAccepted)
243 (constructor CheckpointEnvelopeVerdict CheckpointEnvelopeChunkMismatch bad))
244 (constructor CheckpointEnvelopeVerdict CheckpointEnvelopeForeignSeed))
245 (constructor CheckpointEnvelopeVerdict CheckpointEnvelopeForeignData))
246 (constructor CheckpointEnvelopeVerdict CheckpointEnvelopeForeignLearner))
247 (constructor CheckpointEnvelopeVerdict CheckpointEnvelopeForeignSchema))
248 (constructor CheckpointEnvelopeVerdict CheckpointEnvelopeLayout))
249 (constructor CheckpointEnvelopeVerdict CheckpointEnvelopeTruncated)))))))
250
251-- the checkpoint a contract's host writes: the header -- fixed bytes,
252-- `updates` and `invocations`, the identity digests, each chunk's digest
253-- -- and the payload
254def checkpointEnvelopeWrite =
255 (lambda unrestricted contract : (family CheckpointEnvelopeContract) .
256 (lambda unrestricted updates : Nat .
257 (lambda unrestricted invocations : Nat .
258 (lambda unrestricted payload : Bytes .
259 (let unrestricted chunk = (checkpointEnvelopeContractChunk contract) in
260 (let unrestricted digests =
261 (bytes-builder-build
262 (app (nat-eliminate
263 (lambda unrestricted remaining : Nat . (pi unrestricted index : Nat . BytesBuilder))
264 (lambda unrestricted index : Nat . (bytes-builder-chunk b""))
265 (lambda unrestricted p : Nat .
266 (lambda unrestricted induction : (pi unrestricted index : Nat . BytesBuilder) .
267 (lambda unrestricted index : Nat .
268 (bytes-builder-append
269 (bytes-builder-chunk (sha256RawDigestOrEmpty (checkpointEnvelopeSlice payload (naturalMultiply chunk index) (checkpointEnvelopeChunkExtent contract index))))
270 (induction (succ index))))))
271 (checkpointEnvelopeChunkCount contract))
272 0)) in
273 (let unrestricted body =
274 (bytes-append (checkpointEnvelopeFixed contract)
275 (bytes-append (checkpointEnvelopeWord updates)
276 (bytes-append (checkpointEnvelopeWord invocations)
277 (bytes-append (checkpointEnvelopeIdentityDigests contract) digests)))) in
278 (bytes-append body
279 (bytes-append (checkpointEnvelopeZeros (naturalSaturatingSubtract (checkpointEnvelopeHeaderBytes contract) (bytes-length body)))
280 payload)))))))))The compiler supplied declaration spans and resolved links from this source snapshot. This page does not assert that this file belongs to a checked closure.