fix the benches

Signed-off-by: Yu Chin Fabian Lim <flim@sg.ibm.com>
foundation-model-stack · Aug 21, 2024 · 51499ff · 51499ff
1 parent 263086b
commit 51499ff
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 3 deletions.
diff --git a/plugins/accelerated-moe/README.md b/plugins/accelerated-moe/README.md
@@ -15,5 +15,5 @@ Known Issues
 Currently databricks megablocks does not have a PyPi repository and does not have a proper release, so we have to install from the github repository as below. Please note that installing from github will require CUDA Toolkit to build.
 
 ```
-pip install git+https://github.com/databricks/megablocks.git@bce5d7b2aaf5038bc93b36f76c2baf51c2939bd2
+pip install -r requirements_mb.txt
 ```
diff --git a/plugins/accelerated-moe/src/fms_acceleration_moe/requirements-mb.txt b/plugins/accelerated-moe/src/fms_acceleration_moe/requirements-mb.txt
@@ -0,0 +1 @@
+pip install git+https://github.com/databricks/megablocks.git@bce5d7b2aaf5038bc93b36f76c2baf51c2939bd2
diff --git a/scripts/benchmarks/accelerator-config.json b/scripts/benchmarks/accelerator-config.json
@@ -0,0 +1,5 @@
+{
+    "gradient_accumulation_kwargs": {
+        "sync_each_batch": true
+    }
+}
diff --git a/scripts/benchmarks/scenarios.yaml b/scripts/benchmarks/scenarios.yaml
@@ -98,8 +98,13 @@ scenarios:
         framework_config: 
             - moe-megablocks
         arguments:
-            learning_rate: 2e-4
-            bf16: True
+            learning_rate: 5e-5
             torch_dtype: bfloat16
+            accelerator_config: scripts/benchmarks/accelerator-config.json
+            gradient_accumulation_steps: 16
+            logging_steps: 1
+            packing: False
+            adam_epsilon: 1e-8
+
             model_name_or_path: 
                 - 'mistralai/Mixtral-8x7B-Instruct-v0.1'
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1 @@
		pip install git+https://github.com/databricks/megablocks.git@bce5d7b2aaf5038bc93b36f76c2baf51c2939bd2