ColossalAI/applications/Chat/benchmarks
Wenhao Chen 7b9b86441f
[chat]: update rm, add wandb and fix bugs (#4471)
* feat: modify forward fn of critic and reward model

* feat: modify calc_action_log_probs

* to: add wandb in sft and rm trainer

* feat: update train_sft

* feat: update train_rm

* style: modify type annotation and add warning

* feat: pass tokenizer to ppo trainer

* to: modify trainer base and maker base

* feat: add wandb in ppo trainer

* feat: pass tokenizer to generate

* test: update generate fn tests

* test: update train tests

* fix: remove action_mask

* feat: remove unused code

* fix: fix wrong ignore_index

* fix: fix mock tokenizer

* chore: update requirements

* revert: modify make_experience

* fix: fix inference

* fix: add padding side

* style: modify _on_learn_batch_end

* test: use mock tokenizer

* fix: use bf16 to avoid overflow

* fix: fix workflow

* [chat] fix gemini strategy

* [chat] fix

* sync: update colossalai strategy

* fix: fix args and model dtype

* fix: fix checkpoint test

* fix: fix requirements

* fix: fix missing import and wrong arg

* fix: temporarily skip gemini test in stage 3

* style: apply pre-commit

* fix: temporarily skip gemini test in stage 1&2

---------

Co-authored-by: Mingyan Jiang <1829166702@qq.com>
2023-09-20 15:53:58 +08:00
..
ray [misc] update pre-commit and run all files (#4752) 2023-09-19 14:20:26 +08:00
README.md [doc] update Coati README (#4405) 2023-08-14 15:26:27 +08:00
benchmark_opt_lora_dummy.py [chat]: update rm, add wandb and fix bugs (#4471) 2023-09-20 15:53:58 +08:00

README.md

Benchmarks

Benchmark OPT with LoRA on dummy prompt data

We provide various OPT models (string in parentheses is the corresponding model name used in this script):

  • OPT-125M (125m)
  • OPT-350M (350m)
  • OPT-700M (700m)
  • OPT-1.3B (1.3b)
  • OPT-2.7B (2.7b)
  • OPT-3.5B (3.5b)
  • OPT-5.5B (5.5b)
  • OPT-6.7B (6.7b)
  • OPT-10B (10b)
  • OPT-13B (13b)

We also provide various training strategies:

  • ddp: torch DDP
  • colossalai_gemini: ColossalAI GeminiDDP with placement_policy="cuda", like zero3
  • colossalai_gemini_cpu: ColossalAI GeminiDDP with placement_policy="cpu", like zero3-offload
  • colossalai_zero2: ColossalAI zero2
  • colossalai_zero2_cpu: ColossalAI zero2-offload
  • colossalai_zero1: ColossalAI zero1
  • colossalai_zero1_cpu: ColossalAI zero1-offload

We only support torchrun to launch now. E.g.

# run OPT-125M with no lora (lora_rank=0) on single-node single-GPU with min batch size
torchrun --standalone --nproc_per_node 1 benchmark_opt_lora_dummy.py \
    --model 125m --critic_model 125m --strategy ddp \
    --experience_batch_size 1 --train_batch_size 1 --lora_rank 0
# run Actor (OPT-1.3B) and Critic (OPT-350M) with lora_rank=4 on single-node 4-GPU
torchrun --standalone --nproc_per_node 4 benchmark_opt_lora_dummy.py \
    --model 1.3b --critic_model 350m --strategy colossalai_zero2 --lora_rank 4