Add DistributedDataParallel training

Add the most basic parallelism to the framework, through PyTorch's DDP.
Adjust the dataloaders to also use distributed samplers.

Add other goodies for distributed logging + distributed processing.

optimus/trainer.py

@@ -8,6 +8,8 @@ import torch.nn as nn
 from torch.utils.data import DataLoader
 from transformers.tokenization_utils_base import PreTrainedTokenizerBase
 
+from optimus.utils import dist_utils, logging_utils
+
 
 logger = logging.getLogger(__name__)
 
@@ -22,6 +24,8 @@ class TrainingArguments():
         log_steps (int): Log training progress each number of steps. This
             considers number of updates, so gradient_accumulation_steps
             influences this.
+        log_on_main_rank_only (bool, defaults to True): Whether the main rank
+            should be the only one that logs training information.
         show_progress (bool, defaults to True): Whether to show progress during
             training.
         seed (int): Seed used for reproducibility purposes.
@@ -57,10 +61,12 @@ class TrainingArguments():
         checkpoints_dir: Path,
         save_steps: int,
         save_limit: int,
+        log_on_main_rank_only: bool = True,
         show_progress: bool = True,
     ):
         self.device = device
         self.log_steps = log_steps
+        self.log_on_main_rank_only = log_on_main_rank_only
         self.show_progress = show_progress
         self.seed = seed
         self.optimizer = optimizer
@@ -106,6 +112,9 @@ class Trainer():
         self.eval_dataloader = eval_dataloader
         self.tokenizer = tokenizer
 
