Skip to content

Latest commit

 

History

History
131 lines (100 loc) · 4.97 KB

README.md

File metadata and controls

131 lines (100 loc) · 4.97 KB
Raw

Activation Compression with Guarantees

We explore how to fine-tune language models over slow networks using activation compression with guarantees. This is a research project developed by DS3Lab@ETH Zurich and HazyResearch@Stanford.

Cite Our Paper

@article{jue2022fine,
  title={Fine-tuning Language Models over Slow Networks using Activation Compression with Guarantees}, 
  author={Jue Wang, Binhang Yuan, Luka Rimanic, Yongjun He, Tri Dao, Beidi Chen, Christopher Re, Ce Zhang},
  year={2022},
  eprint={2206.01299},
  archivePrefix={arXiv},
  primaryClass={cs.DC}
}

Setup:

  • Create environment:

    conda create -n acsgd python=3.8
conda activate acsgd

  • Install PyTorch env:

    pip3 install torch==1.9.0+cu111 torchtext -f https://download.pytorch.org/whl/torch_stable.html

# Magic, not sure why cupy-cuda111 would not work, it seems that cupy-cuda111 will use different PTX from torch.
pip3 install cupy-cuda110==8.6.0

    Other dependencies:

    pip3 install datasets==2.2.2
pip3 install transformers==4.19.2
pip3 install sentencepiece==0.1.96 # required by deberta

  • Setup network configuration:

    export GLOO_SOCKET_IFNAME=ens3
export NCCL_SOCKET_IFNAME=ens3

Run Distributed Gpipe:

  • Partition the pre-trained model:

    # gpt2
python convert_gpt2_checkpoint --model-name gpt2-xl --save-dir checkpoints/
    
# or deberta 
python convert_deberta_checkpoint --model-name deberta-v2-xxl --save-dir checkpoints/

  • On each node, run:

    # gpt2
python dist_lm_runner.py $(echo ${ARGS}) --cuda-id 0 --rank i # (i=0,...,N-1)
    
# or deberta
python dist_deberta_runner.py $(echo ${ARGS}) --cuda-id 0 --rank i # (i=0,...,N-1)

    where "ARGS" contains training-related configurations, which should remain the same across nodes. An example could be:

    ARGS="--model-name checkpoints/gpt2-xl \
  --tokenizer-name gpt2-xl \
  --load-pretrained-model true \
  --task-name wikitext --n-epochs 10 --warmup-epochs 1 \
  --num-layers 6 --num-heads 25 --embedding-dim 1600 \
  --num-iters 10000000 --lr 5e-5 --seq-length 1024 --batch-size 32 --micro-batch-size 1 \
  --forward-compress-method delta \
  --forward-bits 4 \
  --backward-compress-method fixpoint \
  --backward-bits 8 \
  --dist-url tcp://XXX.XXX.XXX.XXX:9000 \
  --world-size N --pipeline-group-size N \
  --pp-mode gpipe --profiling no-profiling --do-evaluation true"

    Modify "--dist-url", "--world-size" and "--pipeline-group-size" before running.

    Complete examples can be found "./run_lm.sh" and "./run_deberta.sh".

Arguments

Distributed Related

  • "--dist-url": tcp://XXX.XXX.XXX.XXX:9000
  • "--world-size": number of nodes that participate in the training.
  • "--pipeline-group-size": number of nodes that perform pipeline parallelism.
  • "--data-group-size": number of nodes that perform data parallelism.
  • "--rank": the rank of the current node. (0, ..., world_size-1)
  • "--profiling": "no-profiling" or "tidy_profiling". If "tidy_profiling", a trace file will be generated in "./trace_json/", which can be visualized with "chrome://tracing/".

Compression Related

  • "--forward-compress-method": "none", "fixpoint", "delta", or "delta-lowbits".
    • "none": do not compress.
    • "fixpoint": direct compress the activations. need to specify `"--forward-bits".
    • "delta": compress and communicate the delta of activations. need to specify "--forward-bits" and "--max-activation-cache-size".
    • "delta-lowbits": in addition to "delta", it also compresses the local cache (previous activations). need to specify "--forward-bits", "--forward-bits-act", and "--max-activation-cache-size".
  • "--backward-compress-method": "none" or "fixpoint".
    • "none": do not compress.
    • "fixpoint": direct compress the gradients. need to specify "--backward-bits".

Training Related

  • "--batch-size": macro-batch size.
  • "--micro-batch-size ": micro-batch-size. The macro-batch size should be divisible by micro-batch-size.
  • "--lr": the peak learning rate.
  • "--n-epochs": number of training epochs.
  • "--warmup-epochs": number of epochs for uncompressed training (transfer full-precision activations and gradients).
  • "--warmup-steps": number of training steps where the learning rate grows from 0 to "--lr". Default to be one training epoch.
  • "--do-evaluation": whether do evaluation during training.

Model Related

  • "--model-name": Name or path of the pretrained checkpoint. Usually should be a path to the checkpoint generated by "convert_xxx_checkpoint.py".
  • "--tokenizer-name": Name or path of the tokenizer.
  • "--load-pretrained-model": whether to load the pretrained checkpoint. The checkpoint should be generated by "convert_xxx_checkpoint.py".
  • "--num-layers", "--num-heads", "--embedding-dim" should be inline with the configuration of "--model-name".