[ColossalChat] Add PP support

.github/workflows/run_chatgpt_examples.yml

@@ -31,18 +31,18 @@ jobs:
 
       - name: Install Colossal-AI
         run: |
-          BUILD_EXT=1 pip install -v -e .
+          BUILD_EXT=1 pip install --no-cache-dir -v -e .
 
       - name: Install ChatGPT
         run: |
           cd applications/ColossalChat
-          pip install -v .
+          pip install --no-cache-dir -v .
           export BUILD_EXT=1
-          pip install -r examples/requirements.txt
+          pip install --no-cache-dir -r examples/requirements.txt
 
       - name: Install Transformers
         run: |
-          pip install transformers==4.36.2
+          pip install --no-cache-dir transformers==4.36.2
 
       - name: Execute Examples
         run: |

applications/ColossalChat/.gitignore

@@ -161,3 +161,9 @@ applications/ColossalChat/sft_data
 applications/ColossalChat/prompt_data
 applications/ColossalChat/preference_data
 applications/ColossalChat/temp
+
+# Testing data
+/kto_data/
+/preference_data/
+/prompt_data/
+/sft_data/

applications/ColossalChat/coati/trainer/base.py

@@ -16,7 +16,7 @@
 from coati.experience_maker import Experience
 from torch.optim import Optimizer
 
-from colossalai.booster import Booster
+from colossalai.booster import Booster, Plugin
 
 from .utils import is_rank_0
 
@@ -38,13 +38,15 @@ def __init__(
         max_epochs: int,
         model: nn.Module,
         optimizer: Optimizer,
+        plugin: Plugin,
         start_epoch: int = 0,
     ) -> None:
         super().__init__()
         self.booster = booster
         self.max_epochs = max_epochs
         self.model = model
         self.optimizer = optimizer
+        self.plugin = plugin
         self.start_epoch = start_epoch
 
     @abstractmethod

applications/ColossalChat/coati/trainer/dpo.py

@@ -16,7 +16,7 @@
 from tqdm import trange
 from transformers import PreTrainedTokenizerBase
 
-from colossalai.booster import Booster
+from colossalai.booster import Booster, Plugin
 from colossalai.cluster import DistCoordinator
 from colossalai.utils import get_current_device
 
@@ -50,6 +50,7 @@ def __init__(
         ref_model: Any,
         booster: Booster,
         actor_optim: Optimizer,
+        plugin: Plugin,
         actor_lr_scheduler: _LRScheduler,
         tokenizer: PreTrainedTokenizerBase,
         max_epochs: int = 1,
@@ -63,7 +64,9 @@ def __init__(
         save_dir: str = None,
         coordinator: DistCoordinator = None,
     ) -> None:
-        super().__init__(booster, max_epochs=max_epochs, model=actor, optimizer=actor_optim, start_epoch=start_epoch)
+        super().__init__(
+            booster, max_epochs=max_epochs, model=actor, optimizer=actor_optim, plugin=plugin, start_epoch=start_epoch
+        )
         self.ref_model = ref_model
         self.actor_scheduler = actor_lr_scheduler
         self.tokenizer = tokenizer

applications/ColossalChat/coati/trainer/kto.py

@@ -17,7 +17,7 @@
 from tqdm import trange
 from transformers import PreTrainedTokenizerBase
 
-from colossalai.booster import Booster
+from colossalai.booster import Booster, Plugin
 from colossalai.cluster import DistCoordinator
 from colossalai.utils import get_current_device
 
@@ -53,6 +53,7 @@ def __init__(
         ref_model: Any,
         booster: Booster,
         actor_optim: Optimizer,
+        plugin: Plugin,
         actor_lr_scheduler: _LRScheduler,
         tokenizer: PreTrainedTokenizerBase,
         max_epochs: int = 1,
@@ -66,7 +67,9 @@ def __init__(
         save_dir: str = None,
         coordinator: DistCoordinator = None,
     ) -> None:
-        super().__init__(booster, max_epochs=max_epochs, model=actor, optimizer=actor_optim, start_epoch=start_epoch)
+        super().__init__(
+            booster, max_epochs=max_epochs, model=actor, optimizer=actor_optim, plugin=plugin, start_epoch=start_epoch
+        )
         self.ref_model = ref_model
         self.actor_scheduler = actor_lr_scheduler
         self.tokenizer = tokenizer

applications/ColossalChat/coati/trainer/orpo.py

@@ -16,7 +16,7 @@
 from tqdm import trange
 from transformers import PreTrainedTokenizerBase
 
-from colossalai.booster import Booster
+from colossalai.booster import Booster, Plugin
 from colossalai.cluster import DistCoordinator
 from colossalai.utils import get_current_device
 
@@ -48,6 +48,7 @@ def __init__(
         actor: Any,
         booster: Booster,
         actor_optim: Optimizer,
+        plugin: Plugin,
         actor_lr_scheduler: _LRScheduler,
         tokenizer: PreTrainedTokenizerBase,
         max_epochs: int = 1,
@@ -59,7 +60,9 @@ def __init__(
         save_dir: str = None,
         coordinator: DistCoordinator = None,
     ) -> None:
-        super().__init__(booster, max_epochs=max_epochs, model=actor, optimizer=actor_optim, start_epoch=start_epoch)
+        super().__init__(
+            booster, max_epochs=max_epochs, model=actor, optimizer=actor_optim, plugin=plugin, start_epoch=start_epoch
+        )
         self.actor_scheduler = actor_lr_scheduler
         self.tokenizer = tokenizer
         self.odds_ratio_loss_fn = OddsRatioLoss()

applications/ColossalChat/coati/trainer/rm.py

@@ -15,7 +15,7 @@
 from torch.utils.data import DataLoader
 from transformers import PreTrainedTokenizerBase
 
-from colossalai.booster import Booster
+from colossalai.booster import Booster, Plugin
 from colossalai.cluster import DistCoordinator
 from colossalai.utils import get_current_device
 
@@ -48,6 +48,7 @@ def __init__(
         model: Any,
         booster: Booster,
         optimizer: Optimizer,
+        plugin: Plugin,
         lr_scheduler: _LRScheduler,
         tokenizer: PreTrainedTokenizerBase,
         loss_fn: Optional[Callable] = None,
@@ -59,7 +60,9 @@ def __init__(
         save_dir: str = None,
         coordinator: DistCoordinator = None,
     ) -> None:
-        super().__init__(booster, max_epochs=max_epochs, model=model, optimizer=optimizer, start_epoch=start_epoch)
+        super().__init__(
+            booster, max_epochs=max_epochs, model=model, optimizer=optimizer, plugin=plugin, start_epoch=start_epoch
+        )
         self.actor_scheduler = lr_scheduler
         self.tokenizer = tokenizer
         self.loss_fn = loss_fn if loss_fn is not None else LogSigLoss(beta=beta)