+        if self.args.log_on_main_rank_only:
+            logger.addFilter(logging_utils.FilterMainRankOnly())
+
     def train(self) -> None:
         """
         Training loop of the trainer.
@@ -114,13 +123,13 @@ class Trainer():
         num_examples = len(self.train_dataloader) * self.args.per_device_batch_size
         num_update_steps_per_epoch = len(self.train_dataloader) // self.args.gradient_accumulation_steps
         max_steps = self.args.num_train_epochs * num_update_steps_per_epoch
-        global_batch_size = self.args.per_device_batch_size * self.args.gradient_accumulation_steps * 1
+        global_batch_size = self.args.per_device_batch_size * self.args.gradient_accumulation_steps * dist_utils.get_world_size()
 
         fp16_dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
 
         loss_fn = torch.nn.CrossEntropyLoss()
 
-        self.progress = tqdm(range(max_steps), disable=(not self.args.show_progress))
+        self.progress = tqdm(range(max_steps), disable=(not self.args.show_progress) or dist_utils.get_rank() != 0)
 
         # scaler used for mixed precision fp16 training on GPU
         self.scaler = torch.cuda.amp.GradScaler(enabled=self.args.use_fp16)
@@ -152,7 +161,7 @@ class Trainer():
                     logits = self.model(inputs)
 
                     labels = inputs[..., 1:].contiguous().view(-1)
-                    logits = logits[..., :-1, :].contiguous().view(-1, self.model.vocab_size)
+                    logits = logits[..., :-1, :].contiguous().view(-1, self.model.module.vocab_size)
 
                     loss = loss_fn(logits, labels)
 

optimus/utils/dist_utils.py

+import os
+import logging
+import contextlib
+
+import torch.distributed as dist
+
+
+logger = logging.getLogger(__name__)
+
+
+def init_process_group():
+    if 'RANK' not in os.environ or \
+        'WORLD_SIZE' not in os.environ or \
+        'MASTER_ADDR' not in os.environ or \
+        'MASTER_PORT' not in os.environ:
+        # at this point, we're pretty sure we are not in a distributed
+        # environment
+        pass
+    else:
+        dist.init_process_group('nccl')
+
+
+def destroy_process_group():
+    if 'RANK' not in os.environ or \
+        'WORLD_SIZE' not in os.environ or \
+        'MASTER_ADDR' not in os.environ or \
+        'MASTER_PORT' not in os.environ:
+        # at this point, we're pretty sure we are not in a distributed
+        # environment
+        pass
+    else:
+        dist.destroy_process_group()
+
+
+def is_dist_available_and_initialized():
+    if not dist.is_available():
+        return False
+
+    if not dist.is_initialized():
+        return False
+
+    return True
+
+
+def get_rank():
+    if not is_dist_available_and_initialized():
+        return 0
+
+    return dist.get_rank()
+
+
+def get_local_rank():
+    try:
+        # this should be automatically detected when running with torchrun
+        local_rank = int(os.environ['LOCAL_RANK'])
+    except:
+        # if you run by hand, make sure to correctly set ONLY ONE OF:
+        # - LOCAL_RANK=<x>
+        #   OR
+        # - CUDA_VISIBLE_DEVICES=<x>
+        # , where x is one of the GPU's available on the system
+        local_rank = 0
+    return local_rank
+
+
+def get_world_size():
+    if not is_dist_available_and_initialized():
+        return 1
+
+    return dist.get_world_size()
+
+
+def is_main_rank():
+    return get_rank() == 0
+
+
+@contextlib.contextmanager
+def main_process_first():
+    rank = get_rank()
+    if rank == 0:
+        try:
+            logger.info(f"Rank 0 start work")
+            yield
+        finally:
+            logger.info(f"Rank 0 done with work")
+            if is_dist_available_and_initialized():
+                dist.barrier(device_ids=[get_local_rank()])
+    else:
+        try:
+            logger.info(f"Rank {rank} waiting for rank 0 to process work")
+            if is_dist_available_and_initialized():
+                dist.barrier(device_ids=[get_local_rank()])
+        finally:
+            logger.info(f"Rank {rank} done waiting for rank 0")
+            yield

optimus/utils/logging_utils.py

+import os
+import logging
+
+import torch.distributed as dist
+
+
+class FilterMainRankOnly(logging.Filter):
+    def filter(self, record):
+        if not dist.is_available():
+            return True
+        if not dist.is_initialized():
+            return True
+        return dist.get_rank() == 0

optimus/utils/setup_utils.py

 import logging
 
+from torch.nn.parallel import DistributedDataParallel as DDP
 from transformers import AutoTokenizer
 from datasets import load_from_disk
 
-from optimus.models.optimus import OptimusTransformer, OptimusConfig
+from optimus.models.optimus import OptimusConfig, OptimusTransformer
+from optimus.utils import dist_utils
 
 
 logger = logging.getLogger(__name__)
@@ -71,6 +73,13 @@ def create_model(config_file, device):
     model = OptimusTransformer(config)
     model.to(device)
 
+    model = DDP(
+        model,
+        device_ids=[dist_utils.get_local_rank()],
+        output_device=dist_utils.get_local_rank(),
+        gradient_as_bucket_view=True,
+    )
+
     logger.info(model)
 
     _total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)

training.py

@@ -4,10 +4,12 @@ from pathlib import Path
 
 import fire
 import torch
+import torch.distributed as dist
 from torch.utils.data import DataLoader
+from torch.utils.data.distributed import DistributedSampler
 
 from optimus.trainer import Trainer, TrainingArguments
-from optimus.utils import setup_utils
+from optimus.utils import dist_utils, logging_utils, setup_utils
 
 
 logger = logging.getLogger(__name__)
@@ -25,6 +27,7 @@ def main(
     grad_clip_norm: float = 1.0,
     use_fp16: bool = True,
     seed: int = 42,
+    log_on_main_rank_only: bool = True,
 
     # model + tokenizer
     model_config_path: Path = Path('config.json'),
@@ -63,6 +66,8 @@ def main(
             precision. Bfloat16 is used if available, otherwise fp16.
         seed (int): Seed used for reproducibility purposes. Each process will
             have its initial PyTorch seed set to `seed + process_rank`.
+        log_on_main_rank_only (bool): Whether output should only be printed by
+            the main rank (rank 0).
         model_config_path (Path): Path to a config file which describes the
             model to be trained.
         hf_tokenizer_name (str): HuggingFace tokenizer name.
@@ -85,6 +90,13 @@ def main(
             `data_dir`.
 
     """
