Co

applications/Chat/.gitignore

@@ -145,4 +145,4 @@ docs/.build
 # wandb log
 example/wandb/
 
-examples/awesome-chatgpt-prompts/
+examples/awesome-chatgpt-prompts/

applications/Chat/coati/dataset/sft_dataset.py

@@ -74,15 +74,10 @@ def __getitem__(self, idx):
         return dict(input_ids=self.input_ids[idx], labels=self.labels[idx])
 
 
-def _tokenize_fn(strings: Sequence[str],
-                 tokenizer: transformers.PreTrainedTokenizer,
-                 max_length: int
-                 ) -> Dict[str, torch.Tensor]:
+def _tokenize_fn(strings: Sequence[str], tokenizer: transformers.PreTrainedTokenizer,
+                 max_length: int) -> Dict[str, torch.Tensor]:
     """Tokenize a list of strings."""
-    tokenized_list = tokenizer(
-        strings, return_tensors="pt", padding="longest",
-        max_length=max_length, truncation=True
-    )
+    tokenized_list = tokenizer(strings, return_tensors="pt", padding="longest", max_length=max_length, truncation=True)
     input_ids = labels = tokenized_list["input_ids"]
     input_ids_lens = labels_lens = \
         tokenized_list["input_ids"].ne(tokenizer.pad_token_id).sum(dim=-1)
