diff --git a/megatron/checkpointing.py b/megatron/checkpointing.py index 2c08f46c2..8bcc01f3b 100644 --- a/megatron/checkpointing.py +++ b/megatron/checkpointing.py @@ -27,8 +27,14 @@ import sys import numpy as np -import boto3 -import hf_transfer +try: + import boto3 +except ModuleNotFoundError: + print("For s3 checkpointing, please install boto3 either using requirements/requirements-s3.txt or https://github.com/boto/boto3") +try: + import hf_transfer +except ModuleNotFoundError: + print("For s3 checkpointing, please install hf_transfer either using requirements/requirements-s3.txt or https://github.com/huggingface/hf_transfer") import torch from glob import glob diff --git a/requirements/requirements-s3.txt b/requirements/requirements-s3.txt new file mode 100644 index 000000000..7a2924ccd --- /dev/null +++ b/requirements/requirements-s3.txt @@ -0,0 +1,2 @@ +hf-transfer>=0.1.3 +boto3 \ No newline at end of file diff --git a/requirements/requirements.txt b/requirements/requirements.txt index 03bf47a45..88e49f073 100644 --- a/requirements/requirements.txt +++ b/requirements/requirements.txt @@ -3,7 +3,6 @@ git+https://github.com/EleutherAI/DeeperSpeed.git#egg=deepspeed ftfy>=6.0.1 git+https://github.com/EleutherAI/lm_dataformat.git@4eec05349977071bf67fc072290b95e31c8dd836 huggingface_hub>=0.11.0 -hf-transfer>=0.1.3 lm_eval>=0.3.0 mpi4py>=3.0.3 numpy>=1.22.0