deepseekmirror/DeepSeek-V3
1
0
Fork 0
mirror of https://github.com/deepseek-ai/DeepSeek-V3.git synced 2025-02-23 06:08:58 -05:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

Memory management update

Browse Source
This commit is contained in:
Nripesh Niketan 2025-02-02 10:41:20 +00:00
parent b6e3910fd0
commit 73efe7c631
1 changed files with 7 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
inference/fp8_cast_bf16.py
View File

@ -3,6 +3,7 @@ import json
from argparse import ArgumentParser
from glob import glob
from tqdm import tqdm
import gc
import torch
from safetensors.torch import load_file, save_file
@ -97,7 +98,12 @@ def main(fp8_path, bf16_path):
if len(loaded_files) > 2:
oldest_file = next(iter(loaded_files))
del loaded_files[oldest_file]
torch.cuda.empty_cache()
if torch.cuda.is_available():
torch.cuda.empty_cache()
elif torch.mps.is_available():
torch.mps.empty_cache()
else:
gc.collect()
# Update model index
new_model_index_file = os.path.join(bf16_path, "model.safetensors.index.json")