Learning.DataParallel

1module Learning.DataParallel
2
3import Component.Pointwise
4import Component.PointwiseIdentity
5import Model.Parameter
6import Data.Bytes
7import Std.Fin
8import Std.List
9import Std.Natural
10import Std.Word
11
12-- A data-parallel learner: one parameter vector and its optimizer state
13-- (for instance momentum), both replicated on every owner; immutable samples
14-- partitioned among owners in possibly unequal shares; a per-sample
15-- contribution program; and an update program every owner applies
16-- identically to the mean contribution.
17--
18-- One update, as a mathematical definition independent of placement:
19--   numerator = sum over all samples of contribution(parameters, sample)
20--   mass      = number of samples
21--   mean      = numerator / mass
22--   (parameters, state) := update(parameters, state, mean, rate)
23-- The sum's evaluation order is not part of this definition. A realization
24-- fixes it explicitly (an ordered combine tree, Runtime.OrderedCombine) and
25-- its numerical contract names that order; a different order is a different
26-- floating-point program, never a silent reassociation.
27--
28-- Slots, by position, with k parameters, m optimizer-state words and samples
29-- of width s. The contribution program reads parameters [0, k) and one
30-- sample [k, k + s) and writes only its contribution [k + s, 2k + s). The
31-- update program reads parameters [0, k), the state [k, k + m), the mean
32-- [k + m, 2k + m) and the rate at 2k + m, and writes only parameters and
33-- state [0, k + m). The state is persistent learner state: checkpointed with
34-- the parameters and required to agree across owners. The same typed
35-- assignment semantics (Component.Pointwise) serve gradient and non-gradient
36-- rules: a contribution need not be a gradient.
37family DataParallelSystem : Type 0
38parameter erased dataParallelContributionSlots : Nat
39parameter erased dataParallelUpdateSlots : Nat
40constructor DataParallelSystemOf
41field unrestricted dataParallelParameterCount : Nat
42field unrestricted dataParallelSampleWidth : Nat
43field unrestricted dataParallelContribution : (family StdList (family PointwiseAssignment StdU64 dataParallelContributionSlots))
44field unrestricted dataParallelUpdate : (family StdList (family PointwiseAssignment StdU64 dataParallelUpdateSlots))
45field unrestricted dataParallelInitial : (family StdList StdU64)
46-- the optimizer state's initial words (m = its length; none for plain steps)
47field unrestricted dataParallelInitialState : (family StdList StdU64)
48field unrestricted dataParallelRate : StdU64
49-- per owner, per sample, the sample's s words
50field unrestricted dataParallelSamples : (family StdList (family StdList (family StdList StdU64)))
51
52end-family
53
54def dataParallelParameters =
55  (lambda erased contribution : Nat .
56    (lambda erased update : Nat .
57      (lambda unrestricted system : (family DataParallelSystem contribution update) .
58        (eliminate
59          DataParallelSystem
60          (lambda unrestricted current : (family DataParallelSystem contribution update) . Nat)
61          system
62          (branch
63            DataParallelSystemOf
64            parameters
65            width
66            program
67            rule
68            initial
69            state
70            rate
71            samples
72            .
73            parameters)))))
74
75def dataParallelWidth =
76  (lambda erased contribution : Nat .
77    (lambda erased update : Nat .
78      (lambda unrestricted system : (family DataParallelSystem contribution update) .
79        (eliminate
80          DataParallelSystem
81          (lambda unrestricted current : (family DataParallelSystem contribution update) . Nat)
82          system
83          (branch
84            DataParallelSystemOf
85            parameters
86            width
87            program
88            rule
89            initial
90            state
91            rate
92            samples
93            .
94            width)))))
95
96def dataParallelContributionProgram =
97  (lambda erased contribution : Nat .
98    (lambda erased update : Nat .
99      (lambda unrestricted system : (family DataParallelSystem contribution update) .
100        (eliminate
101          DataParallelSystem
102          (lambda unrestricted current : (family DataParallelSystem contribution update) .
103            (family StdList (family PointwiseAssignment StdU64 contribution)))
104          system
105          (branch
106            DataParallelSystemOf
107            parameters
108            width
109            program
110            rule
111            initial
112            state
113            rate
114            samples
115            .
116            program)))))
117
118def dataParallelUpdateProgram =
119  (lambda erased contribution : Nat .
120    (lambda erased update : Nat .
121      (lambda unrestricted system : (family DataParallelSystem contribution update) .
122        (eliminate
123          DataParallelSystem
124          (lambda unrestricted current : (family DataParallelSystem contribution update) .
125            (family StdList (family PointwiseAssignment StdU64 update)))
126          system
127          (branch
128            DataParallelSystemOf
129            parameters
130            width
131            program
132            rule
133            initial
134            state
135            rate
136            samples
137            .
138            rule)))))
139
140def dataParallelInitialParameters =
141  (lambda erased contribution : Nat .
142    (lambda erased update : Nat .
143      (lambda unrestricted system : (family DataParallelSystem contribution update) .
144        (eliminate
145          DataParallelSystem
146          (lambda unrestricted current : (family DataParallelSystem contribution update) .
147            (family StdList StdU64))
148          system
149          (branch
150            DataParallelSystemOf
151            parameters
152            width
153            program
154            rule
155            initial
156            state
157            rate
158            samples
159            .
160            initial)))))
161
162def dataParallelInitialStateWords =
163  (lambda erased contribution : Nat .
164    (lambda erased update : Nat .
165      (lambda unrestricted system : (family DataParallelSystem contribution update) .
166        (eliminate
167          DataParallelSystem
168          (lambda unrestricted current : (family DataParallelSystem contribution update) .
169            (family StdList StdU64))
170          system
171          (branch
172            DataParallelSystemOf
173            parameters
174            width
175            program
176            rule
177            initial
178            state
179            rate
180            samples
181            .
182            state)))))
183
184-- m, the optimizer state's word count.
185def dataParallelStateCount =
186  (lambda erased contribution : Nat .
187    (lambda erased update : Nat .
188      (lambda unrestricted system : (family DataParallelSystem contribution update) .
189        (stdListLength StdU64 (dataParallelInitialStateWords contribution update system)))))
190
191def dataParallelRateWord =
192  (lambda erased contribution : Nat .
193    (lambda erased update : Nat .
194      (lambda unrestricted system : (family DataParallelSystem contribution update) .
195        (eliminate
196          DataParallelSystem
197          (lambda unrestricted current : (family DataParallelSystem contribution update) . StdU64)
198          system
199          (branch
200            DataParallelSystemOf
201            parameters
202            width
203            program
204            rule
205            initial
206            state
207            rate
208            samples
209            .
210            rate)))))
211
212def dataParallelOwnerSamples =
213  (lambda erased contribution : Nat .
214    (lambda erased update : Nat .
215      (lambda unrestricted system : (family DataParallelSystem contribution update) .
216        (eliminate
217          DataParallelSystem
218          (lambda unrestricted current : (family DataParallelSystem contribution update) .
219            (family StdList (family StdList (family StdList StdU64))))
220          system
221          (branch
222            DataParallelSystemOf
223            parameters
224            width
225            program
226            rule
227            initial
228            state
229            rate
230            samples
231            .
232            samples)))))
233
234def dataParallelOwners =
235  (lambda erased contribution : Nat .
236    (lambda erased update : Nat .
237      (lambda unrestricted system : (family DataParallelSystem contribution update) .
238        (stdListLength
239          (family StdList (family StdList StdU64))
240          (dataParallelOwnerSamples contribution update system)))))
241
242-- The samples one owner holds, in order (none for an owner past the list).
243def dataParallelSamplesOf =
244  (lambda erased contribution : Nat .
245    (lambda erased update : Nat .
246      (lambda unrestricted system : (family DataParallelSystem contribution update) .
247        (lambda unrestricted owner : Nat .
248          (stdListHeadOr
249            (family StdList (family StdList StdU64))
250            (constructor StdList StdListEmpty (family StdList StdU64))
251            (stdListDrop
252              (family StdList (family StdList StdU64))
253              owner
254              (dataParallelOwnerSamples contribution update system)))))))
255
256def dataParallelTotalMass =
257  (lambda erased contribution : Nat .
258    (lambda erased update : Nat .
259      (lambda unrestricted system : (family DataParallelSystem contribution update) .
260        (stdListFold
261          (family StdList (family StdList StdU64))
262          Nat
263          (lambda unrestricted owned : (family StdList (family StdList StdU64)) .
264            (lambda unrestricted rest : Nat .
265              (naturalAdd (stdListLength (family StdList StdU64) owned) rest)))
266          0
267          (dataParallelOwnerSamples contribution update system)))))
268
269-- Every assignment of `program` writes a slot in [low, high).
270def dataParallelWritesWithin =
271  (lambda unrestricted slots : Nat .
272    (lambda unrestricted program : (family StdList (family PointwiseAssignment StdU64 slots)) .
273      (lambda unrestricted low : Nat .
274        (lambda unrestricted high : Nat .
275          (stdListFold
276            (family PointwiseAssignment StdU64 slots)
277            Nat
278            (lambda unrestricted assignment : (family PointwiseAssignment StdU64 slots) .
279              (lambda unrestricted rest : Nat .
280                (eliminate
281                  PointwiseAssignment
282                  (lambda unrestricted current : (family PointwiseAssignment StdU64 slots) . Nat)
283                  assignment
284                  (branch
285                    PointwiseAssign
286                    destination
287                    expression
288                    .
289                    (naturalAnd
290                      rest
291                      (naturalAnd
292                        (naturalLessOrEqual low (stdFinToNatural slots destination))
293                        (naturalLess (stdFinToNatural slots destination) high)))))))
294            1
295            program)))))
296
297-- Admitted when the slot spaces match the declared shapes (with the state's m
298-- words), each program writes only its outputs, every sample has the declared width, at least two
299-- owners exist and the total mass is positive and exactly representable as a
300-- binary64 integer (below 2^53), since realizations carry it in binary64.
301def dataParallelAdmitted =
302  (lambda unrestricted contribution : Nat .
303    (lambda unrestricted update : Nat .
304      (lambda unrestricted system : (family DataParallelSystem contribution update) .
305        (let unrestricted k =
306          (dataParallelParameters contribution update system)
307          in
308          (let unrestricted s =
309            (dataParallelWidth contribution update system)
310            in
311            (let unrestricted m =
312              (dataParallelStateCount contribution update system)
313              in
314              (naturalAnd
315                (naturalNonzero k)
316                (naturalAnd
317                  (naturalEqual contribution (naturalAdd (naturalMultiply 2 k) s))
318                  (naturalAnd
319                    (naturalEqual update (succ (naturalAdd (naturalMultiply 2 k) m)))
320                    (naturalAnd
321                      (naturalEqual
322                        (stdListLength
323                          StdU64
324                          (dataParallelInitialParameters contribution update system))
325                        k)
326                      (naturalAnd
327                        (dataParallelWritesWithin
328                          contribution
329                          (dataParallelContributionProgram contribution update system)
330                          (naturalAdd k s)
331                          contribution)
332                        (naturalAnd
333                          (dataParallelWritesWithin
334                            update
335                            (dataParallelUpdateProgram contribution update system)
336                            0
337                            (naturalAdd k m))
338                          (naturalAnd
339                            (naturalLessOrEqual 2 (dataParallelOwners contribution update system))
340                            (naturalAnd
341                              (naturalNonzero (dataParallelTotalMass contribution update system))
342                              (naturalAnd
343                                (naturalBelowTwoPower
344                                  53
345                                  (dataParallelTotalMass contribution update system))
346                                (stdListFold
347                                  (family StdList (family StdList StdU64))
348                                  Nat
349                                  (lambda unrestricted owned : (family StdList (family StdList StdU64)) .
350                                    (lambda unrestricted rest : Nat .
351                                      (naturalAnd
352                                        rest
353                                        (stdListFold
354                                        (family StdList StdU64)
355                                        Nat
356                                        (lambda unrestricted sample : (family StdList StdU64) .
357                                        (lambda unrestricted more : Nat .
358                                        (naturalAnd
359                                        more
360                                        (naturalEqual (stdListLength StdU64 sample) s))))
361                                        1
362                                        owned))))
363                                  1
364                                  (dataParallelOwnerSamples contribution update system)))))))))))))))))
365
366def dataParallelWordsIdentity =
367  (lambda unrestricted words : (family StdList StdU64) .
368    (bytes-append
369      (pointwiseIdentityWord (stdListLength StdU64 words))
370      (stdListFold
371        StdU64
372        Bytes
373        (lambda unrestricted word : StdU64 .
374          (lambda unrestricted tail : Bytes . (bytes-append (dataBytesWord64LE word) tail)))
375        b""
376        words)))
377
378-- The learner's rule: both programs, initial parameters and state, rate.
379def dataParallelRuleIdentity =
380  (lambda unrestricted contribution : Nat .
381    (lambda unrestricted update : Nat .
382      (lambda unrestricted system : (family DataParallelSystem contribution update) .
383        (bytes-append
384          (pointwiseIdentityWord (dataParallelParameters contribution update system))
385          (bytes-append
386            (pointwiseIdentityWord (dataParallelWidth contribution update system))
387            (bytes-append
388              (pointwiseProgramIdentity
389                contribution
390                (dataParallelContributionProgram contribution update system))
391              (bytes-append
392                (pointwiseProgramIdentity
393                  update
394                  (dataParallelUpdateProgram contribution update system))
395                (bytes-append
396                  (dataParallelWordsIdentity
397                    (dataParallelInitialParameters contribution update system))
398                  (bytes-append
399                    (dataParallelWordsIdentity
400                      (dataParallelInitialStateWords contribution update system))
401                    (dataBytesWord64LE (dataParallelRateWord contribution update system)))))))))))
402
403def dataParallelSamplesIdentity =
404  (lambda unrestricted samples : (family StdList (family StdList StdU64)) .
405    (stdListFold
406      (family StdList StdU64)
407      Bytes
408      (lambda unrestricted sample : (family StdList StdU64) .
409        (lambda unrestricted more : Bytes . (bytes-append (dataParallelWordsIdentity sample) more)))
410      b""
411      samples))
412
413-- The learner's identity: its rule and every owner's samples in order.
414-- Evaluation order is not part of it; a realization adds its numerical
415-- contract beside it.
416def dataParallelIdentity =
417  (lambda unrestricted contribution : Nat .
418    (lambda unrestricted update : Nat .
419      (lambda unrestricted system : (family DataParallelSystem contribution update) .
420        (bytes-append
421          b"alpha-data-parallel-v1;"
422          (bytes-append
423            (dataParallelRuleIdentity contribution update system)
424            (stdListFold
425              (family StdList (family StdList StdU64))
426              Bytes
427              (lambda unrestricted owned : (family StdList (family StdList StdU64)) .
428                (lambda unrestricted rest : Bytes .
429                  (bytes-append
430                    (pointwiseIdentityWord (stdListLength (family StdList StdU64) owned))
431                    (bytes-append (dataParallelSamplesIdentity owned) rest))))
432              b""
433              (dataParallelOwnerSamples contribution update system)))))))
434
435-- Every sample in owner order: the learner's input sequence without its
436-- placement among owners.
437def dataParallelAllSamples =
438  (lambda erased contribution : Nat .
439    (lambda erased update : Nat .
440      (lambda unrestricted system : (family DataParallelSystem contribution update) .
441        (stdListFold
442          (family StdList (family StdList StdU64))
443          (family StdList (family StdList StdU64))
444          (lambda unrestricted owned : (family StdList (family StdList StdU64)) .
445            (lambda unrestricted rest : (family StdList (family StdList StdU64)) .
446              (stdListAppend (family StdList StdU64) owned rest)))
447          (constructor StdList StdListEmpty (family StdList StdU64))
448          (dataParallelOwnerSamples contribution update system)))))
449
450-- The placement-independent identity: the rule and the input sequence, not
451-- how the samples are split among owners. State saved under one split is a
452-- state of the same learner under any other split of the same sequence.
453def dataParallelLearnerIdentity =
454  (lambda unrestricted contribution : Nat .
455    (lambda unrestricted update : Nat .
456      (lambda unrestricted system : (family DataParallelSystem contribution update) .
457        (let unrestricted samples =
458          (dataParallelAllSamples contribution update system)
459          in
460          (bytes-append
461            b"alpha-data-parallel-learner-v1;"
462            (bytes-append
463              (dataParallelRuleIdentity contribution update system)
464              (bytes-append
465                (pointwiseIdentityWord (stdListLength (family StdList StdU64) samples))
466                (dataParallelSamplesIdentity samples))))))))
467
468-- The same learner with its input sequence split into `owners` contiguous
469-- shares whose sizes differ by at most one, larger first.
470def dataParallelRepartitioned =
471  (lambda erased contribution : Nat .
472    (lambda erased update : Nat .
473      (lambda unrestricted system : (family DataParallelSystem contribution update) .
474        (lambda unrestricted owners : Nat .
475          (let unrestricted samples =
476            (dataParallelAllSamples contribution update system)
477            in
478            (let unrestricted total =
479              (stdListLength (family StdList StdU64) samples)
480              in
481              (let unrestricted base =
482                (naturalDivideUnchecked total owners)
483                in
484                (let unrestricted extra =
485                  (naturalModuloUnchecked total owners)
486                  in
487                  (eliminate
488                    DataParallelSystem
489                    (lambda unrestricted current : (family DataParallelSystem contribution update) .
490                      (family DataParallelSystem contribution update))
491                    system
492                    (branch
493                      DataParallelSystemOf
494                      parameters
495                      width
496                      program
497                      rule
498                      initial
499                      state
500                      rate
501                      held
502                      .
503                      (constructor
504                        DataParallelSystem
505                        DataParallelSystemOf
506                        contribution
507                        update
508                        parameters
509                        width
510                        program
511                        rule
512                        initial
513                        state
514                        rate
515                        (stdListMap
516                          Nat
517                          (family StdList (family StdList StdU64))
518                          (lambda unrestricted share : Nat .
519                            (stdListTake
520                              (family StdList StdU64)
521                              (naturalAdd base (naturalLess share extra))
522                              (stdListDrop
523                                (family StdList StdU64)
524                                (naturalAdd
525                                  (naturalMultiply share base)
526                                  (naturalMinimum share extra))
527                                samples)))
528                          (stdListNaturalsBelow owners)))))))))))))
529
530-- ---- coordinate separability ----
531-- An update is separable when each coordinate's results depend only on that
532-- coordinate and on shared values, so different owners can each update the
533-- coordinates they own (optimizer-state partitioning). With q per-coordinate
534-- state words (word j of coordinate i at state index j k + i) and the other
535-- m - q k words shared, an update slot belongs to coordinate i (parameter i,
536-- a state word of coordinate i, mean i) or is shared (a shared state word,
537-- the rate). `dataParallelSlotClass` is succ i for coordinate i and 0 for
538-- shared.
539def dataParallelSlotClass =
540  (lambda unrestricted k : Nat .
541    (lambda unrestricted m : Nat .
542      (lambda unrestricted q : Nat .
543        (lambda unrestricted slot : Nat .
544          (naturalSelect
545            (naturalLess slot k)
546            (succ slot)
547            (naturalSelect
548              (naturalLess slot (naturalAdd k (naturalMultiply q k)))
549              (succ (naturalModuloUnchecked (naturalSaturatingSubtract slot k) k))
550              (naturalSelect
551                (naturalLess slot (naturalAdd k m))
552                0
553                (naturalSelect
554                  (naturalLess slot (naturalAdd (naturalMultiply 2 k) m))
555                  (succ (naturalSaturatingSubtract slot (naturalAdd k m)))
556                  0))))))))
557
558-- Every input `expression` reads is shared or of class `allowed`.
559def dataParallelReadsClass =
560  (lambda unrestricted slots : Nat .
561    (lambda unrestricted k : Nat .
562      (lambda unrestricted m : Nat .
563        (lambda unrestricted q : Nat .
564          (lambda unrestricted allowed : Nat .
565            (lambda unrestricted expression : (family PointwiseExpression StdU64 slots) .
566              (pointwiseFold
567                StdU64
568                Nat
569                slots
570                (lambda unrestricted operation : (family PointwiseUnary) .
571                  (lambda unrestricted child : Nat . child))
572                (lambda unrestricted operation : (family PointwiseBinary) .
573                  (lambda unrestricted left : Nat .
574                    (lambda unrestricted right : Nat . (naturalAnd left right))))
575                (lambda unrestricted word : StdU64 . 1)
576                expression
577                (lambda unrestricted index : (family StdFin slots) .
578                  (let unrestricted class =
579                    (dataParallelSlotClass k m q (stdFinToNatural slots index))
580                    in
581                    (naturalOr (naturalIsZero class) (naturalEqual class allowed)))))))))))
582
583-- The update is separable under q per-coordinate state words: a coordinate's
584-- assignment reads only its coordinate and shared values, and a shared
585-- assignment reads only shared values.
586def dataParallelSeparable =
587  (lambda unrestricted contribution : Nat .
588    (lambda unrestricted update : Nat .
589      (lambda unrestricted system : (family DataParallelSystem contribution update) .
590        (lambda unrestricted q : Nat .
591          (let unrestricted k =
592            (dataParallelParameters contribution update system)
593            in
594            (let unrestricted m =
595              (dataParallelStateCount contribution update system)
596              in
597              (naturalAnd
598                (naturalLessOrEqual (naturalMultiply q k) m)
599                (stdListFold
600                  (family PointwiseAssignment StdU64 update)
601                  Nat
602                  (lambda unrestricted assignment : (family PointwiseAssignment StdU64 update) .
603                    (lambda unrestricted rest : Nat .
604                      (eliminate
605                        PointwiseAssignment
606                        (lambda unrestricted current : (family PointwiseAssignment StdU64 update) .
607                          Nat)
608                        assignment
609                        (branch
610                          PointwiseAssign
611                          destination
612                          expression
613                          .
614                          (naturalAnd
615                            rest
616                            (dataParallelReadsClass
617                              update
618                              k
619                              m
620                              q
621                              (dataParallelSlotClass k m q (stdFinToNatural update destination))
622                              expression))))))
623                  1
624                  (dataParallelUpdateProgram contribution update system)))))))))
625
626-- The assignments one owner runs: every shared assignment and those of the
627-- coordinates `owns` names (owns i = 1), in program order.
628def dataParallelOwnedProgram =
629  (lambda unrestricted contribution : Nat .
630    (lambda unrestricted update : Nat .
631      (lambda unrestricted system : (family DataParallelSystem contribution update) .
632        (lambda unrestricted q : Nat .
633          (lambda unrestricted owns : (pi unrestricted coordinate : Nat . Nat) .
634            (let unrestricted k =
635              (dataParallelParameters contribution update system)
636              in
637              (let unrestricted m =
638                (dataParallelStateCount contribution update system)
639                in
640                (stdListFold
641                  (family PointwiseAssignment StdU64 update)
642                  (family StdList (family PointwiseAssignment StdU64 update))
643                  (lambda unrestricted assignment : (family PointwiseAssignment StdU64 update) .
644                    (lambda unrestricted rest : (family StdList (family PointwiseAssignment StdU64 update)) .
645                      (eliminate
646                        PointwiseAssignment
647                        (lambda unrestricted current : (family PointwiseAssignment StdU64 update) .
648                          (family StdList (family PointwiseAssignment StdU64 update)))
649                        assignment
650                        (branch
651                          PointwiseAssign
652                          destination
653                          expression
654                          .
655                          (let unrestricted class =
656                            (dataParallelSlotClass k m q (stdFinToNatural update destination))
657                            in
658                            (nat-eliminate
659                              (lambda unrestricted current : Nat .
660                                (family StdList (family PointwiseAssignment StdU64 update)))
661                              rest
662                              (lambda unrestricted previous : Nat .
663                                (lambda unrestricted ignored : (family StdList (family PointwiseAssignment StdU64 update)) .
664                                  (constructor
665                                    StdList
666                                    StdListCons
667                                    (family PointwiseAssignment StdU64 update)
668                                    assignment
669                                    rest)))
670                              (naturalOr (naturalIsZero class) (owns (naturalPredecessor class)))))))))
671                  (constructor StdList StdListEmpty (family PointwiseAssignment StdU64 update))
672                  (dataParallelUpdateProgram contribution update system)))))))))