DeepSeek-V3/fuzz/fuzz_target.py
2024-12-31 16:10:31 +05:30

29 lines
928 B
Python

import atheris
import sys
import os
import json
import torch
from safetensors.torch import load_file, save_file
from kernel import weight_dequant
def TestOneInput(data):
fdp = atheris.FuzzedDataProvider(data)
try:
fp8_path = fdp.ConsumeUnicodeNoSurrogates(50)
bf16_path = fdp.ConsumeUnicodeNoSurrogates(50)
torch.set_default_dtype(torch.bfloat16)
os.makedirs(bf16_path, exist_ok=True)
model_index_file = os.path.join(fp8_path, "model.safetensors.index.json")
with open(model_index_file, "r") as f:
model_index = json.load(f)
weight_map = model_index["weight_map"]
for tensor_name in weight_map:
file_name = weight_map[tensor_name]
file_path = os.path.join(fp8_path, file_name)
load_file(file_path, device="cuda")
except Exception as e:
pass
atheris.Setup(sys.argv, TestOneInput)
atheris.Fuzz()