Repository: hpcaitech/ColossalAI

Issue:
[tensor] fix some unittests

[tensor] fix some unittests

[tensor] fix some unittests



Discussion:

Resolving patch:
diff --git a/applications/Chat/coati/models/gpt/gpt_actor.py b/applications/Chat/coati/models/gpt/gpt_actor.py
--- a/applications/Chat/coati/models/gpt/gpt_actor.py
+++ b/applications/Chat/coati/models/gpt/gpt_actor.py
@@ -23,7 +23,8 @@ def __init__(self,
                  config: Optional[GPT2Config] = None,
                  checkpoint: bool = False,
                  lora_rank: int = 0,
-                 lora_train_bias: str = 'none') -> None:
+                 lora_train_bias: str = 'none',
+                 **kwargs) -> None:
         if pretrained is not None:
             model = GPT2LMHeadModel.from_pretrained(pretrained)
         elif config is not None:
@@ -32,4 +33,4 @@ def __init__(self,
             model = GPT2LMHeadModel(GPT2Config())
         if checkpoint:
             model.gradient_checkpointing_enable()
-        super().__init__(model, lora_rank, lora_train_bias)
+        super().__init__(model, lora_rank, lora_train_bias, **kwargs)
diff --git a/applications/Chat/coati/models/gpt/gpt_critic.py b/applications/Chat/coati/models/gpt/gpt_critic.py
--- a/applications/Chat/coati/models/gpt/gpt_critic.py
+++ b/applications/Chat/coati/models/gpt/gpt_critic.py
@@ -24,7 +24,8 @@ def __init__(self,
                  config: Optional[GPT2Config] = None,
                  checkpoint: bool = False,
                  lora_rank: int = 0,
-                 lora_train_bias: str = 'none') -> None:
+                 lora_train_bias: str = 'none',
+                 **kwargs) -> None:
         if pretrained is not None:
             model = GPT2Model.from_pretrained(pretrained)
         elif config is not None:
@@ -34,4 +35,4 @@ def __init__(self,
         if checkpoint:
             model.gradient_checkpointing_enable()
         value_head = nn.Linear(model.config.n_embd, 1)
-        super().__init__(model, value_head, lora_rank, lora_train_bias)
+        super().__init__(model, value_head, lora_rank, lora_train_bias, **kwargs)


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
@@ -2,11 +2,21 @@
 
 set -xue
 
+if [ -z "$SFT_DATASET" ]; then
+    echo "Please set \$SFT_DATASET to the path to sft dataset."
+    exit 1
+fi
+
 if [ -z "$PROMPT_PATH" ]; then
     echo "Please set \$PROMPT_PATH to the path to prompts csv."
     exit 1
 fi
 
+if [ -z "$PRETRAIN_DATASET" ]; then
+    echo "Please set \$PRETRAIN_DATASET to the path to alpaca data."
+    exit 1
+fi
+
 BASE=$(realpath $(dirname $0))
 
 export OMP_NUM_THREADS=8
@@ -14,104 +24,97 @@ export OMP_NUM_THREADS=8
 # install requirements
 pip install -r ${BASE}/requirements.txt
 
-# train dummy
-python ${BASE}/train_dummy.py --strategy naive --num_episodes 1 \
-                              --max_timesteps 2 --update_timesteps 2 \
-                              --max_epochs 1 --train_batch_size 2 --lora_rank 4
-
-torchrun --standalone --nproc_per_node=2 ${BASE}/train_dummy.py \
-         --strategy colossalai_gemini --num_episodes 1 --max_timesteps 2 \
-         --update_timesteps 2 --max_epochs 1 --train_batch_size 2\
-         --pretrain 'facebook/opt-350m' --model opt --lora_rank 4\
-         --save_path ${BASE}/actor_checkpoint_dummy.pt
-python ${BASE}/inference.py --model_path ${BASE}/actor_checkpoint_dummy.pt --pretrain 'facebook/opt-350m' --model opt
-
-torchrun --standalone --nproc_per_node=2 ${BASE}/train_dummy.py \
-         --strategy ddp --num_episodes 1 --max_timesteps 2 \
-         --update_timesteps 2 --max_epochs 1 --train_batch_size 2\
-         --pretrain 'facebook/opt-350m' --model opt --lora_rank 4\
-         --save_path ${BASE}/actor_checkpoint_dummy.pt
-python ${BASE}/inference.py --model_path ${BASE}/actor_checkpoint_dummy.pt --pretrain 'facebook/opt-350m' --model opt
-
-torchrun --standalone --nproc_per_node=2 ${BASE}/train_dummy.py \
-         --strategy colossalai_zero2 --num_episodes 1 --max_timesteps 2 \
-         --update_timesteps 2 --max_epochs 1 --train_batch_size 2\
-         --pretrain 'gpt2' --model gpt2 --lora_rank 4\
-         --save_path ${BASE}/actor_checkpoint_dummy.pt
-python ${BASE}/inference.py --model_path ${BASE}/actor_checkpoint_dummy.pt --pretrain 'gpt2' --model gpt2
+wandb init -m offline
 
-torchrun --standalone --nproc_per_node=2 ${BASE}/train_dummy.py \
-         --strategy colossalai_zero2 --num_episodes 1 --max_timesteps 2 \
-         --update_timesteps 2 --max_epochs 1 --train_batch_size 2\
-         --pretrain 'roberta-base' --model roberta --lora_rank 4\
-         --save_path ${BASE}/actor_checkpoint_dummy.pt
-python ${BASE}/inference.py --model_path ${BASE}/actor_checkpoint_dummy.pt --pretrain 'roberta-base' --model roberta
+# train sft
+torchrun --standalone --nproc_per_node=4 ${BASE}/train_sft.py --pretrain 'bigscience/bloom-560m' \
+        --model 'bloom' --strategy colossalai_zero2 --lora_rank 4\
+        --dataset $SFT_DATASET --max_datasets_size 512 --max_epochs 1 \
+        --save_path ${BASE}/output
 
-rm -rf ${BASE}/actor_checkpoint_dummy.pt
+torchrun --standalone --nproc_per_node=4 ${BASE}/train_sft.py --pretrain 'gpt2' \
+        --model 'gpt2' --strategy colossalai_zero2 \
+        --dataset $SFT_DATASET --max_datasets_size 512 --max_epochs 1 \
+        --save_path ${BASE}/output
 
-# train prompts
-python ${BASE}/train_prompts.py $PROMPT_PATH --strategy naive --num_episodes 1 \
-                                             --max_timesteps 2 --update_timesteps 2 \
-                                             --max_epochs 1 --train_batch_size 2 --lora_rank 4
+torchrun --standalone --nproc_per_node=4 ${BASE}/train_sft.py --pretrain 'facebook/opt-350m' \
+        --model 'opt' --strategy colossalai_zero2 --lora_rank 4\
+        --dataset $SFT_DATASET --max_datasets_size 512 --max_epochs 1 \
+        --save_path ${BASE}/output
 
-torchrun --standalone --nproc_per_node=2 ${BASE}/train_prompts.py $PROMPT_PATH \
-         --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 --lora_rank 4\
-         --save_path ${BASE}/actor_checkpoint_prompts.pt
-python ${BASE}/inference.py --model_path ${BASE}/actor_checkpoint_prompts.pt --pretrain 'facebook/opt-350m' --model opt
+torchrun --standalone --nproc_per_node=4 ${BASE}/train_sft.py --pretrain 'gpt2' \
+        --model 'gpt2' --strategy ddp --lora_rank 4\
+        --dataset $SFT_DATASET --max_datasets_size 512 --max_epochs 1 \
+        --save_path ${BASE}/output
 
-torchrun --standalone --nproc_per_node=2 ${BASE}/train_prompts.py $PROMPT_PATH \
-         --strategy ddp --num_episodes 1 --max_timesteps 2 \
-         --update_timesteps 2 --max_epochs 1 --train_batch_size 2\
-         --pretrain 'gpt2' --model gpt2 --lora_rank 4\
-         --save_path ${BASE}/actor_checkpoint_prompts.pt
-python ${BASE}/inference.py --model_path ${BASE}/actor_checkpoint_prompts.pt --pretrain 'gpt2' --model gpt2
+#torchrun --standalone --nproc_per_node=4 ${BASE}/train_sft.py --pretrain 'facebook/opt-350m' \
+#        --model 'opt' --strategy naive \
+#        --dataset $SFT_DATASET --max_datasets_size 512 --max_epochs 1 \
+#        --save_path ${BASE}/output
 
-torchrun --standalone --nproc_per_node=2 ${BASE}/train_prompts.py $PROMPT_PATH \
-         --strategy colossalai_gemini --num_episodes 1 --max_timesteps 2 \
-         --update_timesteps 2 --max_epochs 1 --train_batch_size 2\
-         --pretrain 'gpt2' --model gpt2 --lora_rank 4\
-         --save_path ${BASE}/actor_checkpoint_prompts.pt
-python ${BASE}/inference.py --model_path ${BASE}/actor_checkpoint_prompts.pt --pretrain 'gpt2' --model gpt2
-
-torchrun --standalone --nproc_per_node=2 ${BASE}/train_prompts.py $PROMPT_PATH \
-         --strategy colossalai_zero2 --num_episodes 1 --max_timesteps 2 \
-         --update_timesteps 2 --max_epochs 1 --train_batch_size 2\
-         --pretrain 'roberta-base' --model roberta --lora_rank 4\
-         --save_path ${BASE}/actor_checkpoint_prompts.pt
-python ${BASE}/inference.py --model_path ${BASE}/actor_checkpoint_prompts.pt --pretrain 'roberta-base' --model roberta
-
-rm -rf ${BASE}/actor_checkpoint_prompts.pt
+rm -rf ${BASE}/output
 
 # train rm
 torchrun --standalone --nproc_per_node=2 ${BASE}/train_reward_model.py \
-                             --pretrain 'facebook/opt-350m' --model 'opt' \
-                             --strategy colossalai_zero2 --loss_fn 'log_sig'\
-                             --dataset 'Anthropic/hh-rlhf' --subset 'harmless-base'\
-                             --test True --lora_rank 4
+                            --pretrain 'facebook/opt-350m' --model 'opt' \
+                            --strategy colossalai_zero2 --loss_fn 'log_sig'\
+                            --dataset 'Anthropic/hh-rlhf' --subset 'harmless-base' \
+                            --test True --lora_rank 4 \
+                            --save_path ${BASE}/rm_ckpt_opt.pt
+
+torchrun --standalone --nproc_per_node=2 ${BASE}/train_reward_model.py \
+                            --pretrain 'gpt2' --model 'gpt2' \
+                            --strategy colossalai_zero2 --loss_fn 'log_exp' \
+                            --dataset 'Dahoas/rm-static' \
+                            --test True  --lora_rank 4 \
+                            --save_path ${BASE}/rm_ckpt_gpt.pt
 
 torchrun --standalone --nproc_per_node=2 ${BASE}/train_reward_model.py \
-                             --pretrain 'gpt2' --model 'gpt2' \
-                             --strategy colossalai_gemini --loss_fn 'log_exp'\
-                             --dataset 'Dahoas/rm-static' --test True --lora_rank 4
+                            --pretrain 'gpt2' --model 'gpt2' \
+                            --strategy ddp --loss_fn 'log_exp' \
+                            --dataset 'Dahoas/rm-static' \
+                            --test True --lora_rank 4 \
+                            --save_path ${BASE}/rm_ckpt.pt
 
 torchrun --standalone --nproc_per_node=2 ${BASE}/train_reward_model.py \
-                             --pretrain 'bigscience/bloom-560m' --model 'bloom' \
-                             --strategy colossalai_zero2 --loss_fn 'log_sig'\
-                             --dataset 'Anthropic/hh-rlhf' --subset 'harmless-base'\
-                             --test True --lora_rank 4
+                            --pretrain 'bigscience/bloom-560m' --model 'bloom' \
+                            --strategy colossalai_zero2 --loss_fn 'log_sig' \
+                            --dataset 'Anthropic/hh-rlhf' --subset 'harmless-base' \
+                            --test True --lora_rank 4 \
+                            --save_path ${BASE}/rm_ckpt.pt
 
 torchrun --standalone --nproc_per_node=2 ${BASE}/train_reward_model.py \
-                             --pretrain 'microsoft/deberta-v3-large' --model 'deberta' \
-                             --strategy colossalai_zero2 --loss_fn 'log_sig'\
-                             --dataset 'Anthropic/hh-rlhf' --subset 'harmless-base'\
-                             --test True --lora_rank 4
+                            --pretrain 'microsoft/deberta-v3-large' --model 'deberta' \
+                            --strategy colossalai_zero2 --loss_fn 'log_sig' \
+                            --dataset 'Anthropic/hh-rlhf' --subset 'harmless-base' \
+                            --test True --lora_rank 4 \
+                            --save_path ${BASE}/rm_ckpt.pt
 
 torchrun --standalone --nproc_per_node=2 ${BASE}/train_reward_model.py \
-                             --pretrain 'roberta-base' --model 'roberta' \
-                             --strategy colossalai_zero2 --loss_fn 'log_exp'\
-                             --dataset 'Anthropic/hh-rlhf' --subset 'harmless-base'\
-                             --test True --lora_rank 4
+                            --pretrain 'roberta-base' --model 'roberta' \
+                            --strategy colossalai_zero2 --loss_fn 'log_exp'\
+                            --dataset 'Anthropic/hh-rlhf' --subset 'harmless-base'\
+                            --test True --lora_rank 4 \
+                            --save_path ${BASE}/rm_ckpt.pt
 
 rm -rf ${BASE}/rm_ckpt.pt
+
+torchrun --standalone --nproc_per_node=2 ${BASE}/train_prompts.py --prompt_path $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 \
+        --rm_pretrain 'facebook/opt-350m' \
+        --rm_path ${BASE}/rm_ckpt_opt.pt \
+        --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 \
+         --strategy colossalai_zero2 --num_episodes 1 --max_timesteps 2 \
+         --update_timesteps 2 --max_epochs 1 --train_batch_size 2 \
+         --pretrain 'gpt2' --model gpt2 \
+         --rm_pretrain 'gpt2' \
+         --rm_path ${BASE}/rm_ckpt_gpt.pt \
+         --save_path ${BASE}/actor_checkpoint_prompts.pt
+rm -rf ${BASE}/rm_ckpt_gpt.pt
+
+rm -rf ${BASE}/actor_checkpoint_prompts.pt
\ No newline at end of file