[squash] Big science v1

README.md

@@ -96,6 +96,12 @@ If you would like to pre-install any of the DeepSpeed extensions/ops (instead
 of JIT compiling) or install pre-compiled ops via PyPI please see our [advanced
 installation instructions](https://www.deepspeed.ai/tutorials/advanced-install/).
 
+On Windows you can build wheel with following steps, currently only inference mode is supported.
+1. Install pytorch, such as pytorch 1.8 + cuda 11.1
+2. Install visual cpp build tools, such as VS2019 C++ x64/x86 build tools
+3. Launch cmd console with Administrator privilege for creating required symlink folders
+4. Run `python setup.py bdist_wheel` to build wheel in `dist` folder
+
 # Features
 Below we provide a brief feature list, see our detailed [feature
 overview](https://www.deepspeed.ai/features/) for descriptions and usage.

csrc/adam/cpu_adam.cpp

@@ -322,42 +322,37 @@ int create_adam_optimizer(int optimizer_id,
                           float betta2 = 0.999,
                           float eps = 1e-8,
                           float weight_decay = 0,
-                          bool adamw_mode = true)
+                          bool adamw_mode = true,
+                          bool should_log = false)
 {
     auto opt =
         std::make_shared<Adam_Optimizer>(alpha, betta1, betta2, eps, weight_decay, adamw_mode);
 
     s_optimizers[optimizer_id] = opt;
+
+    if (should_log) {
+        std::string avx_type = "";
 #if defined(__AVX512__)
-    std::cout << "Adam Optimizer #" << optimizer_id
-              << " is created with AVX512 arithmetic capability." << std::endl;
-    printf("Config: alpha=%f, betas=(%f, %f), weight_decay=%f, adam_w=%d\n",
-           alpha,
-           betta1,
-           betta2,
-           weight_decay,
-           (int)adamw_mode);
+        avx_type = "AVX512";
 #else
 #if defined(__AVX256__)
-    std::cout << "Adam Optimizer #" << optimizer_id
-              << " is created with AVX2 arithmetic capability." << std::endl;
-    printf("Config: alpha=%f, betas=(%f, %f), weight_decay=%f, adam_w=%d\n",
-           alpha,
-           betta1,
-           betta2,
-           weight_decay,
-           (int)adamw_mode);
+        avx_type = "AVX2";
 #else
-    std::cout << "Adam Optimizer #" << optimizer_id
-              << " is created with scalar arithmetic capability." << std::endl;
-    printf("Config: alpha=%f, betas=(%f, %f), weight_decay=%f, adam_w=%d\n",
-           alpha,
-           betta1,
-           betta2,
-           weight_decay,
-           (int)adamw_mode);
+        avx_type = "scalar";
 #endif
 #endif
+
+        printf("Adam Optimizer #%d is created with %s arithmetic capability.\n",
+               optimizer_id,
+               avx_type.c_str());
+        printf("Config: alpha=%f, betas=(%f, %f), weight_decay=%f, adam_w=%d\n",
+               alpha,
+               betta1,
+               betta2,
+               weight_decay,
+               (int)adamw_mode);
+    }
+
     return 0;
 }
 

deepspeed/inference/engine.py

@@ -189,12 +189,16 @@ def _pre_forward_hook(self, module, *inputs, **kwargs):
             if torch.is_tensor(input):
                 input = input.to(torch.cuda.current_device())
                 if self.mp_world_size > 1:
+                    if not input.is_contiguous():
+                        input = input.contiguous()
                     dist.broadcast(input, 0)
 
         for k in kwargs:
             if torch.is_tensor(kwargs[k]):
                 kwargs[k] = kwargs[k].to(torch.cuda.current_device())
                 if self.mp_world_size > 1:
+                    if not kwargs[k].is_contiguous():
+                        kwargs[k] = kwargs[k].contiguous()
                     dist.broadcast(kwargs[k], 0)
 
     def forward(self, *inputs, **kwargs):
@@ -210,12 +214,16 @@ def forward(self, *inputs, **kwargs):
                     if torch.is_tensor(input):
                         input = input.to(torch.cuda.current_device())
                         if self.mp_world_size > 1:
+                            if not input.is_contiguous():
+                                input = input.contiguous()
                             dist.broadcast(input, 0)
 
                 for k in kwargs:
                     if torch.is_tensor(kwargs[k]):
                         kwargs[k] = kwargs[k].to(torch.cuda.current_device())
                         if self.mp_world_size > 1:
+                            if not kwargs[k].is_contiguous():
+                                kwargs[k] = kwargs[k].contiguous()
                             dist.broadcast(kwargs[k], 0)
 
             return self.model_orig_fwd(*inputs, **kwargs)

deepspeed/module_inject/replace_policy.py

@@ -19,7 +19,7 @@ def attention(self):
 
     def get_hidden_heads(self):
         """
-        retun hidden_size and number of heads
+        return hidden_size and number of heads
         """
         raise NotImplementedError
 

deepspeed/ops/adam/cpu_adam.py

@@ -7,6 +7,7 @@
 import time
 from pathlib import Path
 from ..op_builder import CPUAdamBuilder
+from deepspeed.utils.logging import should_log_le
 
 
 class DeepSpeedCPUAdam(torch.optim.Optimizer):
@@ -83,7 +84,8 @@ def __init__(self,
                                      betas[1],
                                      eps,
                                      weight_decay,
-                                     adamw_mode)
+                                     adamw_mode,
+                                     should_log_le("info"))
 
     def __del__(self):
         # need to destroy the C++ object explicitly to avoid a memory leak when deepspeed.initialize

