AntCPLab / puma_benchmarks

Puma_benchmarks

This repo is to reproduce the results in our paper Puma

0. Compile and Launch SecretFlow-SPU

git clone SecretFlow-SPU

git clone https://github.com/secretflow/spu.git & cd spu

and follow SecretFlow-SPU README.md to build spu from source.

1. Launch SPU-backends:

In Puma, we launch 5 nodes, 3 for ABY3 computing nodes, 1 for model provider, and the last one for input provider. On each node, run the following

bazel run -c opt //examples/python/utils:nodectl -- --config `pwd`/examples/python/ml/puma_bert_benchmarks/3pc.json up

To profile the costs of Puma, turn on the following in 3pc.json

"enable_pphlo_profile": true,
"enable_hal_profile": true,

and you will get the time and communication costs on each ABY3 node.

2. Install huggingface transformers library

pip install 'transformers[flax]'

To hijack, we need to modify the activation function of Bert and GPT2 as jax.nn.gelu.

3. Config and Run Puma

After launching SPU, move the directory of puma_bert_benchmarks and puma_gpt2_benchmarks to spu/examples/python/ml and run the scripts. For example, run

bazel run -c opt //examples/python/ml/puma_bert_benchmarks:puma_bert_benchmarks

to re-produce our results for Bert on GLUE. Also, you can modify the model path and task_name easily in the python scripts to run Puma on more model and dataset. For example, Puma achieves the following performance on Bert-base on GLUE-CoLA:

MCC	Time	Comm.
0.613	42s	3.59GB

To run Puma on LLaMA-7B, please refer to: https://github.com/secretflow/spu/tree/main/examples/python/ml/flax_llama7b