1	# A PyTorch checkpoint: the ZIP archive torch.save writes, whose data.pkl
2	# holds the saved object and whose data/<key> entries hold the storages. The
3	# saved module tree, or a state dict, gives one tensor per state-dict entry
4	# under its PyTorch name, read from its storage in the stored format.
5	import Base
6	import ./storage_bytes.bend as Bytes
7	import ./checkpoint_zip.bend as Archive
8	import ./checkpoint_pickle.bend as Pickle
9	import ./storage_element_format.bend as Formats
10	import ./tensor.bend as Tensors
11	import ./tensor_file.bend as Files
12	import ./tensor_model.bend as Models
13	import ./tensor_view.bend as Views
14	import ./tensor_shape.bend as Shapes
15	
16	def Tensor() -> Type:
17	  Tensors.Tensor<F32>
18	
19	def operation() -> String:
20	  "load_checkpoint"
21	
22	# What the pickle says about one state-dict entry: the storage class, the key
23	# of its entry, the element offset in the storage, the shape and the axes of
24	# the view of the storage, one extent and stride each.
25	type Described is Data:
26	  Described{name: String,kind: String,key: String,offset: U32,shape: List<&2,U32>,axes: List<&2,Views.Axis>}
27	  Unusable{name: String,detail: String}
28	
29	# A module still to walk, with the prefix of its tensors' names.
30	type Task is Data:
31	  Task{prefix: String,module: Pickle.Value}
32	
33	type Loading is Type:
34	  Loading{file: File,store: Models.Store,remaining: List<&2,Described>}
35	
36	# Tensors as torch._utils rebuilds them: a parameter wraps a tensor, and a tensor
37	# is a storage reference, an offset, a size and a stride.
38	
39	def numbers(items: List<&2,Pickle.Value>) -> List<&2,U32>:
40	  match items:
41	    case Nil{}: []
42	    case Con{Pickle.Integer{word},rest}: word <> numbers(rest)
43	    case Con{other,rest}: 0 <> numbers(rest)
44	
45	def all_numbers(items: List<&2,Pickle.Value>) -> Bool:
46	  match items:
47	    case Nil{}: True{}
48	    case Con{Pickle.Integer{word},rest}: all_numbers(rest)
49	    case Con{other,rest}: False{}
50	
51	# The axes of the tensor as the library views them.
52	def prepended(+extent: U32,+stride: U32,axes: Maybe<&2,List<&2,Views.Axis>>) -> Maybe<&2,List<&2,Views.Axis>>:
53	  match axes:
54	    case None{}: None{}
55	    case Some{axes}: Some{Views.Axis{extent,stride} <> axes}
56	
57	def axes_of(extents: List<&2,U32>,strides: List<&2,U32>) -> Maybe<&2,List<&2,Views.Axis>>:
58	  match extents strides:
59	    case Nil{} Nil{}: Some{[]}
60	    case Nil{} Con{stride,more}: None{}
61	    case Con{extent,rest} Nil{}: None{}
62	    case Con{extent,rest} Con{stride,more}: prepended(extent,stride,axes_of(rest,more))
63	
64	def described_layout(axes: Maybe<&2,List<&2,Views.Axis>>,+name: String,+kind: String,+key: String,+offset: U32,+shape: List<&2,U32>) -> Described:
65	  match axes:
66	    case Some{axes}: Described{name,kind,key,offset,shape,axes}
67	    case None{}: Unusable{name,"has a different number of sizes and strides"}
68	
69	def described_numbers(plain: Bool,+name: String,+kind: String,+key: String,+offset: U32,+sizes: List<&2,Pickle.Value>,+strides: List<&2,Pickle.Value>) -> Described:
70	  match plain:
71	    case False{}: Unusable{name,"has a size or stride that is not an integer"}
72	    case True{}: described_layout(axes_of(numbers(sizes),numbers(strides)),name,kind,key,offset,numbers(sizes))
73	
74	# The persistent id of a storage: ('storage', torch.<Kind>Storage, key, location, count).
75	def described_storage(+name: String,id: Pickle.Value,+offset: U32,+sizes: List<&2,Pickle.Value>,+strides: List<&2,Pickle.Value>) -> Described:
76	  match id:
77	    case Pickle.Sequence{Con{Pickle.Text{tag},Con{Pickle.Global{module,kind},Con{Pickle.Text{key},others}}}}:
78	      described_numbers(all_numbers(sizes) && all_numbers(strides),name,kind,key,offset,sizes,strides)
79	    case id: Unusable{name,"has a storage reference of an unknown form"}
80	
81	def described_arguments(+name: String,arguments: Pickle.Value) -> Described:
82	  match arguments:
83	    case Pickle.Sequence{Con{Pickle.Persistent{id},Con{Pickle.Integer{offset},Con{Pickle.Sequence{sizes},Con{Pickle.Sequence{strides},others}}}}}:
84	      described_storage(name,id,offset,sizes,strides)
85	    case arguments: Unusable{name,"is rebuilt with unexpected arguments"}
86	
87	def described_call(rebuilt: Bool,+name: String,arguments: Pickle.Value) -> Described:
88	  match rebuilt:
89	    case True{}: described_arguments(name,arguments)
90	    case False{}: Unusable{name,"is not rebuilt by torch._utils._rebuild_tensor_v2"}
91	
92	def tensor_of(+name: String,value: Pickle.Value) -> Described:
93	  match value:
94	    case Pickle.Reduced{Pickle.Global{module,+function},arguments}:
95	      described_call(String.eq(function,"_rebuild_tensor_v2"),name,arguments)
96	    case value: Unusable{name,"is not a tensor"}
97	
98	def unwrapped(parameter: Bool,+name: String,inner: Pickle.Value,+whole: Pickle.Value) -> Described:
99	  match parameter:
100	    case True{}: tensor_of(name,inner)
101	    case False{}: tensor_of(name,whole)
102	
103	# A parameter is _rebuild_parameter(tensor, requires_grad, hooks).
104	def described(+name: String,+value: Pickle.Value) -> Described:
105	  match value:
106	    case Pickle.Reduced{Pickle.Global{module,+function},Pickle.Sequence{Con{+inner,others}}}:
107	      unwrapped(String.eq(function,"_rebuild_parameter"),name,inner,Pickle.Reduced{Pickle.Global{module,function},Pickle.Sequence{inner <> others}})
108	    case value: tensor_of(name,value)
109	
110	# The module tree. A module's state holds _parameters and _buffers, names to
111	# tensors or None, _non_persistent_buffers_set, and _modules, names to child
112	# modules. The state dict lists parameters, then persistent buffers, then each
113	# child in order, so a walk takes a module's own tensors and then its children.
114	
115	def listed(set: Pickle.Value) -> List<&2,Pickle.Value>:
116	  match set:
117	    case Pickle.Reduced{callable,Pickle.Sequence{Con{Pickle.Items{names},others}}}: names
118	    case Pickle.Items{names}: names
119	    case set: []
120	
121	def named(names: List<&2,Pickle.Value>,+name: String) -> Bool:
122	  match names:
123	    case Nil{}: False{}
124	    case Con{Pickle.Text{+text},rest}: String.eq(text,name) || named(rest,name)
125	    case Con{other,rest}: named(rest,name)
126	
127	def entries_of(value: Pickle.Value) -> List<&2,Pickle.Value>:
128	  match value:
129	    case Pickle.Dict{entries}: entries
130	    case value: []
131	
132	# The tensors of one dictionary, newest first onto the found list; None and the
133	# excluded names are skipped.
134	def own(entries: List<&2,Pickle.Value>,+prefix: String,+excluded: List<&2,Pickle.Value>,+found: List<&2,Described>) -> List<&2,Described>:
135	  match entries:
136	    case Nil{}: found
137	    case Con{Pickle.Sequence{Con{Pickle.Text{+name},Con{Pickle.Absent{},others}}},rest}: own(rest,prefix,excluded,found)
138	    case Con{Pickle.Sequence{Con{Pickle.Text{+name},Con{value,others}}},rest}:
139	      own(rest,prefix,excluded,Bool.pick(List<&2,Described>,named(excluded,name),found,described(prefix ++ name,value) <> found))
140	    case Con{entry,rest}: own(rest,prefix,excluded,found)
141	
142	def children(entries: List<&2,Pickle.Value>,+prefix: String) -> List<&2,Task>:
143	  match entries:
144	    case Nil{}: []
145	    case Con{Pickle.Sequence{Con{Pickle.Text{name},Con{Pickle.Absent{},others}}},rest}: children(rest,prefix)
146	    case Con{Pickle.Sequence{Con{Pickle.Text{name},Con{module,others}}},rest}: Task{prefix ++ name ++ ".",module} <> children(rest,prefix)
147	    case Con{entry,rest}: children(rest,prefix)
148	
149	def state_of(module: Pickle.Value) -> List<&2,Pickle.Value>:
150	  match module:
151	    case Pickle.Object{class,arguments,state}: entries_of(state)
152	    case module: []
153	
154	# Every module is one step, so the byte count of the pickle bounds the walk.
155	def walk(fuel: Nat,todo: List<&2,Task>,found: List<&2,Described>) -> List<&2,Described>:
156	  match fuel todo:
157	    case fuel Nil{}: List.reverse(&2,Described,found)
158	    case 0n todo: List.reverse(&2,Described,found)
159	    case 1n+rest Con{Task{+prefix,module},more}:
160	      +state = state_of(module)
161	      walk(rest,List.append(&2,Task,children(entries_of(Pickle.get(state,"_modules")),prefix),more),
162	        own(entries_of(Pickle.get(state,"_buffers")),prefix,listed(Pickle.get(state,"_non_persistent_buffers_set")),
163	          own(entries_of(Pickle.get(state,"_parameters")),prefix,[],found)))
164	
165	# A state dict: names to tensors, in order.
166	def plain(entries: List<&2,Pickle.Value>,found: List<&2,Described>) -> List<&2,Described>:
167	  match entries:
168	    case Nil{}: List.reverse(&2,Described,found)
169	    case Con{Pickle.Sequence{Con{Pickle.Text{+name},Con{value,others}}},rest}: plain(rest,described(name,value) <> found)
170	    case Con{entry,rest}: plain(rest,found)
171	
172	def is_module(value: Pickle.Value) -> Bool:
173	  match value:
174	    case Pickle.Object{class,arguments,Pickle.Dict{entries}}: True{}
175	    case value: False{}
176	
177	def rooted(module: Bool,+fuel: Nat,+value: Pickle.Value) -> List<&2,Described>:
178	  match module:
179	    case True{}: walk(fuel,[Task{"",value}],[])
180	    case False{}: plain(entries_of(value),[])
181	
182	# The saved object: a module, a checkpoint dictionary whose "model" is one, or
183	# a state dict.
184	def tensors_of(+fuel: Nat,+value: Pickle.Value) -> List<&2,Described>:
185	  match value:
186	    case Pickle.Dict{+entries}: rooted(is_module(Pickle.get(entries,"model")),fuel,Bool.pick(Pickle.Value,is_module(Pickle.get(entries,"model")),Pickle.get(entries,"model"),Pickle.Dict{entries}))
187	    case value: rooted(is_module(value),fuel,value)
188	
189	# Reading the storages
190	
191	def format_classified(half: Bool,bfloat: Bool,float: Bool,double: Bool,long: Bool,int: Bool) -> Maybe<&2,Formats.Format>:
192	  match half bfloat float double long int:
193	    case True{} bfloat float double long int: Some{Formats.Fp16{}}
194	    case False{} True{} float double long int: Some{Formats.BFloat16{}}
195	    case False{} False{} True{} double long int: Some{Formats.Fp32{}}
196	    case False{} False{} False{} True{} long int: Some{Formats.Fp64{}}
197	    case False{} False{} False{} False{} True{} int: Some{Formats.Int64{}}
198	    case False{} False{} False{} False{} False{} True{}: Some{Formats.Int32{}}
199	    case False{} False{} False{} False{} False{} False{}: None{}
200	
201	def format_of(+kind: String) -> Maybe<&2,Formats.Format>:
202	  format_classified(String.eq(kind,"HalfStorage"),String.eq(kind,"BFloat16Storage"),String.eq(kind,"FloatStorage"),String.eq(kind,"DoubleStorage"),
203	    String.eq(kind,"LongStorage"),String.eq(kind,"IntStorage"))
204	
205	def fault(+path: String,+name: String,+detail: String,shape: List<&2,U32>) -> Tensor():
206	  Tensors.failed(~F32,operation(),Tensors.FileFault{path,"holds " ++ name ++ ", which " ++ detail},shape,0.0,1)
207	
208	def storage_name(+prefix: List<&2,U32>,+key: String) -> List<&2,U32>:
209	  List.append(&2,U32,prefix,List.append(&2,U32,Bytes.codes("data/"),Bytes.codes(key)))
210	
211	def bounded(+count: U32,+width: U32,+offset: U32,+length: U32) -> Bool:
212	  +first = (offset * width : U32)
213	  +bytes = (count * width : U32)
214	  U32.is_eq((first / width : U32),offset) && U32.is_eq((bytes / width : U32),count) && U32.is_le(first,length) && U32.is_le(bytes,(length - first : U32))
215	
216	# The elements a view reaches from its offset: one past its last address, or
217	# none when an extent is zero. None when that does not fit a word.
218	def joined(product: Maybe<&2,U32>,rest: Maybe<&2,U32>) -> Maybe<&2,U32>:
219	  match product:
220	    case Some{product}: Shapes.checked_add_to(rest,product)
221	    case None{}: None{}
222	
223	def last_address(axes: List<&2,Views.Axis>) -> Maybe<&2,U32>:
224	  match axes:
225	    case []: Some{0}
226	    case Views.Axis{extent,stride} <> rest: joined(Shapes.checked_multiply((extent - 1 : U32),stride),last_address(rest))
227	
228	def reach(empty: Bool,axes: List<&2,Views.Axis>) -> Maybe<&2,U32>:
229	  match empty:
230	    case True{}: Some{0}
231	    case False{}: Shapes.checked_add_to(last_address(axes),1)
232	
233	# A row-major view is read as the tensor itself. Any other view is read as the
234	# elements it reaches and copied through the view into row-major storage.
235	def strided(tensor: Tensor(),+axes: List<&2,Views.Axis>,+shape: List<&2,U32>) -> Tensor():
236	  Tensors.with_whole(~F32,~Tensor(),tensor,storage => reached => fill => status => workers =>
237	    Tensors.result_of(~F32,Tensors.gather_status(~F32,operation(),status,Some{Views.View{axes,0}},[],shape,shape,storage,reached,fill,workers)))
238	
239	def rearranged(pair: File & Tensor(),+axes: List<&2,Views.Axis>,+shape: List<&2,U32>) -> File & Tensor():
240	  (file,tensor) = pair
241	  (file,strided(tensor,axes,shape))
242	
243	def arranged(row_major: Bool,file: File,+path: String,+shape: List<&2,U32>,+axes: List<&2,Views.Axis>,+span: U32,+start: U32,+format: Formats.Format) -> IO(File & Tensor()):
244	  match row_major:
245	    case True{}: Files.load_range(file,operation(),path,shape,start,format)
246	    case False{}:
247	      do IO<File & Tensor()>:
248	        pair : File & Tensor() <- Files.load_range(file,operation(),path,[span],start,format)
249	        IO.pure(File & Tensor(),rearranged(pair,axes,shape))
250	
251	# The elements the view reaches lie within the entry.
252	def within(fits: Bool,file: File,+path: String,+name: String,+shape: List<&2,U32>,+axes: List<&2,Views.Axis>,+span: U32,+start: U32,+format: Formats.Format) -> IO(File & Tensor()):
253	  match fits:
254	    case False{}: IO.pure(File & Tensor(),(file,fault(path,name,"reaches beyond its storage",shape)))
255	    case True{}: arranged(Maybe.is_some(&2,U32,Views.contiguous_size(axes)),file,path,shape,axes,span,start,format)
256	
257	def spanned(span: Maybe<&2,U32>,file: File,+path: String,+name: String,+shape: List<&2,U32>,+axes: List<&2,Views.Axis>,+data: U32,+offset: U32,+length: U32,
258	  +format: Formats.Format) -> IO(File & Tensor()):
259	  match span:
260	    case None{}: IO.pure(File & Tensor(),(file,fault(path,name,"reaches beyond any storage",shape)))
261	    case Some{+span}:
262	      within(bounded(span,Formats.width(format),offset,length),file,path,name,shape,axes,span,(data + offset * Formats.width(format) : U32),format)
263	
264	def located(data: Maybe<&2,U32>,file: File,+path: String,+name: String,+shape: List<&2,U32>,+axes: List<&2,Views.Axis>,+offset: U32,+length: U32,+format: Formats.Format) -> IO(File & Tensor()):
265	  match data:
266	    case None{}: IO.pure(File & Tensor(),(file,fault(path,name,"has a storage entry that cannot be read",shape)))
267	    case Some{+data}: spanned(reach(Shapes.has_zero(shape),axes),file,path,name,shape,axes,data,offset,length,format)
268	
269	def stored(pair: File & Maybe<&2,U32>,+path: String,+name: String,+shape: List<&2,U32>,+axes: List<&2,Views.Axis>,+offset: U32,+length: U32,+format: Formats.Format) -> IO(File & Tensor()):
270	  (file,data) = pair
271	  located(data,file,path,name,shape,axes,offset,length,format)
272	
273	def recorded(record: Maybe<&2,Archive.Record>,file: File,+path: String,+name: String,+shape: List<&2,U32>,+axes: List<&2,Views.Axis>,+offset: U32,+format: Formats.Format) -> IO(File & Tensor()):
274	  match record:
275	    case None{}: IO.pure(File & Tensor(),(file,fault(path,name,"has no storage entry",shape)))
276	    case Some{Archive.Record{entry,header,+length,method}}:
277	      do IO<File & Tensor()>:
278	        pair : File & Maybe<&2,U32> <- Archive.data_offset(file,Archive.Record{entry,header,length,method})
279	        stored(pair,path,name,shape,axes,offset,length,format)
280	
281	def formatted(format: Maybe<&2,Formats.Format>,file: File,+path: String,+name: String,+kind: String,+key: String,+offset: U32,+shape: List<&2,U32>,+axes: List<&2,Views.Axis>,
282	  +records: List<&2,Archive.Record>,+prefix: List<&2,U32>) -> IO(File & Tensor()):
283	  match format:
284	    case None{}: IO.pure(File & Tensor(),(file,fault(path,name,"is stored as " ++ kind ++ ", a format that is not supported",shape)))
285	    case Some{format}: recorded(Archive.find(records,storage_name(prefix,key)),file,path,name,shape,axes,offset,format)
286	
287	def load_one(described: Described,file: File,+path: String,+records: List<&2,Archive.Record>,+prefix: List<&2,U32>) -> IO(File & Tensor()):
288	  match described:
289	    case Unusable{+name,+detail}: IO.pure(File & Tensor(),(file,fault(path,name,detail,[])))
290	    case Described{+name,+kind,+key,+offset,+shape,+axes}: formatted(format_of(kind),file,path,name,kind,key,offset,shape,axes,records,prefix)
291	
292	def name_of(described: Described) -> String:
293	  match described:
294	    case Unusable{name,detail}: name
295	    case Described{name,kind,key,offset,shape,axes}: name
296	
297	def entered(pair: File & Tensor(),+name: String,store: Models.Store,remaining: List<&2,Described>) -> Loading:
298	  (file,tensor) = pair
299	  Loading{file,Models.Entry{name,tensor,store},remaining}
300	
301	# The described tensors are visited last first, so the store lists them in order.
302	def load_all(fuel: Nat,loading: Loading,+path: String,+records: List<&2,Archive.Record>,+prefix: List<&2,U32>) -> IO(Loading):
303	  match fuel loading:
304	    case 0n loading: IO.pure(Loading,loading)
305	    case 1n+rest Loading{file,store,Nil{}}: IO.pure(Loading,Loading{file,store,[]})
306	    case 1n+rest Loading{file,store,Con{+first,more}}:
307	      do IO<Loading>:
308	        pair : File & Tensor() <- load_one(first,file,path,records,prefix)
309	        load_all(rest,entered(pair,name_of(first),store,more),path,records,prefix)
310	
311	def closed(loading: Loading) -> IO(Models.Store):
312	  Loading{file,store,remaining} = loading
313	  do IO<Models.Store>:
314	    File.close(file)
315	    IO.pure(Models.Store,store)
316	
317	def refused(file: File,+path: String,+detail: String) -> IO(Models.Store):
318	  do IO<Models.Store>:
319	    File.close(file)
320	    IO.pure(Models.Store,Models.Unavailable{Tensors.Problem{operation(),Tensors.FileFault{path,detail}}})
321	
322	def parsed(outcome: Pickle.Parsed,file: File,+path: String,+records: List<&2,Archive.Record>,+prefix: List<&2,U32>,+count: Nat) -> IO(Models.Store):
323	  match outcome:
324	    case Pickle.Rejected{detail}: refused(file,path,"has a data.pkl that " ++ detail)
325	    case Pickle.Parsed{value}:
326	      +described = List.reverse(&2,Described,tensors_of(count,value))
327	      do IO<Models.Store>:
328	        loading : Loading <- load_all(Nat.add(List.length(&2,Described,described),1n),Loading{file,Models.Empty{},described},path,records,prefix)
329	        closed(loading)
330	
331	def fetched(outcome: Archive.Fetched,+path: String,+records: List<&2,Archive.Record>,+prefix: List<&2,U32>) -> IO(Models.Store):
332	  Archive.Fetched{file,bytes} = outcome
333	  match bytes:
334	    case None{}: refused(file,path,"has a data.pkl that cannot be read")
335	    case Some{+bytes}: parsed(Pickle.parse(bytes),file,path,records,prefix,List.length(&2,U32,bytes))
336	
337	def pickled(record: Maybe<&2,Archive.Record>,file: File,+path: String,+records: List<&2,Archive.Record>) -> IO(Models.Store):
338	  match record:
339	    case None{}: refused(file,path,"has no data.pkl entry")
340	    case Some{Archive.Record{+entry,header,length,method}}:
341	      do IO<Models.Store>:
342	        read : Archive.Fetched <- Archive.entry(file,Archive.Record{entry,header,length,method})
343	        fetched(read,path,records,List.take(&2,U32,entry,Nat.sub(List.length(&2,U32,entry),8n)))
344	
345	def listed_archive(listing: Archive.Listing,file: File,+path: String) -> IO(Models.Store):
346	  match listing:
347	    case Archive.Unreadable{detail}: refused(file,path,detail)
348	    case Archive.Records{+records}: pickled(Archive.first_matching(records,Bytes.codes("data.pkl")),file,path,records)
349	
350	def directory(pair: File & Archive.Listing,+path: String) -> IO(Models.Store):
351	  (file,listing) = pair
352	  listed_archive(listing,file,path)
353	
354	def opened(result: Result<&1,&1,U32 & String,File>,+path: String) -> IO(Models.Store):
355	  match result:
356	    case Fail{error}: IO.pure(Models.Store,Models.Unavailable{Tensors.Problem{operation(),Tensors.FileFault{path,"cannot be opened"}}})
357	    case Done{file}:
358	      do IO<Models.Store>:
359	        pair : File & Archive.Listing <- Archive.directory(file)
360	        directory(pair,path)
361	
362	# One tensor per state-dict entry under its PyTorch name. FP16 and BFloat16
363	# values widen exactly; integers round to the nearest FP32. A file that is not
364	# a checkpoint gives a store whose every take reports why.
365	def load(+path: String) -> IO(Models.Store):
366	  do IO<Models.Store>:
367	    result : Result<&1,&1,U32 & String,File> <- File.open(path,"r")
368	    opened(result,path)