foundation-model-stack / fastsafetensors Goto Github PK

huggingface-cli download HuggingFaceH4/zephyr-7b-beta --include "*.safetensors"

import time

import torch
from safetensors.torch import safe_open
from fastsafetensors.loader import fastsafe_open

ZEPHYR_MODEL_PATHS = [
    f"/mnt/local_storage/data/cache/huggingface/hub/models--HuggingFaceH4--zephyr-7b-beta/snapshots/b70e0c9a2d9e14bd1e812d3c398e5f313e93b473/model-0000{i}-of-00008.safetensors" for i in range(1, 9)
]


device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
state_dict = {}
cloned_state_dict = {}
start = time.perf_counter()
with fastsafe_open(ZEPHYR_MODEL_PATHS, device="cuda:0", nogds=False, debug_log=True) as f:
    for k in f.get_keys():
        state_dict[k] = f.get_tensor(k)

    end = time.perf_counter()
    
    f.bufs.close() # this releases memory... so my call to the safetensor run works or else OOM
print(f"fastsafe_open: {end - start} seconds")

start = time.perf_counter()
state_dict = {}
for model_path in ZEPHYR_MODEL_PATHS:
    with safe_open(model_path, framework="pt", device=0) as f:
        for k in f.keys():
            state_dict[k] = f.get_tensor(k)
end = time.perf_counter()
print(f"safe_open: {end - start} seconds")

fastsafe_open: 14.059521514000153 seconds
safe_open: 2.213277554999877 seconds