Use hf transfer as default (pytorch#2046)

Co-authored-by: Felipe Mello <[email protected]>
vrvrv · Nov 22, 2024 · a9aadf5 · a9aadf5
1 parent 2d258f4
commit a9aadf5
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 1 deletion.
diff --git a/pyproject.toml b/pyproject.toml
@@ -13,7 +13,7 @@ dependencies = [
 
     # Hugging Face integrations
     "datasets",
-    "huggingface_hub",
+    "huggingface_hub[hf_transfer]",
     "safetensors",
 
     # Kaggle Integrations

diff --git a/torchtune/__init__.py b/torchtune/__init__.py
@@ -23,6 +23,18 @@
         """
     ) from e
 
+# Enables faster downloading. For more info: https://huggingface.co/docs/huggingface_hub/en/guides/download#faster-downloads
+# To disable, run `HF_HUB_ENABLE_HF_TRANSFER=0 tune download <model_config>`
+try:
+    import os
+
+    import hf_transfer  # noqa
+
+    if os.environ.get("HF_HUB_ENABLE_HF_TRANSFER") is None:
+        os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
+except ImportError:
+    pass
+
 from torchtune import datasets, generation, models, modules, utils
 
 __all__ = [datasets, models, modules, utils, generation]