Skip to content

Commit

Permalink
Improve HF download speed (Lightning-AI#1899)
Browse files Browse the repository at this point in the history
  • Loading branch information
rasbt authored Jan 8, 2025
1 parent 40c08dc commit a439338
Showing 1 changed file with 3 additions and 0 deletions.
3 changes: 3 additions & 0 deletions litgpt/scripts/download.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
import os
from concurrent.futures import ProcessPoolExecutor
from contextlib import contextmanager
import importlib.util
from pathlib import Path
from typing import List, Optional, Tuple

Expand Down Expand Up @@ -56,6 +57,8 @@ def download_from_hub(
return

from huggingface_hub import snapshot_download
if importlib.util.find_spec("hf_transfer") is None:
print("It is recommended to install hf_transfer for faster checkpoint download speeds: `pip install hf_transfer`")

download_files = ["tokenizer*", "generation_config.json", "config.json"]
if not tokenizer_only:
Expand Down

0 comments on commit a439338

Please sign in to comment.