Make s3 imports try-except and separate requirements to s3 file

EleutherAI · Sep 23, 2023 · 3d76d4f · 3d76d4f
1 parent 4e548a4
commit 3d76d4f
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 3 deletions.
diff --git 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
@@ -0,0 +1,2 @@
+hf-transfer>=0.1.3
+boto3
diff --git 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