[MoE/ZeRO] fix .github conflict with main branch.

.cuda_ext.json

@@ -7,10 +7,6 @@
     {
       "torch_command": "pip install torch==2.1.0 torchvision==0.16.0 torchaudio==2.1.0 --index-url https://download.pytorch.org/whl/cu118",
       "cuda_image": "hpcaitech/cuda-conda:11.8"
-    },
-    {
-      "torch_command": "pip install torch==2.0.0 torchvision==0.15.1 torchaudio==2.0.1",
-      "cuda_image": "hpcaitech/cuda-conda:11.7"
     }
   ]
 }

.github/workflows/compatiblity_test_on_dispatch.yml

@@ -51,11 +51,11 @@ jobs:
     container:
       image: ${{ matrix.container }}
       options: --gpus all --rm -v /dev/shm -v /data/scratch/:/data/scratch/
-    timeout-minutes: 120
+    timeout-minutes: 200
     steps:
       - name: Install dependencies
         run: |
-          pip install -U pip setuptools wheel --user
+          pip install -U pip setuptools==68.2.2 wheel --user
       - uses: actions/checkout@v2
         with:
           repository: hpcaitech/TensorNVMe

.github/workflows/compatiblity_test_on_pr.yml

@@ -42,14 +42,14 @@ jobs:
     container:
       image: ${{ matrix.container }}
       options: --gpus all --rm -v /dev/shm -v /data/scratch/:/data/scratch/
-    timeout-minutes: 120
+    timeout-minutes: 200
     concurrency:
       group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}-run-test-${{ matrix.container }}
       cancel-in-progress: true
     steps:
       - name: Install dependencies
         run: |
-          pip install -U pip setuptools wheel --user
+          pip install -U pip setuptools==68.2.2 wheel --user
       - uses: actions/checkout@v2
         with:
           repository: hpcaitech/TensorNVMe

.github/workflows/compatiblity_test_on_schedule.yml

@@ -39,11 +39,11 @@ jobs:
     container:
       image: ${{ matrix.container }}
       options: --gpus all --rm -v /dev/shm -v /data/scratch/:/data/scratch/
-    timeout-minutes: 120
+    timeout-minutes: 200
     steps:
       - name: Install dependencies
         run: |
-          pip install -U pip setuptools wheel --user
+          pip install -U pip setuptools==68.2.2 wheel --user
 
       - uses: actions/checkout@v2
         with:

.github/workflows/release_docker_after_publish.yml

@@ -28,6 +28,8 @@ jobs:
           docker tag $tag $latest
           echo "tag=${tag}" >> $GITHUB_OUTPUT
           echo "latest=${latest}" >> $GITHUB_OUTPUT
+        env:
+          DOCKER_BUILDKIT: 0
 
       - name: Log in to Docker Hub
         uses: docker/login-action@f054a8b539a109f9f41c372932f1ae047eff08c9

.github/workflows/run_chatgpt_examples.yml

@@ -4,10 +4,11 @@ on:
   pull_request:
     types: [synchronize, opened, reopened]
     paths:
-      - "applications/Chat/coati/**"
-      - "applications/Chat/requirements.txt"
-      - "applications/Chat/setup.py"
-      - "applications/Chat/examples/**"
+      - "applications/ColossalChat/coati/**"
+      - "applications/ColossalChat/requirements.txt"
+      - "applications/ColossalChat/setup.py"
+      - "applications/ColossalChat/examples/**"
+      - "applications/ColossalChat/tests/**"
 
 jobs:
   tests:
@@ -41,7 +42,7 @@ jobs:
 
       - name: Install Transformers
         run: |
-          pip install transformers==4.34.1
+          pip install transformers==4.36.2
 
       - name: Execute Examples
         run: |

.github/workflows/run_chatgpt_unit_tests.yml

@@ -4,12 +4,11 @@ on:
   pull_request:
     types: [synchronize, opened, reopened]
     paths:
-      - 'applications/Chat/coati/**'
-      - 'applications/Chat/requirements.txt'
-      - 'applications/Chat/setup.py'
-      - 'applications/Chat/requirements-test.txt'
-      - 'applications/Chat/tests/**'
-      - 'applications/Chat/pytest.ini'
+      - 'applications/ColossalChat/coati/**'
+      - 'applications/ColossalChat/requirements.txt'
+      - 'applications/ColossalChat/setup.py'
+      - 'applications/ColossalChat/tests/**'
+      - 'applications/ColossalChat/pytest.ini'
 
 jobs:
   tests:

