From ed1dbb26c7f200cc12f9618a6cafc23c8c1c04c6 Mon Sep 17 00:00:00 2001
From: CWHer <cwher@outlook.com>
Date: Sun, 25 Jun 2023 18:09:32 +0800
Subject: [PATCH 01/21] feat: remove on_learn_epoch fn as not used

---
 applications/Chat/coati/trainer/callbacks/base.py | 6 ------
 1 file changed, 6 deletions(-)

diff --git a/applications/Chat/coati/trainer/callbacks/base.py b/applications/Chat/coati/trainer/callbacks/base.py
index f5616048855b..4351c0079b8d 100644
--- a/applications/Chat/coati/trainer/callbacks/base.py
+++ b/applications/Chat/coati/trainer/callbacks/base.py
@@ -26,12 +26,6 @@ def on_make_experience_start(self) -> None:
     def on_make_experience_end(self, experience: Experience) -> None:
         pass
 
-    def on_learn_epoch_start(self, epoch: int) -> None:
-        pass
-
-    def on_learn_epoch_end(self, epoch: int) -> None:
-        pass
-
     def on_learn_batch_start(self) -> None:
         pass
 

From a7edd3e1387458b094f3718e0df1a2f1ad5e1b14 Mon Sep 17 00:00:00 2001
From: CWHer <cwher@outlook.com>
Date: Mon, 26 Jun 2023 11:50:09 +0800
Subject: [PATCH 02/21] revert: add _on_learn_epoch fn

---
 applications/Chat/coati/trainer/callbacks/base.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/applications/Chat/coati/trainer/callbacks/base.py b/applications/Chat/coati/trainer/callbacks/base.py
index 4351c0079b8d..f5616048855b 100644
--- a/applications/Chat/coati/trainer/callbacks/base.py
+++ b/applications/Chat/coati/trainer/callbacks/base.py
@@ -26,6 +26,12 @@ def on_make_experience_start(self) -> None:
     def on_make_experience_end(self, experience: Experience) -> None:
         pass
 
+    def on_learn_epoch_start(self, epoch: int) -> None:
+        pass
+
+    def on_learn_epoch_end(self, epoch: int) -> None:
+        pass
+
     def on_learn_batch_start(self) -> None:
         pass
 

From 9069341b66b49bfe2545b8ab141f751862553313 Mon Sep 17 00:00:00 2001
From: CWHer <cwher@outlook.com>
Date: Mon, 26 Jun 2023 15:15:50 +0800
Subject: [PATCH 03/21] feat: remove NaiveStrategy

---
 applications/Chat/coati/trainer/strategies/ddp.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/applications/Chat/coati/trainer/strategies/ddp.py b/applications/Chat/coati/trainer/strategies/ddp.py
index e1c1bbf19f35..2fdfebd2a6a8 100644
--- a/applications/Chat/coati/trainer/strategies/ddp.py
+++ b/applications/Chat/coati/trainer/strategies/ddp.py
@@ -63,7 +63,7 @@ def _post_init(self) -> None:
             f'{type(self).__name__}\'s plugin is not initialized properly.'
 
     def setup_distributed(self) -> None:
-        self._try_init_dist(force=True)
+        self._try_init_dist(force=False)
         self.set_seed(self.seed)
 
     def set_seed(self, seed: int) -> None:

From 436e52a509096a922cd70b64ee81a8b70d75bfb9 Mon Sep 17 00:00:00 2001
From: CWHer <cwher@outlook.com>
Date: Wed, 28 Jun 2023 11:49:59 +0800
Subject: [PATCH 04/21] test: update train_prompts tests

---
 applications/Chat/examples/test_ci.sh | 17 +++++++++++------
 1 file changed, 11 insertions(+), 6 deletions(-)

diff --git a/applications/Chat/examples/test_ci.sh b/applications/Chat/examples/test_ci.sh
index dec1f7c036c8..b498b4e79236 100755
--- a/applications/Chat/examples/test_ci.sh
+++ b/applications/Chat/examples/test_ci.sh
@@ -43,18 +43,23 @@ pip install -r ${BASE}/requirements.txt
 
 wandb init -m offline
 
-# FIXME: This is a hack to skip tests that are not working (tested at commit b3ab7fbabf)
+# FIXME: This is a hack to skip tests that are not working
 #  - gpt2-ddp: RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation
-#  - llama-*: Repository Not Found for url: https://huggingface.co/{...}/resolve/main/tokenizer.model.
-#  - roberta-*: RuntimeError: CUDA error: CUBLAS_STATUS_NOT_INITIALIZED when calling `cublasCreate(handle)`
+#  - llama-ddp, llama-gemini: Out of memory
+#  - roberta-ddp: RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation
+#  - roberta-gemini: RuntimeError: ZERO DDP error: the synchronization of gradients doesn't exit properly.
+#  - roberta-zero2: "zero/low_level/low_level_optim.py", line 498, assert param_shape == flat_master_avg_grads.shape
 SKIPPED_TESTS=(
     "gpt2-ddp"
-    "llama-ddp" "llama-colossalai_gemini" "llama-colossalai_zero2"
-    "roberta-ddp" "roberta-colossalai_gemini" "roberta-colossalai_zero2"
+    "llama-ddp"
+    "llama-gemini"
+    "roberta-ddp"
+    "roberta-colossalai_gemini"
+    "roberta-colossalai_zero2"
 )
 
 # These tests are quick and do not have any dependencies
-for model in 'gpt2' 'bloom' 'opt' 'llama' 'roberta'; do
+for model in 'gpt2' 'bloom' 'opt' 'roberta' 'llama'; do
     for strategy in 'ddp' 'colossalai_gemini' 'colossalai_zero2'; do
         if [[ " ${SKIPPED_TESTS[*]} " =~ " ${model}-${strategy} " ]]; then
             echo "[Test]: Skipped $model-$strategy"

From 20e704c5a90117e914de2f68a295df328008adca Mon Sep 17 00:00:00 2001
From: CWHer <cwher@outlook.com>
Date: Wed, 28 Jun 2023 11:55:36 +0800
Subject: [PATCH 05/21] fix: remove prepare_llama_tokenizer_and_embedding

---
 applications/Chat/examples/train_prompts.py   |  7 ++---
 .../Chat/examples/train_reward_model.py       |  7 ++---
 applications/Chat/examples/train_sft.py       | 26 ++++++++-----------
 3 files changed, 15 insertions(+), 25 deletions(-)

diff --git a/applications/Chat/examples/train_prompts.py b/applications/Chat/examples/train_prompts.py
index c748eeb21065..13e6cf00c0f7 100644
--- a/applications/Chat/examples/train_prompts.py
+++ b/applications/Chat/examples/train_prompts.py
@@ -122,17 +122,14 @@ def main(args):
     elif args.model == 'opt':
         tokenizer = AutoTokenizer.from_pretrained("facebook/opt-350m")
     elif args.model == 'llama':
-        tokenizer = LlamaTokenizer.from_pretrained(args.pretrain)
+        tokenizer = LlamaTokenizer.from_pretrained("hf-internal-testing/llama-tokenizer")
         tokenizer.eos_token = '<\s>'
     elif args.model == 'roberta':
         tokenizer = RobertaTokenizer.from_pretrained("roberta-base")
     else:
         raise ValueError(f'Unsupported model "{args.model}"')
 
-    if args.model == 'llama':
-        tokenizer = prepare_llama_tokenizer_and_embedding(tokenizer, actor)
-    else:
-        tokenizer.pad_token = tokenizer.eos_token
+    tokenizer.pad_token = tokenizer.eos_token
 
     data_collator = DataCollatorForSupervisedDataset(tokenizer=tokenizer)
 
diff --git a/applications/Chat/examples/train_reward_model.py b/applications/Chat/examples/train_reward_model.py
index e9618e0c1d5e..933139db134f 100644
--- a/applications/Chat/examples/train_reward_model.py
+++ b/applications/Chat/examples/train_reward_model.py
@@ -71,17 +71,14 @@ def train(args):
     elif args.model == 'deberta':
         tokenizer = DebertaV2Tokenizer.from_pretrained('microsoft/deberta-v3-large')
     elif args.model == 'llama':
-        tokenizer = LlamaTokenizer.from_pretrained(args.pretrain)
+        tokenizer = LlamaTokenizer.from_pretrained("hf-internal-testing/llama-tokenizer")
     elif args.model == 'roberta':
         tokenizer = RobertaTokenizer.from_pretrained("roberta-base")
     else:
         raise ValueError(f'Unsupported model "{args.model}"')
     max_len = args.max_len
 
-    if args.model == 'llama':
-        tokenizer = prepare_llama_tokenizer_and_embedding(tokenizer, model)
-    else:
-        tokenizer.pad_token = tokenizer.eos_token
+    tokenizer.pad_token = tokenizer.eos_token
 
     # configure optimizer
     if args.strategy.startswith('colossalai'):
