deepseekmirror/DeepSeek-V3
1
0
Fork 0
mirror of https://github.com/deepseek-ai/DeepSeek-V3.git synced 2025-04-20 02:28:57 -04:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

fix(fp8_cast): Add robust memory management and error handling

Browse Source
This commit is contained in:
ajwise9 2025-02-04 16:36:07 +00:00
parent b5d872ead0
commit dca08f2cfd
1 changed files with 8 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
inference/fp8_cast_bf16.py
View File

@ -88,10 +88,14 @@ def main(fp8_path, bf16_path):
save_file(new_state_dict, new_safetensor_file) save_file(new_state_dict, new_safetensor_file)
# Memory management: keep only the 2 most recently used files # Memory management: keep only the 2 most recently used files
if len(loaded_files) > 2: try:
oldest_file = next(iter(loaded_files)) if len(loaded_files) > 2:
del loaded_files[oldest_file] oldest_file = next(iter(loaded_files))
torch.cuda.empty_cache() del loaded_files[oldest_file]
torch.cuda.empty_cache()
except RuntimeError as e:
print(f"Memory error: {e}")
# Implement fallback strategy or graceful exit
# Update model index # Update model index
new_model_index_file = os.path.join(bf16_path, "model.safetensors.index.json") new_model_index_file = os.path.join(bf16_path, "model.safetensors.index.json")