[chat]fix readme (#3429)

* fix stage 2

fix stage 2

* add torch
pull/3368/head
kingkingofall 2023-04-06 10:58:53 +08:00 committed by GitHub
parent 7d8d825681
commit 57a3c4db6d
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 2 additions and 1 deletions

View File

@ -57,7 +57,7 @@ You can run the `examples/train_rm.sh` to start a reward model training.
You can also use the following cmd to start training a reward model.
```
torchrun --standalone --nproc_per_node=4 train_reward_model.py
torchrun --standalone --nproc_per_node=4 train_reward_model.py \
--pretrain "/path/to/LLaMa-7B/" \
--model 'llama' \
--strategy colossalai_zero2 \

View File

@ -51,6 +51,7 @@ Please ensure you have downloaded HF-format model weights of LLaMA models.
Usage:
```python
import torch
from transformers import LlamaForCausalLM
USE_8BIT = True # use 8-bit quantization; otherwise, use fp16