fixes to readme and tox

Signed-off-by: Yu Chin Fabian Lim <flim@sg.ibm.com>
foundation-model-stack · Aug 23, 2024 · 9cf7008 · 9cf7008
1 parent e00fcd0
commit 9cf7008
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 2 deletions.
diff --git a/plugins/accelerated-moe/README.md b/plugins/accelerated-moe/README.md
@@ -12,8 +12,15 @@ Plugin | Description | Depends | Loading | Augmentation | Callbacks
 
 ## Running Benchmarks
 
+Run the below in the top-level directory of this repo:
+- the `megablocks` dep is not included by default, so the `-x` switch installs it.
+
 ```
-tox -e run-benches -- 8 8 scenarios.yaml accelerated-moe-megablocks
+tox -e run-benches \
+    -x testenv:run-benches.deps+="-r plugins/accelerated-moe/requirements-mb.txt" \
+    -- \
+    8 8 benchmark_outputs scenarios.yaml accelerated-moe-megablocks
+
 ```
 
 ## Expert-Parallel MoE with Megablocks

diff --git a/plugins/accelerated-moe/requirements-mb.txt b/plugins/accelerated-moe/requirements-mb.txt
@@ -0,0 +1,3 @@
+megablocks @ git+https://github.com/databricks/megablocks.git@bce5d7b2aaf5038bc93b36f76c2baf51c2939bd2
+
+# auto_gptq @ git+https://github.com/AutoGPTQ/AutoGPTQ.git@ea829c7bbe83561c2b1de26795b6592992373ef7
diff --git a/plugins/accelerated-moe/src/fms_acceleration_moe/requirements-mb.txt b/plugins/accelerated-moe/src/fms_acceleration_moe/requirements-mb.txt
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1,3 @@
		megablocks @ git+https://github.com/databricks/megablocks.git@bce5d7b2aaf5038bc93b36f76c2baf51c2939bd2

		# auto_gptq @ git+https://github.com/AutoGPTQ/AutoGPTQ.git@ea829c7bbe83561c2b1de26795b6592992373ef7