Bob.TrainingRun

1module Bob.TrainingRun
2
3import Bob.Model
4import Bob.Learner
5import Checkpoint.Envelope
6import Data.Float32Bits
7import Model.Config
8import Model.Parameter
9import Model.Word64
10import Std.Natural
11import Std.Physical
12import Std.Word
13
14-- The physical transport is raw contiguous FP32 P/M/V.  The optimizer step
15-- remains an external receipt, matching the currently qualified native path.
16family BobCheckpointTransportSchema : Type 0
17constructor BobCheckpointTransportSchemaValue
18field unrestricted bobCheckpointSchemaIdentity : Bytes
19field unrestricted bobCheckpointParameterCount : Nat
20field unrestricted bobCheckpointPlaneCount : Nat
21field unrestricted bobCheckpointBytesPerScalar : Nat
22field unrestricted bobCheckpointPayloadBytes : Nat
23field unrestricted bobCheckpointOrderingIdentity : Bytes
24field unrestricted bobCheckpointInBandStepMetadata : Nat
25field unrestricted bobCheckpointExternalStepReceiptRequired : Nat
26
27end-family
28
29family BobTrainingRunSemantics : Type 0
30constructor BobTrainingRunSemanticsValue
31field unrestricted bobRunUpdatesPerInvocation : Nat
32field unrestricted bobRunDeterministicInitializationIdentity : Bytes
33field unrestricted bobRunFreshCheckpointInputBytes : ByteCount
34field unrestricted bobRunResumeCheckpointRequired : Nat
35field unrestricted bobRunFreshProcessResumeRequired : Nat
36field unrestricted bobRunNativeOnlyRequired : Nat
37field unrestricted bobRunAllowedHostFallbacks : Nat
38
39end-family
40
41def bobCheckpointPlaneCount : Nat =
42  3
43
44def bobCheckpointBytesPerScalar : Nat =
45  4
46
47-- Derived from the model's table extent, optimizer planes and representation
48-- width: 256 * 256 * 3 * 4 = 786432 bytes.
49def bobCheckpointPayloadByteCount : Nat =
50  (nat-multiply
51    bobParameterCount
52    (nat-multiply bobCheckpointPlaneCount bobCheckpointBytesPerScalar))
53
54def bobCheckpointPayloadBytes : Nat =
55  bobCheckpointPayloadByteCount
56
57def bobCheckpointSchema : (family BobCheckpointTransportSchema) =
58  (record
59    BobCheckpointTransportSchema
60    (bobCheckpointSchemaIdentity = b"bob-raw-pmv-f32-v1")
61    (bobCheckpointParameterCount = bobParameterCount)
62    (bobCheckpointPlaneCount = bobCheckpointPlaneCount)
63    (bobCheckpointBytesPerScalar = bobCheckpointBytesPerScalar)
64    (bobCheckpointPayloadBytes = bobCheckpointPayloadBytes)
65    (bobCheckpointOrderingIdentity = b"parameter-order:p-then-m-then-v")
66    (bobCheckpointInBandStepMetadata = 1)
67    (bobCheckpointExternalStepReceiptRequired = 0))
68
69def bobUpdatesPerInvocation : Nat = 10
70
71def bobTrainingRunSemantics : (family BobTrainingRunSemantics) =
72  (record
73    BobTrainingRunSemantics
74    (bobRunUpdatesPerInvocation = bobUpdatesPerInvocation)
75    (bobRunDeterministicInitializationIdentity = b"bob-seed-v1")
76    (bobRunFreshCheckpointInputBytes = (stdByteCount 0))
77    (bobRunResumeCheckpointRequired = 1)
78    (bobRunFreshProcessResumeRequired = 1)
79    (bobRunNativeOnlyRequired = 1)
80    (bobRunAllowedHostFallbacks = 0))
81
82-- The input record one invocation reads (systems/bob/tools/build_fixture.py
83-- writes it): the magic and version (8 bytes each), the corpus's 256
84-- (byte, target) pairs, each byte's target as a 32-bit word, and each
85-- prompt's last byte as a 32-bit word.
86def bobInputHeaderBytes : Nat = 16
87def bobInputWordBytes : Nat = 4
88def bobPromptCount : Nat = 3
89def bobInputTargetsOffset : Nat =
90  (naturalAdd bobInputHeaderBytes (naturalMultiply 2 bobVocabularySize))
91def bobInputPromptsOffset : Nat =
92  (naturalAdd bobInputTargetsOffset (naturalMultiply bobInputWordBytes bobVocabularySize))
93def bobInputRecordBytes : Nat =
94  (naturalAdd bobInputPromptsOffset (naturalMultiply bobInputWordBytes bobPromptCount))
95
96-- What an invocation returns per prompt: its token and its validity, a
97-- 32-bit word each
98def bobPredictionBytes : Nat = 8
99
100-- The checkpoint's envelope (Checkpoint.Envelope): the state's schema, the
101-- learner (objective, optimizer, its binary32 hyperparameters' bytes), the
102-- data (the input record: the model's input, the vocabulary, the prompts,
103-- the extent), the initialization's seed, one chunk per plane, the updates
104-- an invocation makes.  A checkpoint of any other is refused by name.
105def bobIdentityPart =
106  (lambda unrestricted name : Bytes .
107    (lambda unrestricted value : Nat .
108      (bytes-append b"|" (bytes-append name (bytes-append b"=" (naturalDecimalBytesWithin 24 value))))))
109
110def bobCheckpointContract : (family CheckpointEnvelopeContract) =
111  (constructor CheckpointEnvelopeContract CheckpointEnvelopeContractValue
112    (bytes-append b"bob-raw-pmv-f32-v1|parameter-order:p-then-m-then-v" (bobIdentityPart b"payload" bobCheckpointPayloadBytes))
113    (bytes-append bobObjectiveIdentity
114      (bytes-append b"|" (bytes-append bobLearnerIdentity
115        (bytes-append b"|" (bytes-append (dataFloat32EncodeLE bobLearningRate) (bytes-append (dataFloat32EncodeLE bobBeta1)
116          (bytes-append (dataFloat32EncodeLE bobBeta2) (bytes-append (dataFloat32EncodeLE bobEpsilon) (dataFloat32EncodeLE bobWeightDecay)))))))))
117    (bytes-append b"raw-byte-sequence"
118      (bytes-append (bobIdentityPart b"vocabulary" bobVocabularySize)
119        (bytes-append (bobIdentityPart b"prompts" bobPromptCount) (bobIdentityPart b"record" bobInputRecordBytes))))
120    b"bob-seed-v1"
121    bobCheckpointPayloadBytes
122    (naturalMultiply bobCheckpointBytesPerScalar bobParameterCount)
123    bobUpdatesPerInvocation)