deepspeed/runtime/activation_checkpointing/checkpointing.py

@@ -24,7 +24,7 @@
 
 from deepspeed.runtime.config import DeepSpeedConfig
 from deepspeed.utils import logger
-from deepspeed.runtime.utils import move_to_device, see_memory_usage
+from deepspeed.runtime.utils import move_to_device, see_memory_usage, bwc_tensor_model_parallel_rank
 from deepspeed.utils.timer import SynchronizedWallClockTimer as Timers
 
 # DeepSpeed Checkpointing Enabled or Disabled
@@ -213,9 +213,12 @@ def model_parallel_cuda_manual_seed(seed):
                               model parallel regions.
     """
     global mpu
+
+    tp_rank = bwc_tensor_model_parallel_rank(mpu)
+
     # 2718 is just for fun and any POSITIVE value will work.
     offset = seed + 2718
-    model_parallel_seed = offset + mpu.get_model_parallel_rank()
+    model_parallel_seed = offset + tp_rank
     # Data parallel gets the original sedd.
     data_parallel_seed = seed
 
@@ -225,7 +228,7 @@ def model_parallel_cuda_manual_seed(seed):
             'model parallel rank {}, and data parallel rank {} with '
             'model parallel seed: {} and data parallel seed: {}'.format(
                 torch.distributed.get_rank(),
-                mpu.get_model_parallel_rank(),
+                tp_rank,
                 mpu.get_data_parallel_rank(),
                 model_parallel_seed,
                 data_parallel_seed),
@@ -384,9 +387,14 @@ def save_args_for_backward(*all_args):
         global data_offsets, size_offsets
         if mp_rank is None:
             if mpu is not None:
-                mp_rank = mpu.get_model_parallel_rank()
-                mp_size = mpu.get_model_parallel_world_size()
-                mp_group = mpu.get_model_parallel_group()
+                if hasattr(mpu, 'get_tensor_model_parallel_rank'):
+                    mp_rank = mpu.get_tensor_model_parallel_rank()
+                    mp_size = mpu.get_tensor_model_parallel_world_size()
+                    mp_group = mpu.get_tensor_model_parallel_group()
+                else:
+                    mp_rank = mpu.get_model_parallel_rank()
+                    mp_size = mpu.get_model_parallel_world_size()
+                    mp_group = mpu.get_model_parallel_group()
             else:
                 mp_rank = 0
                 mp_size = 1
@@ -416,7 +424,7 @@ def save_args_for_backward(*all_args):
 
             inputs = []
             for i, item in enumerate(args[:-1]):
-                if not torch.is_tensor(item):
+                if not torch.is_tensor(item) or mp_size > item.numel():
                     inputs.append(item)
                     continue