+    # set up logging before anything else
+    logging.basicConfig(
+        level=os.environ.get('LOG_LEVEL', logging.INFO),
+        format='{asctime} - {filename} | {levelname} | {message}',
+        style='{',
+    )
+
     # create paths
     if isinstance(data_dir, str):
         data_dir = Path(data_dir)
@@ -94,14 +106,6 @@ def main(
     else:
         dataset_dir = data_dir / dataset_dir
 
-    # set up logging before anything else
-    logging.basicConfig(
-        level=os.environ.get('LOG_LEVEL', logging.INFO),
-        format='{asctime} - {filename} | {levelname} | {message}',
-        style='{',
-    )
-
-    # create paths
     if isinstance(model_config_path, str):
         model_config_path = Path(model_config_path)
 
@@ -139,6 +143,7 @@ def main(
         f"\t- gradient clipping norm: {grad_clip_norm}\n"
         f"\t- 16-bit floating-point training (fp16): {use_fp16}\n"
         f"\t- seed: {seed}\n"
+        f"\t- only main rank logs: {log_on_main_rank_only}\n"
         f"\t- model config file: {model_config_path}\n"
         f"\t- huggingface tokenizer: {hf_tokenizer_name}\n"
         f"\t- training data directory: {str(data_dir)}\n"
@@ -148,12 +153,26 @@ def main(
         f"\t- saved model directory: {str(save_dir)}\n"
         f"Please seek '--help' if you want to change any of these settings")
 
+    if log_on_main_rank_only:
+        logger.addFilter(logging_utils.FilterMainRankOnly())
+
+    # init distributed process group
+    dist_utils.init_process_group()
+    torch.manual_seed(seed + dist_utils.get_rank())
+
     # set device
-    device = f'cuda'
+    device = f'cuda:{dist_utils.get_local_rank()}'
+
+    logger.info(f'Rank {dist_utils.get_rank()}, world size {dist_utils.get_world_size()}, device {device}')
+
+    # set each process's GPU
+    torch.cuda.set_device(dist_utils.get_local_rank())
+    torch.cuda.empty_cache()
 
     # load dataset and split into batches
-    dataset = setup_utils.load_and_chunk_dataset(str(dataset_dir), seq_len)
-    dataset.set_format('torch')
+    with dist_utils.main_process_first():
+        dataset = setup_utils.load_and_chunk_dataset(str(dataset_dir), seq_len)
+        dataset.set_format('torch')
 
     # load tokenizer
     tokenizer = setup_utils.create_tokenizer(hf_tokenizer_name)
@@ -161,13 +180,30 @@ def main(
     # create model and move to device
     model = setup_utils.create_model(model_config_path, device)
 
-    # create dataloaders
+    # create samplers + dataloaders
+    train_dist_sampler = DistributedSampler(
+        dataset['train'],
+        num_replicas=dist_utils.get_world_size(),
+        rank=dist_utils.get_rank(),
+        shuffle=True,
+        seed=seed,
+        drop_last=True,
+    )
+
     train_dataloader = DataLoader(
         dataset['train'],
         batch_size=batch_size, # per GPU
         num_workers=6, # allow pre-fetching data through multi-process
         pin_memory=True, # fast CPU-GPU transfer
+        sampler=train_dist_sampler,
+    )
+
+    eval_dist_sampler = DistributedSampler(
+        dataset['test'],
+        num_replicas=dist_utils.get_world_size(),
+        rank=dist_utils.get_rank(),
         shuffle=True,
+        seed=seed,
         drop_last=True,
     )
 
@@ -176,7 +212,7 @@ def main(
         batch_size=batch_size, # per GPU
         num_workers=6, # allow pre-fetching data through multi-process
         pin_memory=True, # fast CPU-GPU transfer
-        drop_last=True,
+        sampler=eval_dist_sampler,
     )
 
     # create optimizer
@@ -201,6 +237,7 @@ def main(
 
         # logging
         log_steps=100,
+        log_on_main_rank_only=log_on_main_rank_only,
 
         # core training
         seed=seed,
@@ -239,6 +276,9 @@ def main(
     # save log data
     trainer.save_logs(log_dir)
 
+    # clean distributed process group
+    dist_utils.destroy_process_group()
+
 
 if __name__=='__main__':
     fire.Fire(main)