artin/llama2.c
1
0
Fork 0
mirror of https://github.com/trholding/llama2.c.git synced 2026-02-06 11:26:53 +00:00
Code Issues Packages Projects Releases Wiki Activity

use out_file.tell() instead of nbytes += arithmetic

Browse Source
This commit is contained in:
Andrej Karpathy 2023-08-20 16:51:35 +00:00
parent fa8dfd854e
commit f3db92a2dc
1 changed files with 2 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
export.py
View File

@ -124,13 +124,10 @@ def version1_export(model, filepath):
out_file = open(filepath, 'wb')
# first write out the header. the header will be 256 bytes
nbytes = 0
# 1) write magic, which will be uint32 of "ak42" in ASCII
out_file.write(struct.pack('I', 0x616b3432))
nbytes += 4
# 2) write version, which will be int
out_file.write(struct.pack('i', version))
nbytes += 4
# 3) write the params, which will be 7 ints
p = model.params
hidden_dim = model.layers[0].feed_forward.w1.weight.shape[0]
@ -138,12 +135,10 @@ def version1_export(model, filepath):
header = struct.pack('iiiiiii', p.dim, hidden_dim, p.n_layers, p.n_heads,
n_kv_heads, p.vocab_size, p.max_seq_len)
out_file.write(header)
nbytes += 7*4
# 4) write some other flags
shared_classifier = 1 # we do share a classifier, write flag as a byte
out_file.write(struct.pack('B', shared_classifier))
nbytes += 1
pad = 256 - nbytes # pad the rest with zeros
pad = 256 - out_file.tell() # pad rest with zeros; tell returns current pos
assert pad >= 0
out_file.write(b'\0' * pad)
@ -198,13 +193,10 @@ def version2_export(model, filepath, group_size=64):
# write
out_file = open(filepath, 'wb')
# first write out the header. the header will be 256 bytes
nbytes = 0
# 1) write magic, which will be uint32 of "ak42" in ASCII
out_file.write(struct.pack('I', 0x616b3432))
nbytes += 4
# 2) write version, which will be int
out_file.write(struct.pack('i', version))
nbytes += 4
# 3) write the params, which will be 7 ints
p = model.params
hidden_dim = model.layers[0].feed_forward.w1.weight.shape[0]
@ -212,14 +204,11 @@ def version2_export(model, filepath, group_size=64):
header = struct.pack('iiiiiii', p.dim, hidden_dim, p.n_layers, p.n_heads,
n_kv_heads, p.vocab_size, p.max_seq_len)
out_file.write(header)
nbytes += 7*4
# 4) write some other flags
shared_classifier = 1 # we do share a classifier, write flag as a byte
out_file.write(struct.pack('B', shared_classifier))
nbytes += 1
out_file.write(struct.pack('i', group_size)) # group size used for quantization
nbytes += 4
pad = 256 - nbytes # pad the rest with zeros
pad = 256 - out_file.tell() # pad rest with zeros; tell returns current pos
assert pad >= 0
out_file.write(b'\0' * pad)
# now that the header is done, let's write out the model