diff --git a/applications/Chat/examples/train_sft.py b/applications/Chat/examples/train_sft.py
index 30becd8a68a1..0e8f34554536 100644
--- a/applications/Chat/examples/train_sft.py
+++ b/applications/Chat/examples/train_sft.py
@@ -70,7 +70,7 @@ def train(args):
         tokenizer = AutoTokenizer.from_pretrained("facebook/opt-350m")
     elif args.model == 'llama':
         tokenizer = AutoTokenizer.from_pretrained(
-            args.pretrain,
+            "hf-internal-testing/llama-tokenizer",
             padding_side="right",
             use_fast=False,
         )
@@ -79,20 +79,16 @@ def train(args):
         raise ValueError(f'Unsupported model "{args.model}"')
     tokenizer.pad_token = tokenizer.eos_token
     max_len = args.max_len
-    if args.model == 'llama':
-        tokenizer = prepare_llama_tokenizer_and_embedding(tokenizer, model)
-
-        if args.strategy == 'colossalai_gemini':
-            # this is a hack to deal with the resized embedding
-            # to make sure all parameters are ColoParameter for Colossal-AI Gemini Compatibility
-            for name, param in model.named_parameters():
-                if not isinstance(param, ColoParameter):
-                    sub_module_name = '.'.join(name.split('.')[:-1])
-                    weight_name = name.split('.')[-1]
-                    sub_module = model.get_submodule(sub_module_name)
-                    setattr(sub_module, weight_name, ColoParameter(param))
-    else:
-        tokenizer.pad_token = tokenizer.eos_token
+
+    if args.model == 'llama' and args.strategy == 'colossalai_gemini':
+        # this is a hack to deal with the resized embedding
+        # to make sure all parameters are ColoParameter for Colossal-AI Gemini Compatibility
+        for name, param in model.named_parameters():
+            if not isinstance(param, ColoParameter):
+                sub_module_name = '.'.join(name.split('.')[:-1])
+                weight_name = name.split('.')[-1]
+                sub_module = model.get_submodule(sub_module_name)
+                setattr(sub_module, weight_name, ColoParameter(param))
 
     # configure optimizer
     if args.strategy.startswith('colossalai'):

From 5ae6da1542487a5d28b0512920b2c0fcdb014d3d Mon Sep 17 00:00:00 2001
From: CWHer <cwher@outlook.com>
Date: Wed, 28 Jun 2023 12:01:57 +0800
Subject: [PATCH 06/21] test: add lora arg

---
 applications/Chat/examples/test_ci.sh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/applications/Chat/examples/test_ci.sh b/applications/Chat/examples/test_ci.sh
index b498b4e79236..12a12019837f 100755
--- a/applications/Chat/examples/test_ci.sh
+++ b/applications/Chat/examples/test_ci.sh
@@ -69,7 +69,7 @@ for model in 'gpt2' 'bloom' 'opt' 'roberta' 'llama'; do
             --prompt_dataset $PROMPT_PATH --pretrain_dataset $PRETRAIN_DATASET \
             --strategy $strategy --model $model \
             --num_episodes 1 --num_collect_steps 2 --num_update_steps 1 \
-            --train_batch_size 2
+            --train_batch_size 2 --lora_rank 4
     done
 done
 

From 9fe4ac4215cb0f641cfc50ebd93ca229f85da20b Mon Sep 17 00:00:00 2001
From: CWHer <cwher@outlook.com>
Date: Thu, 29 Jun 2023 11:27:58 +0800
Subject: [PATCH 07/21] feat: remove roberta support in train_prompts due to
 runtime errs

---
 applications/Chat/examples/train_prompts.py | 17 +++--------------
 1 file changed, 3 insertions(+), 14 deletions(-)

diff --git a/applications/Chat/examples/train_prompts.py b/applications/Chat/examples/train_prompts.py
index 13e6cf00c0f7..7a3440b2d3c8 100644
--- a/applications/Chat/examples/train_prompts.py
+++ b/applications/Chat/examples/train_prompts.py
@@ -7,14 +7,13 @@
 from coati.models.gpt import GPTRM, GPTActor, GPTCritic
 from coati.models.llama import LlamaActor, LlamaCritic, LlamaRM
 from coati.models.opt import OPTRM, OPTActor, OPTCritic
-from coati.models.roberta import RoBERTaActor, RoBERTaCritic, RoBERTaRM
 from coati.trainer import PPOTrainer
 from coati.trainer.strategies import DDPStrategy, GeminiStrategy, LowLevelZeroStrategy
 from coati.utils import prepare_llama_tokenizer_and_embedding
 from torch.optim import Adam
 from torch.utils.data import DataLoader
 from torch.utils.data.distributed import DistributedSampler
-from transformers import AutoTokenizer, BloomTokenizerFast, GPT2Tokenizer, LlamaTokenizer, RobertaTokenizer
+from transformers import AutoTokenizer, BloomTokenizerFast, GPT2Tokenizer, LlamaTokenizer
 
 from colossalai.nn.optimizer import HybridAdam
 
@@ -43,8 +42,6 @@ def main(args):
             initial_model = OPTActor(pretrained=args.pretrain)
         elif args.model == 'llama':
             initial_model = LlamaActor(pretrained=args.pretrain)
-        elif args.model == 'roberta':
-            initial_model = RoBERTaActor(pretrained=args.pretrain)
         else:
             raise ValueError(f'Unsupported actor model "{args.model}"')
 
@@ -61,8 +58,6 @@ def main(args):
             reward_model = OPTRM(pretrained=args.rm_pretrain)
         elif rm_model_name == 'llama':
             reward_model = LlamaRM(pretrained=args.rm_pretrain)
-        elif rm_model_name == 'roberta':
-            reward_model = RoBERTaRM(pretrained=args.rm_pretrain)
         else:
             raise ValueError(f'Unsupported reward model "{rm_model_name}"')
 
@@ -80,8 +75,6 @@ def main(args):
             actor = OPTActor(pretrained=args.pretrain, lora_rank=args.lora_rank)
         elif args.model == 'llama':
             actor = LlamaActor(pretrained=args.pretrain, lora_rank=args.lora_rank)
-        elif args.model == 'roberta':
-            actor = RoBERTaActor(pretrained=args.pretrain, lora_rank=args.lora_rank)
         else:
             raise ValueError(f'Unsupported actor model "{args.model}"')
 
@@ -93,8 +86,6 @@ def main(args):
             critic = OPTCritic(pretrained=args.rm_pretrain, lora_rank=args.lora_rank, use_action_mask=True)
         elif rm_model_name == 'llama':
             critic = LlamaCritic(pretrained=args.rm_pretrain, lora_rank=args.lora_rank, use_action_mask=True)
-        elif rm_model_name == 'roberta':
-            critic = RoBERTaCritic(pretrained=args.rm_pretrain, lora_rank=args.lora_rank, use_action_mask=True)
         else:
             raise ValueError(f'Unsupported reward model "{rm_model_name}"')
 
@@ -124,8 +115,6 @@ def main(args):
     elif args.model == 'llama':
         tokenizer = LlamaTokenizer.from_pretrained("hf-internal-testing/llama-tokenizer")
         tokenizer.eos_token = '<\s>'
-    elif args.model == 'roberta':
-        tokenizer = RobertaTokenizer.from_pretrained("roberta-base")
     else:
         raise ValueError(f'Unsupported model "{args.model}"')
 
@@ -206,9 +195,9 @@ def main(args):
                         choices=['ddp', 'colossalai_gemini', 'colossalai_zero2'],
                         default='colossalai_zero2',
                         help='strategy to use')
-    parser.add_argument('--model', default='gpt2', choices=['gpt2', 'bloom', 'opt', 'llama', 'roberta'])
+    parser.add_argument('--model', default='gpt2', choices=['gpt2', 'bloom', 'opt', 'llama'])
     parser.add_argument('--pretrain', type=str, default=None)
-    parser.add_argument('--rm_model', default=None, choices=['gpt2', 'bloom', 'opt', 'llama', 'roberta'])
+    parser.add_argument('--rm_model', default=None, choices=['gpt2', 'bloom', 'opt', 'llama'])
     parser.add_argument('--rm_path', type=str, default=None)
     parser.add_argument('--rm_pretrain', type=str, default=None)
     parser.add_argument('--save_path', type=str, default='actor_checkpoint_prompts')

From 45fe095422e653c6d93fdef12bdf78be0d2000e8 Mon Sep 17 00:00:00 2001
From: CWHer <cwher@outlook.com>
Date: Thu, 29 Jun 2023 11:32:01 +0800
Subject: [PATCH 08/21] feat: remove deberta & roberta in rm as not used

---
 applications/Chat/examples/train_reward_model.py | 16 ++--------------
 1 file changed, 2 insertions(+), 14 deletions(-)

diff --git a/applications/Chat/examples/train_reward_model.py b/applications/Chat/examples/train_reward_model.py
index 933139db134f..15fea9fd8440 100644
--- a/applications/Chat/examples/train_reward_model.py
+++ b/applications/Chat/examples/train_reward_model.py
@@ -1,18 +1,14 @@
 import argparse
 from random import randint
 
-import loralib as lora
 import torch
 import torch.distributed as dist
 from coati.dataset import HhRlhfDataset, RmStaticDataset
 from coati.models import LogExpLoss, LogSigLoss
-from coati.models.base import RewardModel
 from coati.models.bloom import BLOOMRM
-from coati.models.deberta import DebertaRM
 from coati.models.gpt import GPTRM
 from coati.models.llama import LlamaRM
 from coati.models.opt import OPTRM
-from coati.models.roberta import RoBERTaRM
 from coati.trainer import RewardModelTrainer
 from coati.trainer.strategies import DDPStrategy, GeminiStrategy, LowLevelZeroStrategy
 from coati.utils import prepare_llama_tokenizer_and_embedding
@@ -21,7 +17,7 @@
 from torch.optim.lr_scheduler import CosineAnnealingLR
 from torch.utils.data import DataLoader
 from torch.utils.data.distributed import DistributedSampler
-from transformers import AutoTokenizer, BloomTokenizerFast, DebertaV2Tokenizer, LlamaTokenizer, RobertaTokenizer
+from transformers import AutoTokenizer, BloomTokenizerFast, LlamaTokenizer
 from transformers.models.gpt2.tokenization_gpt2 import GPT2Tokenizer
 
 from colossalai.nn.optimizer import HybridAdam
@@ -46,12 +42,8 @@ def train(args):
             model = OPTRM(pretrained=args.pretrain, lora_rank=args.lora_rank).to(torch.cuda.current_device())
         elif args.model == 'gpt2':
             model = GPTRM(pretrained=args.pretrain, lora_rank=args.lora_rank).to(torch.cuda.current_device())
-        elif args.model == 'deberta':
-            model = DebertaRM(pretrained=args.pretrain, lora_rank=args.lora_rank).to(torch.cuda.current_device())
         elif args.model == 'llama':
             model = LlamaRM(pretrained=args.pretrain, lora_rank=args.lora_rank).to(torch.cuda.current_device())
-        elif args.model == 'roberta':
-            model = RoBERTaRM(pretrained=args.pretrain, lora_rank=args.lora_rank).to(torch.cuda.current_device())
         else:
             raise ValueError(f'Unsupported model "{args.model}"')
 
@@ -68,12 +60,8 @@ def train(args):
         tokenizer = BloomTokenizerFast.from_pretrained('bigscience/bloom-560m')
     elif args.model == 'opt':
         tokenizer = AutoTokenizer.from_pretrained("facebook/opt-350m")
-    elif args.model == 'deberta':
-        tokenizer = DebertaV2Tokenizer.from_pretrained('microsoft/deberta-v3-large')
     elif args.model == 'llama':
         tokenizer = LlamaTokenizer.from_pretrained("hf-internal-testing/llama-tokenizer")
-    elif args.model == 'roberta':
-        tokenizer = RobertaTokenizer.from_pretrained("roberta-base")
     else:
         raise ValueError(f'Unsupported model "{args.model}"')
     max_len = args.max_len
@@ -192,7 +180,7 @@ def train(args):
     parser.add_argument('--strategy',
                         choices=['ddp', 'colossalai_gemini', 'colossalai_zero2'],
                         default='colossalai_zero2')
-    parser.add_argument('--model', choices=['gpt2', 'bloom', 'opt', 'deberta', 'llama', 'roberta'], default='bloom')
+    parser.add_argument('--model', choices=['gpt2', 'bloom', 'opt', 'llama'], default='bloom')
     parser.add_argument('--pretrain', type=str, default=None)
     parser.add_argument('--model_path', type=str, default=None)
     parser.add_argument('--need_optim_ckpt', type=bool, default=False)

From 2c54f310345305511c9394d0deed164b8909f8d7 Mon Sep 17 00:00:00 2001
From: CWHer <cwher@outlook.com>
Date: Thu, 29 Jun 2023 11:32:53 +0800
Subject: [PATCH 09/21] test: remove deberta and roberta tests

---
 applications/Chat/examples/test_ci.sh | 29 ++++-----------------------
 1 file changed, 4 insertions(+), 25 deletions(-)

diff --git a/applications/Chat/examples/test_ci.sh b/applications/Chat/examples/test_ci.sh
index 12a12019837f..fe2af471017e 100755
--- a/applications/Chat/examples/test_ci.sh
+++ b/applications/Chat/examples/test_ci.sh
@@ -45,21 +45,16 @@ wandb init -m offline
 
 # FIXME: This is a hack to skip tests that are not working
 #  - gpt2-ddp: RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation
-#  - llama-ddp, llama-gemini: Out of memory
-#  - roberta-ddp: RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation
-#  - roberta-gemini: RuntimeError: ZERO DDP error: the synchronization of gradients doesn't exit properly.
-#  - roberta-zero2: "zero/low_level/low_level_optim.py", line 498, assert param_shape == flat_master_avg_grads.shape
+#  - llama-*: These tests can be passed locally, skipped for long execution time
 SKIPPED_TESTS=(
     "gpt2-ddp"
     "llama-ddp"
-    "llama-gemini"
-    "roberta-ddp"
-    "roberta-colossalai_gemini"
-    "roberta-colossalai_zero2"
+    "llama-colossalai_gemini"
+    "llama-colossalai_zero2"
 )
 
 # These tests are quick and do not have any dependencies
-for model in 'gpt2' 'bloom' 'opt' 'roberta' 'llama'; do
+for model in 'gpt2' 'bloom' 'opt' 'llama'; do
     for strategy in 'ddp' 'colossalai_gemini' 'colossalai_zero2'; do
         if [[ " ${SKIPPED_TESTS[*]} " =~ " ${model}-${strategy} " ]]; then
             echo "[Test]: Skipped $model-$strategy"
@@ -129,22 +124,6 @@ torchrun --standalone --nproc_per_node=2 ${BASE}/train_reward_model.py \
     --save_path ${BASE}/rm_ckpt.pt
 rm -rf ${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 \
-    --save_path ${BASE}/rm_ckpt.pt
-rm -rf ${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 \
-    --save_path ${BASE}/rm_ckpt.pt
-rm -rf ${BASE}/rm_ckpt.pt
-
 # train rl
 torchrun --standalone --nproc_per_node=2 ${BASE}/train_prompts.py \
     --prompt_dataset $PROMPT_PATH --pretrain_dataset $PRETRAIN_DATASET \

From b907e3ea67e5b1f06a96822d6b40528d76a380c2 Mon Sep 17 00:00:00 2001
From: CWHer <cwher@outlook.com>
Date: Thu, 29 Jun 2023 11:35:18 +0800
Subject: [PATCH 10/21] feat: remove deberta and roberta models as not used

---
 .../Chat/coati/models/deberta/__init__.py     |  4 --
 .../coati/models/deberta/deberta_critic.py    | 36 -----------------
 .../Chat/coati/models/deberta/deberta_rm.py   | 37 ------------------
 .../Chat/coati/models/roberta/__init__.py     |  5 ---
 .../coati/models/roberta/roberta_actor.py     | 35 -----------------
 .../coati/models/roberta/roberta_critic.py    | 38 ------------------
 .../Chat/coati/models/roberta/roberta_rm.py   | 39 -------------------
 7 files changed, 194 deletions(-)
 delete mode 100644 applications/Chat/coati/models/deberta/__init__.py
 delete mode 100644 applications/Chat/coati/models/deberta/deberta_critic.py
 delete mode 100644 applications/Chat/coati/models/deberta/deberta_rm.py
 delete mode 100644 applications/Chat/coati/models/roberta/__init__.py
 delete mode 100644 applications/Chat/coati/models/roberta/roberta_actor.py
 delete mode 100644 applications/Chat/coati/models/roberta/roberta_critic.py
 delete mode 100644 applications/Chat/coati/models/roberta/roberta_rm.py

diff --git a/applications/Chat/coati/models/deberta/__init__.py b/applications/Chat/coati/models/deberta/__init__.py
deleted file mode 100644
index b66888f34fd0..000000000000
--- a/applications/Chat/coati/models/deberta/__init__.py
+++ /dev/null
@@ -1,4 +0,0 @@
-from .deberta_critic import DebertaCritic
-from .deberta_rm import DebertaRM
-
-__all__ = ['DebertaCritic', 'DebertaRM']
diff --git a/applications/Chat/coati/models/deberta/deberta_critic.py b/applications/Chat/coati/models/deberta/deberta_critic.py
deleted file mode 100644
index e84c1dbd8380..000000000000
--- a/applications/Chat/coati/models/deberta/deberta_critic.py
+++ /dev/null
@@ -1,36 +0,0 @@
-from typing import Optional
-
-import torch.nn as nn
-from transformers import DebertaV2Config, DebertaV2Model
-
-from ..base import Critic
-
-
-class DebertaCritic(Critic):
-    """
-    Deberta Critic model.
-
-    Args:
-        pretrained (str): Pretrained model name or path.
-        config (DebertaV2Config): Model config.
-        checkpoint (bool): Enable gradient checkpointing.
-        lora_rank (int): Rank of the LO-RA decomposition.
-        lora_train_bias (str): LoRA bias training mode.
-    """
-
-    def __init__(self,
-                 pretrained: Optional[str] = None,
-                 config: Optional[DebertaV2Config] = None,
-                 checkpoint: bool = False,
-                 lora_rank: int = 0,
-                 lora_train_bias: str = 'none') -> None:
-        if pretrained is not None:
-            model = DebertaV2Model.from_pretrained(pretrained)
-        elif config is not None:
-            model = DebertaV2Model(config)
-        else:
-            model = DebertaV2Model(DebertaV2Config())
-        if checkpoint:
-            model.gradient_checkpointing_enable()
-        value_head = nn.Linear(model.config.hidden_size, 1)
-        super().__init__(model, value_head, lora_rank, lora_train_bias)
diff --git a/applications/Chat/coati/models/deberta/deberta_rm.py b/applications/Chat/coati/models/deberta/deberta_rm.py
deleted file mode 100644
index 2448c879ec85..000000000000
--- a/applications/Chat/coati/models/deberta/deberta_rm.py
+++ /dev/null
@@ -1,37 +0,0 @@
-from typing import Optional
-
-import torch.nn as nn
-from transformers import DebertaV2Config, DebertaV2Model
-
-from ..base import RewardModel
-
-
-class DebertaRM(RewardModel):
-    """
-    Deberta Reward model.
-
-    Args:
-        pretrained (str): Pretrained model name or path.
-        config (DebertaV2Config): Model config.
-        checkpoint (bool): Enable gradient checkpointing.
-        lora_rank (int): Rank of the LO-RA decomposition.
-        lora_train_bias (str): LoRA bias training mode.
-    """
-
-    def __init__(self,
-                 pretrained: str = None,
-                 config: Optional[DebertaV2Config] = None,
-                 checkpoint: bool = False,
-                 lora_rank: int = 0,
-                 lora_train_bias: str = 'none') -> None:
-        if pretrained is not None:
-            model = DebertaV2Model.from_pretrained(pretrained)
-        elif config is not None:
-            model = DebertaV2Model(config)
-        else:
-            model = DebertaV2Model(DebertaV2Config())
-        if checkpoint:
-            model.gradient_checkpointing_enable()
-        value_head = nn.Linear(model.config.hidden_size, 1)
-        value_head.weight.data.normal_(mean=0.0, std=1 / (model.config.hidden_size + 1))
-        super().__init__(model, value_head, lora_rank, lora_train_bias)
diff --git a/applications/Chat/coati/models/roberta/__init__.py b/applications/Chat/coati/models/roberta/__init__.py
deleted file mode 100644
index 0f4a8de067b1..000000000000
--- a/applications/Chat/coati/models/roberta/__init__.py
+++ /dev/null
@@ -1,5 +0,0 @@
-from .roberta_actor import RoBERTaActor
-from .roberta_critic import RoBERTaCritic
-from .roberta_rm import RoBERTaRM
-
-__all__ = ['RoBERTaActor', 'RoBERTaCritic', 'RoBERTaRM']
\ No newline at end of file
diff --git a/applications/Chat/coati/models/roberta/roberta_actor.py b/applications/Chat/coati/models/roberta/roberta_actor.py
deleted file mode 100644
index e35fa6eb19a8..000000000000
--- a/applications/Chat/coati/models/roberta/roberta_actor.py
+++ /dev/null
@@ -1,35 +0,0 @@
-from typing import Optional
-
-from transformers.models.roberta.configuration_roberta import RobertaConfig
-from transformers.models.roberta.modeling_roberta import RobertaForCausalLM
-
-from ..base import Actor
-
-class RoBERTaActor(Actor):
-    """
-    RoBERTa Actor model.
-
-    Args:
-        pretrained (str): Pretrained model name or path.
-        config (RoBERTaConfig): Model config.
-        checkpoint (bool): Enable gradient checkpointing.
-        lora_rank (int): Rank of the low-rank approximation.
-        lora_train_bias (str): LoRA bias training mode.
-    """
-
-
-    def __init__(self,
-                 pretrained: Optional[str] = None,
-                 config: Optional[RobertaConfig] = None,
-                 checkpoint: bool = False,
-                 lora_rank: int = 0,
-                 lora_train_bias: str = 'none') -> None:
-        if pretrained is not None:
-            model = RobertaForCausalLM.from_pretrained(pretrained)
-        elif config is not None:
-            model = RobertaForCausalLM(config)
-        else:
-            model = RobertaForCausalLM(RobertaConfig())
-        if checkpoint:
-            model.gradient_checkpointing_enable()
-        super().__init__(model, lora_rank, lora_train_bias)
diff --git a/applications/Chat/coati/models/roberta/roberta_critic.py b/applications/Chat/coati/models/roberta/roberta_critic.py
deleted file mode 100644
index c8dc0d9e14f2..000000000000
--- a/applications/Chat/coati/models/roberta/roberta_critic.py
+++ /dev/null
@@ -1,38 +0,0 @@
-from typing import Optional
-
-import torch.nn as nn
-from transformers.models.roberta.configuration_roberta import RobertaConfig
-from transformers.models.roberta.modeling_roberta import RobertaModel
-
-from ..base import Critic
-
-
-class RoBERTaCritic(Critic):
-    """
-    RoBERTa Critic model.
-
-    Args:
-        pretrained (str): Pretrained model name or path.
-        config (RoBERTa Config): Model config.
-        checkpoint (bool): Enable gradient checkpointing.
-        lora_rank (int): Rank of the low-rank approximation.
-        lora_train_bias (str): LoRA bias training mode.
-    """
-
-    def __init__(self,
-                 pretrained: Optional[str] = None,
-                 config: Optional[RobertaConfig] = None,
-                 checkpoint: bool = False,
-                 lora_rank: int = 0,
-                 lora_train_bias: str = 'none',
-                 **kwargs) -> None:
-        if pretrained is not None:
-            model = RobertaModel.from_pretrained(pretrained, add_pooling_layer=False)
-        elif config is not None:
-            model = RobertaModel(config)
-        else:
-            model = RobertaModel(RobertaConfig())
-        if checkpoint:
-            model.gradient_checkpointing_enable()
-        value_head = nn.Linear(model.config.hidden_size, 1)
-        super().__init__(model, value_head, lora_rank, lora_train_bias, **kwargs)
diff --git a/applications/Chat/coati/models/roberta/roberta_rm.py b/applications/Chat/coati/models/roberta/roberta_rm.py
deleted file mode 100644
index 77075052978b..000000000000
--- a/applications/Chat/coati/models/roberta/roberta_rm.py
+++ /dev/null
@@ -1,39 +0,0 @@
-from typing import Optional
-
-import torch.nn as nn
-from transformers import RobertaConfig, RobertaModel
-
-
-from ..base import RewardModel
-
-
-class RoBERTaRM(RewardModel):
-    """
-    RoBERTa Reward model.
-
-    Args:
-        pretrained (str): Pretrained model name or path.
-        config (RoBERTaConfig): Model config.
-        checkpoint (bool): Enable gradient checkpointing.
-        lora_rank (int): Rank of the low-rank approximation.
-        lora_train_bias (str): LoRA bias training mode.
-    """
-
-    def __init__(self,
-                 pretrained: Optional[str] = None,
-                 config: Optional[RobertaConfig] = None,
-                 checkpoint: bool = False,
-                 lora_rank: int = 0,
-                 lora_train_bias: str = 'none') -> None:
-        if pretrained is not None:
-            model = RobertaModel.from_pretrained(pretrained, add_pooling_layer=False)
-        elif config is not None:
-            model = RobertaModel(config)
-        else:
-            model = RobertaModel(RobertaConfig())
-        if checkpoint:
-            model.gradient_checkpointing_enable()
-
-        value_head = nn.Linear(model.config.hidden_size, 1)
-        value_head.weight.data.normal_(mean=0.0, std=1/(model.config.hidden_size + 1))
-        super().__init__(model, value_head, lora_rank, lora_train_bias)
\ No newline at end of file

From bd0da7d2045f546381bfb9092e516dfe07760934 Mon Sep 17 00:00:00 2001
From: CWHer <cwher@outlook.com>
Date: Thu, 29 Jun 2023 11:41:52 +0800
Subject: [PATCH 11/21] fix: remove calls to roberta

---
 applications/Chat/coati/ray/utils.py    | 11 +----------
 applications/Chat/examples/inference.py |  9 ++-------
 2 files changed, 3 insertions(+), 17 deletions(-)

diff --git a/applications/Chat/coati/ray/utils.py b/applications/Chat/coati/ray/utils.py
index 4f8e0b8a87e9..e0ba1046d2c7 100644
--- a/applications/Chat/coati/ray/utils.py
+++ b/applications/Chat/coati/ray/utils.py
@@ -9,10 +9,9 @@
 from coati.models.gpt import GPTRM, GPTActor, GPTCritic
 from coati.models.llama import LlamaActor, LlamaCritic, LlamaRM
 from coati.models.opt import OPTRM, OPTActor, OPTCritic
-from coati.models.roberta import RoBERTaActor, RoBERTaCritic, RoBERTaRM
 from coati.trainer.strategies import DDPStrategy, GeminiStrategy, LowLevelZeroStrategy
 from coati.utils import prepare_llama_tokenizer_and_embedding
-from transformers import AutoTokenizer, BloomTokenizerFast, GPT2Tokenizer, LlamaTokenizer, RobertaTokenizer
+from transformers import AutoTokenizer, BloomTokenizerFast, GPT2Tokenizer, LlamaTokenizer
 
 
 def is_rank_0() -> bool:
@@ -36,8 +35,6 @@ def get_actor_from_args(model: str, pretrained: str = None, config=None, lora_ra
         actor = OPTActor(pretrained=pretrained, config=config, lora_rank=lora_rank)
     elif model == 'llama':
         actor = LlamaActor(pretrained=pretrained, config=config, lora_rank=lora_rank)
-    elif model == 'roberta':
-        actor = RoBERTaActor(pretrained=pretrained, config=config, lora_rank=lora_rank)
     else:
         raise ValueError(f'Unsupported actor model "{model}"')
     return actor
@@ -52,8 +49,6 @@ def get_critic_from_args(model: str, pretrained: str = None, config=None, lora_r
         critic = OPTCritic(pretrained=pretrained, lora_rank=lora_rank, config=config, use_action_mask=True)
     elif model == 'llama':
         critic = LlamaCritic(pretrained=pretrained, lora_rank=lora_rank, config=config, use_action_mask=True)
-    elif model == 'roberta':
-        critic = RoBERTaCritic(pretrained=pretrained, lora_rank=lora_rank, config=config, use_action_mask=True)
     else:
         raise ValueError(f'Unsupported reward model "{model}"')
     return critic
@@ -68,8 +63,6 @@ def get_reward_model_from_args(model: str, pretrained: str = None, config=None):
         reward_model = OPTRM(pretrained=pretrained, config=config)
     elif model == 'llama':
         reward_model = LlamaRM(pretrained=pretrained, config=config)
-    elif model == 'roberta':
-        reward_model = RoBERTaRM(pretrained=pretrained, config=config)
     else:
         raise ValueError(f'Unsupported reward model "{model}"')
     return reward_model
@@ -101,8 +94,6 @@ def get_tokenizer_from_args(model: str, **kwargs):
     elif model == 'llama':
         pretrain_path = kwargs["pretrain"]
         tokenizer = AutoTokenizer.from_pretrained(pretrain_path)
-    elif model == 'roberta':
-        tokenizer = RobertaTokenizer.from_pretrained("roberta-base")
     else:
         raise ValueError(f'Unsupported model "{model}"')
 
diff --git a/applications/Chat/examples/inference.py b/applications/Chat/examples/inference.py
index ae59d91c1822..f75950804d2e 100644
--- a/applications/Chat/examples/inference.py
+++ b/applications/Chat/examples/inference.py
@@ -4,8 +4,7 @@
 from coati.models.bloom import BLOOMActor
 from coati.models.gpt import GPTActor
 from coati.models.opt import OPTActor
-from coati.models.roberta import RoBERTaActor
-from transformers import AutoTokenizer, RobertaTokenizer
+from transformers import AutoTokenizer
 from transformers.models.gpt2.tokenization_gpt2 import GPT2Tokenizer
 
 
@@ -17,8 +16,6 @@ def eval(args):
         actor = BLOOMActor(pretrained=args.pretrain).to(torch.cuda.current_device())
     elif args.model == 'opt':
         actor = OPTActor(pretrained=args.pretrain).to(torch.cuda.current_device())
-    elif args.model == 'roberta':
-        actor = RoBERTaActor(pretrained=args.pretrain).to(torch.cuda.current_device())
     else:
         raise ValueError(f'Unsupported model "{args.model}"')
 
@@ -34,8 +31,6 @@ def eval(args):
         tokenizer.pad_token = tokenizer.eos_token
     elif args.model == 'opt':
         tokenizer = AutoTokenizer.from_pretrained('facebook/opt-350m')
-    elif args.model == 'roberta':
-        tokenizer = RobertaTokenizer.from_pretrained("roberta-base")
     else:
         raise ValueError(f'Unsupported model "{args.model}"')
 
@@ -54,7 +49,7 @@ def eval(args):
 
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
-    parser.add_argument('--model', default='gpt2', choices=['gpt2', 'bloom', 'opt', 'roberta'])
+    parser.add_argument('--model', default='gpt2', choices=['gpt2', 'bloom', 'opt'])
     # We suggest to use the pretrained model from HuggingFace, use pretrain to configure model
     parser.add_argument('--pretrain', type=str, default=None)
     parser.add_argument('--model_path', type=str, default=None)

From e2499651477c0212810f143d9efa7ea02471a8fd Mon Sep 17 00:00:00 2001
From: CWHer <cwher@outlook.com>
Date: Thu, 29 Jun 2023 11:48:26 +0800
Subject: [PATCH 12/21] fix: remove prepare_llama_tokenizer_and_embedding

---
 applications/Chat/coati/ray/utils.py          |  1 -
 applications/Chat/coati/utils/__init__.py     |  3 -
 .../Chat/coati/utils/tokenizer_utils.py       | 73 -------------------
 .../community/peft/train_peft_prompts.py      |  6 +-
 .../examples/community/peft/train_peft_sft.py | 25 +++----
 applications/Chat/examples/train_prompts.py   |  1 -
 .../Chat/examples/train_reward_model.py       |  1 -
 applications/Chat/examples/train_sft.py       |  1 -
 8 files changed, 11 insertions(+), 100 deletions(-)
 delete mode 100644 applications/Chat/coati/utils/__init__.py
 delete mode 100644 applications/Chat/coati/utils/tokenizer_utils.py

diff --git a/applications/Chat/coati/ray/utils.py b/applications/Chat/coati/ray/utils.py
index e0ba1046d2c7..761186b95ee5 100644
--- a/applications/Chat/coati/ray/utils.py
+++ b/applications/Chat/coati/ray/utils.py
@@ -10,7 +10,6 @@
 from coati.models.llama import LlamaActor, LlamaCritic, LlamaRM
 from coati.models.opt import OPTRM, OPTActor, OPTCritic
 from coati.trainer.strategies import DDPStrategy, GeminiStrategy, LowLevelZeroStrategy
-from coati.utils import prepare_llama_tokenizer_and_embedding
 from transformers import AutoTokenizer, BloomTokenizerFast, GPT2Tokenizer, LlamaTokenizer
 
 
diff --git a/applications/Chat/coati/utils/__init__.py b/applications/Chat/coati/utils/__init__.py
deleted file mode 100644
index 112b82b97064..000000000000
--- a/applications/Chat/coati/utils/__init__.py
+++ /dev/null
@@ -1,3 +0,0 @@
-from .tokenizer_utils import prepare_llama_tokenizer_and_embedding, smart_tokenizer_and_embedding_resize
-
-__all__ = ['smart_tokenizer_and_embedding_resize', 'prepare_llama_tokenizer_and_embedding']
\ No newline at end of file
diff --git a/applications/Chat/coati/utils/tokenizer_utils.py b/applications/Chat/coati/utils/tokenizer_utils.py
deleted file mode 100644
index e0d96cfc8be2..000000000000
--- a/applications/Chat/coati/utils/tokenizer_utils.py
+++ /dev/null
@@ -1,73 +0,0 @@
-#    Copyright 2023 Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li
-#
-#    Licensed under the Apache License, Version 2.0 (the "License");
-#    you may not use this file except in compliance with the License.
-#    You may obtain a copy of the License at
-#
-#        http://www.apache.org/licenses/LICENSE-2.0
-#
-#    Unless required by applicable law or agreed to in writing, software
-#    distributed under the License is distributed on an "AS IS" BASIS,
-#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-#    See the License for the specific language governing permissions and
-#    limitations under the License.
-
-from typing import Dict
-
-import transformers
-
-DEFAULT_PAD_TOKEN = "[PAD]"
-DEFAULT_EOS_TOKEN = "</s>"
-DEFAULT_BOS_TOKEN = "</s>"
-DEFAULT_UNK_TOKEN = "</s>"
-
-
-def prepare_llama_tokenizer_and_embedding(
-        tokenizer: transformers.PreTrainedTokenizer,
-        model: transformers.PreTrainedModel,
-        special_tokens_dict: Dict = dict(pad_token=DEFAULT_PAD_TOKEN),
-):
-    """prepare llama tokenizer and embedding.
-
-    """
-
-    if tokenizer.pad_token is None:
-        smart_tokenizer_and_embedding_resize(
-            special_tokens_dict=dict(pad_token=DEFAULT_PAD_TOKEN),
-            tokenizer=tokenizer,
-            model=model,
-        )
-
-    tokenizer.add_special_tokens({
-        "eos_token": DEFAULT_EOS_TOKEN,
-        "bos_token": DEFAULT_BOS_TOKEN,
-        "unk_token": DEFAULT_UNK_TOKEN,
-    })
-
-    return tokenizer
-
-
-def smart_tokenizer_and_embedding_resize(
-        tokenizer: transformers.PreTrainedTokenizer,
-        model: transformers.PreTrainedModel,
-        special_tokens_dict: Dict = dict(pad_token=DEFAULT_PAD_TOKEN),
-):
-    """Resize tokenizer and embedding.
-
-    Note: This is the unoptimized version that may make your embedding size not be divisible by 64.
-    """
-
-    if tokenizer.pad_token is None:
-        num_new_tokens = tokenizer.add_special_tokens(special_tokens_dict)
-
-        model.resize_token_embeddings(len(tokenizer))
-
-        if num_new_tokens > 0:
-            input_embeddings = model.get_input_embeddings().weight.data
-            output_embeddings = model.get_output_embeddings().weight.data
-
-            input_embeddings_avg = input_embeddings[:-num_new_tokens].mean(dim=0, keepdim=True)
-            output_embeddings_avg = output_embeddings[:-num_new_tokens].mean(dim=0, keepdim=True)
-
-            input_embeddings[-num_new_tokens:] = input_embeddings_avg
-            output_embeddings[-num_new_tokens:] = output_embeddings_avg
diff --git a/applications/Chat/examples/community/peft/train_peft_prompts.py b/applications/Chat/examples/community/peft/train_peft_prompts.py
index 9d8dbb38ae5d..fc179769ebd0 100644
--- a/applications/Chat/examples/community/peft/train_peft_prompts.py
+++ b/applications/Chat/examples/community/peft/train_peft_prompts.py
@@ -10,7 +10,6 @@
 from coati.models.opt import OPTRM, OPTActor, OPTCritic
 from coati.trainer import PPOTrainer
 from coati.trainer.strategies import DDPStrategy, GeminiStrategy, LowLevelZeroStrategy
-from coati.utils import prepare_llama_tokenizer_and_embedding
 from easy_dataset import EasyPromptsDataset, EasySupervisedDataset
 from easy_models import BLOOMActor
 from peft import PeftModel
@@ -122,10 +121,7 @@ def main(args):
     else:
         raise ValueError(f'Unsupported model "{args.model}"')
 
-    if args.model == 'llama':
-        tokenizer = prepare_llama_tokenizer_and_embedding(tokenizer, actor)
-    else:
-        tokenizer.pad_token = tokenizer.eos_token
+    tokenizer.pad_token = tokenizer.eos_token
 
     data_collator = DataCollatorForSupervisedDataset(tokenizer=tokenizer)
 
diff --git a/applications/Chat/examples/community/peft/train_peft_sft.py b/applications/Chat/examples/community/peft/train_peft_sft.py
index 54fe0ad55049..924bf7fb7a57 100644
--- a/applications/Chat/examples/community/peft/train_peft_sft.py
+++ b/applications/Chat/examples/community/peft/train_peft_sft.py
@@ -12,7 +12,6 @@
 from coati.models.opt import OPTLM
 from coati.trainer import SFTTrainer
 from coati.trainer.strategies import DDPStrategy, GeminiStrategy, LowLevelZeroStrategy
-from coati.utils import prepare_llama_tokenizer_and_embedding
 from datasets import load_dataset
 from easy_dataset import EasyDataset
 from peft import LoraConfig, PeftModel, TaskType, get_peft_model
@@ -79,20 +78,16 @@ def train(args):
     else:
         raise ValueError(f'Unsupported model "{args.model}"')
     tokenizer.pad_token = tokenizer.eos_token
-    if args.model == 'llama':
-        tokenizer = prepare_llama_tokenizer_and_embedding(tokenizer, model)
-
-        if args.strategy == 'colossalai_gemini':
-            # this is a hack to deal with the resized embedding
-            # to make sure all parameters are ColoParameter for Colossal-AI Gemini Compatibility
-            for name, param in model.named_parameters():
-                if not isinstance(param, ColoParameter):
-                    sub_module_name = '.'.join(name.split('.')[:-1])
-                    weight_name = name.split('.')[-1]
-                    sub_module = model.get_submodule(sub_module_name)
-                    setattr(sub_module, weight_name, ColoParameter(param))
-    else:
-        tokenizer.pad_token = tokenizer.eos_token
+
+    if args.model == 'llama' and args.strategy == 'colossalai_gemini':
+        # this is a hack to deal with the resized embedding
+        # to make sure all parameters are ColoParameter for Colossal-AI Gemini Compatibility
+        for name, param in model.named_parameters():
+            if not isinstance(param, ColoParameter):
+                sub_module_name = '.'.join(name.split('.')[:-1])
+                weight_name = name.split('.')[-1]
+                sub_module = model.get_submodule(sub_module_name)
+                setattr(sub_module, weight_name, ColoParameter(param))
 
     # configure optimizer
     if args.strategy.startswith('colossalai'):
diff --git a/applications/Chat/examples/train_prompts.py b/applications/Chat/examples/train_prompts.py
index 7a3440b2d3c8..2b3ce423f8e7 100644
--- a/applications/Chat/examples/train_prompts.py
+++ b/applications/Chat/examples/train_prompts.py
@@ -9,7 +9,6 @@
 from coati.models.opt import OPTRM, OPTActor, OPTCritic
 from coati.trainer import PPOTrainer
 from coati.trainer.strategies import DDPStrategy, GeminiStrategy, LowLevelZeroStrategy
-from coati.utils import prepare_llama_tokenizer_and_embedding
 from torch.optim import Adam
 from torch.utils.data import DataLoader
 from torch.utils.data.distributed import DistributedSampler
diff --git a/applications/Chat/examples/train_reward_model.py b/applications/Chat/examples/train_reward_model.py
index 15fea9fd8440..837519cca766 100644
--- a/applications/Chat/examples/train_reward_model.py
+++ b/applications/Chat/examples/train_reward_model.py
@@ -11,7 +11,6 @@
 from coati.models.opt import OPTRM
 from coati.trainer import RewardModelTrainer
 from coati.trainer.strategies import DDPStrategy, GeminiStrategy, LowLevelZeroStrategy
-from coati.utils import prepare_llama_tokenizer_and_embedding
 from datasets import load_dataset
 from torch.optim import Adam
 from torch.optim.lr_scheduler import CosineAnnealingLR
diff --git a/applications/Chat/examples/train_sft.py b/applications/Chat/examples/train_sft.py
index 0e8f34554536..dee049662c60 100644
--- a/applications/Chat/examples/train_sft.py
+++ b/applications/Chat/examples/train_sft.py
@@ -9,7 +9,6 @@
 from coati.models import convert_to_lora_module
 from coati.trainer import SFTTrainer
 from coati.trainer.strategies import DDPStrategy, GeminiStrategy, LowLevelZeroStrategy
-from coati.utils import prepare_llama_tokenizer_and_embedding
 from datasets import load_dataset
 from torch.optim import Adam
 from torch.utils.data import DataLoader

From 9d34193ec486a1b686dac684c90a5081106396ec Mon Sep 17 00:00:00 2001
From: CWHer <cwher@outlook.com>
Date: Thu, 29 Jun 2023 12:08:33 +0800
Subject: [PATCH 13/21] chore: update transformers version

---
 .github/workflows/run_chatgpt_examples.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/run_chatgpt_examples.yml b/.github/workflows/run_chatgpt_examples.yml
index 129bf7ed3270..e8b1b6fa24a7 100644
--- a/.github/workflows/run_chatgpt_examples.yml
+++ b/.github/workflows/run_chatgpt_examples.yml
@@ -39,7 +39,7 @@ jobs:
       - name: Install Transformers
         run: |
           cd applications/Chat
-          git clone https://github.com/hpcaitech/transformers
+          git clone --depth 1 --branch v4.30.2 https://github.com/huggingface/transformers.git
           cd transformers
           pip install -v .
 

From 1356ba5ac6aed5992445d48ded8fd85298079fe8 Mon Sep 17 00:00:00 2001
From: CWHer <cwher@outlook.com>
Date: Thu, 29 Jun 2023 12:08:54 +0800
Subject: [PATCH 14/21] docs: update transformers version

---
 applications/Chat/README.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/applications/Chat/README.md b/applications/Chat/README.md
index 016272ed8c89..2d36d44a03b6 100644
--- a/applications/Chat/README.md
+++ b/applications/Chat/README.md
@@ -101,7 +101,7 @@ pip install .
 Given Hugging Face hasn't officially supported the LLaMA models, We fork a branch of Transformers that can be compatible with our code
 
 ```shell
-git clone https://github.com/hpcaitech/transformers
+git clone --depth 1 --branch v4.30.2 https://github.com/huggingface/transformers.git
 cd transformers
 pip install .
 ```

From ef110b589bd9e8065ce0126f985f4ded99691482 Mon Sep 17 00:00:00 2001
From: CWHer <cwher@outlook.com>
Date: Thu, 29 Jun 2023 14:54:20 +0800
Subject: [PATCH 15/21] fix: fix actor inference

---
 applications/Chat/examples/inference.py | 16 +++++++++-------
 1 file changed, 9 insertions(+), 7 deletions(-)

diff --git a/applications/Chat/examples/inference.py b/applications/Chat/examples/inference.py
index f75950804d2e..4b49e76088bc 100644
--- a/applications/Chat/examples/inference.py
+++ b/applications/Chat/examples/inference.py
@@ -2,6 +2,7 @@
 
 import torch
 from coati.models.bloom import BLOOMActor
+from coati.models.generation import generate
 from coati.models.gpt import GPTActor
 from coati.models.opt import OPTActor
 from transformers import AutoTokenizer
@@ -20,7 +21,7 @@ def eval(args):
         raise ValueError(f'Unsupported model "{args.model}"')
 
     state_dict = torch.load(args.model_path)
-    actor.model.load_state_dict(state_dict)
+    actor.load_state_dict(state_dict)
 
     # configure tokenizer
     if args.model == 'gpt2':
@@ -37,12 +38,13 @@ def eval(args):
     actor.eval()
     input = args.input
     input_ids = tokenizer.encode(input, return_tensors='pt').to(torch.cuda.current_device())
-    outputs = actor.generate(input_ids,
-                             max_length=args.max_length,
-                             do_sample=True,
-                             top_k=50,
-                             top_p=0.95,
-                             num_return_sequences=1)
+    outputs = generate(actor,
+                       input_ids,
+                       max_length=args.max_length,
+                       do_sample=True,
+                       top_k=50,
+                       top_p=0.95,
+                       num_return_sequences=1)
     output = tokenizer.batch_decode(outputs[0], skip_special_tokens=True)
     print(output)
 

From 82ee53273612b72101a0b2589b08e7eca7642b51 Mon Sep 17 00:00:00 2001
From: CWHer <cwher@outlook.com>
Date: Fri, 30 Jun 2023 10:32:45 +0800
Subject: [PATCH 16/21] fix: fix ci

---
 .github/workflows/run_chatgpt_examples.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/run_chatgpt_examples.yml b/.github/workflows/run_chatgpt_examples.yml
index e8b1b6fa24a7..5b1e9383fe17 100644
--- a/.github/workflows/run_chatgpt_examples.yml
+++ b/.github/workflows/run_chatgpt_examples.yml
@@ -39,7 +39,7 @@ jobs:
       - name: Install Transformers
         run: |
           cd applications/Chat
-          git clone --depth 1 --branch v4.30.2 https://github.com/huggingface/transformers.git
+          git clone --depth 1 --branch v4.30.2 https://github.com/huggingface/transformers
           cd transformers
           pip install -v .
 

From baadbb138110417ff8b4e41c45c7189b235858c3 Mon Sep 17 00:00:00 2001
From: CWHer <cwher@outlook.com>
Date: Mon, 3 Jul 2023 14:39:07 +0800
Subject: [PATCH 17/21] feat: change llama pad token to unk

---
 applications/Chat/examples/train_prompts.py   |  6 ++++--
 .../Chat/examples/train_reward_model.py       | 19 ++++++++++---------
 applications/Chat/examples/train_sft.py       | 10 +++++-----
 3 files changed, 19 insertions(+), 16 deletions(-)

diff --git a/applications/Chat/examples/train_prompts.py b/applications/Chat/examples/train_prompts.py
index 2b3ce423f8e7..791ba52e2f4c 100644
--- a/applications/Chat/examples/train_prompts.py
+++ b/applications/Chat/examples/train_prompts.py
@@ -107,18 +107,20 @@ def main(args):
     # configure tokenizer
     if args.model == 'gpt2':
         tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
+        tokenizer.pad_token = tokenizer.eos_token
     elif args.model == 'bloom':
         tokenizer = BloomTokenizerFast.from_pretrained('bigscience/bloom-560m')
+        tokenizer.pad_token = tokenizer.eos_token
     elif args.model == 'opt':
         tokenizer = AutoTokenizer.from_pretrained("facebook/opt-350m")
+        tokenizer.pad_token = tokenizer.eos_token
     elif args.model == 'llama':
         tokenizer = LlamaTokenizer.from_pretrained("hf-internal-testing/llama-tokenizer")
         tokenizer.eos_token = '<\s>'
+        tokenizer.pad_token = tokenizer.unk_token
     else:
         raise ValueError(f'Unsupported model "{args.model}"')
 
-    tokenizer.pad_token = tokenizer.eos_token
-
     data_collator = DataCollatorForSupervisedDataset(tokenizer=tokenizer)
 
     prompt_dataset = PromptDataset(tokenizer=tokenizer, data_path=args.prompt_dataset, max_datasets_size=16384)
diff --git a/applications/Chat/examples/train_reward_model.py b/applications/Chat/examples/train_reward_model.py
index 837519cca766..003f2a0f3597 100644
--- a/applications/Chat/examples/train_reward_model.py
+++ b/applications/Chat/examples/train_reward_model.py
@@ -55,17 +55,18 @@ def train(args):
     # configure tokenizer
     if args.model == 'gpt2':
         tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
+        tokenizer.pad_token = tokenizer.eos_token
     elif args.model == 'bloom':
         tokenizer = BloomTokenizerFast.from_pretrained('bigscience/bloom-560m')
+        tokenizer.pad_token = tokenizer.eos_token
     elif args.model == 'opt':
         tokenizer = AutoTokenizer.from_pretrained("facebook/opt-350m")
+        tokenizer.pad_token = tokenizer.eos_token
     elif args.model == 'llama':
         tokenizer = LlamaTokenizer.from_pretrained("hf-internal-testing/llama-tokenizer")
+        tokenizer.pad_token = tokenizer.unk_token
     else:
         raise ValueError(f'Unsupported model "{args.model}"')
-    max_len = args.max_len
-
-    tokenizer.pad_token = tokenizer.eos_token
 
     # configure optimizer
     if args.strategy.startswith('colossalai'):
@@ -96,13 +97,13 @@ def train(args):
     valid_data = data['test'].select((randint(0, len(eval_data) - 1) for _ in range(len(eval_data) // 5)))
 
     if args.dataset == 'Dahoas/rm-static':
-        train_dataset = RmStaticDataset(train_data, tokenizer, max_len)
-        valid_dataset = RmStaticDataset(valid_data, tokenizer, max_len)
-        eval_dataset = RmStaticDataset(eval_data, tokenizer, max_len)
+        train_dataset = RmStaticDataset(train_data, tokenizer, args.max_len)
+        valid_dataset = RmStaticDataset(valid_data, tokenizer, args.max_len)
+        eval_dataset = RmStaticDataset(eval_data, tokenizer, args.max_len)
     elif args.dataset == 'Anthropic/hh-rlhf':
-        train_dataset = HhRlhfDataset(train_data, tokenizer, max_len)
-        valid_dataset = HhRlhfDataset(valid_data, tokenizer, max_len)
-        eval_dataset = HhRlhfDataset(eval_data, tokenizer, max_len)
+        train_dataset = HhRlhfDataset(train_data, tokenizer, args.max_len)
+        valid_dataset = HhRlhfDataset(valid_data, tokenizer, args.max_len)
+        eval_dataset = HhRlhfDataset(eval_data, tokenizer, args.max_len)
     else:
         raise ValueError(f'Unsupported dataset "{args.dataset}"')
 
diff --git a/applications/Chat/examples/train_sft.py b/applications/Chat/examples/train_sft.py
index dee049662c60..acc17769133c 100644
--- a/applications/Chat/examples/train_sft.py
+++ b/applications/Chat/examples/train_sft.py
@@ -67,6 +67,7 @@ def train(args):
         tokenizer.pad_token = tokenizer.eos_token
     elif args.model == 'opt':
         tokenizer = AutoTokenizer.from_pretrained("facebook/opt-350m")
+        tokenizer.pad_token = tokenizer.eos_token
     elif args.model == 'llama':
         tokenizer = AutoTokenizer.from_pretrained(
             "hf-internal-testing/llama-tokenizer",
@@ -74,10 +75,9 @@ def train(args):
             use_fast=False,
         )
         tokenizer.eos_token = '<\s>'
+        tokenizer.pad_token = tokenizer.unk_token
     else:
         raise ValueError(f'Unsupported model "{args.model}"')
-    tokenizer.pad_token = tokenizer.eos_token
-    max_len = args.max_len
 
     if args.model == 'llama' and args.strategy == 'colossalai_gemini':
         # this is a hack to deal with the resized embedding
@@ -102,14 +102,14 @@ def train(args):
         train_data = load_dataset(args.dataset, 'super_natural_instructions', split='train')
         eval_data = load_dataset(args.dataset, 'super_natural_instructions', split='test')
 
-        train_dataset = SFTDataset(train_data, tokenizer, max_len)
-        eval_dataset = SFTDataset(eval_data, tokenizer, max_len)
+        train_dataset = SFTDataset(train_data, tokenizer, args.max_len)
+        eval_dataset = SFTDataset(eval_data, tokenizer, args.max_len)
 
     else:
         train_dataset = SupervisedDataset(tokenizer=tokenizer,
                                           data_path=args.dataset,
                                           max_datasets_size=args.max_datasets_size,
-                                          max_length=max_len)
+                                          max_length=args.max_len)
         eval_dataset = None
     data_collator = DataCollatorForSupervisedDataset(tokenizer=tokenizer)
 

From 78f8d792d8d4538be7d64b58282aea010856b1d6 Mon Sep 17 00:00:00 2001
From: CWHer <cwher@outlook.com>
Date: Mon, 3 Jul 2023 14:47:07 +0800
Subject: [PATCH 18/21] revert: revert ddp setup_distributed

---
 applications/Chat/coati/trainer/strategies/ddp.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/applications/Chat/coati/trainer/strategies/ddp.py b/applications/Chat/coati/trainer/strategies/ddp.py
index 2fdfebd2a6a8..e1c1bbf19f35 100644
--- a/applications/Chat/coati/trainer/strategies/ddp.py
+++ b/applications/Chat/coati/trainer/strategies/ddp.py
@@ -63,7 +63,7 @@ def _post_init(self) -> None:
             f'{type(self).__name__}\'s plugin is not initialized properly.'
 
     def setup_distributed(self) -> None:
-        self._try_init_dist(force=False)
+        self._try_init_dist(force=True)
         self.set_seed(self.seed)
 
     def set_seed(self, seed: int) -> None:

From 1a0f74c401ca080606cf93ed65b67f8bb3c0dad2 Mon Sep 17 00:00:00 2001
From: CWHer <cwher@outlook.com>
Date: Mon, 3 Jul 2023 15:01:55 +0800
Subject: [PATCH 19/21] fix: change llama pad token to unk

---
 .../Chat/examples/community/peft/train_peft_prompts.py      | 6 ++++--
 applications/Chat/examples/community/peft/train_peft_sft.py | 5 +++--
 2 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/applications/Chat/examples/community/peft/train_peft_prompts.py b/applications/Chat/examples/community/peft/train_peft_prompts.py
index fc179769ebd0..9385e457d852 100644
--- a/applications/Chat/examples/community/peft/train_peft_prompts.py
+++ b/applications/Chat/examples/community/peft/train_peft_prompts.py
@@ -111,18 +111,20 @@ def main(args):
     # configure tokenizer
     if args.model == 'gpt2':
         tokenizer = GPT2Tokenizer.from_pretrained(args.rm_pretrain)
+        tokenizer.pad_token = tokenizer.eos_token
     elif args.model == 'bloom':
         tokenizer = BloomTokenizerFast.from_pretrained(args.rm_pretrain)
+        tokenizer.pad_token = tokenizer.eos_token
     elif args.model == 'opt':
         tokenizer = AutoTokenizer.from_pretrained(args.rm_pretrain)
+        tokenizer.pad_token = tokenizer.eos_token
     elif args.model == 'llama':
         tokenizer = LlamaTokenizer.from_pretrained(args.pretrain)
         tokenizer.eos_token = '<\s>'
+        tokenizer.pad_token = tokenizer.unk_token
     else:
         raise ValueError(f'Unsupported model "{args.model}"')
 
-    tokenizer.pad_token = tokenizer.eos_token
-
     data_collator = DataCollatorForSupervisedDataset(tokenizer=tokenizer)
 
     prompt_dataset = EasyPromptsDataset(args.prompt_path, tokenizer)
diff --git a/applications/Chat/examples/community/peft/train_peft_sft.py b/applications/Chat/examples/community/peft/train_peft_sft.py
index 924bf7fb7a57..4af08e6d0141 100644
--- a/applications/Chat/examples/community/peft/train_peft_sft.py
+++ b/applications/Chat/examples/community/peft/train_peft_sft.py
@@ -64,10 +64,11 @@ def train(args):
         tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
         tokenizer.pad_token = tokenizer.eos_token
     elif args.model == 'bloom':
-        tokenizer = BloomTokenizerFast.from_pretrained(args.pretrain)
+        tokenizer = BloomTokenizerFast.from_pretrained("bigscience/bloom-560m")
         tokenizer.pad_token = tokenizer.eos_token
     elif args.model == 'opt':
         tokenizer = AutoTokenizer.from_pretrained("facebook/opt-350m")
+        tokenizer.pad_token = tokenizer.eos_token
     elif args.model == 'llama':
         tokenizer = AutoTokenizer.from_pretrained(
             args.pretrain,
@@ -75,9 +76,9 @@ def train(args):
             use_fast=False,
         )
         tokenizer.eos_token = '<\s>'
+        tokenizer.pad_token = tokenizer.unk_token
     else:
         raise ValueError(f'Unsupported model "{args.model}"')
-    tokenizer.pad_token = tokenizer.eos_token
 
     if args.model == 'llama' and args.strategy == 'colossalai_gemini':
         # this is a hack to deal with the resized embedding

From 2ceba11c8f79746cef6b20f70b4309cabb3a3e1b Mon Sep 17 00:00:00 2001
From: CWHer <cwher@outlook.com>
Date: Tue, 4 Jul 2023 10:57:13 +0800
Subject: [PATCH 20/21] revert: undo unnecessary changes

---
 applications/Chat/examples/train_prompts.py      | 2 +-
 applications/Chat/examples/train_reward_model.py | 2 +-
 applications/Chat/examples/train_sft.py          | 2 +-
 3 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/applications/Chat/examples/train_prompts.py b/applications/Chat/examples/train_prompts.py
index 791ba52e2f4c..7338a6d51142 100644
--- a/applications/Chat/examples/train_prompts.py
+++ b/applications/Chat/examples/train_prompts.py
@@ -115,7 +115,7 @@ def main(args):
         tokenizer = AutoTokenizer.from_pretrained("facebook/opt-350m")
         tokenizer.pad_token = tokenizer.eos_token
     elif args.model == 'llama':
-        tokenizer = LlamaTokenizer.from_pretrained("hf-internal-testing/llama-tokenizer")
+        tokenizer = LlamaTokenizer.from_pretrained(args.pretrain)
         tokenizer.eos_token = '<\s>'
         tokenizer.pad_token = tokenizer.unk_token
     else:
diff --git a/applications/Chat/examples/train_reward_model.py b/applications/Chat/examples/train_reward_model.py
index 003f2a0f3597..5b1b8d3d16b2 100644
--- a/applications/Chat/examples/train_reward_model.py
+++ b/applications/Chat/examples/train_reward_model.py
@@ -63,7 +63,7 @@ def train(args):
         tokenizer = AutoTokenizer.from_pretrained("facebook/opt-350m")
         tokenizer.pad_token = tokenizer.eos_token
     elif args.model == 'llama':
-        tokenizer = LlamaTokenizer.from_pretrained("hf-internal-testing/llama-tokenizer")
+        tokenizer = LlamaTokenizer.from_pretrained(args.pretrain)
         tokenizer.pad_token = tokenizer.unk_token
     else:
         raise ValueError(f'Unsupported model "{args.model}"')
diff --git a/applications/Chat/examples/train_sft.py b/applications/Chat/examples/train_sft.py
index acc17769133c..cb3eb649d76c 100644
--- a/applications/Chat/examples/train_sft.py
+++ b/applications/Chat/examples/train_sft.py
@@ -70,7 +70,7 @@ def train(args):
         tokenizer.pad_token = tokenizer.eos_token
     elif args.model == 'llama':
         tokenizer = AutoTokenizer.from_pretrained(
-            "hf-internal-testing/llama-tokenizer",
+            args.pretrain,
             padding_side="right",
             use_fast=False,
         )

From 3d5d4ec3f473fd7edb35dc2fdf00a7fcdd641c54 Mon Sep 17 00:00:00 2001
From: CWHer <cwher@outlook.com>
Date: Tue, 4 Jul 2023 11:34:45 +0800
Subject: [PATCH 21/21] fix: use pip to install transformers

---
 .github/workflows/run_chatgpt_examples.yml | 14 +++++---------
 applications/Chat/README.md                |  5 +----
 2 files changed, 6 insertions(+), 13 deletions(-)

diff --git a/.github/workflows/run_chatgpt_examples.yml b/.github/workflows/run_chatgpt_examples.yml
index 5b1e9383fe17..510f6b6f0985 100644
--- a/.github/workflows/run_chatgpt_examples.yml
+++ b/.github/workflows/run_chatgpt_examples.yml
@@ -4,11 +4,10 @@ on:
   pull_request:
     types: [synchronize, opened, reopened]
     paths:
-      - 'applications/Chat/coati/**'
-      - 'applications/Chat/requirements.txt'
-      - 'applications/Chat/setup.py'
-      - 'applications/Chat/examples/**'
-
+      - "applications/Chat/coati/**"
+      - "applications/Chat/requirements.txt"
+      - "applications/Chat/setup.py"
+      - "applications/Chat/examples/**"
 
 jobs:
   tests:
@@ -38,10 +37,7 @@ jobs:
 
       - name: Install Transformers
         run: |
-          cd applications/Chat
-          git clone --depth 1 --branch v4.30.2 https://github.com/huggingface/transformers
-          cd transformers
-          pip install -v .
+          pip install transformers==4.30.2
 
       - name: Execute Examples
         run: |
diff --git a/applications/Chat/README.md b/applications/Chat/README.md
index 2d36d44a03b6..162528cee414 100644
--- a/applications/Chat/README.md
+++ b/applications/Chat/README.md
@@ -98,12 +98,9 @@ pip install .
 ```
 
 ### Install the Transformers
-Given Hugging Face hasn't officially supported the LLaMA models, We fork a branch of Transformers that can be compatible with our code
 
 ```shell
-git clone --depth 1 --branch v4.30.2 https://github.com/huggingface/transformers.git
-cd transformers
-pip install .
+pip install transformers==4.30.2
 ```
 
 ## How to use?