1module Platform.Linux.Nvidia.PlanHostAdamW
2
3import Checkpoint.Envelope
4import Data.Bytes
5import Data.Float32Bits
6import Model.Config
7import Model.Parameter
8import Model.Word64
9import Platform.Linux.Nvidia.PlanHost
10import Platform.Linux.Nvidia.PlanHostRequest
11import Runtime.NativePhysicalProgram
12import Std.Float
13import Std.List
14import Std.Natural
15
16-- AdamW's step-dependent scalars, supplied by the host at run time. The
17-- bias-corrected step size lr sqrt(1 - beta2^t) / (1 - beta1^t) and the
18-- scaled epsilon eps sqrt(1 - beta2^t) depend on the update's number t,
19-- which counts on across invocations: the checkpoint's header carries the
20-- updates taken (Checkpoint.Envelope), so an invocation's k-th update is
21-- t = that count + k + 1. After the recipes are expanded, and before
22-- anything is submitted, the host computes both for each of its updates
23-- and writes them into the parameter word of every AdamW launch of that
24-- update (the realization reports where it placed each launch's parameter
25-- block: NvidiaWholeProgramParameterAddresses). A fresh start reads a
26-- zero count, so it takes t = 1, 2, ...
27--
28-- The arithmetic is the learner's: binary64 throughout, rounded once to
29-- binary32 at the end, or binary32 at every step (the binary64 operation
30-- rounded to binary32, which is exactly the binary32 operation). The
31-- powers are running products, taken first over the updates already done.
32
33family NvidiaPlanHostArithmetic : Type 0
34constructor NvidiaPlanHostBinary64
35constructor NvidiaPlanHostBinary32
36end-family
37
38-- a hyperparameter: a ratio of naturals (the nearest value of the
39-- arithmetic's precision), or a binary32's bits
40family NvidiaPlanHostConstant : Type 0
41constructor NvidiaPlanHostRatio
42field unrestricted nvidiaPlanHostRatioNumerator : Nat
43field unrestricted nvidiaPlanHostRatioDenominator : Nat
44constructor NvidiaPlanHostBinary32Bits
45field unrestricted nvidiaPlanHostBinary32BitsValue : Nat
46end-family
47
48family NvidiaPlanHostAdamW : Type 0
49constructor NvidiaPlanHostAdamWValue
50field unrestricted nvidiaPlanHostAdamWArithmetic : (family NvidiaPlanHostArithmetic)
51field unrestricted nvidiaPlanHostAdamWLearningRate : (family NvidiaPlanHostConstant)
52field unrestricted nvidiaPlanHostAdamWBeta1 : (family NvidiaPlanHostConstant)
53field unrestricted nvidiaPlanHostAdamWBeta2 : (family NvidiaPlanHostConstant)
54field unrestricted nvidiaPlanHostAdamWEpsilon : (family NvidiaPlanHostConstant)
55end-family
56
57-- a scheduled learner (the warmup-cosine schedule, at the end of this module)
58family NvidiaPlanHostAdamWSchedule : Type 0
59constructor NvidiaPlanHostAdamWWarmupCosine
60field unrestricted nvidiaPlanHostSchedulePeak : (family NvidiaPlanHostConstant)
61field unrestricted nvidiaPlanHostScheduleFloor : (family NvidiaPlanHostConstant)
62field unrestricted nvidiaPlanHostScheduleWarmup : Nat
63field unrestricted nvidiaPlanHostScheduleHorizon : Nat
64field unrestricted nvidiaPlanHostScheduleBeta1 : (family NvidiaPlanHostConstant)
65field unrestricted nvidiaPlanHostScheduleBeta2 : (family NvidiaPlanHostConstant)
66field unrestricted nvidiaPlanHostScheduleEpsilon : (family NvidiaPlanHostConstant)
67field unrestricted nvidiaPlanHostScheduleWeightDecay : (family NvidiaPlanHostConstant)
68field unrestricted nvidiaPlanHostScheduleDecayLow : Nat
69field unrestricted nvidiaPlanHostScheduleOrigin : Nat
70end-family
71
72-- a binary32's bits, for NvidiaPlanHostBinary32Bits
73def nvidiaPlanHostBinary32Of =
74 (lambda unrestricted value : F32 . (nvidiaPlanHostLittleNatural 4 (dataFloat32EncodeLE value)))
75
76-- ---- where the words go ----
77-- entry `ordinal` of a realization's per-launch table of 8-byte words: a
78-- launch's parameter block's device address
79-- (NvidiaWholeProgramParameterAddresses), or a word in it
80-- (NvidiaWholeProgramParameterWords)
81def nvidiaPlanHostParameterBlock =
82 (lambda unrestricted addresses : Bytes .
83 (lambda unrestricted ordinal : Nat .
84 (nvidiaPlanHostLittleNatural 8 (dataBytesDropValidated (naturalMultiply 8 ordinal) addresses))))
85
86-- the entries at `ordinals` (ascending) of such a table, in one walk
87-- (the table is dropped through once, not once per entry)
88def nvidiaPlanHostTableEntries =
89 (lambda unrestricted table : Bytes .
90 (lambda unrestricted ordinals : (family StdList Nat) .
91 (app
92 (app
93 (eliminate StdList
94 (lambda unrestricted current : (family StdList Nat) .
95 (pi unrestricted at : Nat . (pi unrestricted rest : Bytes . (family StdList Nat))))
96 ordinals
97 (branch StdListEmpty .
98 (lambda unrestricted at : Nat . (lambda unrestricted rest : Bytes . (constructor StdList StdListEmpty Nat))))
99 (branch StdListCons ordinal tail induction .
100 (lambda unrestricted at : Nat .
101 (lambda unrestricted rest : Bytes .
102 (let unrestricted here = (dataBytesDropValidated (naturalMultiply 8 (naturalSaturatingSubtract ordinal at)) rest) in
103 (constructor StdList StdListCons Nat (nvidiaPlanHostLittleNatural 8 here) (induction ordinal here)))))))
104 0)
105 table)))
106
107-- every update's launches, in order (the updates' launches ascend), and
108-- the stretch of such a list (or of its entries) that is update k's
109def adAllLaunches =
110 (lambda unrestricted updates : Nat .
111 (lambda unrestricted launches : (pi unrestricted k : Nat . (family StdList Nat)) .
112 (nat-eliminate
113 (lambda unrestricted current : Nat . (family StdList Nat))
114 (constructor StdList StdListEmpty Nat)
115 (lambda unrestricted k : Nat .
116 (lambda unrestricted earlier : (family StdList Nat) . (stdListAppend Nat earlier (launches k))))
117 updates)))
118
119def adUpdateStretch =
120 (lambda unrestricted launches : (pi unrestricted k : Nat . (family StdList Nat)) .
121 (lambda unrestricted all : (family StdList Nat) .
122 (lambda unrestricted k : Nat .
123 (stdListTake Nat (stdListLength Nat (launches k))
124 (stdListDrop Nat (stdListLength Nat (adAllLaunches k launches)) all)))))
125
126-- whether a buffer maps a device address
127def nvidiaPlanHostBufferMaps =
128 (lambda unrestricted buffer : (family NvidiaPlanHostBuffer) .
129 (lambda unrestricted device : Nat .
130 (naturalAnd (naturalLessOrEqual (phBufferGPU buffer) device)
131 (naturalLess device (naturalAdd (phBufferGPU buffer) (phBufferExtent buffer))))))
132
133def adMappedHost =
134 (lambda unrestricted buffer : (family NvidiaPlanHostBuffer) .
135 (lambda unrestricted device : Nat .
136 (naturalSelect (nvidiaPlanHostBufferMaps buffer device)
137 (naturalAdd (phBufferHost buffer) (naturalSaturatingSubtract device (phBufferGPU buffer)))
138 0)))
139
140-- a device address's host address in a host-mapped data buffer (else 0)
141def adMappedDataHost =
142 (lambda unrestricted buffers : (family NvidiaPlanHostBuffers) .
143 (lambda unrestricted device : Nat .
144 (eliminate NvidiaPlanHostBuffers (lambda unrestricted current : (family NvidiaPlanHostBuffers) . Nat) buffers
145 (branch NvidiaPlanHostBuffersEnd . 0)
146 (branch NvidiaPlanHostBuffersNext head rest induction .
147 (naturalAdd (naturalSelect (naturalNonzero (phBufferHost head)) (adMappedHost head device) 0) induction)))))
148
149-- Parameter blocks live in the program buffer or past the QMD table -- in
150-- the QMD buffer, or in a mapped data buffer after it when the table and
151-- its spill outgrow it (Coppelius's QMD overflow, under the GB10's 512-byte
152-- QMD records): a device address's host address there (else 0).
153def nvidiaPlanHostDeviceToHost =
154 (lambda unrestricted layout : (family NvidiaPlanHostLayout) .
155 (lambda unrestricted device : Nat .
156 (eliminate NvidiaPlanHostLayout (lambda unrestricted current : (family NvidiaPlanHostLayout) . Nat) layout
157 (branch NvidiaPlanHostLayoutValue gpfifo push sem program qmd userd data abi lifecycle expectedName errorHost .
158 (naturalAdd (adMappedHost program device)
159 (naturalAdd (adMappedHost qmd device) (adMappedDataHost data device)))))))
160
161-- update k's sites, from its launches (ordinals): the host address of the
162-- word at `offset` in each launch's parameter block
163def nvidiaPlanHostAdamWHostSites =
164 (lambda unrestricted layout : (family NvidiaPlanHostLayout) .
165 (lambda unrestricted addresses : Bytes .
166 (lambda unrestricted offset : Nat .
167 (lambda unrestricted updates : Nat .
168 (lambda unrestricted launches : (pi unrestricted k : Nat . (family StdList Nat)) .
169 (let unrestricted sites =
170 (stdListMap Nat Nat
171 (lambda unrestricted block : Nat . (nvidiaPlanHostDeviceToHost layout (naturalAdd block offset)))
172 (nvidiaPlanHostTableEntries addresses (adAllLaunches updates launches))) in
173 (adUpdateStretch launches sites)))))))
174
175-- What admits the schedule (1): every update has launches (ascending); each one's
176-- site is mapped; and each holds, as realized, the word `baked k` the plan
177-- put there for its update (the first invocation's; `words` is the
178-- realization's NvidiaWholeProgramParameterWords at the site's offset) --
179-- so a site is the word the launch's AdamW reads, not a neighbour.
180def nvidiaPlanHostAdamWSitesAdmitted =
181 (lambda unrestricted layout : (family NvidiaPlanHostLayout) .
182 (lambda unrestricted addresses : Bytes .
183 (lambda unrestricted words : Bytes .
184 (lambda unrestricted offset : Nat .
185 (lambda unrestricted updates : Nat .
186 (lambda unrestricted launches : (pi unrestricted k : Nat . (family StdList Nat)) .
187 (lambda unrestricted baked : (pi unrestricted k : Nat . Nat) .
188 (let unrestricted all = (adAllLaunches updates launches) in
189 (let unrestricted blocks = (nvidiaPlanHostTableEntries addresses all) in
190 (let unrestricted realized = (nvidiaPlanHostTableEntries words all) in
191 (nat-eliminate
192 (lambda unrestricted current : Nat . Nat)
193 1
194 (lambda unrestricted k : Nat .
195 (lambda unrestricted induction : Nat .
196 (naturalAnd induction
197 (naturalAnd (naturalNonzero (stdListLength Nat (launches k)))
198 (naturalAnd
199 (stdListFold Nat Nat
200 (lambda unrestricted block : Nat .
201 (lambda unrestricted rest : Nat .
202 (naturalAnd (naturalNonzero (nvidiaPlanHostDeviceToHost layout (naturalAdd block offset))) rest)))
203 1
204 (adUpdateStretch launches blocks k))
205 (stdListFold Nat Nat
206 (lambda unrestricted word : Nat .
207 (lambda unrestricted rest : Nat . (naturalAnd (naturalEqual word (baked k)) rest)))
208 1
209 (adUpdateStretch launches realized k)))))))
210 updates)))))))))))
211
212-- ---- the computation ----
213-- the scratch words (PlanHostRequest.prLearnerScratch)
214def adWord = (lambda unrestricted index : Nat . (naturalAdd prLearnerScratch (naturalMultiply 8 index)))
215def adOne : Nat = (adWord 0)
216def adBeta1 : Nat = (adWord 1)
217def adBeta2 : Nat = (adWord 2)
218def adRate : Nat = (adWord 3)
219def adEpsilon : Nat = (adWord 4)
220def adPower1 : Nat = (adWord 5)
221def adPower2 : Nat = (adWord 6)
222def adGap1 : Nat = (adWord 7)
223def adRoot : Nat = (adWord 8)
224def adStep : Nat = (adWord 9)
225def adScaled : Nat = (adWord 10)
226def adDenominator : Nat = (adWord 11)
227-- the parameter word: the step size's binary32, then the epsilon's (the
228-- second store's high half lands in the word after)
229def adOut : Nat = (adWord 12)
230
231def adOp =
232 (lambda unrestricted kind : (family NativePhysicalFloat64Operation) .
233 (lambda unrestricted destination : Nat .
234 (lambda unrestricted left : (family NativePhysicalOperand) .
235 (lambda unrestricted right : (family NativePhysicalOperand) .
236 (lambda unrestricted tail : (family NativePhysicalCommands) .
237 (phNext (constructor NativePhysicalOperation NativePhysicalFloat64 kind (phState destination) left right) tail))))))
238
239
240-- the arithmetic's rounding of a word just computed
241def adRound =
242 (lambda unrestricted arithmetic : (family NvidiaPlanHostArithmetic) .
243 (lambda unrestricted at : Nat .
244 (lambda unrestricted tail : (family NativePhysicalCommands) .
245 (eliminate NvidiaPlanHostArithmetic (lambda unrestricted current : (family NvidiaPlanHostArithmetic) . (family NativePhysicalCommands)) arithmetic
246 (branch NvidiaPlanHostBinary64 . tail)
247 (branch NvidiaPlanHostBinary32 .
248 (adOp (constructor NativePhysicalFloat64Operation NativePhysicalFloat64ToBinary32) at (phLoad at) (phImm 0)
249 (adOp (constructor NativePhysicalFloat64Operation NativePhysicalFloat64FromBinary32) at (phLoad at) (phImm 0) tail)))))))
250
251def adArith =
252 (lambda unrestricted arithmetic : (family NvidiaPlanHostArithmetic) .
253 (lambda unrestricted kind : (family NativePhysicalFloat64Operation) .
254 (lambda unrestricted destination : Nat .
255 (lambda unrestricted left : Nat .
256 (lambda unrestricted right : Nat .
257 (lambda unrestricted tail : (family NativePhysicalCommands) .
258 (adOp kind destination (phLoad left) (phLoad right) (adRound arithmetic destination tail))))))))
259
260def adConstant =
261 (lambda unrestricted arithmetic : (family NvidiaPlanHostArithmetic) .
262 (lambda unrestricted constant : (family NvidiaPlanHostConstant) .
263 (lambda unrestricted destination : Nat .
264 (lambda unrestricted tail : (family NativePhysicalCommands) .
265 (eliminate NvidiaPlanHostConstant (lambda unrestricted current : (family NvidiaPlanHostConstant) . (family NativePhysicalCommands)) constant
266 (branch NvidiaPlanHostRatio numerator denominator .
267 (adOp (constructor NativePhysicalFloat64Operation NativePhysicalFloat64FromNatural) destination (phImm numerator) (phImm 0)
268 (adOp (constructor NativePhysicalFloat64Operation NativePhysicalFloat64FromNatural) adDenominator (phImm denominator) (phImm 0)
269 (adArith arithmetic (constructor NativePhysicalFloat64Operation NativePhysicalFloat64Divide) destination destination adDenominator tail))))
270 (branch NvidiaPlanHostBinary32Bits bits .
271 (phNext (constructor NativePhysicalOperation NativePhysicalStoreWord64 (phState destination) (phImm bits))
272 (adOp (constructor NativePhysicalFloat64Operation NativePhysicalFloat64FromBinary32) destination (phLoad destination) (phImm 0) tail))))))))
273
274def adMultiply = (constructor NativePhysicalFloat64Operation NativePhysicalFloat64Multiply)
275def adSubtract = (constructor NativePhysicalFloat64Operation NativePhysicalFloat64Subtract)
276def adDivide = (constructor NativePhysicalFloat64Operation NativePhysicalFloat64Divide)
277
278-- one more update's powers
279def adAdvance =
280 (lambda unrestricted arithmetic : (family NvidiaPlanHostArithmetic) .
281 (lambda unrestricted tail : (family NativePhysicalCommands) .
282 (adArith arithmetic adMultiply adPower1 adPower1 adBeta1
283 (adArith arithmetic adMultiply adPower2 adPower2 adBeta2 tail))))
284
285-- update k: its powers, its two scalars, their parameter word into every
286-- site of the update
287def adUpdate =
288 (lambda unrestricted arithmetic : (family NvidiaPlanHostArithmetic) .
289 (lambda unrestricted sites : (family StdList Nat) .
290 (lambda unrestricted tail : (family NativePhysicalCommands) .
291 (adAdvance arithmetic
292 (adArith arithmetic adSubtract adRoot adOne adPower2
293 (adOp (constructor NativePhysicalFloat64Operation NativePhysicalFloat64SquareRoot) adRoot (phImm 0) (phLoad adRoot)
294 (adRound arithmetic adRoot
295 (adArith arithmetic adSubtract adGap1 adOne adPower1
296 (adArith arithmetic adMultiply adStep adRate adRoot
297 (adArith arithmetic adDivide adStep adStep adGap1
298 (adArith arithmetic adMultiply adScaled adEpsilon adRoot
299 (adOp (constructor NativePhysicalFloat64Operation NativePhysicalFloat64ToBinary32) adOut (phLoad adStep) (phImm 0)
300 (adOp (constructor NativePhysicalFloat64Operation NativePhysicalFloat64ToBinary32) (naturalAdd adOut 4) (phLoad adScaled) (phImm 0)
301 (stdListFold Nat (family NativePhysicalCommands)
302 (lambda unrestricted site : Nat .
303 (lambda unrestricted rest : (family NativePhysicalCommands) .
304 (phNext (constructor NativePhysicalOperation NativePhysicalStoreWord64 (phImm site) (phLoad adOut)) rest)))
305 tail
306 sites))))))))))))))
307
308-- The commands: the constants, the powers after the updates the
309-- checkpoint's header counts, then each of the invocation's `updates` in
310-- order, update k's word written to `sites k` (host addresses).
311def nvidiaPlanHostAdamWCommands =
312 (lambda unrestricted adamW : (family NvidiaPlanHostAdamW) .
313 (lambda unrestricted updates : Nat .
314 (lambda unrestricted sites : (pi unrestricted k : Nat . (family StdList Nat)) .
315 (lambda unrestricted tail : (family NativePhysicalCommands) .
316 (eliminate NvidiaPlanHostAdamW (lambda unrestricted current : (family NvidiaPlanHostAdamW) . (family NativePhysicalCommands)) adamW
317 (branch NvidiaPlanHostAdamWValue arithmetic rate beta1 beta2 epsilon .
318 (adOp (constructor NativePhysicalFloat64Operation NativePhysicalFloat64FromNatural) adOne (phImm 1) (phImm 0)
319 (adConstant arithmetic beta1 adBeta1
320 (adConstant arithmetic beta2 adBeta2
321 (adConstant arithmetic rate adRate
322 (adConstant arithmetic epsilon adEpsilon
323 (phNext (constructor NativePhysicalOperation NativePhysicalStoreWord64 (phState adPower1) (phLoad adOne))
324 (phNext (constructor NativePhysicalOperation NativePhysicalStoreWord64 (phState adPower2) (phLoad adOne))
325 (phNext (constructor NativePhysicalOperation NativePhysicalRepeatBeginCounted
326 (phLoad (naturalAdd prCheckpointHeaderIn checkpointEnvelopeUpdatesAt)))
327 (adAdvance arithmetic
328 (phNext (constructor NativePhysicalOperation NativePhysicalRepeatEnd)
329 (app
330 (nat-eliminate
331 (lambda unrestricted current : Nat .
332 (pi unrestricted after : (family NativePhysicalCommands) . (family NativePhysicalCommands)))
333 (lambda unrestricted after : (family NativePhysicalCommands) . after)
334 (lambda unrestricted k : Nat .
335 (lambda unrestricted induction : (pi unrestricted after : (family NativePhysicalCommands) . (family NativePhysicalCommands)) .
336 (lambda unrestricted after : (family NativePhysicalCommands) .
337 (induction (adUpdate arithmetic (sites k) after)))))
338 updates)
339 tail)))))))))))))))))
340
341-- One more update's commands, for a run loop's body
342-- (PlanHostRequest.NvidiaPlanHostStepCommands): the powers advanced, the
343-- two scalars computed, their word written into every one of `sites`. The
344-- constants and the powers the checkpoint counts come first, from
345-- `nvidiaPlanHostAdamWCommands` with no updates of its own.
346def nvidiaPlanHostAdamWStepCommands =
347 (lambda unrestricted adamW : (family NvidiaPlanHostAdamW) .
348 (lambda unrestricted sites : (family StdList Nat) .
349 (eliminate NvidiaPlanHostAdamW (lambda unrestricted current : (family NvidiaPlanHostAdamW) . (family NativePhysicalCommands)) adamW
350 (branch NvidiaPlanHostAdamWValue arithmetic rate beta1 beta2 epsilon .
351 (adUpdate arithmetic sites (constructor NativePhysicalCommands NativePhysicalCommandsEnd))))))
352
353-- ---- a scheduled learning rate: warmup, then a cosine decay ----
354-- llm.c's schedule over the update's number s (the updates taken before
355-- it since `origin`, the checkpoint count the schedule starts from: the
356-- first update has s = 0; a run continued from another learner's
357-- checkpoint starts its moments, and so their bias correction, there too):
358-- peak (s + 1) / W over the W warmup
359-- updates, then floor + (peak - floor) (1 + cos (pi p)) / 2 with
360-- p = (s - W) / (T - W), T the horizon (p held at 1 past it). The weight
361-- decay's factor follows the rate: 1 - lr wd, rewritten every update (its
362-- binary32 the high half of the 8-byte parameter word whose low half is
363-- `decayLow`'s bits: AdamWHalfSM86 pairs 1 - beta2 with the decay).
364--
365-- Branch-free over the host's binary64 operations: min(a, b) =
366-- (a + b - |a - b|) / 2 and max(d, 0) = (d + |d|) / 2 with |x| = sqrt(x x),
367-- exact on these whole numbers (each below 2^26: the horizon is admitted
368-- below it). With theta = pi (p - 1/2), (1 + cos (pi p)) / 2 =
369-- (1 - sin theta) / 2, sin by its Taylor series through theta^15 in Horner
370-- form (|theta| <= pi / 2: the first term left out is below 7e-12). pi is
371-- the ratio 245850922 / 78256779, a convergent within 1e-17 of it, whose
372-- nearest binary64 is pi's. The arithmetic is binary64, each word rounded
373-- once to binary32; scripts/alpha/gpt2-oracle.py repeats these operations in
374-- this order.
375
376-- warmup at least one update, the horizon past it and below 2^26 (the
377-- branch-free minima stay exact)
378def nvidiaPlanHostAdamWScheduleAdmitted =
379 (lambda unrestricted schedule : (family NvidiaPlanHostAdamWSchedule) .
380 (eliminate NvidiaPlanHostAdamWSchedule (lambda unrestricted current : (family NvidiaPlanHostAdamWSchedule) . Nat) schedule
381 (branch NvidiaPlanHostAdamWWarmupCosine peak floor warmup horizon beta1 beta2 epsilon decay decayLow origin .
382 (naturalAnd (naturalLess 0 warmup)
383 (naturalAnd (naturalLess warmup horizon)
384 (naturalAnd (naturalLess horizon (naturalPowerOfTwo 26))
385 (naturalAnd (naturalLess origin (naturalPowerOfTwo 40))
386 (naturalLess decayLow (naturalPowerOfTwo 32)))))))))
387
388-- the schedule's scratch words (PlanHostRequest.prScheduleScratch); the
389-- step size, epsilon, powers and rate are adUpdate's
390def asWord = (lambda unrestricted index : Nat . (naturalAdd prScheduleScratch (naturalMultiply 8 index)))
391-- s, the updates taken
392def asUpdates : Nat = (asWord 0)
393def asPeak : Nat = (asWord 1)
394def asFloor : Nat = (asWord 2)
395def asSpan : Nat = (asWord 3)
396def asWarmup : Nat = (asWord 4)
397-- T - W
398def asReach : Nat = (asWord 5)
399def asWeightDecay : Nat = (asWord 6)
400def asHalf : Nat = (asWord 7)
401def asPi : Nat = (asWord 8)
402def asA : Nat = (asWord 9)
403def asB : Nat = (asWord 10)
404def asX : Nat = (asWord 11)
405def asSquare : Nat = (asWord 12)
406def asSeries : Nat = (asWord 13)
407def asWarm : Nat = (asWord 14)
408def asDecay : Nat = (asWord 15)
409-- the decay's parameter word (its binary32 store reaches word 17)
410def asDecayOut : Nat = (asWord 16)
411-- the Horner steps' divisors, innermost first: 14 x 15, 12 x 13, ..., 2 x 3
412def asSeriesTerms : Nat = 7
413def asDivisorOf = (lambda unrestricted k : Nat .
414 (let unrestricted n = (naturalMultiply 2 (naturalSaturatingSubtract asSeriesTerms k)) in
415 (naturalMultiply n (succ n))))
416def asDivisor = (lambda unrestricted k : Nat . (asWord (naturalAdd 18 k)))
417-- s as a whole number (the prelude's count of the powers' advances)
418def asCount : Nat = (asWord 25)
419def asScratchFits : (equal Nat (naturalLessOrEqual (naturalAdd 18 asSeriesTerms) 25) 1) = (refl Nat 1)
420-- 2^64 - 1, the word whose sum with x is x - 1 (the host's words wrap)
421def asWordMaximum : Nat = 18446744073709551615
422
423def asAdd = (constructor NativePhysicalFloat64Operation NativePhysicalFloat64Add)
424def asOp =
425 (lambda unrestricted kind : (family NativePhysicalFloat64Operation) .
426 (lambda unrestricted destination : Nat . (lambda unrestricted left : Nat . (lambda unrestricted right : Nat .
427 (lambda unrestricted tail : (family NativePhysicalCommands) .
428 (adOp kind destination (phLoad left) (phLoad right) tail))))))
429-- |source| into destination: sqrt(source source)
430def asAbsolute =
431 (lambda unrestricted destination : Nat . (lambda unrestricted source : Nat .
432 (lambda unrestricted tail : (family NativePhysicalCommands) .
433 (asOp adMultiply destination source source
434 (adOp (constructor NativePhysicalFloat64Operation NativePhysicalFloat64SquareRoot) destination (phImm 0) (phLoad destination)
435 tail)))))
436
437-- the Horner steps k .. asSeriesTerms - 1: series = 1 - theta^2 series / divisor
438def asSeriesSteps =
439 (lambda unrestricted tail : (family NativePhysicalCommands) .
440 (app
441 (nat-eliminate
442 (lambda unrestricted current : Nat . (pi unrestricted k : Nat . (family NativePhysicalCommands)))
443 (lambda unrestricted k : Nat . tail)
444 (lambda unrestricted predecessor : Nat .
445 (lambda unrestricted rest : (pi unrestricted k : Nat . (family NativePhysicalCommands)) .
446 (lambda unrestricted k : Nat .
447 (asOp adMultiply asA asSquare asSeries
448 (asOp adDivide asA asA (asDivisor k)
449 (asOp adSubtract asSeries adOne asA
450 (rest (succ k))))))))
451 asSeriesTerms)
452 0))
453
454-- the update's rate into adRate and its decay factor into asDecay, from s
455def asRateCommands =
456 (lambda unrestricted tail : (family NativePhysicalCommands) .
457 -- warm = min(s + 1, W) / W
458 (asOp asAdd asA asUpdates adOne
459 (asOp adSubtract asB asA asWarmup
460 (asAbsolute asB asB
461 (asOp asAdd asX asA asWarmup
462 (asOp adSubtract asX asX asB
463 (asOp adMultiply asX asX asHalf
464 (asOp adDivide asWarm asX asWarmup
465 -- q = min(max(s - W, 0), T - W)
466 (asOp adSubtract asB asUpdates asWarmup
467 (asAbsolute asA asB
468 (asOp asAdd asB asB asA
469 (asOp adMultiply asB asB asHalf
470 (asOp adSubtract asA asB asReach
471 (asAbsolute asA asA
472 (asOp asAdd asX asB asReach
473 (asOp adSubtract asX asX asA
474 (asOp adMultiply asX asX asHalf
475 -- theta = pi (q / (T - W) - 1/2), and its square
476 (asOp adDivide asX asX asReach
477 (asOp adSubtract asX asX asHalf
478 (asOp adMultiply asX asX asPi
479 (asOp adMultiply asSquare asX asX
480 -- sin theta
481 (phNext (constructor NativePhysicalOperation NativePhysicalStoreWord64 (phState asSeries) (phLoad adOne))
482 (asSeriesSteps
483 (asOp adMultiply asSeries asSeries asX
484 -- rate = warm (floor + (peak - floor) (1 - sin theta) / 2)
485 (asOp adSubtract asSeries adOne asSeries
486 (asOp adMultiply asSeries asSeries asHalf
487 (asOp adMultiply asSeries asSeries asSpan
488 (asOp asAdd asSeries asSeries asFloor
489 (asOp adMultiply adRate asSeries asWarm
490 -- the decay factor
491 (asOp adMultiply asDecay adRate asWeightDecay
492 (asOp adSubtract asDecay adOne asDecay
493 tail)))))))))))))))))))))))))))))))
494
495-- the decay's word into every one of `sites`
496def asDecayStores =
497 (lambda unrestricted decayLow : Nat . (lambda unrestricted sites : (family StdList Nat) .
498 (lambda unrestricted tail : (family NativePhysicalCommands) .
499 (phNext (constructor NativePhysicalOperation NativePhysicalStoreWord64 (phState asDecayOut) (phImm decayLow))
500 (adOp (constructor NativePhysicalFloat64Operation NativePhysicalFloat64ToBinary32) (naturalAdd asDecayOut 4) (phLoad asDecay) (phImm 0)
501 (stdListFold Nat (family NativePhysicalCommands)
502 (lambda unrestricted site : Nat .
503 (lambda unrestricted rest : (family NativePhysicalCommands) .
504 (phNext (constructor NativePhysicalOperation NativePhysicalStoreWord64 (phImm site) (phLoad asDecayOut)) rest)))
505 tail
506 sites))))))
507
508-- the schedule's divisors, k from `first`
509def asDivisorCommands =
510 (lambda unrestricted tail : (family NativePhysicalCommands) .
511 (app
512 (nat-eliminate
513 (lambda unrestricted current : Nat . (pi unrestricted k : Nat . (family NativePhysicalCommands)))
514 (lambda unrestricted k : Nat . tail)
515 (lambda unrestricted predecessor : Nat .
516 (lambda unrestricted rest : (pi unrestricted k : Nat . (family NativePhysicalCommands)) .
517 (lambda unrestricted k : Nat .
518 (adOp (constructor NativePhysicalFloat64Operation NativePhysicalFloat64FromNatural) (asDivisor k)
519 (phImm (asDivisorOf k)) (phImm 0)
520 (rest (succ k))))))
521 asSeriesTerms)
522 0))
523
524-- An invocation's prelude: the constants, s = the updates the checkpoint's
525-- header counts less the origin (two wrapping word sums: + (2^64 - 1 -
526-- origin), + 1), and the powers advanced s times.
527def nvidiaPlanHostAdamWScheduleCommands =
528 (lambda unrestricted schedule : (family NvidiaPlanHostAdamWSchedule) .
529 (lambda unrestricted tail : (family NativePhysicalCommands) .
530 (let unrestricted binary64 = (constructor NvidiaPlanHostArithmetic NvidiaPlanHostBinary64) in
531 (eliminate NvidiaPlanHostAdamWSchedule (lambda unrestricted current : (family NvidiaPlanHostAdamWSchedule) . (family NativePhysicalCommands)) schedule
532 (branch NvidiaPlanHostAdamWWarmupCosine peak floor warmup horizon beta1 beta2 epsilon decay decayLow origin .
533 (adOp (constructor NativePhysicalFloat64Operation NativePhysicalFloat64FromNatural) adOne (phImm 1) (phImm 0)
534 (adConstant binary64 beta1 adBeta1
535 (adConstant binary64 beta2 adBeta2
536 (adConstant binary64 epsilon adEpsilon
537 (adConstant binary64 peak asPeak
538 (adConstant binary64 floor asFloor
539 (adConstant binary64 decay asWeightDecay
540 (adConstant binary64 (constructor NvidiaPlanHostConstant NvidiaPlanHostRatio 1 2) asHalf
541 (adConstant binary64 (constructor NvidiaPlanHostConstant NvidiaPlanHostRatio 245850922 78256779) asPi
542 (adOp (constructor NativePhysicalFloat64Operation NativePhysicalFloat64FromNatural) asWarmup (phImm warmup) (phImm 0)
543 (adOp (constructor NativePhysicalFloat64Operation NativePhysicalFloat64FromNatural) asReach
544 (phImm (naturalSaturatingSubtract horizon warmup)) (phImm 0)
545 (asOp adSubtract asSpan asPeak asFloor
546 (asDivisorCommands
547 (phNext (constructor NativePhysicalOperation NativePhysicalAddWord64 (phState asCount)
548 (phLoad (naturalAdd prCheckpointHeaderIn checkpointEnvelopeUpdatesAt))
549 (phImm (naturalSaturatingSubtract asWordMaximum origin)))
550 (phNext (constructor NativePhysicalOperation NativePhysicalAddWord64 (phState asCount) (phLoad asCount) (phImm 1))
551 (phNext (constructor NativePhysicalOperation NativePhysicalStoreWord64 (phState adPower1) (phLoad adOne))
552 (phNext (constructor NativePhysicalOperation NativePhysicalStoreWord64 (phState adPower2) (phLoad adOne))
553 (phNext (constructor NativePhysicalOperation NativePhysicalRepeatBeginCounted (phLoad asCount))
554 (adAdvance binary64
555 (phNext (constructor NativePhysicalOperation NativePhysicalRepeatEnd)
556 (adOp (constructor NativePhysicalFloat64Operation NativePhysicalFloat64FromNatural) asUpdates
557 (phLoad asCount) (phImm 0)
558 tail))))))))))))))))))))))))))
559
560-- One update's commands, for a run loop's body: its rate and decay from s,
561-- the powers advanced, the step size and epsilon into every one of
562-- `stepSites` (adUpdate's word), the decay's word into every one of
563-- `decaySites`, then s + 1.
564def nvidiaPlanHostAdamWScheduleStepCommands =
565 (lambda unrestricted schedule : (family NvidiaPlanHostAdamWSchedule) .
566 (lambda unrestricted stepSites : (family StdList Nat) .
567 (lambda unrestricted decaySites : (family StdList Nat) .
568 (eliminate NvidiaPlanHostAdamWSchedule (lambda unrestricted current : (family NvidiaPlanHostAdamWSchedule) . (family NativePhysicalCommands)) schedule
569 (branch NvidiaPlanHostAdamWWarmupCosine peak floor warmup horizon beta1 beta2 epsilon decay decayLow origin .
570 (asRateCommands
571 (adUpdate (constructor NvidiaPlanHostArithmetic NvidiaPlanHostBinary64) stepSites
572 (asDecayStores decayLow decaySites
573 (asOp asAdd asUpdates asUpdates adOne
574 (constructor NativePhysicalCommands NativePhysicalCommandsEnd))))))))))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.