add benches

Signed-off-by: Yu Chin Fabian Lim <flim@sg.ibm.com>

plugins/accelerated-moe/README.md

@@ -12,6 +12,9 @@ Plugin | Description | Depends | Loading | Augmentation | Callbacks
 
 ## Running Benchmarks
 
+See the benchmarks [a100_80gb_mb.csv](../../scripts/benchmarks/refs/a100_80gb_mb.csv)
+
+
 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.
 

scripts/benchmarks/refs/a100_80gb_mb.csv

@@ -0,0 +1,3 @@
+framework_config,mem_nvidia_mem_reserved,mem_peak_torch_mem_alloc_in_bytes,mem_torch_mem_alloc_in_bytes,torch_dtype,train_loss,train_runtime,train_samples_per_second,train_steps_per_second,train_tokens_per_second
+none,65598.5,58936741888,47259259904,bfloat16,0.859970542192459,4170.0391,3.07,0.024,80.575
+moe-megablocks,52284.0,48874301952,35987686400,bfloat16,0.8570401281118393,1404.3938,9.114,0.071,239.249

scripts/benchmarks/refs/requirements_mb.txt

@@ -0,0 +1,88 @@
+accelerate==0.33.0
+aiohappyeyeballs==2.4.0
+aiohttp==3.10.5
+aiosignal==1.3.1
+async-timeout==4.0.3
+attrs==24.2.0
+bitsandbytes==0.43.3
+certifi==2024.7.4
+charset-normalizer==3.3.2
+contourpy==1.2.1
+cycler==0.12.1
+datasets==2.21.0
+dill==0.3.8
+docstring_parser==0.16
+einops==0.8.0
+filelock==3.15.4
+fire==0.6.0
+flash-attn==2.6.3
+-e git+https://github.com/foundation-model-stack/fms-acceleration.git@9cf70081f3bfc00e84331102e7d13b333f17ee26#egg=fms_acceleration&subdirectory=plugins/framework
+-e git+https://github.com/foundation-model-stack/fms-acceleration.git@9cf70081f3bfc00e84331102e7d13b333f17ee26#egg=fms_acceleration_foak&subdirectory=plugins/fused-ops-and-kernels
+-e git+https://github.com/foundation-model-stack/fms-acceleration.git@9cf70081f3bfc00e84331102e7d13b333f17ee26#egg=fms_acceleration_moe&subdirectory=plugins/accelerated-moe
+-e git+https://github.com/foundation-model-stack/fms-acceleration.git@9cf70081f3bfc00e84331102e7d13b333f17ee26#egg=fms_acceleration_peft&subdirectory=plugins/accelerated-peft
+fms-hf-tuning @ git+https://github.com/foundation-model-stack/fms-hf-tuning.git@daca5510ab76cc8ecf0283fd31fc220697a75040
+fonttools==4.53.1
+frozenlist==1.4.1
+fsspec==2024.6.1
+huggingface-hub==0.24.6
+idna==3.7
+Jinja2==3.1.4
+kiwisolver==1.4.5
+markdown-it-py==3.0.0
+MarkupSafe==2.1.5
+matplotlib==3.9.2
+mdurl==0.1.2
+megablocks @ git+https://github.com/databricks/megablocks.git@bce5d7b2aaf5038bc93b36f76c2baf51c2939bd2
+mpmath==1.3.0
+multidict==6.0.5
+multiprocess==0.70.16
+networkx==3.3
+numpy==1.26.4
+nvidia-cublas-cu12==12.1.3.1
+nvidia-cuda-cupti-cu12==12.1.105
+nvidia-cuda-nvrtc-cu12==12.1.105
+nvidia-cuda-runtime-cu12==12.1.105
+nvidia-cudnn-cu12==8.9.2.26
+nvidia-cufft-cu12==11.0.2.54
+nvidia-curand-cu12==10.3.2.106
+nvidia-cusolver-cu12==11.4.5.107
+nvidia-cusparse-cu12==12.1.0.106
+nvidia-nccl-cu12==2.20.5
+nvidia-nvjitlink-cu12==12.6.20
+nvidia-nvtx-cu12==12.1.105
+packaging==24.1
+pandas==2.2.2
+peft==0.12.0
+pillow==10.4.0
+protobuf==5.27.3
+psutil==6.0.0
+pyarrow==17.0.0
+Pygments==2.18.0
+pyparsing==3.1.2
+python-dateutil==2.9.0.post0
+pytz==2024.1
+PyYAML==6.0.2
+regex==2024.7.24
+requests==2.32.3
+rich==13.7.1
+safetensors==0.4.4
+sentencepiece==0.2.0
+shtab==1.7.1
+simpleeval==0.9.13
+six==1.16.0
+stanford-stk @ git+https://git@github.com/stanford-futuredata/stk.git@a1ddf98466730b88a2988860a9d8000fd1833301
+sympy==1.13.2
+termcolor==2.4.0
+threadpoolctl==3.5.0
+tokenizers==0.19.1
+torch==2.3.1
+tqdm==4.66.5
+transformers==4.44.2
+triton==2.3.1
+trl==0.9.6
+typing_extensions==4.12.2
+tyro==0.8.8
+tzdata==2024.1
+urllib3==2.2.2
+xxhash==3.5.0
+yarl==1.9.4