Deepspeed can be used outside of fbcode (facebookresearch#1727)

Summary: Pull Request resolved: fairinternal/fairseq-py#1727

Reviewed By: myleott

Differential Revision: D27213955

Pulled By: sshleifer

fbshipit-source-id: be84e7f7c1c55c407ee7445fad9b3026a79763fb

fairseq/optim/cpu_adam.py

@@ -16,13 +16,21 @@
 
 
 try:
-    import deepspeed.op_extensions.cpu_adam as ds_opt_adam
-    has_deepspeed_cpu_adam = True
+    import deepspeed
+    has_deepspeed = True
 except ImportError as e:
-    logging.warning(e)
-    has_deepspeed_cpu_adam = False
+    has_deepspeed = False
 
 
+def _get_cpu_adam():
+    try:
+        from deepspeed.ops.op_builder import CPUAdamBuilder
+        return CPUAdamBuilder().load()
+    except ImportError:
+        # fbcode
+        from deepspeed.ops.adam import DeepSpeedCPUAdam as ds_opt_adam
+        return ds_opt_adam
+
 @dataclass
 class FairseqCPUAdamConfig(FairseqDataclass):
     adam_betas: str = field(
@@ -97,13 +105,13 @@ def __init__(
         self.use_fp16_stats = use_fp16_stats
         self.FLOAT16_MAX = 65504.0
 
-        if not has_deepspeed_cpu_adam:
+        if not has_deepspeed:
             raise ImportError("Please install DeepSpeed: pip install deepspeed")
 
         self.opt_id = CPUAdam.optimizer_id
         CPUAdam.optimizer_id = CPUAdam.optimizer_id + 1
 
-        self.ds_opt_adam = ds_opt_adam
+        self.ds_opt_adam = _get_cpu_adam()
         adamw_mode = True
         self.ds_opt_adam.create_adam(
             self.opt_id, lr, betas[0], betas[1], eps, weight_decay, adamw_mode

scripts/test_fsdp.sh

@@ -0,0 +1,13 @@
+#!/usr/bin/env bash
+rm -rf fsdp_dummy
+mkdir -p fsdp_dummy
+fairseq-train /private/home/sshleifer/data-bin/stories_mmap \
+    --ddp-backend fully_sharded --fp16 --fp16-init-scale 4 \
+    --cpu-offload --checkpoint-activations \
+    --task language_modeling --tokens-per-sample 256 --batch-size 8 \
+    --arch transformer_lm_gpt2_tiny \
+    --optimizer cpu_adam --adam-betas "(0.9,0.98)" \
+    --lr 0.0001 --lr-scheduler polynomial_decay --warmup-updates 5 --total-num-update 10 \
+    --max-update 10 --log-format json --log-interval 1 \
+    --save-interval-updates 10 --save-dir fsdp_dummy \
+    --restore-file x.pt "$@"