@@ -103,8 +98,7 @@ def preprocess(
     """Preprocess the data by tokenizing."""
     examples = [s + t for s, t in zip(sources, targets)]
     examples_tokenized, sources_tokenized = [
-        _tokenize_fn(strings, tokenizer, max_length)
-        for strings in (examples, sources)
+        _tokenize_fn(strings, tokenizer, max_length) for strings in (examples, sources)
     ]
     input_ids = examples_tokenized["input_ids"]
     labels = copy.deepcopy(input_ids)
@@ -116,7 +110,11 @@ def preprocess(
 class SupervisedDataset(Dataset):
     """Dataset for supervised fine-tuning."""
 
-    def __init__(self, data_path: str, tokenizer: transformers.PreTrainedTokenizer, max_datasets_size: int = None, max_length: int = 512):
+    def __init__(self,
+                 data_path: str,
+                 tokenizer: transformers.PreTrainedTokenizer,
+                 max_datasets_size: int = None,
+                 max_length: int = 512):
         super(SupervisedDataset, self).__init__()
         logger.info("Loading data...")
         list_data_dict = jload(data_path)

applications/Chat/coati/models/base/actor.py

@@ -21,16 +21,13 @@ def __init__(self, model: nn.Module, lora_rank: int = 0, lora_train_bias: str =
         self.model = model
         self.convert_to_lora()
 
-    def forward(self,
-                input_ids: torch.LongTensor,
-                attention_mask: Optional[torch.Tensor] = None,
-                **model_kwargs,  # HACK: `generate` method may pass more kwargs
-                ) -> torch.Tensor:
+    def forward(
+            self,
+            input_ids: torch.LongTensor,
+            attention_mask: Optional[torch.Tensor] = None,
+            **model_kwargs,    # HACK: `generate` method may pass more kwargs
+    ) -> torch.Tensor:
         """Returns model output.
         """
-        output = self.model(
-            input_ids,
-            attention_mask=attention_mask,
-            **model_kwargs
-        )
+        output = self.model(input_ids, attention_mask=attention_mask, **model_kwargs)
         return output

applications/Chat/coati/models/generation.py

@@ -5,7 +5,6 @@
 import torch.nn as nn
 import torch.nn.functional as F
 
-
 try:
     from transformers.generation_logits_process import (
         LogitsProcessorList,
@@ -148,12 +147,12 @@ def generate(model: nn.Module,
 
 
 @torch.no_grad()
-def generate_with_actor(actor_model: nn.Module,
-                        input_ids: torch.Tensor,
-                        return_action_mask: bool = True,
-                        **kwargs
-                        ) -> Union[Tuple[torch.LongTensor, torch.LongTensor],
-                                   Tuple[torch.LongTensor, torch.LongTensor, torch.BoolTensor]]:
+def generate_with_actor(
+    actor_model: nn.Module,
+    input_ids: torch.Tensor,
+    return_action_mask: bool = True,
+    **kwargs
+) -> Union[Tuple[torch.LongTensor, torch.LongTensor], Tuple[torch.LongTensor, torch.LongTensor, torch.BoolTensor]]:
     """Generate token sequence with actor model. Refer to `generate` for more details.
     """
     # generate sequences

applications/Chat/coati/models/utils.py

@@ -46,10 +46,7 @@ def log_probs_from_logits(logits: torch.Tensor, labels: torch.Tensor) -> torch.T
     return log_probs_labels.squeeze(-1)
 
 
-def calc_action_log_probs(output: torch.Tensor,
-                          sequences: torch.LongTensor,
-                          num_actions: int
-                          ) -> torch.Tensor:
+def calc_action_log_probs(output: torch.Tensor, sequences: torch.LongTensor, num_actions: int) -> torch.Tensor:
     """Calculate action log probs.
 
     Args:

applications/Chat/coati/trainer/base.py

@@ -25,12 +25,13 @@ class SLTrainer(ABC):
         optim (Optimizer): the optimizer to use for training
     """
 
-    def __init__(self,
-                 strategy: Strategy,
-                 max_epochs: int,
-                 model: nn.Module,
-                 optimizer: Optimizer,
-                 ) -> None:
+    def __init__(
+        self,
+        strategy: Strategy,
+        max_epochs: int,
+        model: nn.Module,
+        optimizer: Optimizer,
+    ) -> None:
         super().__init__()
         self.strategy = strategy
         self.max_epochs = max_epochs
@@ -50,10 +51,7 @@ def _before_fit(self):
 
     def fit(self, *args, **kwargs):
         self._before_fit(*args, **kwargs)
-        for epoch in tqdm.trange(self.max_epochs,
-                                 desc="Epochs",
-                                 disable=not is_rank_0() or self.no_epoch_bar
-                                 ):
+        for epoch in tqdm.trange(self.max_epochs, desc="Epochs", disable=not is_rank_0() or self.no_epoch_bar):
             self._train(epoch)
             self._eval(epoch)
 
@@ -75,8 +73,7 @@ def __init__(self,
                  buffer: NaiveReplayBuffer,
                  sample_buffer: bool,
                  dataloader_pin_memory: bool,
-                 callbacks: List[Callback] = []
-                 ) -> None:
+                 callbacks: List[Callback] = []) -> None:
         super().__init__()
         self.strategy = strategy
         self.buffer = buffer
@@ -138,7 +135,7 @@ def _make_experience(self, collect_step: int):
     @abstractmethod
     def _learn(self, update_step: int):
         """
-        Implement this method to learn from experience, either 
+        Implement this method to learn from experience, either
         sample from buffer or transform buffer into dataloader.
         """
         raise NotImplementedError()
@@ -154,13 +151,14 @@ def _update_phase(self, update_step: int):
         self._learn(update_step)
         self._on_learn_epoch_end(update_step)
 
-    def fit(self,
-            prompt_dataloader: DataLoader,
-            pretrain_dataloader: DataLoader,
-            num_episodes: int,
-            num_collect_steps: int,
-            num_update_steps: int,
-            ):
+    def fit(
+        self,
+        prompt_dataloader: DataLoader,
+        pretrain_dataloader: DataLoader,
+        num_episodes: int,
+        num_collect_steps: int,
+        num_update_steps: int,
+    ):
         """
         The main training loop of on-policy rl trainers.
 
@@ -175,23 +173,16 @@ def fit(self,
         self.pretrain_dataloader = CycledDataLoader(pretrain_dataloader)
 
         with self._fit_ctx():
-            for episode in tqdm.trange(num_episodes,
-                                       desc="Episodes",
-                                       disable=not is_rank_0()):
+            for episode in tqdm.trange(num_episodes, desc="Episodes", disable=not is_rank_0()):
                 with self._episode_ctx(episode):
-                    for collect_step in tqdm.trange(num_collect_steps,
-                                                    desc="Collect steps",
-                                                    disable=not is_rank_0()):
+                    for collect_step in tqdm.trange(num_collect_steps, desc="Collect steps", disable=not is_rank_0()):
                         self._collect_phase(collect_step)
                     if not self.sample_buffer:
                         # HACK(cwher): according to the design of boost API, dataloader should also be boosted,
                         #  but it is impractical to adapt this pattern in RL training. Thus, I left dataloader unboosted.
                         #  I only call strategy.setup_dataloader() to setup dataloader.
-                        self.dataloader = self.strategy.setup_dataloader(self.buffer,
-                                                                         self.dataloader_pin_memory)
-                    for update_step in tqdm.trange(num_update_steps,
-                                                   desc="Update steps",
-                                                   disable=not is_rank_0()):
+                        self.dataloader = self.strategy.setup_dataloader(self.buffer, self.dataloader_pin_memory)
+                    for update_step in tqdm.trange(num_update_steps, desc="Update steps", disable=not is_rank_0()):
                         self._update_phase(update_step)
                     # NOTE: this is for on-policy algorithms
                     self.buffer.clear()

applications/Chat/coati/trainer/strategies/base.py

@@ -79,8 +79,7 @@ def prepare(self, *boost_args: _BoostArgSpec) -> Union[List[_BoostArgSpec], _Boo
                     model, optimizer = arg
                 except ValueError:
                     raise RuntimeError(f'Expect (model, optimizer) pair, got a tuple with size "{len(arg)}"')
-                model, optimizer, *_ = self.booster.boost(model=model,
-                                                          optimizer=optimizer)
+                model, optimizer, *_ = self.booster.boost(model=model, optimizer=optimizer)
                 rets.append((model, optimizer))
             elif isinstance(arg, Dict):
                 model, optimizer, criterion, dataloader, lr_scheduler = self.booster.boost(**arg)
@@ -90,10 +89,7 @@ def prepare(self, *boost_args: _BoostArgSpec) -> Union[List[_BoostArgSpec], _Boo
                                     dataloader=dataloader,
                                     lr_scheduler=lr_scheduler)
                 # remove None values
-                boost_result = {
-                    key: value
-                    for key, value in boost_result.items() if value is not None
-                }
+                boost_result = {key: value for key, value in boost_result.items() if value is not None}
                 rets.append(boost_result)
             else:
                 raise RuntimeError(f'Type {type(arg)} is not supported')
@@ -112,23 +108,13 @@ def unwrap_model(model: nn.Module) -> nn.Module:
         """
         return model
 
-    def save_model(self,
-                   model: nn.Module,
-                   path: str,
-                   only_rank0: bool = True,
-                   **kwargs
-                   ) -> None:
+    def save_model(self, model: nn.Module, path: str, only_rank0: bool = True, **kwargs) -> None:
         self.booster.save_model(model, path, shard=not only_rank0, **kwargs)
 
     def load_model(self, model: nn.Module, path: str, strict: bool = True) -> None:
         self.booster.load_model(model, path, strict)
 
-    def save_optimizer(self,
-                       optimizer: Optimizer,
-                       path: str,
-                       only_rank0: bool = False,
-                       **kwargs
-                       ) -> None:
+    def save_optimizer(self, optimizer: Optimizer, path: str, only_rank0: bool = False, **kwargs) -> None:
         self.booster.save_optimizer(optimizer, path, shard=not only_rank0, **kwargs)
 
     def load_optimizer(self, optimizer: Optimizer, path: str) -> None:

applications/Chat/coati/trainer/utils.py

@@ -14,9 +14,10 @@ class CycledDataLoader:
     NOTE: next(iter(dataloader)) is not equivalent to for batch in dataloader: break, it causes slightly different behavior.
     """
 
-    def __init__(self,
-                 dataloader: DataLoader,
-                 ) -> None:
+    def __init__(
+        self,
+        dataloader: DataLoader,
+    ) -> None:
         self.dataloader = dataloader
 
         self.count = 0

applications/Chat/examples/ray/mmmt_prompt.py

@@ -87,8 +87,8 @@ def model_fn():
             kl_coef=0.1,
             debug=args.debug,
             update_lora_weights=not (args.lora_rank == 0),
-            # sync_models_from_trainers=True,
-            # generation kwargs:
+    # sync_models_from_trainers=True,
+    # generation kwargs:
             max_length=512,
             do_sample=True,
             temperature=1.0,
@@ -161,12 +161,10 @@ def tokenize_fn(texts):
     parser.add_argument('--prompt_path', type=str, default=None)
     parser.add_argument('--num_makers', type=int, default=1)
     parser.add_argument('--num_trainers', type=int, default=1)
-    parser.add_argument('--trainer_strategy',
-                        choices=[
-                            'ddp', 'colossalai_gemini', 'colossalai_zero2', 'colossalai_gemini_cpu',
-                            'colossalai_zero2_cpu'
-                        ],
-                        default='ddp')
+    parser.add_argument(
+        '--trainer_strategy',
+        choices=['ddp', 'colossalai_gemini', 'colossalai_zero2', 'colossalai_gemini_cpu', 'colossalai_zero2_cpu'],
+        default='ddp')
     parser.add_argument('--maker_strategy', choices=['naive'], default='naive')
     parser.add_argument('--model', default='gpt2', choices=['gpt2', 'bloom', 'opt', 'llama'])
     parser.add_argument('--critic_model', default='gpt2', choices=['gpt2', 'bloom', 'opt', 'llama'])

applications/Chat/examples/train_reward_model.py

@@ -150,9 +150,7 @@ def train(args):
                                  pin_memory=True)
 
     lr_scheduler = CosineAnnealingLR(optim, train_dataloader.__len__() // 100)
-    strategy_dict = strategy.prepare(
-        dict(model=model, optimizer=optim, lr_scheduler=lr_scheduler)
-    )
+    strategy_dict = strategy.prepare(dict(model=model, optimizer=optim, lr_scheduler=lr_scheduler))
     model = strategy_dict['model']
     optim = strategy_dict['optimizer']
     lr_scheduler = strategy_dict['lr_scheduler']
@@ -163,9 +161,7 @@ def train(args):
                                  loss_fn=loss_fn,
                                  max_epochs=args.max_epochs)
 
-    trainer.fit(train_dataloader=train_dataloader,
-                valid_dataloader=valid_dataloader,
-                eval_dataloader=eval_dataloader)
+    trainer.fit(train_dataloader=train_dataloader, valid_dataloader=valid_dataloader, eval_dataloader=eval_dataloader)
     # save model checkpoint after fitting on only rank0
     strategy.save_model(model, args.save_path, only_rank0=True)
     # save optimizer checkpoint on all ranks

applications/Chat/inference/requirements.txt

@@ -10,4 +10,4 @@ uvicorn
 git+https://github.com/huggingface/transformers
 accelerate
 bitsandbytes
-jieba
+jieba

applications/Chat/inference/server.py

@@ -14,7 +14,7 @@
 from slowapi.util import get_remote_address
 from sse_starlette.sse import EventSourceResponse
 from transformers import AutoTokenizer, GenerationConfig, LlamaForCausalLM
-from utils import ChatPromptProcessor, Dialogue, LockedIterator, sample_streamingly, update_model_kwargs_fn, load_json
+from utils import ChatPromptProcessor, Dialogue, LockedIterator, load_json, sample_streamingly, update_model_kwargs_fn
 
 CONTEXT = 'Below is an instruction that describes a task. Write a response that appropriately completes the request. Do not generate new instructions.'
 MAX_LEN = 512
@@ -145,7 +145,9 @@ def generate_no_stream(data: GenerationTaskReq, request: Request):
                         help='Group size for GPTQ. This is only useful when quantization mode is 4bit. Default: 128.')
     parser.add_argument('--http_host', default='0.0.0.0')
     parser.add_argument('--http_port', type=int, default=7070)
-    parser.add_argument('--profanity_file', default=None, help='Path to profanity words list. It should be a JSON file containing a list of words.')
+    parser.add_argument('--profanity_file',
+                        default=None,
+                        help='Path to profanity words list. It should be a JSON file containing a list of words.')
     args = parser.parse_args()
 
     if args.quant == '4bit':

colossalai/auto_parallel/offload/amp_optimizer.py

@@ -1,24 +1,25 @@
-from typing import Dict, Tuple
 from enum import Enum
+from typing import Dict, Tuple
+
 import torch
 from torch.optim import Optimizer
 
+from colossalai.amp.naive_amp.grad_scaler import DynamicGradScaler
 from colossalai.logging import get_dist_logger
 from colossalai.nn.optimizer import ColossalaiOptimizer
-from colossalai.amp.naive_amp.grad_scaler import DynamicGradScaler
 from colossalai.utils import get_current_device
 
 from .base_offload_module import BaseOffloadModule
-from .region_manager import RegionManager
 from .region import Region
+from .region_manager import RegionManager
 
 
 class OptimState(Enum):
     SCALED = 0
     UNSCALED = 1
 
-class AMPOptimizer(ColossalaiOptimizer):
 
+class AMPOptimizer(ColossalaiOptimizer):
     """
     A wrapper for Optimizer.
     Code reference: https://github.com/hpcaitech/ColossalAI/blob/main/colossalai/nn/optimizer/zero_optimizer.py
@@ -174,4 +175,4 @@ def __init__optimizer(self):
 
         # Leverage state_dict() and load_state_dict() to
         # recast preexisting per-param state tensors
-        self.optim.load_state_dict(self.optim.state_dict())
+        self.optim.load_state_dict(self.optim.state_dict())

colossalai/auto_parallel/passes/runtime_preparation_pass.py

@@ -55,7 +55,7 @@ def size_processing(size: Union[int, torch.Size],
 
 
 def solution_annotation_pass(gm: torch.fx.GraphModule, solution: List[int],
-                               strategies_constructor: StrategiesConstructor):
+                             strategies_constructor: StrategiesConstructor):
     """
     This method is used to stick the solution strategy to the nodes and add the information
     required in runtime into graph as placeholder nodes.