Merge pull request #26 from EleutherAI/deeperspeed_and_wandb

Add wandb and deeperspeed

Dockerfile

@@ -23,6 +23,7 @@ RUN pip install --upgrade tensorflow
 RUN pip install -r requirements.txt
 RUN pip install -v --disable-pip-version-check --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" git+https://github.com/NVIDIA/apex.git
 RUN echo 'deb http:https://archive.ubuntu.com/ubuntu/ focal main restricted' >> /etc/apt/sources.list && apt-get install --upgrade libpython3-dev
+RUN sudo apt-get update -y && sudo apt-get install -y libpython3-dev
 
 WORKDIR /app
 

deepy.py

@@ -0,0 +1,38 @@
+#!/usr/bin/env python
+
+"""
+Helper script to enable wandb with deepspeed. Drop in replacement for `deepspeed` command.
+
+i.e. `$ deepy.py --hosts ....` == `$ deepspeed --hosts ....`
+"""
+
+import shortuuid
+import sys
+import os
+import deepspeed
+from deepspeed.launcher.runner import main
+import requests
+
+def get_wandb_api_key():
+ """ Get Weights and Biases API key from ENV or .netrc file. Otherwise return None """
+ if 'WANDB_API_KEY' in os.environ:
+ return os.environ['WANDB_API_KEY']
+
+ wandb_token = requests.utils.get_netrc_auth('https://api.wandb.ai')
+
+ if wandb_token is not None:
+ return wandb_token[1]
+
+# Generate unique run group name
+wandb_group = shortuuid.uuid()
+sys.argv.extend(['--group_name', wandb_group])
+
+# Extract wandb API key and inject into worker environments
+wandb_token = get_wandb_api_key()
+if wandb_token is not None:
+ deepspeed.launcher.runner.EXPORT_ENVS.append('WANDB_API_KEY')
+ os.environ['WANDB_API_KEY'] = wandb_token
+
+if __name__ == '__main__':
+ main()
+

requirements.txt

@@ -4,8 +4,9 @@ six
 regex
 numpy
 nltk
-deepspeed==0.3.10 # bug in DeepSpeed, see https:https://github.com/huggingface/transformers/issues/9996#issuecomment-773725303
+-e git+git:https://github.com/EleutherAI/DeeperSpeed.git@5a454114886709987a954ed315bf9f6b14cc4efc#egg=deepspeed
 autopep8
 zstandard
 cupy-cuda102
-mpi4py
+mpi4py
+wandb