applications/ColossalChat/coati/dataset/__init__.py

@@ -5,7 +5,6 @@
     DataCollatorForSupervisedDataset,
     StatefulDistributedSampler,
     load_tokenized_dataset,
-    setup_distributed_dataloader,
 )
 from .tokenization_utils import supervised_tokenize_sft, tokenize_prompt_dataset, tokenize_rlhf
 
@@ -17,7 +16,6 @@
     "DataCollatorForSupervisedDataset",
     "StatefulDistributedSampler",
     "load_tokenized_dataset",
-    "setup_distributed_dataloader",
     "supervised_tokenize_pretrain",
     "supervised_tokenize_sft",
     "tokenize_rlhf",

applications/ColossalChat/coati/dataset/conversation.py

@@ -17,14 +17,17 @@ class Conversation:
     system_message: str
     chat_template: str
     stop_ids: List[int]
+    end_of_assistant: str
 
     @classmethod
     def from_config(cls, tokenizer: PreTrainedTokenizer, config: Dict):
         """
         Setup the conversation template from config
         """
         tokenizer.chat_template = config["chat_template"]
-        conv = cls(tokenizer, config["system_message"], config["chat_template"], config["stop_ids"])
+        conv = cls(
+            tokenizer, config["system_message"], config["chat_template"], config["stop_ids"], config["end_of_assistant"]
+        )
         conv.clear()
         return conv
 
