Support faster checkpoint downloads in the cloud (Lightning-AI#918)

murthyrudra · Feb 6, 2024 · b8032e8 · b8032e8
1 parent 99e9f4e
commit b8032e8
Show file tree

Hide file tree

Showing 18 changed files with 35 additions and 20 deletions.
diff --git a/scripts/download.py b/scripts/download.py
@@ -13,6 +13,7 @@
 sys.path.append(str(wd))
 
 _SAFETENSORS_AVAILABLE = RequirementCache("safetensors")
+_HF_TRANSFER_AVAILABLE = RequirementCache("hf_transfer")
 
 
 def download_from_hub(
@@ -51,6 +52,15 @@ def download_from_hub(
     elif from_safetensors:
         raise ValueError("`--from_safetensors=True` won't have an effect with `--tokenizer_only=True`")
 
+    import huggingface_hub._snapshot_download as download
+    import huggingface_hub.constants as constants
+
+    previous = constants.HF_HUB_ENABLE_HF_TRANSFER
+    if _HF_TRANSFER_AVAILABLE and not previous:
+        print("Setting HF_HUB_ENABLE_HF_TRANSFER=1")
+        constants.HF_HUB_ENABLE_HF_TRANSFER = True
+        download.HF_HUB_ENABLE_HF_TRANSFER = True
+
     directory = checkpoint_dir / repo_id
     snapshot_download(
         repo_id,
@@ -61,6 +71,9 @@ def download_from_hub(
         token=access_token,
     )
 
+    constants.HF_HUB_ENABLE_HF_TRANSFER = previous
+    download.HF_HUB_ENABLE_HF_TRANSFER = previous
+
     # convert safetensors to PyTorch binaries
     if from_safetensors:
         from safetensors import SafetensorError

diff --git a/tutorials/download_code_llama.md b/tutorials/download_code_llama.md
@@ -36,7 +36,7 @@ codellama/CodeLlama-70b-Instruct-hf
 In order to use a specific checkpoint, for instance [CodeLlama-7b-Python-hf](https://huggingface.co/codellama/CodeLlama-7b-Python-hf), download the weights and convert the checkpoint to the lit-gpt format.
 
 ```bash
-pip install huggingface_hub
+pip install 'huggingface_hub[hf_transfer] @ git+https://github.com/huggingface/huggingface_hub'
 
 python scripts/download.py --repo_id codellama/CodeLlama-7b-Python-hf
 

diff --git a/tutorials/download_dolly.md b/tutorials/download_dolly.md
@@ -26,7 +26,7 @@ databricks/dolly-v2-12b
 In order to use a specific Dolly checkpoint, for instance [dolly-v2-3b](https://huggingface.co/databricks/dolly-v2-3b), download the weights and convert the checkpoint to the lit-gpt format:
 
 ```bash
-pip install huggingface_hub
+pip install 'huggingface_hub[hf_transfer] @ git+https://github.com/huggingface/huggingface_hub'
 
 python scripts/download.py --repo_id databricks/dolly-v2-3b
 

diff --git a/tutorials/download_falcon.md b/tutorials/download_falcon.md
@@ -26,7 +26,7 @@ tiiuae/falcon-180B-chat
 In order to use a specific Falcon checkpoint, for instance [falcon-7b](https://huggingface.co/tiiuae/falcon-7b), download the weights and convert the checkpoint to the lit-gpt format:
 
 ```bash
-pip install huggingface_hub
+pip install 'huggingface_hub[hf_transfer] @ git+https://github.com/huggingface/huggingface_hub'
 
 python scripts/download.py --repo_id tiiuae/falcon-7b
 

diff --git a/tutorials/download_freewilly_2.md b/tutorials/download_freewilly_2.md
@@ -5,7 +5,7 @@ Stability AI announced FreeWilly inspired by the methodology pioneered by Micros
 FreeWilly2 leverages the Llama 2 70B foundation model to reach a performance that compares favorably with GPT-3.5 for some tasks.
 
 ```bash
-pip install huggingface_hub
+pip install 'huggingface_hub[hf_transfer] @ git+https://github.com/huggingface/huggingface_hub'
 
 python scripts/download.py --repo_id stabilityai/FreeWilly2
 

diff --git a/tutorials/download_function_calling_llama_2.md b/tutorials/download_function_calling_llama_2.md
@@ -8,7 +8,7 @@ The model responds with a structured json argument with the function name and ar
 In order to use the checkpoint, download the weights and convert the checkpoint to the lit-gpt format.
 
 ```bash
-pip install huggingface_hub
+pip install 'huggingface_hub[hf_transfer] @ git+https://github.com/huggingface/huggingface_hub'
 
 python scripts/download.py --repo_id Trelis/Llama-2-7b-chat-hf-function-calling-v2 --from_safetensors true
 

diff --git a/tutorials/download_llama_2.md b/tutorials/download_llama_2.md
@@ -33,7 +33,7 @@ This requires that you've been granted access to the weights on the HuggingFace
 After access is granted, you can find your HF hub token in <https://huggingface.co/settings/tokens>.
 
 ```bash
-pip install huggingface_hub
+pip install 'huggingface_hub[hf_transfer] @ git+https://github.com/huggingface/huggingface_hub'
 
 python scripts/download.py --repo_id meta-llama/Llama-2-7b-chat-hf --access_token your_hf_token
 

diff --git a/tutorials/download_longchat.md b/tutorials/download_longchat.md
@@ -19,7 +19,7 @@ lmsys/longchat-13b-16k
 In order to use a specific checkpoint, for instance [longchat-7b-16k](https://huggingface.co/lmsys/longchat-7b-16k), download the weights and convert the checkpoint to the lit-gpt format:
 
 ```bash
-pip install huggingface_hub
+pip install 'huggingface_hub[hf_transfer] @ git+https://github.com/huggingface/huggingface_hub'
 
 python scripts/download.py --repo_id lmsys/longchat-7b-16k
 

diff --git a/tutorials/download_mistral.md b/tutorials/download_mistral.md
@@ -32,7 +32,7 @@ mistralai/Mistral-7B-Instruct-v0.2
 In order to use the Mistral 7B model checkpoint, which requires about 14 GB of disk space, download the weights and convert the checkpoint to the lit-gpt format:
 
 ```bash
-pip install huggingface_hub
+pip install 'huggingface_hub[hf_transfer] @ git+https://github.com/huggingface/huggingface_hub'
 
 python scripts/download.py --repo_id mistralai/Mistral-7B-Instruct-v0.2
 
@@ -57,7 +57,7 @@ Details about the data used to train the model or training procedure have not be
 In order to use the Mixtral 7B model checkpoint, which requires about 94 GB of disk space, download the weights and convert the checkpoint to the lit-gpt format:
 
 ```bash
-pip install huggingface_hub
+pip install 'huggingface_hub[hf_transfer] @ git+https://github.com/huggingface/huggingface_hub'
 
 python scripts/download.py --repo_id mistralai/Mixtral-8x7B-Instruct-v0.1 --from_safetensors true
 

diff --git a/tutorials/download_openllama.md b/tutorials/download_openllama.md
@@ -21,7 +21,7 @@ openlm-research/open_llama_13b
 In order to use a specific OpenLLaMA checkpoint, for instance [open_llama_3b](https://huggingface.co/openlm-research/open_llama_3b), download the weights and convert the checkpoint to the lit-gpt format:
 
 ```bash
-pip install huggingface_hub
+pip install 'huggingface_hub[hf_transfer] @ git+https://github.com/huggingface/huggingface_hub'
 
 python scripts/download.py --repo_id openlm-research/open_llama_3b
 

diff --git a/tutorials/download_phi.md b/tutorials/download_phi.md
@@ -8,8 +8,10 @@ The model weights are released under [*Microsoft Research license*](https://hugg
 To download the model weights and convert them to the lit-gpt format, run
 
 ```bash
-pip install huggingface_hub
+pip install 'huggingface_hub[hf_transfer] @ git+https://github.com/huggingface/huggingface_hub'
+
 python scripts/download.py --repo_id microsoft/phi-2 --from_safetensors True
+
 python scripts/convert_hf_checkpoint.py --checkpoint_dir checkpoints/microsoft/phi-2
 ```
 
@@ -59,7 +61,7 @@ The model weights are released under a [*Microsoft Research license*](https://hu
 In order to use the phi-1.5 model checkpoint, which requires about 3 Gb of disk space, download the weights and convert the checkpoint to the lit-gpt format:
 
 ```bash
-pip install huggingface_hub
+pip install 'huggingface_hub[hf_transfer] @ git+https://github.com/huggingface/huggingface_hub'
 
 python scripts/download.py --repo_id microsoft/phi-1_5
 

diff --git a/tutorials/download_pythia.md b/tutorials/download_pythia.md
@@ -37,7 +37,7 @@ EleutherAI/pythia-12b-deduped
 In order to use a specific Pythia checkpoint, for instance [pythia-1b](https://huggingface.co/EleutherAI/pythia-1b), download the weights and convert the checkpoint to the lit-gpt format:
 
 ```bash
-pip install huggingface_hub
+pip install 'huggingface_hub[hf_transfer] @ git+https://github.com/huggingface/huggingface_hub'
 
 python scripts/download.py --repo_id EleutherAI/pythia-1b
 

diff --git a/tutorials/download_redpajama_incite.md b/tutorials/download_redpajama_incite.md
@@ -27,7 +27,7 @@ togethercomputer/RedPajama-INCITE-Instruct-7B-v0.1
 In order to use a specific RedPajama-INCITE checkpoint, for instance [RedPajama-INCITE-Base-3B-v1](https://huggingface.co/togethercomputer/RedPajama-INCITE-Base-3B-v1), download the weights and convert the checkpoint to the lit-gpt format:
 
 ```bash
-pip install huggingface_hub
+pip install 'huggingface_hub[hf_transfer] @ git+https://github.com/huggingface/huggingface_hub'
 
 python scripts/download.py --repo_id togethercomputer/RedPajama-INCITE-Base-3B-v1
 

diff --git a/tutorials/download_stablelm.md b/tutorials/download_stablelm.md
@@ -23,7 +23,7 @@ stabilityai/stablelm-zephyr-3b
 In order to use a specific StableLM checkpoint, for instance [stablelm-base-alpha-3b](http://huggingface.co/stabilityai/stablelm-base-alpha-3b), download the weights and convert the checkpoint to the lit-gpt format:
 
 ```bash
-pip install huggingface_hub
+pip install 'huggingface_hub[hf_transfer] @ git+https://github.com/huggingface/huggingface_hub'
 
 python scripts/download.py --repo_id stabilityai/stablelm-base-alpha-3b
 
@@ -51,7 +51,7 @@ More details can be found in the [announcement](https://stability.ai/news/stable
 In order to use this model, download the weights and convert the checkpoint to the lit-gpt format. As this version of the model is in `safetensor` format, to download it an additional flag is required:
 
 ```bash
-pip install huggingface_hub
+pip install 'huggingface_hub[hf_transfer] @ git+https://github.com/huggingface/huggingface_hub'
 
 python scripts/download.py --repo_id stabilityai/stablelm-zephyr-3b --from_safetensors=True
 ```

diff --git a/tutorials/download_tinyllama.md b/tutorials/download_tinyllama.md
@@ -21,7 +21,7 @@ TinyLlama/TinyLlama-1.1B-Chat-v1.0
 In order to use a specific checkpoint, for instance [TinyLlama 1.1B base model](https://huggingface.co/TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T), which requires about 5 GB of disk space, download the weights and convert the checkpoint to the lit-gpt format:
 
 ```bash
-pip install huggingface_hub
+pip install 'huggingface_hub[hf_transfer] @ git+https://github.com/huggingface/huggingface_hub'
 
 python scripts/download.py --repo_id TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T
 

diff --git a/tutorials/download_vicuna.md b/tutorials/download_vicuna.md
@@ -23,7 +23,7 @@ lmsys/vicuna-13b-v1.5-16k
 In order to use a specific Vicuna checkpoint, for instance [vicuna-7b-v1.5](https://huggingface.co/lmsys/vicuna-7b-v1.5), download the weights and convert the checkpoint to the lit-gpt format:
 
 ```bash
-pip install huggingface_hub
+pip install 'huggingface_hub[hf_transfer] @ git+https://github.com/huggingface/huggingface_hub'
 
 python scripts/download.py --repo_id lmsys/vicuna-7b-v1.5
 

diff --git a/tutorials/pretrain_redpajama.md b/tutorials/pretrain_redpajama.md
@@ -51,7 +51,7 @@ on it, you need to read, tokenize, and write the data in binary chunks. This wil
 streaming dataset that comes with lit-gpt. You will need to have the tokenizer config available:
 
 ```bash
-pip install huggingface_hub sentencepiece
+pip install 'huggingface_hub[hf_transfer] @ git+https://github.com/huggingface/huggingface_hub' sentencepiece
 
 python scripts/download.py \
    --repo_id meta-llama/Llama-2-7b-chat-hf \

diff --git a/tutorials/pretrain_tinyllama.md b/tutorials/pretrain_tinyllama.md
@@ -49,7 +49,7 @@ In order to start pretraining lit-gpt on it, you need to read, tokenize, and wri
 First, install additional dependencies for preprocessing:
 
 ```bash
-pip install lightning[data] torchmetrics tensorboard sentencepiece zstandard pandas pyarrow huggingface_hub
+pip install 'lightning[data]' torchmetrics tensorboard sentencepiece zstandard pandas pyarrow 'huggingface_hub[hf_transfer] @ git+https://github.com/huggingface/huggingface_hub'
 ```
 
 You will need to have the tokenizer config available: