Repository: hpcaitech/ColossalAI

Issue:
[DOC]: Clarify More about Datasets for Chat Training Stage 3
### 📚 The doc issue

As raised by some users, the current guide to training Stages 3 regarding the required datasets is yet clear. Some user has composed a [guide](https://juejin.cn/post/7222250219570757690) which still confuses the datasets. 

Some clarification is needed in this [section](https://github.com/hpcaitech/ColossalAI/tree/main/applications/Chat/examples#stage3---training-model-using-prompts-with-rl), regarding how to download and set `prompt_path` and `pretrain_dataset`.

Below is an example user query I received on Slack.

```
torchrun --standalone --nproc_per_node=4 train_prompts.py \
         --pretrain "bigscience/bloom-560m" \
         --model 'bloom' \
         --strategy colossalai_zero2 \
         --prompt_path /data/chenhao/train/ColossalAI/prompt_dataset/data.json \   -------Where is the data.json data
         --pretrain_dataset /data/chenhao/train/ColossalAI/pretrain_dataset/data.json \    -----------Where is the data.json data
         --rm_pretrain /data/chenhao/train/ColossalAI/Coati-7B \
         --rm_path /data/chenhao/train/ColossalAI/rmstatic.pt \
         --train_batch_size 4 \
         --experience_batch_size 4 \
         --max_epochs 1 \
         --num_episodes 1
```



[tensor] fix some unittests

[tensor] fix some unittests



Discussion:

Resolving patch:
diff --git a/applications/Chat/examples/train_prompts.py b/applications/Chat/examples/train_prompts.py
--- a/applications/Chat/examples/train_prompts.py
+++ b/applications/Chat/examples/train_prompts.py
@@ -139,7 +139,7 @@ def main(args):
 
     data_collator = DataCollatorForSupervisedDataset(tokenizer=tokenizer)
 
-    prompt_dataset = PromptDataset(tokenizer=tokenizer, data_path=args.prompt_path, max_datasets_size=16384)
+    prompt_dataset = PromptDataset(tokenizer=tokenizer, data_path=args.prompt_dataset, max_datasets_size=16384)
     if dist.is_initialized() and dist.get_world_size() > 1:
         prompt_sampler = DistributedSampler(prompt_dataset, shuffle=True, seed=42, drop_last=True)
     else:
@@ -204,7 +204,7 @@ def main(args):
 
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
-    parser.add_argument('--prompt_path', type=str, default=None, help='path to the prompt dataset')
+    parser.add_argument('--prompt_dataset', type=str, default=None, help='path to the prompt dataset')
     parser.add_argument('--pretrain_dataset', type=str, default=None, help='path to the pretrained dataset')
     parser.add_argument('--strategy',
                         choices=['naive', 'ddp', 'colossalai_gemini', 'colossalai_zero2'],


Test patch:
diff --git a/applications/Chat/examples/test_ci.sh b/applications/Chat/examples/test_ci.sh
--- a/applications/Chat/examples/test_ci.sh
+++ b/applications/Chat/examples/test_ci.sh
@@ -99,7 +99,7 @@ torchrun --standalone --nproc_per_node=2 ${BASE}/train_reward_model.py \
 
 rm -rf ${BASE}/rm_ckpt.pt
 
-torchrun --standalone --nproc_per_node=2 ${BASE}/train_prompts.py --prompt_path $PROMPT_PATH --pretrain_dataset $PRETRAIN_DATASET \
+torchrun --standalone --nproc_per_node=2 ${BASE}/train_prompts.py --prompt_dataset $PROMPT_PATH --pretrain_dataset $PRETRAIN_DATASET \
         --strategy colossalai_zero2 --num_episodes 1 --max_timesteps 2 \
         --update_timesteps 2 --max_epochs 1 --train_batch_size 2 \
         --pretrain 'facebook/opt-350m' --model opt \
@@ -108,7 +108,7 @@ torchrun --standalone --nproc_per_node=2 ${BASE}/train_prompts.py --prompt_path
         --save_path ${BASE}/actor_checkpoint_prompts.pt
 rm -rf ${BASE}/rm_ckpt_opt.pt
 
-torchrun --standalone --nproc_per_node=2 ${BASE}/train_prompts.py --prompt_path $PROMPT_PATH --pretrain_dataset $PRETRAIN_DATASET \
+torchrun --standalone --nproc_per_node=2 ${BASE}/train_prompts.py --prompt_dataset $PROMPT_PATH --pretrain_dataset $PRETRAIN_DATASET \
          --strategy colossalai_zero2 --num_episodes 1 --max_timesteps 2 \
          --update_timesteps 2 --max_epochs 1 --train_batch_size 2 \
          --pretrain 'gpt2' --model gpt2 \