support ddp (#26)

Author: helloyongyang

llmc/__main__.py

@@ -9,6 +9,7 @@
 import yaml
 from easydict import EasyDict
 from loguru import logger
+from torch.distributed import destroy_process_group, init_process_group
 
 from llmc.compression.quantization import *
 from llmc.compression.sparsification import *
@@ -111,20 +112,29 @@ def main(config):
     llmc_start_time = time.time()
     parser = argparse.ArgumentParser()
     parser.add_argument('--config', type=str, required=True)
+    parser.add_argument('--task_id', type=str, required=True)
     args = parser.parse_args()
 
     with open(args.config, 'r') as file:
         config = yaml.safe_load(file)
     config = EasyDict(config)
 
+    init_process_group(backend='nccl')
+    torch.cuda.set_device(int(os.environ['LOCAL_RANK']))
+
+    if int(os.environ['RANK']) != 0:
+        logger.remove()
+
     check_config(config)
 
     logger.info(f'args: {args}')
     logger.info(f'config:\n{json.dumps(config, ensure_ascii=False, indent=4)}')
 
     print_important_package_version()
 
-    seed_all(config.base.seed)
+    logger.info(f'WORLD_SIZE : {int(os.environ["WORLD_SIZE"])}')
+
+    seed_all(config.base.seed + int(os.environ['RANK']))
 
     # mkdirs
     if 'save' in config:
@@ -149,6 +159,8 @@ def main(config):
 
     main(config)
 
+    destroy_process_group()
+
     llmc_end_time = time.time()
     llmc_duration_time = llmc_end_time - llmc_start_time
     logger.info(f'llmc_duration_time: {llmc_duration_time} s')

llmc/compression/quantization/awq.py

@@ -1,6 +1,8 @@
 import gc
+import os
 
 import torch
+import torch.distributed as dist
 import torch.nn as nn
 from loguru import logger
 
@@ -136,6 +138,8 @@ def search_scale_subset(self, layers_dict, input, inspect_module, subset_kwargs)
                 best_error = loss_mean
                 best_scales = scales_mean
         best_scales = best_scales.view(-1)
+        dist.all_reduce(best_scales, op=dist.ReduceOp.SUM)
+        best_scales /= int(os.environ['WORLD_SIZE'])
         del org_out_dict
         gc.collect()
         torch.cuda.empty_cache()

llmc/compression/quantization/base_blockwise_quantization.py

@@ -1,10 +1,12 @@
 import functools
 import gc
 import json
+import os
 from collections import defaultdict
 from functools import partial
 
 import torch
+import torch.distributed as dist
 import torch.nn as nn
 from loguru import logger
 
@@ -487,6 +489,12 @@ def auto_clip(self, block, input_feat, n_sample_token):
                     n_sample_token=n_sample_token,
                 )
 
+                dist.all_reduce(max_val, op=dist.ReduceOp.SUM)
+                max_val /= int(os.environ['WORLD_SIZE'])
+
+                dist.all_reduce(min_val, op=dist.ReduceOp.SUM)
+                min_val /= int(os.environ['WORLD_SIZE'])
+
                 self.apply_clip(m, min_val, max_val, n)
 
     @torch.no_grad()
@@ -802,6 +810,8 @@ def contiguous_params(self):
 
     @torch.no_grad()
     def save_model(self, path):
+        if int(os.environ['RANK']) != 0:
+            return
         if self.online_rotate:
             self.contiguous_params()
         if self.config.model.type == 'Llava':

llmc/data/dataset/base_dataset.py

@@ -1,3 +1,4 @@
+import os
 from abc import ABCMeta
 
 import torch
@@ -84,6 +85,10 @@ def get_calib_samples(self):
 
     def get_calib_dataset(self):
         samples = self.get_calib_samples()
+        logger.info(f'len(samples) all : {len(samples)}')
+        assert len(samples) % int(os.environ['WORLD_SIZE']) == 0
+        samples = samples[int(os.environ['RANK'])::int(os.environ['WORLD_SIZE'])]
+        logger.info(f'len(samples) rank : {len(samples)}')
         calib_samples = []
         if self.calib_bs < 0:
             batch = torch.cat(samples, dim=0)