Convert back to AF2 scheduler/optimizer in deepspeed config

deepspeed_config.json

@@ -10,13 +10,6 @@
   "bfloat16": {
     "enabled": true
   },
-  "optimizer": {
-    "type": "Adam",
-    "params": {
-      "lr": 1e-3,
-      "eps": 1e-5
-    }
-  },
   "zero_optimization": {
     "stage": 2,
     "offload_optimizer": {
@@ -29,5 +22,6 @@
     "cpu_checkpointing": false,
     "profile": false
   },
-  "gradient_clipping": 0.1
+  "gradient_clipping": 0.1,
+  "zero_force_ds_cpu_optimizer": false
 }

environment.yml

@@ -16,7 +16,6 @@ dependencies:
   - pytorch::pytorch=1.12.*
   - pip:
       - biopython==1.79
-      - deepspeed==0.5.10
       - dm-tree==0.1.6
       - ml-collections==0.1.0
       - numpy==1.21.2
@@ -29,3 +28,8 @@ dependencies:
       - wandb==0.12.21
       - modelcif==0.7
       - git+https://github.com/NVIDIA/dllogger.git
+      - git+https://github.com/NVIDIA/cutlass.git
+      - git+https://github.com/microsoft/DeepSpeed.git
+      # TODO: Replace above when version becomes available
+      # - deepspeed==0.10.3
+