Coppelius.TrainingRun

1module Coppelius.TrainingRun
2
3import Checkpoint.Envelope
4import Coppelius.Learner
5import Coppelius.Model
6import Model.Config
7import Model.Parameter
8import Model.Word64
9import Std.Natural
10import Std.Physical
11import Std.Word
12
13-- Coppelius trains under its run policy (below): one process, a fixed
14-- number of AdamW steps, each on the next context of a token stream.  Loss
15-- is telemetry, never a stopping condition.  A zero-byte checkpoint input
16-- selects fresh deterministic initialization; a full raw P/M/V image
17-- selects continuation.
18family CoppeliusTrainingRunSemantics : Type 0
19constructor CoppeliusTrainingRunSemanticsValue
20field unrestricted coppeliusRunUpdatesPerInvocation : Nat
21field unrestricted coppeliusRunCheckpointInterval : Nat
22field unrestricted coppeliusRunLossControlsTermination : Nat
23field unrestricted coppeliusRunDeterministicInitialization : Nat
24field unrestricted coppeliusRunFreshCheckpointInputBytes : ByteCount
25field unrestricted coppeliusRunResumeCheckpointRequired : Nat
26field unrestricted coppeliusRunCheckpointOutputRequired : Nat
27field unrestricted coppeliusRunFreshProcessResumeRequired : Nat
28field unrestricted coppeliusRunNativeOnlyRequired : Nat
29field unrestricted coppeliusRunAllowedHostFallbacks : Nat
30
31end-family
32
33-- The checkpoint's payload: a contiguous FP32 P/M/V image.  The file wraps
34-- it in the envelope (coppeliusCheckpointContract, below), whose header
35-- carries the update count in band.
36family CoppeliusCheckpointTransportSchema : Type 0
37constructor CoppeliusCheckpointTransportSchemaValue
38field unrestricted coppeliusCheckpointSchemaIdentity : Bytes
39field unrestricted coppeliusCheckpointParameterCount : StdU64
40field unrestricted coppeliusCheckpointPlaneCount : Nat
41field unrestricted coppeliusCheckpointBytesPerScalar : Nat
42field unrestricted coppeliusCheckpointPayloadBytes : ByteCount
43field unrestricted coppeliusCheckpointOrderingIdentity : Bytes
44field unrestricted coppeliusCheckpointInBandStepMetadata : Nat
45field unrestricted coppeliusCheckpointExternalStepReceiptRequired : Nat
46
47end-family
48
49-- ---- the run policy ----
50-- One process takes `coppeliusRunSteps` AdamW steps, step s on context
51-- c0 + s of the token stream the request names (c0 the contexts its
52-- checkpoint has consumed: 0 fresh), and publishes its checkpoint after
53-- every `coppeliusRunCheckpointInterval` steps -- the last publication at
54-- its end.  A checkpoint goes to the request's checkpoint path through its
55-- temporary: written whole, digested, synced, then renamed over the path,
56-- so the path always holds a complete checkpoint (Platform.Linux.Nvidia.
57-- PlanHostRequest).  The loop is compiled: the host issues one update
58-- submission of the plan per step (Coppelius.Build.NativeHost).
59def coppeliusRunSteps : Nat =
60  10000
61
62def coppeliusRunCheckpointInterval : Nat =
63  1000
64
65-- publications a run makes: the interval divides the steps
66def coppeliusRunPublications : Nat =
67  (naturalDivideUnchecked coppeliusRunSteps coppeliusRunCheckpointInterval)
68
69def coppeliusRunIntervalDividesSteps :
70  (equal Nat (naturalMultiply coppeliusRunPublications coppeliusRunCheckpointInterval) coppeliusRunSteps) =
71  (refl Nat coppeliusRunSteps)
72
73-- the plan's update submissions: one, issued once per step
74def coppeliusUpdatePieces : Nat =
75  1
76
77def coppeliusFreshCheckpointInputBytes : ByteCount =
78  (stdByteCount 0)
79
80def coppeliusTrainingRunSemantics : (family CoppeliusTrainingRunSemantics) =
81  (record
82    CoppeliusTrainingRunSemantics
83    (coppeliusRunUpdatesPerInvocation = coppeliusRunSteps)
84    (coppeliusRunCheckpointInterval = coppeliusRunCheckpointInterval)
85    (coppeliusRunLossControlsTermination = 0)
86    (coppeliusRunDeterministicInitialization = 1)
87    (coppeliusRunFreshCheckpointInputBytes = coppeliusFreshCheckpointInputBytes)
88    (coppeliusRunResumeCheckpointRequired = 1)
89    (coppeliusRunCheckpointOutputRequired = 1)
90    (coppeliusRunFreshProcessResumeRequired = 1)
91    (coppeliusRunNativeOnlyRequired = 1)
92    (coppeliusRunAllowedHostFallbacks = 0))
93
94def coppeliusCheckpointParameterCount : StdU64 =
95  57_705_472
96
97def coppeliusCheckpointPayloadBytesNatural : Nat = 692_465_664
98def coppeliusCheckpointPayloadBytes : ByteCount =
99  (stdByteCount (modelWord64FromNaturalTruncated coppeliusCheckpointPayloadBytesNatural))
100
101-- The input record a step consumes, as the kernels' parameter blocks place
102-- it in the staging region: the context's 1024 input tokens, and its 1024
103-- targets at 0x1000.  Nothing else: the rotary tables are computed from the
104-- model (Coppelius.Build.Graph.cgRotaryTables), and a fresh model is drawn
105-- from the run's seed.
106def coppeliusInputSequence : Nat = 1024
107def coppeliusInputTargetsOffset : Nat = 0x1000
108def coppeliusInputRecordBytes : Nat =
109  (naturalAdd coppeliusInputTargetsOffset (naturalMultiply 4 coppeliusInputSequence))
110
111-- The token stream the host reads the records from: a token cache (the
112-- tokenizer's output), a 16-byte header (the source's modification time
113-- and size) then little-endian u32 token ids.  Context c is ids
114-- [1024 c, 1024 c + 1024]: its inputs the first 1024, its targets the last
115-- 1024 -- the same window one token on.
116def coppeliusTokenStreamHeaderBytes : Nat = 16
117def coppeliusTokenBytes : Nat = 4
118def coppeliusContextStreamBytes : Nat = (naturalMultiply coppeliusTokenBytes coppeliusInputSequence)
119
120-- the seed a fresh model is drawn from
121def coppeliusRunSeed : Nat =
122  20260925
123
124-- What the kernels leave in the staging region for the host to return (the
125-- R6 result layout, R9 appended the host timestamp table): the loss block
126-- at 0x180000 (16 KiB, four 4 KiB windows: the last step's per-row
127-- cross-entropy -- the target logit gathered into `correct` makes the
128-- cross-entropy rows the loss; before that they held the log-sum-exp --
129-- and its embedding gradient, then the embedding parameters and their AdamW
130-- first moment after the final step), the three next-token predictions at
131-- 0x1F0000 (three u64), and the result record at 0x200000 -- its first
132-- 147,456 bytes are written after the forward submission, before training
133-- (the marker below is put at its start first), the next 49,152 after the
134-- run.  The result file is, in the order the host writes it: the first
135-- 147,456 bytes of the record (the forward on the run's first context,
136-- before any step); every step's per-row cross-entropy (1024 binary32
137-- words, one 4 KiB window per step, appended as the step completes: the
138-- update's loss copy leaves it at 0x190000); then, after the run, the
139-- checkpoint input's stat (160 bytes, with gpGet, gpPut and the first
140-- semaphore word at 96, 100 and 104), the loss block, 49,152 bytes of the
141-- record (row 1023's logits after the final step), the semaphore region
142-- (32 bytes per submission of the plan: the push segment's final release,
143-- the fence the host waits for, and its pre-launch release, each with the
144-- device's timestamp -- a submission issued again leaves its last issue's),
145-- 64 bytes of the error notifier, the host's 160-byte status block, the
146-- timestamp table (64 bytes per submission the run issues, in issue order).
147-- A step's loss is the forward on a context no earlier step has trained
148-- on: the loss on unseen text.
149def coppeliusResultWindowBytes : Nat = 0x1000
150def coppeliusResultLossRecordOffset : Nat = 0x190000
151def coppeliusResultLossBlockOffset : Nat = 0x180000
152def coppeliusResultLossBlockBytes : Nat = 16384
153def coppeliusResultPredictionsOffset : Nat = 0x1F0000
154def coppeliusResultPredictionsBytes : Nat = 24
155-- the rows whose next token the invocation predicts, in the order their
156-- predictions (token, validity: 8 bytes each) and logits are returned
157def coppeliusPredictionRow0 : Nat = 255
158def coppeliusPredictionRow1 : Nat = 511
159def coppeliusPredictionRow2 : Nat = 1023
160def coppeliusPredictionRows : Nat = 3
161def coppeliusResultPredictionBytes : Nat = 8
162def coppeliusResultRecordOffset : Nat = 0x200000
163def coppeliusResultRecordForwardBytes : Nat = 147456
164def coppeliusResultRecordFinalBytes : Nat = 49152
165def coppeliusResultRecordMarker : Bytes = b"\xbe\xba\xfe\xca\xbe\xba\xfe\xca"
166def coppeliusResultErrorNotifierBytes : Nat = 64
167def coppeliusResultStatGPGetOffset : Nat = 96
168def coppeliusResultStatGPPutOffset : Nat = 100
169def coppeliusResultStatSemaphoreOffset : Nat = 104
170
171def coppeliusRawCheckpointSchemaIdentity : Bytes =
172  b"coppelius-raw-pmv-f32-v1"
173
174def coppeliusParameterOrderingIdentity : Bytes =
175  b"parameter-order:p-then-m-then-v"
176
177def coppeliusCheckpointSchema : (family CoppeliusCheckpointTransportSchema) =
178  (record
179    CoppeliusCheckpointTransportSchema
180    (coppeliusCheckpointSchemaIdentity = coppeliusRawCheckpointSchemaIdentity)
181    (coppeliusCheckpointParameterCount = coppeliusCheckpointParameterCount)
182    (coppeliusCheckpointPlaneCount = 3)
183    (coppeliusCheckpointBytesPerScalar = 4)
184    (coppeliusCheckpointPayloadBytes = coppeliusCheckpointPayloadBytes)
185    (coppeliusCheckpointOrderingIdentity = coppeliusParameterOrderingIdentity)
186    (coppeliusCheckpointInBandStepMetadata = 1)
187    (coppeliusCheckpointExternalStepReceiptRequired = 0))
188
189-- The checkpoint's envelope (Checkpoint.Envelope): the state's schema (the
190-- raw P/M/V layout, its order, its extent), the learner (objective,
191-- optimizer, hyperparameters), the data (the token stream: vocabulary,
192-- context, its layout), the initialization (the generator and its seed),
193-- chunks of the staging window's size, a run's steps.  A checkpoint of any
194-- other is refused by name.
195def coppeliusCheckpointChunkBytes : Nat = 4194304
196
197def coppeliusIdentityPart =
198  (lambda unrestricted name : Bytes .
199    (lambda unrestricted value : Nat .
200      (bytes-append b"|" (bytes-append name (bytes-append b"=" (naturalDecimalBytesWithin 24 value))))))
201
202def coppeliusCheckpointContract : (family CheckpointEnvelopeContract) =
203  (constructor CheckpointEnvelopeContract CheckpointEnvelopeContractValue
204    (bytes-append coppeliusRawCheckpointSchemaIdentity
205      (bytes-append b"|" (bytes-append coppeliusParameterOrderingIdentity (coppeliusIdentityPart b"payload" coppeliusCheckpointPayloadBytesNatural))))
206    (bytes-append coppeliusObjectiveIdentity
207      (bytes-append b"|halves=" (bytes-append coppeliusHalfFormatIdentity
208      (bytes-append b"|" (bytes-append coppeliusLearnerIdentity
209        (bytes-append (coppeliusIdentityPart b"lr" coppeliusLearningRateNumerator) (bytes-append (coppeliusIdentityPart b"/" coppeliusLearningRateDenominator)
210        (bytes-append (coppeliusIdentityPart b"beta1" coppeliusBeta1Numerator) (bytes-append (coppeliusIdentityPart b"/" coppeliusBeta1Denominator)
211        (bytes-append (coppeliusIdentityPart b"beta2" coppeliusBeta2Numerator) (bytes-append (coppeliusIdentityPart b"/" coppeliusBeta2Denominator)
212        (bytes-append (coppeliusIdentityPart b"decay" coppeliusWeightDecayNumerator) (bytes-append (coppeliusIdentityPart b"/" coppeliusWeightDecayDenominator)
213        (bytes-append (coppeliusIdentityPart b"epsilon" coppeliusEpsilonNumerator) (coppeliusIdentityPart b"/" coppeliusEpsilonDenominator)))))))))))))))
214    (bytes-append b"coppelius-token-stream"
215      (bytes-append (coppeliusIdentityPart b"vocabulary" (stdU32ToNatural modelCoppeliusVocab))
216        (bytes-append (coppeliusIdentityPart b"context" coppeliusInputSequence)
217          (bytes-append (coppeliusIdentityPart b"header" coppeliusTokenStreamHeaderBytes) (coppeliusIdentityPart b"token" coppeliusTokenBytes)))))
218    (bytes-append b"coppelius-initialization:random-normal" (coppeliusIdentityPart b"seed" coppeliusRunSeed))
219    coppeliusCheckpointPayloadBytesNatural
220    coppeliusCheckpointChunkBytes
221    coppeliusRunSteps)