Skip to content

perf: optimize safetensors metadata fetching latency by fetching in parallel - #99

Open
saquibsaifee wants to merge 1 commit into
GenAI-Security-Project:mainfrom
saquibsaifee:feat/optimize-safetensors-metadata-fetch-13406851267980656361
Open

saquibsaifee wants to merge 1 commit into
GenAI-Security-Project:mainfrom
saquibsaifee:feat/optimize-safetensors-metadata-fetch-13406851267980656361

Conversation

@saquibsaifee

Copy link
Copy Markdown
Contributor

💡 What: Refactored fetch_safetensors_metadata to use concurrent.futures.ThreadPoolExecutor for fetching config.json, tokenizer_config.json, and safetensors metadata in parallel rather than sequentially.
🎯 Why: To reduce overall network I/O latency when fetching metadata about safetensors. The previous implementation executed three blocking downloads in sequence.
📊 Measured Improvement: Running a basic warmup-based test benchmarking TinyLlama/TinyLlama-1.1B-Chat-v1.0, latency dropped from ~0.46s down to ~0.32s, approximately a 30% reduction in fetch time. The latency improvement applies best for models containing tokenizer_config.json and valid safetensors metadata alongside config.json by combining their network IO time.

@saquibsaifee saquibsaifee changed the title Feat/optimize safetensors metadata fetch perf: optimize safetensors metadata fetching latency by fetching in parallel Aug 31, 2026
@saquibsaifee

Copy link
Copy Markdown
Contributor Author

Benchmark can be reproduced using:

import time
import argparse
from src.models.safetensors_metadata import fetch_safetensors_metadata

def benchmark(repo_id: str):
    # Warmup
    try:
        fetch_safetensors_metadata(repo_id)
    except Exception:
        pass

    start = time.time()
    res = fetch_safetensors_metadata(repo_id)
    end = time.time()

    print(f"Time taken for {repo_id}: {end - start:.4f} seconds")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Benchmark safetensors metadata fetch.")
    parser.add_argument("--repo_id", default="TinyLlama/TinyLlama-1.1B-Chat-v1.0", help="Hugging Face repo ID")
    args = parser.parse_args()
    benchmark(args.repo_id)

@saquibsaifee

Copy link
Copy Markdown
Contributor Author

@eaglei15 this PR is ready to be reviewed.

Refactored fetch_safetensors_metadata to fetch config.json, tokenizer_config.json
and the safetensors metadata concurrently using concurrent.futures.ThreadPoolExecutor.
This reduces metadata network operations latency by running them in parallel rather
than sequentially. Removes benchmark_verify.py script per PR review.

Signed-off-by: saquibsaifee <saquibsaifee2@gmail.com>
@saquibsaifee
saquibsaifee force-pushed the feat/optimize-safetensors-metadata-fetch-13406851267980656361 branch from d6b2b1d to be3e440 Compare August 31, 2026 17:03
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant