Merge branch 'main' into feature/fp8_comm

.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)