tensor_checkpoint.bend
stelliferous@0.0.3.0 · 368 lines · 0 definitions
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.
5import Base
6import stelliferous@0.0.3.0/storage_bytes.bend as Bytes
7import stelliferous@0.0.3.0/checkpoint_zip.bend as Archive
8import stelliferous@0.0.3.0/checkpoint_pickle.bend as Pickle
9import stelliferous@0.0.3.0/storage_element_format.bend as Formats
10import stelliferous@0.0.3.0/tensor.bend as Tensors
11import stelliferous@0.0.3.0/tensor_file.bend as Files
12import stelliferous@0.0.3.0/tensor_model.bend as Models
13import stelliferous@0.0.3.0/tensor_view.bend as Views
14import stelliferous@0.0.3.0/tensor_shape.bend as Shapes
15
16def Tensor() -> Type:
17 Tensors.Tensor<F32>
18
19def 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.
25type 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.
30type Task is Data:
31 Task{prefix: String,module: Pickle.Value}
32
33type 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
39def 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
45def 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.
52def 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
57def 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
64def 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
69def 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).
75def 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
81def 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
87def 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
92def 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
98def 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).
104def 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
115def 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
121def 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
127def 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.
134def 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
142def 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
149def 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.
155def 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.
166def 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
172def is_module(value: Pickle.Value) -> Bool:
173 match value:
174 case Pickle.Object{class,arguments,Pickle.Dict{entries}}: True{}
175 case value: False{}
176
177def 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.
184def 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
191def 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
201def 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
205def 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
208def 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
211def 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.
218def 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
223def 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
228def 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.
235def 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
239def 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
243def 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.
252def 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
257def 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
264def 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
269def 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
273def 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
281def 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
287def 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
292def 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
297def 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.
302def 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
311def 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
317def 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
322def 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
331def 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
337def 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
345def 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
350def directory(pair: File & Archive.Listing,+path: String) -> IO(Models.Store):
351 (file,listing) = pair
352 listed_archive(listing,file,path)
353
354def 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.
365def 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)