Add log util (#230)

* Add log util

* replace all occurrences of print and logging

* address format

* disable propagate to avoid duplicate log

deepspeed/__init__.py

@@ -5,6 +5,7 @@ Copyright 2020 The Microsoft DeepSpeed Team
 from deepspeed.pt.deepspeed_light import DeepSpeedLight
 from deepspeed.pt.deepspeed_light import ADAM_OPTIMIZER, LAMB_OPTIMIZER
 from deepspeed.pt.deepspeed_lr_schedules import add_tuning_arguments
+from deepspeed.pt.log_utils import logger
 from deepspeed.pt.deepspeed_cuda import DeepSpeedTransformerLayer, DeepSpeedTransformerConfig
 from deepspeed.pt.deepspeed_config import DeepSpeedConfig
 
@@ -82,11 +83,12 @@ def initialize(args,
         * ``lr_scheduler``: Wrapped lr scheduler if user ``lr_scheduler`` is passed, or
           if ``lr_scheduler`` specified in JSON configuration. Otherwise ``None``.
     """
-    print("DeepSpeed info: version={}, git-hash={}, git-branch={}".format(
-        __version__,
-        __git_hash__,
-        __git_branch__),
-          flush=True)
+    logger.info(
+        "DeepSpeed info: version={}, git-hash={}, git-branch={}".format(
+            __version__,
+            __git_hash__,
+            __git_branch__),
+    )
 
     engine = DeepSpeedLight(args=args,
                             model=model,

deepspeed/pt/deepspeed_checkpointing.py

@@ -22,6 +22,7 @@ from torch.cuda import _lazy_call, device as device_ctx_manager
 from deepspeed.pt.deepspeed_timer import SynchronizedWallClockTimer as Timers
 import torch.distributed as dist
 from deepspeed.pt.deepspeed_config import DeepSpeedConfig
+from deepspeed.pt.log_utils import logger
 
 #DeepSpeed Checkpointing Enabled or Disabled
 deepspeed_checkpointing_enabled = False
@@ -58,20 +59,23 @@ def see_memory_usage(message, force=False):
         return
     #dist.barrier()
     if dist.get_rank() == 0:
-        print(message)
-        print("Memory Allocated ",
-              torch.cuda.memory_allocated() / (1024 * 1024 * 1024),
-              "GigaBytes")
-        print("Max Memory Allocated ",
-              torch.cuda.max_memory_allocated() / (1024 * 1024 * 1024),
-              "GigaBytes")
-        print("Cache Allocated ",
-              torch.cuda.memory_cached() / (1024 * 1024 * 1024),
-              "GigaBytes")
-        print("Max cache Allocated ",
-              torch.cuda.max_memory_cached() / (1024 * 1024 * 1024),
-              "GigaBytes")
-        print(" ")
+        logger.info(message)
+        logger.info(
+            "Memory Allocated %s GigaBytes",
+            torch.cuda.memory_allocated() / (1024 * 1024 * 1024),
+        )
+        logger.info(
+            "Max Memory Allocated %s GigaBytes",
+            torch.cuda.max_memory_allocated() / (1024 * 1024 * 1024),
+        )
+        logger.info(
+            "Cache Allocated %s GigaBytes",
+            torch.cuda.memory_cached() / (1024 * 1024 * 1024),
+        )
+        logger.info(
+            "Max cache Allocated %s GigaBytes",
+            torch.cuda.max_memory_cached() / (1024 * 1024 * 1024),
+        )
         #input("Press Any Key To Continue ..")
 
 
@@ -240,15 +244,16 @@ def model_parallel_cuda_manual_seed(seed):
     data_parallel_seed = seed
 
     if torch.distributed.get_rank() == 0:
-        print('> initializing model parallel cuda seeds on global rank {}, '
-              '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(),
-                  mpu.get_data_parallel_rank(),
-                  model_parallel_seed,
-                  data_parallel_seed),
-              flush=True)
+        logger.info(
+            '> initializing model parallel cuda seeds on global rank {}, '
+            '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(),
+                mpu.get_data_parallel_rank(),
+                model_parallel_seed,
+                data_parallel_seed),
+        )
     _CUDA_RNG_STATE_TRACKER.reset()
     # Set the default state.
     torch.cuda.manual_seed(data_parallel_seed)
@@ -348,15 +353,15 @@ class CheckpointFunction(torch.autograd.Function):
         if cuda_device is None:
             see_memory_usage("First Forward Begining", force=True)
             if dist.get_rank() == 0:
-                print(f"Activation Checkpointing Information")
-                print(
+                logger.info(f"Activation Checkpointing Information")
+                logger.info(
                     f"----Partition Activations {PARTITION_ACTIVATIONS}, CPU CHECKPOINTING {PA_TO_CPU}"
                 )
-                print(
+                logger.info(
                     f"----contiguous Memory Checkpointing {CONTIGUOUS_CHECKPOINTING} with {num_layers} total layers"
                 )
-                print(f"----Synchronization {SYNCHRONIZE}")
-                print(f"----Profiling {PROFILE_TIME}")
+                logger.info(f"----Synchronization {SYNCHRONIZE}")
+                logger.info(f"----Profiling {PROFILE_TIME}")
 
             cuda_device = torch.cuda.current_device()
             transport_stream = torch.cuda.Stream(device=cuda_device)
@@ -464,7 +469,7 @@ class CheckpointFunction(torch.autograd.Function):
                 else:
                     new_args.append(size)
                 #if dist.get_rank() == 0:
-                #    print (f"The stored tensor is {contiguous_size} and orginal one is {size} ")
+                #    logger.info(f"The stored tensor is {contiguous_size} and orginal one is {size} ")
 
             ctx.save_for_backward(*new_args)
         else:
@@ -562,7 +567,8 @@ def partition_activations_in_checkpoint(partition_activation):
     global PARTITION_ACTIVATIONS
     PARTITION_ACTIVATIONS = partition_activation
     if dist.get_rank() == 0:
-        print(f"**************Partition Activations {PARTITION_ACTIVATIONS}************")
+        logger.info(
+            f"**************Partition Activations {PARTITION_ACTIVATIONS}************")
 
 
 def set_num_layers(nlayers):
@@ -601,7 +607,7 @@ def _configure_using_config_file(deepspeed_config):
             PA_TO_CPU, SYNCHRONIZE, PROFILE_TIME
 
     config = DeepSpeedConfig(deepspeed_config).activation_checkpointing_config
-    print(config.repr())
+    logger.info(config.repr())
     PARTITION_ACTIVATIONS = config.partition_activations
     CONTIGUOUS_CHECKPOINTING = config.contiguous_memory_optimization
     num_layers = config.number_checkpoints

deepspeed/pt/deepspeed_config.py

@@ -4,13 +4,13 @@ Licensed under the MIT license.
 """
 
 import torch
-import logging
 import json
 from deepspeed.pt.deepspeed_constants import *
 from deepspeed.pt.loss_scaler import INITIAL_LOSS_SCALE, SCALE_WINDOW, DELAYED_SHIFT, MIN_LOSS_SCALE
 from deepspeed.pt.deepspeed_config_utils import get_scalar_param, dict_raise_error_on_duplicate_keys
 from deepspeed.pt.deepspeed_zero_config import DeepSpeedZeroConfig
 from deepspeed.pt.deepspeed_checkpointing_config import DeepSpeedActivationCheckpointingConfig
+from deepspeed.pt.log_utils import logger
 
 TENSOR_CORE_ALIGN_SIZE = 8
 ADAM_OPTIMIZER = 'adam'
@@ -407,7 +407,7 @@ class DeepSpeedConfig(object):
             assert False, \
                 'Either train_batch_size or micro_batch_per_gpu needs to be provided'
 
-        print(
+        logger.info(
             f' After Train batch {self.train_batch_size} micro_batch {self.train_micro_batch_size_per_gpu} and grad_acc {self.gradient_accumulation_steps}'
         )
 
@@ -421,13 +421,13 @@ class DeepSpeedConfig(object):
         self._do_warning_check()
 
     def print(self, name):
-        print('{}:'.format(name), flush=True)
+        logger.info('{}:'.format(name))
         for arg in sorted(vars(self)):
             if arg != '_param_dict':
                 dots = '.' * (29 - len(arg))
-                print('  {} {} {}'.format(arg, dots, getattr(self, arg)), flush=True)
+                logger.info('  {} {} {}'.format(arg, dots, getattr(self, arg)))
 
-        print('  json = {}'.format(
+        logger.info('  json = {}'.format(
             json.dumps(self._param_dict,
                        sort_keys=True,
                        indent=4,
@@ -449,7 +449,7 @@ class DeepSpeedConfig(object):
 
         vocabulary_size = self._param_dict.get(VOCABULARY_SIZE, VOCABULARY_SIZE_DEFAULT)
         if vocabulary_size and vocabulary_size % TENSOR_CORE_ALIGN_SIZE != 0:
-            logging.warning(
+            logger.warning(
                 "DeepSpeedConfig: vocabulary size {} is not aligned to {}, may import tensor core utilization."
                 .format(vocabulary_size,
                         TENSOR_CORE_ALIGN_SIZE))
@@ -458,12 +458,12 @@ class DeepSpeedConfig(object):
             MAX_GRAD_NORM in self.optimizer_params.keys() and \
                 self.optimizer_params[MAX_GRAD_NORM] > 0:
             if fp16_enabled:
-                logging.warning(
+                logger.warning(
                     'DeepSpeedConfig: In FP16 mode, DeepSpeed will pass {}:{} to FP16 wrapper'
                     .format(MAX_GRAD_NORM,
                             self.optimizer_params[MAX_GRAD_NORM]))
             else:
-                logging.warning(
+                logger.warning(
                     'DeepSpeedConfig: In FP32 mode, DeepSpeed does not permit MAX_GRAD_NORM ({}) > 0, setting to zero'
                     .format(self.optimizer_params[MAX_GRAD_NORM]))
                 self.optimizer_params[MAX_GRAD_NORM] = 0.0

deepspeed/pt/deepspeed_launch.py

@@ -10,6 +10,8 @@ import base64
 from collections import defaultdict
 from argparse import ArgumentParser, REMAINDER
 
+from deepspeed.pt.log_utils import logger
+
 
 def parse_args():
     parser = ArgumentParser(description="DeepSpeed distributed training launch"
@@ -59,24 +61,24 @@ def main():
 
     for k in current_env.keys():
         if "NCCL" in k:
-            print(args.node_rank, k, current_env[k], flush=True)
+            logger.info("%s %s %s", args.node_rank, k, current_env[k])
 
     world_info = None
     assert args.world_info != "None", "must provide world info dict"
     world_info = base64.urlsafe_b64decode(args.world_info)
     world_info = json.loads(world_info)
 
-    print("WORLD INFO DICT: {}".format(world_info), flush=True)
+    logger.info("WORLD INFO DICT: {}".format(world_info))
     node_list = list(world_info.keys())
     args.nnodes = len(node_list)
     local_node = node_list[args.node_rank]
     local_gpu_ids = world_info[local_node]
     num_local_procs = len(local_gpu_ids)
-    print("nnodes={}, num_local_procs={}, node_rank={}".format(
-        args.nnodes,
-        num_local_procs,
-        args.node_rank),
-          flush=True)
+    logger.info(
+        "nnodes={}, num_local_procs={}, node_rank={}".format(args.nnodes,
+                                                             num_local_procs,
+                                                             args.node_rank),
+    )
 
     global_rank_mapping = defaultdict(list)
     curr_global_rank = 0
@@ -87,11 +89,11 @@ def main():
         for gid in gids:
             global_rank_mapping[node_id].append(curr_global_rank)
             curr_global_rank += 1
-    print("global_rank_mapping={}".format(global_rank_mapping), flush=True)
-    print("dist_world_size={}".format(dist_world_size), flush=True)
+    logger.info("global_rank_mapping={}".format(global_rank_mapping))
+    logger.info("dist_world_size={}".format(dist_world_size))
     current_env["CUDA_VISIBLE_DEVICES"] = ",".join(map(str, local_gpu_ids))
-    print("Setting CUDA_VISIBLE_DEVICES={}".format(current_env["CUDA_VISIBLE_DEVICES"]),
-          flush=True)
+    logger.info("Setting CUDA_VISIBLE_DEVICES={}".format(
+        current_env["CUDA_VISIBLE_DEVICES"]))
     exclusion_counts_per_node = None
 
     # set PyTorch distributed related environmental variables

deepspeed/pt/deepspeed_light.py

@@ -2,7 +2,6 @@
 Copyright 2019 The Microsoft DeepSpeed Team
 '''
 
-import logging
 import torch
 import os
 import warnings
@@ -15,6 +14,7 @@ from tensorboardX import SummaryWriter
 from deepspeed.pt.deepspeed_timer import ThroughputTimer, SynchronizedWallClockTimer
 from deepspeed.pt.deepspeed_zero_optimizer import FP16_DeepSpeedZeroOptimizer
 from deepspeed.pt.zero_optimizer_stage1 import FP16_DeepSpeedZeroOptimizer_Stage1
+from deepspeed.pt.log_utils import logger
 import deepspeed.pt.deepspeed_checkpointing as deepspeed_activation_checkpointing
 
 from deepspeed.pt.fp16_optimizer import FP16_Optimizer
@@ -42,7 +42,7 @@ except ImportError:
     try:
         _ = warned_flatten
     except NameError:
-        print(
+        logger.warning(
             "Warning:  apex was installed without --cpp_ext.  Falling back to Python flatten and unflatten."
         )
         warned_flatten = True
@@ -69,7 +69,9 @@ def _initialize_parameter_parallel_groups(parameter_parallel_size=None):
     data_parallel_size = int(dist.get_world_size())
     if parameter_parallel_size is None:
         parameter_parallel_size = int(data_parallel_size)
-    print(data_parallel_size, parameter_parallel_size)
+    logger.info("data_parallel_size: %s, parameter_parallel_size: %s",
+                data_parallel_size,
+                parameter_parallel_size)
     assert data_parallel_size % parameter_parallel_size == 0, \
         'world size should be divisible by parameter parallel size'
     rank = dist.get_rank()
@@ -83,10 +85,10 @@ def _initialize_parameter_parallel_groups(parameter_parallel_size=None):
 
 
 def print_configuration(args, name):
-    print('{}:'.format(name), flush=True)
+    logger.info('{}:'.format(name))
     for arg in sorted(vars(args)):
         dots = '.' * (29 - len(arg))
-        print('  {} {} {}'.format(arg, dots, getattr(args, arg)), flush=True)
+        logger.info('  {} {} {}'.format(arg, dots, getattr(args, arg)))
 
 
 class DeepSpeedLight(Module):
@@ -105,10 +107,6 @@ class DeepSpeedLight(Module):
                  config_params=None):
         super(DeepSpeedLight, self).__init__()
 
-        logging.basicConfig(level=logging.INFO,
-                            format="[%(levelname)s %(asctime)s] %(message)s",
-                            datefmt="%Y-%m-%d %H:%M:%S")
-
         self.client_optimizer = optimizer
         self.client_model_parameters = model_parameters
         self.client_lr_scheduler = lr_scheduler
@@ -131,11 +129,11 @@ class DeepSpeedLight(Module):
         self.dist_backend = "nccl"
         if dist_init_required:
             if not dist.is_initialized():
-                logging.info("Initializing torch distributed with backend: {}".format(
+                logger.info("Initializing torch distributed with backend: {}".format(
                     self.dist_backend))
                 dist.init_process_group(backend=self.dist_backend)
             else:
-                logging.warning(
+                logger.warning(
                     "Was given dist_init_required=True but detected that torch"
                     "distributed was already initialized, cannot initialize twice.")
 
@@ -178,8 +176,8 @@ class DeepSpeedLight(Module):
             for name, module in self.module.named_modules():
                 if isinstance(module, torch.nn.Embedding):
                     self.csr_tensor_module_names.add(name + ".weight")
-                    logging.info("Will convert {} to sparse (csr) "
-                                 "tensor during training".format(name))
+                    logger.info("Will convert {} to sparse (csr) "
+                                "tensor during training".format(name))
 
         self.save_non_zero_checkpoint = False
         self.save_zero_checkpoint = False
@@ -216,7 +214,7 @@ class DeepSpeedLight(Module):
             os.environ['MASTER_ADDR'] = master_addr
             os.environ['MASTER_PORT'] = TORCH_DISTRIBUTED_DEFAULT_PORT
 
-            logging.info(
+            logger.info(
                 "Discovered MPI settings of world_rank={}, local_rank={}, world_size={}, master_addr={}, master_port={}"
                 .format(os.environ['RANK'],
                         args.local_rank,
@@ -354,13 +352,13 @@ class DeepSpeedLight(Module):
         # First check for scheduler in json configuration
         lr_scheduler = self._scheduler_from_config(self.optimizer)
         if lr_scheduler:
-            logging.info(
+            logger.info(
                 f'DeepSpeed using configured LR scheduler = {self.scheduler_name()}')
             self.lr_scheduler = lr_scheduler
         else:
-            logging.warning('DeepSpeed using client LR scheduler')
+            logger.warning('DeepSpeed using client LR scheduler')
             self.lr_scheduler = client_lr_scheduler
-        logging.info(f'DeepSpeed LR Scheduler = {self.lr_scheduler}')
+        logger.info(f'DeepSpeed LR Scheduler = {self.lr_scheduler}')
 
     def _configure_checkpointing(self, dist_init_required):
 
@@ -401,7 +399,7 @@ class DeepSpeedLight(Module):
             self.device = torch.device("cuda", self.local_rank)
             self.world_size = dist.get_world_size()
             self.global_rank = dist.get_rank()
-            logging.info("Set device to local rank {} within node.".format(
+            logger.info("Set device to local rank {} within node.".format(
                 self.local_rank))
         else:
             self.world_size = 1
@@ -418,7 +416,7 @@ class DeepSpeedLight(Module):
     # Validate command line arguments
     def _do_args_sanity_check(self, args):
         if hasattr(args, 'deepscale_config') and args.deepscale_config is not None:
-            logging.warning(
+            logger.warning(
                 "************ --deepscale_config is deprecated, please use --deepspeed_config ************"
             )
             if hasattr(args, 'deepspeed_config'):
@@ -464,7 +462,7 @@ class DeepSpeedLight(Module):
             self.data_parallel_group = self.mpu.get_data_parallel_group()
             self.dp_world_size = self.mpu.get_data_parallel_world_size()
             src_rank = _get_global_rank(self.mpu.get_data_parallel_group(), 0)
-            print(f"global src_rank={src_rank}")
+            logger.info(f"global src_rank={src_rank}")
         for p in self.module.parameters():
             if torch.is_tensor(p):
                 dist.broadcast(p, src_rank, group=self.data_parallel_group)
@@ -478,21 +476,21 @@ class DeepSpeedLight(Module):
     def _configure_optimizer(self, client_optimizer, model_parameters):
         if client_optimizer is not None:
             basic_optimizer = client_optimizer
-            logging.info('Using client Optimizer as basic optimizer')
+            logger.info('Using client Optimizer as basic optimizer')
         else:
             basic_optimizer = self._configure_basic_optimizer(model_parameters)
-            logging.info(
+            logger.info(
                 'Using DeepSpeed Optimizer param name {} as basic optimizer'.format(
                     self.optimizer_name()))
 
-        logging.info('DeepSpeed Basic Optimizer = {}'.format(basic_optimizer))
+        logger.info('DeepSpeed Basic Optimizer = {}'.format(basic_optimizer))
 
         if self.zero_optimization():
             if self.optimizer_name() != ADAM_OPTIMIZER:
                 assert self.zero_allow_untested_optimizer(), \
                 'You are using an untested ZeRO Optimizer. Please add <"zero_allow_untested_optimizer": true> in the configuration file to use it.'
 
-                logging.warning(
+                logger.warning(
                     "**** You are using ZeRO with an untested optimizer, proceed with caution *****"
                 )
             self.optimizer = self._configure_zero_optimizer(basic_optimizer)
@@ -501,7 +499,7 @@ class DeepSpeedLight(Module):
         else:
             self.optimizer = basic_optimizer
 
-        # logging.info('DeepSpeed Final Optimizer = {}'.format(self.optimizer.state_dict()))
+        # logger.info('DeepSpeed Final Optimizer = {}'.format(self.optimizer.state_dict()))
 
     def _configure_basic_optimizer(self, model_parameters):
         optimizer_parameters = self.optimizer_params()
@@ -525,7 +523,7 @@ class DeepSpeedLight(Module):
         clip_grad = self.gradient_clipping()
         if self.optimizer_name() == ADAM_OPTIMIZER:
             if self.dynamic_loss_scale():
-                logging.info('Creating fp16 optimizer with dynamic loss scale')
+                logger.info('Creating fp16 optimizer with dynamic loss scale')
                 optimizer = FP16_Optimizer(
                     optimizer,
                     dynamic_loss_scale=True,
@@ -535,7 +533,7 @@ class DeepSpeedLight(Module):
                     clip_grad=clip_grad,
                     fused_adam_legacy=self.optimizer_legacy_fusion())
             else:
-                logging.info('Creating fp16 optimizer with static loss scale: {}'.format(
+                logger.info('Creating fp16 optimizer with static loss scale: {}'.format(
                     self.loss_scale()))
                 optimizer = FP16_Optimizer(
                     optimizer,
@@ -544,7 +542,7 @@ class DeepSpeedLight(Module):
                     clip_grad=clip_grad,
                     fused_adam_legacy=self.optimizer_legacy_fusion())
         else:
-            logging.info('Creating fp16 unfused optimizer with dynamic loss scale')
+            logger.info('Creating fp16 unfused optimizer with dynamic loss scale')
             optimizer = FP16_UnfusedOptimizer(
                 optimizer,
                 dynamic_loss_scale=self.dynamic_loss_scale(),
@@ -557,11 +555,11 @@ class DeepSpeedLight(Module):
 
     def _configure_zero_optimizer(self, optimizer):
         zero_stage = self.zero_optimization_stage()
-        logging.info('Creating fp16 ZeRO stage {} optimizer'.format(zero_stage))
+        logger.info('Creating fp16 ZeRO stage {} optimizer'.format(zero_stage))
 
         if zero_stage == ZERO_OPTIMIZATION_OPTIMIZER_STATES:
             assert self.zero_reduce_scatter(), 'Stage 1 only supports reduce scatter mode'
-            logging.info('Creating fp16 ZeRO Optimizer Stage 1')
+            logger.info('Creating fp16 ZeRO Optimizer Stage 1')
             optimizer = FP16_DeepSpeedZeroOptimizer_Stage1(
                 optimizer,
                 static_loss_scale=self.loss_scale(),
@@ -593,7 +591,7 @@ class DeepSpeedLight(Module):
                 gradient_predivide_factor=self.gradient_predivide_factor())
         else:
             raise NotImplementedError("ZeRO stage {} not implemented".format(zero_stage))
-        logging.info('Creating fp16 zero stage {} optimizer'.format(zero_stage))
+        logger.info('Creating fp16 zero stage {} optimizer'.format(zero_stage))
 
         return optimizer
 
@@ -658,7 +656,7 @@ class DeepSpeedLight(Module):
         else:
             scaled_loss = prescaled_loss
             if self.warn_unscaled_loss:
-                logging.warning(
+                logger.warning(
                     f'DeepSpeed unable to scale loss because of type: {type(prescaled_loss)}'
                 )
                 self.warn_unscaled_loss = False
@@ -892,7 +890,7 @@ class DeepSpeedLight(Module):
     def _report_progress(self, step):
         lr = self.get_lr()
         mom = self.get_mom()
-        logging.info('rank:{} step={}, skipped={}, lr={}, mom={}'.format(
+        logger.info('rank:{} step={}, skipped={}, lr={}, mom={}'.format(
             self.global_rank,
             step,
             self.skipped_steps,
@@ -1096,12 +1094,12 @@ class DeepSpeedLight(Module):
         load_path = self._get_ckpt_name(load_dir, tag)
 
         if not os.path.exists(load_path):
-            logging.warn(
+            logger.warn(
                 'Client provided checkpoint load path: {} does not exist ... skip checkpoint load'
                 .format(load_path))
             return None, None
 
-        logging.info('Loading checkpoint: {}'.format(load_path))
+        logger.info('Loading checkpoint: {}'.format(load_path))
         checkpoint = torch.load(load_path, map_location=lambda storage, loc: storage)
 
         self.load_module_state_dict(state_dict=checkpoint['module'],
@@ -1136,7 +1134,7 @@ class DeepSpeedLight(Module):
         zero_checkpoint_name = self._get_zero_ckpt_name(load_dir, tag)
 
         if not os.path.exists(zero_checkpoint_name):
-            logging.warn(
+            logger.warn(
                 'Client provided checkpoint load path: {} does not exist ... skip checkpoint load'
                 .format(zero_checkpoint_name))
             return None
@@ -1144,7 +1142,7 @@ class DeepSpeedLight(Module):
         zero_sd = torch.load(zero_checkpoint_name, map_location='cpu')
         self.optimizer.load_state_dict(zero_sd['optimizer_state_dict'],
                                        load_optimizer_states=load_optimizer_states)
-        logging.info('loading zero checkpoint {}'.format(zero_checkpoint_name))
+        logger.info('loading zero checkpoint {}'.format(zero_checkpoint_name))
 
     def save_checkpoint(self, save_dir, tag, client_state={}):
         r"""Save training checkpoint
@@ -1180,7 +1178,7 @@ class DeepSpeedLight(Module):
                         checkpoint_name = self._get_zero_ckpt_name(save_dir, tag)
                         self._ensure_directory_exists(checkpoint_name)
                 except:
-                    logging.error(
+                    logger.error(
                         f'Failed Saving model checkpoint to {save_dir} with tag {tag}')
                     return False
             dist.barrier()
@@ -1207,7 +1205,7 @@ class DeepSpeedLight(Module):
         }
         state.update(client_state)
 
-        logging.info('Saving model checkpoint: {}'.format(save_path))
+        logger.info('Saving model checkpoint: {}'.format(save_path))
         torch.save(state, save_path)
 
     def _save_zero_checkpoint(self, save_path, tag):
@@ -1215,4 +1213,4 @@ class DeepSpeedLight(Module):
         #self._ensure_directory_exists(zero_checkpoint_name)
         zero_sd = {'optimizer_state_dict': self.optimizer.state_dict()}
         torch.save(zero_sd, zero_checkpoint_name)
-        logging.info('zero checkpoint saved {}'.format(zero_checkpoint_name))
+        logger.info('zero checkpoint saved {}'.format(zero_checkpoint_name))

deepspeed/pt/deepspeed_lr_schedules.py

@@ -13,7 +13,7 @@ from torch.optim import Optimizer
 from typing import Union, List
 import math
 from deepspeed.pt.deepspeed_constants import *
-import logging
+from deepspeed.pt.log_utils import logger
 
 LR_SCHEDULE = 'lr_schedule'
 LR_RANGE_TEST = 'LRRangeTest'
@@ -550,7 +550,7 @@ class OneCycle(object):
                              last_batch_iteration):
         if 'betas' not in optimizer.defaults:
             optimizer_name = type(optimizer).__name__
-            logging.warn(
+            logger.warn(
                 f"cycle_momentum is disabled because optimizer {optimizer_name} does not support momentum, no betas attribute in defaults"
             )
             self.cycle_momentum = False

deepspeed/pt/deepspeed_run.py

@@ -7,7 +7,6 @@ import sys
 import json
 import shutil
 import base64
-import logging
 import argparse
 import subprocess
 import collections
@@ -16,6 +15,7 @@ from copy import deepcopy
 import torch.cuda
 
 from deepspeed.pt.deepspeed_constants import TORCH_DISTRIBUTED_DEFAULT_PORT
+from deepspeed.pt.log_utils import logger
 
 DLTS_HOSTFILE = "/job/hostfile"
 EXPORT_ENVS = ["NCCL", "PYTHON"]
@@ -87,8 +87,8 @@ def parse_args(args=None):
 
 def fetch_hostfile(hostfile_path):
     if not os.path.isfile(hostfile_path):
-        logging.warning("Unable to find hostfile, will proceed with training "
-                        "with local resources only.")
+        logger.warning("Unable to find hostfile, will proceed with training "
+                       "with local resources only.")
         return None
 
     # e.g., worker-0 slots=16
@@ -101,12 +101,12 @@ def fetch_hostfile(hostfile_path):
                 _, slot_count = slots.split("=")
                 slot_count = int(slot_count)
             except ValueError as err:
-                logging.error("Hostfile is not formatted correctly, unable to "
-                              "proceed with training.")
+                logger.error("Hostfile is not formatted correctly, unable to "
+                             "proceed with training.")
                 raise err
             if hostname in resource_pool:
-                logging.error("Hostfile contains duplicate hosts, unable to "
-                              "proceed with training.")
+                logger.error("Hostfile contains duplicate hosts, unable to "
+                             "proceed with training.")
                 raise ValueError("host {} is already defined".format(hostname))
             resource_pool[hostname] = slot_count
 
@@ -169,7 +169,7 @@ def parse_resource_filter(host_info, include_str="", exclude_str=""):
                 filtered_hosts[hostname] = slots
             elif exclude_str:
                 for s in slots:
-                    print('removing {} from {}'.format(s, hostname))
+                    logger.info('removing {} from {}'.format(s, hostname))
                     filtered_hosts[hostname].remove(s)
 
         # User just specified the whole node
@@ -256,7 +256,7 @@ def main(args=None):
         hostname_cmd = ["ssh {} hostname -I".format(first_host)]
         result = subprocess.check_output(hostname_cmd, shell=True)
         args.master_addr = result.decode('utf-8').split()[0]
-        logging.info("Using IP address of {} for node {}".format(
+        logger.info("Using IP address of {} for node {}".format(
             args.master_addr,
             first_host))
 
@@ -292,7 +292,7 @@ def main(args=None):
         env['PDSH_RCMD_TYPE'] = 'ssh'
 
         active_workers = ",".join(active_resources.keys())
-        logging.info("Running on the following workers: %s" % active_workers)
+        logger.info("Running on the following workers: %s" % active_workers)
 
         pdsh_cmd_args = ['pdsh', '-w', active_workers]
 
@@ -330,13 +330,10 @@ def main(args=None):
             "--master_port={}".format(args.master_port)
         ]
         cmd = pdsh_cmd_args + deepspeed_launch + [args.user_script] + args.user_args
-    print("cmd={}".format(cmd), flush=True)
+    logger.info("cmd={}".format(cmd))
     result = subprocess.Popen(cmd, env=env)
     result.wait()
 
 
 if __name__ == "__main__":
-    logging.basicConfig(level=logging.INFO,
-                        format="[%(levelname)s %(asctime)s] %(message)s",
-                        datefmt="%Y-%m-%d %H:%M:%S")
     main()

deepspeed/pt/deepspeed_timer.py

@@ -3,17 +3,18 @@ Copyright 2019 The Microsoft DeepSpeed Team
 '''
 
 import time
-import logging
 import psutil
 import torch
 
+from deepspeed.pt.log_utils import logger
+
 
 def print_rank_0(message):
     if torch.distributed.is_initialized():
         if torch.distributed.get_rank() == 0:
-            print(message, flush=True)
+            logger.info(message)
     else:
-        print(message, flush=True)
+        logger.info(message)
 
 
 class SynchronizedWallClockTimer:
@@ -117,7 +118,7 @@ class ThroughputTimer():
         self.monitor_memory = monitor_memory
         self.logging = logging_fn
         if self.logging is None:
-            self.logging = logging.info
+            self.logging = logger.info
         self.initialized = False
 
     def update_epoch_count(self):

deepspeed/pt/deepspeed_utils.py

@@ -9,6 +9,8 @@ Helper functions and classes from multiple sources.
 import torch
 from torch._six import inf
 
+from deepspeed.pt.log_utils import logger
+
 
 class CheckOverflow(object):
     '''Checks for overflow in gradient across parallel process'''
@@ -112,7 +114,7 @@ def _handle_overflow(cpu_sum, x, i):
             if not math.isfinite(float(v)):
                 t_i = v_i
                 break
-        print(
+        logger.info(
             f"rank {rank} detected overflow {cpu_sum} in tensor {i}:{t_i} shape {x.shape}"
         )
 
@@ -253,21 +255,20 @@ def see_memory_usage(message):
         return
 
     # Print message except when distributed but not rank 0
-    print(message, flush=True)
-    print("Memory Allocated ",
-          torch.cuda.memory_allocated() / (1024 * 1024 * 1024),
-          "GigaBytes",
-          flush=True)
-    print("Max Memory Allocated ",
-          torch.cuda.max_memory_allocated() / (1024 * 1024 * 1024),
-          "GigaBytes",
-          flush=True)
-    print("Cache Allocated ",
-          torch.cuda.memory_cached() / (1024 * 1024 * 1024),
-          "GigaBytes",
-          flush=True)
-    print("Max cache Allocated ",
-          torch.cuda.max_memory_cached() / (1024 * 1024 * 1024),
-          "GigaBytes",
-          flush=True)
-    print(" ", flush=True)
+    logger.info(message)
+    logger.info(
+        "Memory Allocated %s GigaBytes ",
+        torch.cuda.memory_allocated() / (1024 * 1024 * 1024),
+    )
+    logger.info(
+        "Max Memory Allocated %s GigaBytes",
+        torch.cuda.max_memory_allocated() / (1024 * 1024 * 1024),
+    )
+    logger.info(
+        "Cache Allocated %s GigaBytes",
+        torch.cuda.memory_cached() / (1024 * 1024 * 1024),
+    )
+    logger.info(
+        "Max cache Allocated %s GigaBytes",
+        torch.cuda.max_memory_cached() / (1024 * 1024 * 1024),
+    )

deepspeed/pt/deepspeed_zero_config.py

@@ -3,9 +3,9 @@ Copyright (c) Microsoft Corporation
 Licensed under the MIT license.
 """
 
-import logging
 #from deepspeed.pt.deepspeed_constants import *
 from deepspeed.pt.deepspeed_config_utils import get_scalar_param
+from deepspeed.pt.log_utils import logger
 
 #########################################
 # ZeRO optimization
@@ -104,7 +104,7 @@ class DeepSpeedZeroConfig(object):
                 ZERO_OPTIMIZATION_ALLGATHER_BUCKET_SIZE_DEPRECATED,
                 ZERO_OPTIMIZATION_ALLGATHER_BUCKET_SIZE_DEFAULT)
 
-        logging.warning(
+        logger.warning(
             'DeepSpeedConfig: this format of ZeRO optimization setup is deprecated. Please use the following format: {}'
             .format(ZERO_FORMAT))
         return zero_config_dict

deepspeed/pt/deepspeed_zero_optimizer.py

@@ -17,6 +17,8 @@ from deepspeed.pt.deepspeed_utils import see_memory_usage, is_model_parallel_par
 #with gradient partitioning and without
 pg_correctness_test = False
 
+from deepspeed.pt.log_utils import logger
+
 try:
     from apex_C import flatten
     from apex_C import unflatten
@@ -24,8 +26,8 @@ except ImportError:
     try:
         _ = warned_flatten
     except NameError:
-        print(
-            "Warning:  apex was installed without --cpp_ext.  Falling back to Python flatten and unflatten."
+        logger.warning(
+            "apex was installed without --cpp_ext.  Falling back to Python flatten and unflatten."
         )
         warned_flatten = True
     from torch._utils import _flatten_dense_tensors as flatten
@@ -117,8 +119,8 @@ class FP16_DeepSpeedZeroOptimizer(object):
                  gradient_predivide_factor=1.0):
 
         if dist.get_rank() == 0:
-            print(f"Reduce bucket size {reduce_bucket_size}")
-            print(f"Allgather bucket size {allgather_bucket_size}")
+            logger.info(f"Reduce bucket size {reduce_bucket_size}")
+            logger.info(f"Allgather bucket size {allgather_bucket_size}")
         # The fused optimizer does all the work. We need this layer for two reason:
         # 1. maintain same user API from apex.fp16_utils
         # 2. keep common stuff here in case we need to add ne552w fused optimizer later
@@ -217,7 +219,6 @@ class FP16_DeepSpeedZeroOptimizer(object):
             if dist.get_rank(group=self.dp_process_group) == 0:
                 see_memory_usage(
                     f"After Flattening and after emptying param group {i} cache")
-                print("")
 
             # set model fp16 weight to slices of flattened buffer
             updated_params = _unflatten_dense_tensors(self.fp16_groups_flat[i],
@@ -345,11 +346,10 @@ class FP16_DeepSpeedZeroOptimizer(object):
         see_memory_usage("After initializing optimizer states")
 
         if dist.get_rank() == 0:
-            print(f"optimizer state initialized")
+            logger.info(f"optimizer state initialized")
 
         if dist.get_rank(group=self.dp_process_group) == 0:
             see_memory_usage(f"After initializing ZeRO optimizer")
-            print("")
 
     def _release_ipg_buffers(self):
         if self.contiguous_gradients:
@@ -417,7 +417,7 @@ class FP16_DeepSpeedZeroOptimizer(object):
         self.report_ipg_memory_usage(f"In ipg_epilogue after reduce_ipg_grads", 0)
 
         #if dist.get_rank() == 0:
-        #    print("Params already reduced ", self.params_already_reduced)
+        #    logger.info("Params already reduced %s", self.params_already_reduced)
         for i in range(len(self.params_already_reduced)):
             self.params_already_reduced[i] = False
 
@@ -572,7 +572,7 @@ class FP16_DeepSpeedZeroOptimizer(object):
 
     def print_rank_0(self, message):
         if dist.get_rank() == 0:
-            print(message)
+            logger.info(message)
 
     def gradient_reduction_w_predivide(self, tensor):
         dp_world_size = dist.get_world_size(group=self.dp_process_group)
@@ -740,7 +740,7 @@ class FP16_DeepSpeedZeroOptimizer(object):
         flatten_tensor = _flatten_dense_tensors(tensors)
 
         def print_func():
-            print(flatten_tensor.contiguous().view(-1).narrow(0, start, n))
+            logger.info(flatten_tensor.contiguous().view(-1).narrow(0, start, n))
 
         self.sequential_execution(print_func, message)
 
@@ -779,7 +779,7 @@ class FP16_DeepSpeedZeroOptimizer(object):
         if group is None:
             group = self.dp_process_group
         if dist.get_rank(group=group) == 0:
-            print(message)
+            logger.info(message)
         for id in range(dist.get_world_size(group=group)):
             if id == dist.get_rank(group=group):
                 function()
@@ -984,7 +984,7 @@ class FP16_DeepSpeedZeroOptimizer(object):
         else:
             total_norm = 0.0
             #if dist.get_rank() == 0:
-            #    print(f"Total Norm begining {total_norm}")
+            #    logger.info(f"Total Norm begining {total_norm}")
             for g, p in zip(gradients, params):
                 if is_model_parallel_parameter(p) or (self.model_parallel_rank == 0):
                     param_norm = g.data.double().norm(2)
@@ -1080,7 +1080,7 @@ class FP16_DeepSpeedZeroOptimizer(object):
             self.zero_grad()
             see_memory_usage('After overflow after clearing gradients')
 
-            print(
+            logger.info(
                 "[deepscale] OVERFLOW! Rank {} Skipping step. Attempted loss scale: {}, "
                 "reducing to {}".format(dist.get_rank(),
                                         prev_scale,
@@ -1417,6 +1417,6 @@ def _handle_overflow(cpu_sum, x, i):
             if not math.isfinite(float(v)):
                 t_i = v_i
                 break
-        print(
+        logger.info(
             f"rank {rank} detected overflow {cpu_sum} in tensor {i}:{t_i} shape {x.shape}"
         )

deepspeed/pt/fp16_optimizer.py

@@ -6,12 +6,12 @@ This file is adapted from FP16_Optimizer in NVIDIA/apex
 '''
 
 import torch
-import logging
 import math
 from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors
 
 from deepspeed.pt.deepspeed_utils import get_grad_norm, CheckOverflow, get_weight_norm
 from deepspeed.pt.loss_scaler import INITIAL_LOSS_SCALE, SCALE_WINDOW, MIN_LOSS_SCALE
+from deepspeed.pt.log_utils import logger
 
 
 class FP16_Optimizer(object):
@@ -137,9 +137,10 @@ class FP16_Optimizer(object):
 
         if self.overflow:
             if self.verbose:
-                print("[deepspeed] OVERFLOW! Skipping step. Attempted loss "
-                      "scale: {}, reducing to {}".format(prev_scale,
-                                                         self.cur_scale))
+                logger.info("[deepspeed] OVERFLOW! Skipping step. Attempted loss "
+                            "scale: {}, reducing to {}".format(
+                                prev_scale,
+                                self.cur_scale))
             return self.overflow
         combined_scale = self.unscale_and_clip_grads(grads_groups_flat,
                                                      norm_groups,
@@ -190,9 +191,10 @@ class FP16_Optimizer(object):
 
         if self.overflow:
             if self.verbose:
-                print("[deepspeed] OVERFLOW! Skipping step. Attempted loss "
-                      "scale: {}, reducing to {}".format(prev_scale,
-                                                         self.cur_scale))
+                logger.info("[deepspeed] OVERFLOW! Skipping step. Attempted loss "
+                            "scale: {}, reducing to {}".format(
+                                prev_scale,
+                                self.cur_scale))
             return self.overflow
 
         self.unscale_and_clip_grads(grads_groups_flat, norm_groups)
@@ -250,8 +252,8 @@ class FP16_Optimizer(object):
                                      self.min_loss_scale)
                 self.last_overflow_iter = self.cur_iter
                 if self.verbose:
-                    print(f"\nGrad overflow on iteration {self.cur_iter}")
-                    print(
+                    logger.info(f"\nGrad overflow on iteration {self.cur_iter}")
+                    logger.info(
                         f"Reducing dynamic loss scale from {prev_scale} to {self.cur_scale}"
                     )
             else:
@@ -260,14 +262,15 @@ class FP16_Optimizer(object):
                 if (stable_interval > 0) and (stable_interval % self.scale_window == 0):
                     self.cur_scale *= self.scale_factor
                     if self.verbose:
-                        print(f"\nNo Grad overflow for {self.scale_window} iterations")
-                        print(
+                        logger.info(
+                            f"No Grad overflow for {self.scale_window} iterations")
+                        logger.info(
                             f"Increasing dynamic loss scale from {prev_scale} to {self.cur_scale}"
                         )
         else:
             if skip:
-                print("\nGrad overflow on iteration", self.cur_iter)
-                print("Using static loss scale of", self.cur_scale)
+                logger.info("Grad overflow on iteration: %s", self.cur_iter)
+                logger.info("Using static loss scale of: %s", self.cur_scale)
         self.cur_iter += 1
         return
 

deepspeed/pt/fp16_unfused_optimizer.py

@@ -8,10 +8,10 @@ This file is adapted from FP16_Optimizer in NVIDIA/apex
 import torch
 from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors
 import math
-import logging
 
 from deepspeed.pt.deepspeed_utils import get_grad_norm, CheckOverflow, get_weight_norm
 from deepspeed.pt.loss_scaler import INITIAL_LOSS_SCALE, SCALE_WINDOW, MIN_LOSS_SCALE
+from deepspeed.pt.log_utils import logger
 
 
 class FP16_UnfusedOptimizer(object):
@@ -33,7 +33,7 @@ class FP16_UnfusedOptimizer(object):
         self.fused_lamb_legacy = fused_lamb_legacy
 
         if torch.distributed.get_rank() == 0:
-            logging.info(f'Fused Lamb Legacy : {self.fused_lamb_legacy} ')
+            logger.info(f'Fused Lamb Legacy : {self.fused_lamb_legacy} ')
 
         if not torch.cuda.is_available:
             raise SystemError("Cannot use fp16 without CUDA.")
@@ -137,9 +137,10 @@ class FP16_UnfusedOptimizer(object):
         self._update_scale(self.overflow)
         if self.overflow:
             if self.verbose:
-                print("[deepspeed] OVERFLOW! Skipping step. Attempted loss "
-                      "scale: {}, reducing to {}".format(prev_scale,
-                                                         self.cur_scale))
+                logger.info("[deepspeed] OVERFLOW! Skipping step. Attempted loss "
+                            "scale: {}, reducing to {}".format(
+                                prev_scale,
+                                self.cur_scale))
             return self.overflow
 
         combined_scale = self.unscale_and_clip_grads(norm_groups, apply_scale=False)
@@ -162,9 +163,10 @@ class FP16_UnfusedOptimizer(object):
         self._update_scale(self.overflow)
         if self.overflow:
             if self.verbose:
-                print("[deepspeed] OVERFLOW! Skipping step. Attempted loss "
-                      "scale: {}, reducing to {}".format(prev_scale,
-                                                         self.cur_scale))
+                logger.info("[deepspeed] OVERFLOW! Skipping step. Attempted loss "
+                            "scale: {}, reducing to {}".format(
+                                prev_scale,
+                                self.cur_scale))
             return self.overflow
 
         norm_groups = []
@@ -236,8 +238,8 @@ class FP16_UnfusedOptimizer(object):
                                      self.min_loss_scale)
                 self.last_overflow_iter = self.cur_iter
                 if self.verbose:
-                    print("\nGrad overflow on iteration", self.cur_iter)
-                    print(
+                    logger.info("Grad overflow on iteration: %s", self.cur_iter)
+                    logger.info(
                         f"Reducing dynamic loss scale from {prev_scale} to {self.cur_scale}"
                     )
             else:
@@ -246,14 +248,15 @@ class FP16_UnfusedOptimizer(object):
                 if (stable_interval > 0) and (stable_interval % self.scale_window == 0):
                     self.cur_scale *= self.scale_factor
                     if self.verbose:
-                        print(f"\nNo Grad overflow for {self.scale_window} iterations")
-                        print(
+                        logger.info(
+                            f"No Grad overflow for {self.scale_window} iterations")
+                        logger.info(
                             f"Increasing dynamic loss scale from {prev_scale} to {self.cur_scale}"
                         )
         else:
             if skip:
-                print("\nGrad overflow on iteration", self.cur_iter)
-                print("Using static loss scale of", self.cur_scale)
+                logger.info("Grad overflow on iteration %s", self.cur_iter)
+                logger.info("Using static loss scale of %s", self.cur_scale)
         self.cur_iter += 1
         return
 

deepspeed/pt/log_utils.py

+import logging
+import sys
+
+import torch.distributed as dist
+
+
+class LoggerFactory:
+    @staticmethod
+    def create_logger(name=None, level=logging.INFO):
+        """create a logger
+
+        Args:
+            name (str): name of the logger
+            level: level of logger
+
+        Raises:
+            ValueError is name is None
+        """
+
+        if name is None:
+            raise ValueError("name for logger cannot be None")
+
+        formatter = logging.Formatter(
+            "[%(asctime)s] [%(levelname)s] "
+            "[%(filename)s:%(lineno)d:%(funcName)s] %(message)s")
+
+        logger_ = logging.getLogger(name)
+        logger_.setLevel(level)
+        logger_.propagate = False
+        ch = logging.StreamHandler(stream=sys.stdout)
+        ch.setLevel(level)
+        ch.setFormatter(formatter)
+        logger_.addHandler(ch)
+        return logger_
+
+
+logger = LoggerFactory.create_logger(name="DeepSpeed", level=logging.INFO)
+
+
+def log_dist(message, ranks=None, level=logging.INFO):
+    """Log message when one of following condition meets
+
+    + not dist.is_initialized()
+    + dist.get_rank() in ranks if ranks is not None or ranks = [-1]
+
+    Args:
+        message (str)
+        ranks (list)
+        level (int)
+
+    """
+    should_log = not dist.is_initialized()
+    ranks = ranks or []
+    my_rank = dist.get_rank() if dist.is_initialized() else -1
+    if ranks and not should_log:
+        should_log = ranks[0] == -1
+        should_log = should_log or (my_rank in set(ranks))
+    if should_log:
+        final_message = "[Rank {}] {}".format(my_rank, message)
+        logger.log(level, final_message)

deepspeed/pt/zero_optimizer_stage1.py

@@ -4,8 +4,8 @@ import torch.distributed as dist
 from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors
 from collections import defaultdict
 
-from deepspeed.pt.zero_utils import _initialize_parameter_parallel_groups, \
-    pprint
+from deepspeed.pt.zero_utils import _initialize_parameter_parallel_groups
+from deepspeed.pt.log_utils import log_dist, logger
 from deepspeed.pt.loss_scaler import LossScaler, DynamicLossScaler
 from deepspeed.pt.deepspeed_utils import get_grad_norm, CheckOverflow
 
@@ -18,9 +18,10 @@ def flatten_dense_tensors_sub_partition_aligned(tensor_list,
     for tensor in tensor_list:
         num_elements = num_elements + tensor.numel()
 
-    pprint("Total number of elements in model: {}, max elements per com: {}".format(
+    log_dist("Total number of elements in model: {}, max elements per com: {}".format(
         num_elements,
-        max_elements_per_comm))
+        max_elements_per_comm),
+             ranks=[0])
 
     max_elements_per_comm = min(max_elements_per_comm, num_elements)
     sub_partition_size = int(max_elements_per_comm // dp)
@@ -41,10 +42,11 @@ def flatten_dense_tensors_sub_partition_aligned(tensor_list,
     if remaining:
         elements_to_add = alignment - remaining
         # adding padded tensor later after we check comm alignment
-        pprint("adding pad tensor for alignment, {} + {}->{}".format(
+        log_dist("adding pad tensor for alignment, {} + {}->{}".format(
             num_elements,
             elements_to_add,
-            num_elements + elements_to_add))
+            num_elements + elements_to_add),
+                 ranks=[0])
         #num_elements = num_elements + elements_to_add
     else:
         padded_tensor_list = tensor_list
@@ -55,11 +57,12 @@ def flatten_dense_tensors_sub_partition_aligned(tensor_list,
                                                                        "size"
     num_comm_intervals = int(num_partitions // dp)
     partition_remaining = int(num_partitions % dp)
-    pprint("num_comm_intervals={}, partition_remaining={}".format(
+    log_dist("num_comm_intervals={}, partition_remaining={}".format(
         num_comm_intervals,
-        partition_remaining))
+        partition_remaining),
+             ranks=[0])
     if partition_remaining != 0:
-        pprint("adding pad tensor and/or extra sub partition")
+        log_dist("adding pad tensor and/or extra sub partition", ranks=[0])
         # add pad tensor for alignment of comm interval, this overrules previous possibly sub-partition alignment
         num_comm_intervals += 1
         aligned_comm_elements = num_comm_intervals * sub_partition_size * dp
@@ -69,10 +72,11 @@ def flatten_dense_tensors_sub_partition_aligned(tensor_list,
                                  device=tensor_list[0].device,
                                  dtype=tensor_list[0].dtype)
         padded_tensor_list = tensor_list + [pad_tensor]
-        pprint("adding pad tensor and/or extra sub partition, {} + {}->{}".format(
+        log_dist("adding pad tensor and/or extra sub partition, {} + {}->{}".format(
             num_elements,
             elements_to_add,
-            num_elements + elements_to_add))
+            num_elements + elements_to_add),
+                 ranks=[0])
         num_elements += elements_to_add
     elif elements_to_add > 0:
         # add pad tensor for just alignment of sub-partition
@@ -83,7 +87,7 @@ def flatten_dense_tensors_sub_partition_aligned(tensor_list,
         num_elements += elements_to_add
 
     if pg is None or dist.get_rank(group=pg) == 0:
-        print("Number of Elements (w. padding) is ", num_elements)
+        logger.info("Number of Elements (w. padding) is %s", num_elements)
 
     padded_num_elems = 0
     for p in padded_tensor_list:
@@ -161,7 +165,7 @@ class FP16_DeepSpeedZeroOptimizer_Stage1(object):
         self.allgather_size = allgather_size
 
         self.max_elements_per_comm = max_elements_per_comm
-        print("max_elements_per_comm={}".format(max_elements_per_comm))
+        logger.info("max_elements_per_comm={}".format(max_elements_per_comm))
 
         # param flattened by groups
         self.fp16_groups = []
@@ -311,14 +315,14 @@ class FP16_DeepSpeedZeroOptimizer_Stage1(object):
         assert num_sub_partitions % world_size == 0, "{} % {} != 0".format(num_sub_partitions, world_size)
 
         if not dist.is_initialized() or dist.get_rank(group=dp_process_group) == 0:
-            print("**** partition info:")
-            print("\t total_num_elements=", total_num_elements)
-            print("\t world_size=", world_size)
-            print("\t max_elements_per_comm=", max_elements_per_comm)
-            print("\t sub_partition_size=", sub_partition_size)
-            print("\t num_sub_partitions=", num_sub_partitions)
-            print("\t num_comm_intervals=", num_comm_intervals)
-            print("****")
+            logger.info("**** partition info:")
+            logger.info("\t total_num_elements=%s", total_num_elements)
+            logger.info("\t world_size=%s", world_size)
+            logger.info("\t max_elements_per_comm=%s", max_elements_per_comm)
+            logger.info("\t sub_partition_size=%s", sub_partition_size)
+            logger.info("\t num_sub_partitions=%s", num_sub_partitions)
+            logger.info("\t num_comm_intervals=%s", num_comm_intervals)
+            logger.info("****")
 
         # [comm_id] -> [rank]
         comm_partitions = []
@@ -369,14 +373,14 @@ class FP16_DeepSpeedZeroOptimizer_Stage1(object):
             for iii, tensor in enumerate(tensor_list):
                 tensor_size = tensor.numel()
                 #if local_rank == 0:
-                #    #print("rank={}, current_index={}, tensor_size={}, tensor-idx={}".format(rank,
+                #    # logger.info("rank={}, current_index={}, tensor_size={}, tensor-idx={}".format(rank,
                 #        current_index, tensor_size, iii))
                 results_list = _range_check(current_index,
                                             all_element_intervals[rank],
                                             tensor_size)
                 for contained, offset, comm_idx in results_list:
                     #if local_rank == 0:
-                    #    print("rank={}, contained={}, offset={}, comm_idx={}".format(rank, contained,
+                    #    logger.info("rank={}, contained={}, offset={}, comm_idx={}".format(rank, contained,
                     #        offset, comm_idx))
                     if contained:
                         if prev_comm_idx != comm_idx:
@@ -479,7 +483,7 @@ class FP16_DeepSpeedZeroOptimizer_Stage1(object):
             flat_sub_partitions.append(_flatten_dense_tensors(flat_tensor_list))
         if num_comm_intervals is not None and len(
                 flat_sub_partitions) < num_comm_intervals:
-            #print("padding w. sub partitions to ensure uniform communication")
+            # logger.info("padding w. sub partitions to ensure uniform communication")
             device = flat_sub_partitions[0].device
             for _ in range(num_comm_intervals - len(flat_sub_partitions)):
                 flat_sub_partitions.append(
@@ -611,9 +615,10 @@ class FP16_DeepSpeedZeroOptimizer_Stage1(object):
         if self.overflow:
             self.zero_grad()
             if self.verbose:
-                print("[deepspeed] OVERFLOW! Skipping step. Attempted loss "
-                      "scale: {}, reducing to {}".format(prev_scale,
-                                                         self.loss_scale))
+                logger.info("[deepspeed] OVERFLOW! Skipping step. Attempted loss "
+                            "scale: {}, reducing to {}".format(
+                                prev_scale,
+                                self.loss_scale))
             return self.overflow
 
         norm_groups = []
@@ -649,7 +654,7 @@ class FP16_DeepSpeedZeroOptimizer_Stage1(object):
                 default_device=self.local_sub_partitions_of_fp32_groups[i][0].device)
 
             #RS: update all our local params with sub-partition grads
-            #print("self.local_sub_partitions_of_fp32_groups[i]={}, local_grad_sub_partitions={}".format(len(self.local_sub_partitions_of_fp32_groups[i]), len(local_grad_sub_partitions)))
+            #logger. info("self.local_sub_partitions_of_fp32_groups[i]={}, local_grad_sub_partitions={}".format(len(self.local_sub_partitions_of_fp32_groups[i]), len(local_grad_sub_partitions)))
             for idx, sub_partition_param in enumerate(self.local_sub_partitions_of_fp32_groups[i]):
                 sub_partition_param.grad = local_grad_sub_partitions[idx]
             #self.single_partition_of_fp32_groups[i].grad = single_grad_partition

deepspeed/pt/zero_utils.py

 import torch
 import torch.distributed as dist
 
+from deepspeed.pt.log_utils import logger
+
 
 def _initialize_parameter_parallel_groups(parameter_parallel_size=None):
     data_parallel_size = int(dist.get_world_size())
-    if parameter_parallel_size is None:
-        parameter_parallel_size = int(data_parallel_size)
-    print(data_parallel_size, parameter_parallel_size)
+    parameter_parallel_size = parameter_parallel_size or data_parallel_size
+    logger.info("data_parallel_size: %s, parameter_parallel_size: %s",
+                data_parallel_size,
+                parameter_parallel_size)
     assert data_parallel_size % parameter_parallel_size == 0, \
         'world size should be divisible by parameter parallel size'
     rank = dist.get_rank()
     my_group = None
-    for i in range(dist.get_world_size() // parameter_parallel_size):
+    for i in range(data_parallel_size // parameter_parallel_size):
         ranks = range(i * parameter_parallel_size, (i + 1) * parameter_parallel_size)
         group = torch.distributed.new_group(ranks)
         if rank in ranks:
             my_group = group
     return my_group
-
-
-def pprint(msg):
-    if not dist.is_initialized() or dist.get_rank() == 0:
-        print(msg)