@@ -109,6 +112,8 @@ def setup_conversation_template(
     """
     if any([s not in chat_template_config.keys() for s in Conversation.get_conversation_template_keys()]):
         # Try to automatically set up conversation template, if fail, it throws an error that you need to do it manually
+        if "end_of_assistant" not in chat_template_config:
+            raise ValueError("Please set the end of assistant token.")
         if "system_message" not in chat_template_config:
             logger.warning("No system message is provided, will not use system message.")
         if "chat_template" not in chat_template_config:

applications/ColossalChat/coati/dataset/loader.py

@@ -4,22 +4,16 @@
 Dataloader for sft, dpo, ppo
 """
 
-import math
 import os
-import random
 from dataclasses import dataclass
-from typing import Callable, Dict, Iterator, List, Optional, Sequence, Union
+from typing import Dict, Iterator, List, Optional, Sequence, Union
 
-import numpy as np
 import torch
-import torch.distributed as dist
 import torch.nn.functional as F
 from coati.dataset.utils import chuncate_sequence, pad_to_max_len
 from datasets import Dataset as HFDataset
 from datasets import dataset_dict, load_from_disk
-from torch.distributed import ProcessGroup
-from torch.distributed.distributed_c10d import _get_default_group
-from torch.utils.data import ConcatDataset, DataLoader, Dataset, DistributedSampler
+from torch.utils.data import ConcatDataset, Dataset, DistributedSampler
 from transformers.tokenization_utils import PreTrainedTokenizer
 
 DatasetType = Union[Dataset, ConcatDataset, dataset_dict.Dataset]
@@ -236,148 +230,26 @@ def __call__(self, instances: Sequence[Dict[str, List[int]]]) -> Dict[str, torch
 
 
 class StatefulDistributedSampler(DistributedSampler):
-    """
-    Stateful distributed sampler for multi-stage training.
-    """
-
     def __init__(
         self,
-        dataset: DatasetType,
+        dataset: Dataset,
         num_replicas: Optional[int] = None,
         rank: Optional[int] = None,
         shuffle: bool = True,
         seed: int = 0,
         drop_last: bool = False,
-        use_tp: Optional[bool] = False,
     ) -> None:
-        if not use_tp:
-            super().__init__(
-                dataset=dataset,
-                num_replicas=num_replicas,
-                rank=rank,
-                shuffle=shuffle,
-                seed=seed,
-                drop_last=drop_last,
-            )
-        else:
-            # adapted from https://github.com/pytorch/pytorch/blob/4979f9c0d72490970e2019bb1d2284f83d93f76b/torch/utils/data/distributed.py#L62
-            # TODO: support tp_group>1. will fix it later
-            num_replicas = 1
-            if rank is None:
-                rank = dist.get_rank()
-            if rank < 0:
-                raise ValueError(f"Invalid rank {rank}, rank should be in the interval [0, 0]")
-            self.dataset = dataset
-            self.num_replicas = num_replicas
-            self.rank = rank
-            self.epoch = 0
-            self.drop_last = drop_last
-            # If the dataset length is evenly divisible by # of replicas, then there
-            # is no need to drop any data, since the dataset will be split equally.
-            if self.drop_last and len(self.dataset) % self.num_replicas != 0:  # type: ignore[arg-type]
-                # Split to nearest available length that is evenly divisible.
-                # This is to ensure each rank receives the same amount of data when
-                # using this Sampler.
-                self.num_samples = math.ceil(
-                    (len(self.dataset) - self.num_replicas) / self.num_replicas  # type: ignore[arg-type]
-                )
-            else:
-                self.num_samples = math.ceil(len(self.dataset) / self.num_replicas)  # type: ignore[arg-type]
-            self.total_size = self.num_samples * self.num_replicas
-            self.shuffle = shuffle
-            self.seed = seed
-        self.start_index = 0
-        self.use_tp = use_tp
+        super().__init__(dataset, num_replicas, rank, shuffle, seed, drop_last)
+        self.start_index: int = 0
 
     def __iter__(self) -> Iterator:
-        if self.use_tp:
-            # TODO Add support for tp_group not equal to 1
-            pass
-            # adpated from https://github.com/pytorch/pytorch/blob/4979f9c0d72490970e2019bb1d2284f83d93f76b/torch/utils/data/distributed.py#L96
-            if self.shuffle:
-                # deterministically shuffle based on epoch and seed
-                g = torch.Generator()
-                g.manual_seed(self.seed + self.epoch)
-                indices = torch.randperm(len(self.dataset), generator=g).tolist()  # type: ignore[arg-type]
-            else:
-                indices = list(range(len(self.dataset)))  # type: ignore[arg-type]
-
-            if not self.drop_last:
-                # add extra samples to make it evenly divisible
-                padding_size = self.total_size - len(indices)
-                if padding_size <= len(indices):
-                    indices += indices[:padding_size]
-                else:
-                    indices += (indices * math.ceil(padding_size / len(indices)))[:padding_size]
-            else:
-                # remove tail of data to make it evenly divisible.
-                indices = indices[: self.total_size]
-            assert len(indices) == self.total_size
-
-            # subsample
-            indices = indices[
-                : self.total_size : self.num_replicas
-            ]  # num_replicas=tp_group=1, we only support tp_group==1 for now
-            assert len(indices) == self.num_samples
-
-            return iter(indices)
-
-        else:
-            iterator = super().__iter__()
-            indices = list(iterator)
-            indices = indices[self.start_index :]
-            return iter(indices)
+        iterator = super().__iter__()
+        indices = list(iterator)
+        indices = indices[self.start_index :]
+        return iter(indices)
 
     def __len__(self) -> int:
         return self.num_samples - self.start_index
 
     def set_start_index(self, start_index: int) -> None:
         self.start_index = start_index
-
-
-def setup_distributed_dataloader(
-    dataset: DatasetType,
-    batch_size: int = 1,
-    shuffle: bool = False,
-    seed: int = 1024,
-    drop_last: bool = False,
-    pin_memory: bool = False,
-    num_workers: int = 0,
-    collate_fn: Callable[[Sequence[Dict[str, Union[str, List[int]]]]], Dict[str, torch.Tensor]] = None,
-    process_group: Optional[ProcessGroup] = None,
-    use_tp: Optional[bool] = False,
-    **kwargs,
-) -> DataLoader:
-    """
-    Setup dataloader for distributed training.
-    """
-    _kwargs = kwargs.copy()
-    process_group = process_group or _get_default_group()
-    sampler = StatefulDistributedSampler(
-        dataset=dataset,
-        num_replicas=process_group.size() if not use_tp else 1,
-        rank=process_group.rank(),
-        shuffle=shuffle,
-        seed=seed,
-        drop_last=drop_last,
-        use_tp=use_tp,
-    )
-
-    # Deterministic dataloader
-    def seed_worker(worker_id: int) -> None:
-        worker_seed = seed
-        np.random.seed(worker_seed)
-        torch.manual_seed(worker_seed)
-        random.seed(worker_seed)
-
-    return DataLoader(
-        dataset=dataset,
-        batch_size=batch_size,
-        sampler=sampler,
-        num_workers=num_workers,
-        collate_fn=collate_fn,
-        pin_memory=pin_memory,
-        drop_last=drop_last,
-        worker_init_fn=seed_worker,
-        **_kwargs,
-    )