def __init__(self, sess, tensor_shape_ops=None): self.sess = sess self.graph = sess.graph self.step_cnt = 0 self.feed_dict_meta = {} self.tensor_shape_ops = tensor_shape_ops self.trace_dir = os.path.join(os.environ.get("BYTEPS_TRACE_DIR", "."), str(bps.local_rank())) if not os.path.exists(self.trace_dir): os.makedirs(self.trace_dir) if os.environ.get("BYTEPS_TRACE_ON", "") != '1': self._end_trace = True return self._end_trace = False self.end_step = int(os.environ.get("BYTEPS_TRACE_END_STEP", "30")) self.start_step = int(os.environ.get("BYTEPS_TRACE_START_STEP", "20")) if not self._end_trace and self.start_step < 1: raise ValueError("BYTEPS_TRACE_START_STEP must be larger than 1") if not self._end_trace and self.end_step <= self.start_step: raise ValueError( "BYTEPS_TRACE_END_STEP must be larger than BYTEPS_TRACE_START_STEP" ) ### Timeline configuratoin self.run_options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE) self.run_metadata = tf.RunMetadata() self.traces = {"traceEvents": []} self.dag = None
def __init__(self, average=True): """ Args: average (bool): whether to average or sum the gradients across processes. """ import byteps.tensorflow as bps self.hvd = bps # BytePS has the same interface as Horovod self.hvd.allreduce = bps.push_pull # https://github.com/bytedance/byteps/issues/8 # TODO bootstrap env vars bps.init() self.is_chief = bps.rank() == 0 self._local_rank = bps.local_rank() self._rank = bps.rank() self._average = average self._compression = None self._has_compression = False logger.info("[BytePSTrainer] local rank={}".format(self._local_rank)) SingleCostTrainer.__init__(self)
def __init__(self, average=True): """ Args: average (bool): whether to average or sum the gradients across processes. """ import byteps.tensorflow as bps self.hvd = bps # BytePS has the same interface as Horovod self.hvd.allreduce = bps.push_pull # https://github.com/bytedance/byteps/issues/8 assert os.environ.get("DMLC_ROLE", None) == "worker" assert "DMLC_WORKER_ID" in os.environ and "DMLC_NUM_WORKER" in os.environ bps.init() self.is_chief = bps.rank() == 0 self._local_rank = bps.local_rank() self._rank = bps.rank() self._average = average self._compression = None self._has_compression = False logger.info("[BytePSTrainer] local rank={}".format(self._local_rank)) SingleCostTrainer.__init__(self)
def main(_): # BytePS: initialize BytePS. bps.init() # Keras automatically creates a cache directory in ~/.keras/datasets for # storing the downloaded MNIST data. This creates a race # condition among the workers that share the same filesystem. If the # directory already exists by the time this worker gets around to creating # it, ignore the resulting exception and continue. cache_dir = os.path.join(os.path.expanduser('~'), '.keras', 'datasets') if not os.path.exists(cache_dir): try: os.mkdir(cache_dir) except OSError as e: if e.errno == errno.EEXIST and os.path.isdir(cache_dir): pass else: raise # Download and load MNIST dataset. (x_train, y_train), (x_test, y_test) = \ keras.datasets.mnist.load_data('MNIST-data-%d' % bps.rank()) # The shape of downloaded data is (-1, 28, 28), hence we need to reshape it # into (-1, 784) to feed into our network. Also, need to normalize the # features between 0 and 1. x_train = np.reshape(x_train, (-1, 784)) / 255.0 x_test = np.reshape(x_test, (-1, 784)) / 255.0 # Build model... with tf.name_scope('input'): image = tf.placeholder(tf.float32, [None, 784], name='image') label = tf.placeholder(tf.float32, [None], name='label') predict, loss = conv_model(image, label, tf.estimator.ModeKeys.TRAIN) # BytePS: adjust learning rate based on number of GPUs. opt = tf.train.RMSPropOptimizer(0.001 * bps.size()) # BytePS: add BytePS Distributed Optimizer. opt = bps.DistributedOptimizer(opt) global_step = tf.train.get_or_create_global_step() train_op = opt.minimize(loss, global_step=global_step) hooks = [ # BytePS: BroadcastGlobalVariablesHook broadcasts initial variable states # from rank 0 to all other processes. This is necessary to ensure consistent # initialization of all workers when training is started with random weights # or restored from a checkpoint. bps.BroadcastGlobalVariablesHook(0), # BytePS: adjust number of steps based on number of GPUs. tf.train.StopAtStepHook(last_step=200000 // bps.size()), tf.train.LoggingTensorHook(tensors={'step': global_step, 'loss': loss}, every_n_iter=10), ] # BytePS: pin GPU to be used to process local rank (one GPU per process) config = tf.ConfigProto() config.gpu_options.allow_growth = True config.gpu_options.visible_device_list = str(bps.local_rank()) # BytePS: save checkpoints only on worker 0 to prevent other workers from # corrupting them. checkpoint_dir = './checkpoints' if bps.rank() == 0 else None training_batch_generator = train_input_generator(x_train, y_train, batch_size=100) # The MonitoredTrainingSession takes care of session initialization, # restoring from a checkpoint, saving to a checkpoint, and closing when done # or an error occurs. with tf.train.MonitoredTrainingSession(checkpoint_dir=checkpoint_dir, hooks=hooks, config=config) as mon_sess: while not mon_sess.should_stop(): # Run a training step synchronously. image_, label_ = next(training_batch_generator) mon_sess.run(train_op, feed_dict={image: image_, label: label_})
parser.add_argument('--eager', action='store_true', default=False, help='enables eager execution') parser.add_argument('--no-cuda', action='store_true', default=False, help='disables CUDA training') args = parser.parse_args() args.cuda = not args.no_cuda bps.init() # BytePS: pin GPU to be used to process local rank (one GPU per process) config = tf.ConfigProto() if args.cuda: config.gpu_options.allow_growth = True config.gpu_options.visible_device_list = str(bps.local_rank()) else: os.environ["CUDA_VISIBLE_DEVICES"] = "-1" config.gpu_options.allow_growth = False config.gpu_options.visible_device_list = '' if args.eager: tf.enable_eager_execution(config) # Set up standard model. # Check https://github.com/keras-team/keras-applications for all supported models, e.g., ResNet50, VGG16 model = getattr(applications, args.model)(weights=None) opt = tf.train.GradientDescentOptimizer(0.01) # BytePS: (optional) compression algorithm.
import tensorflow as tf import byteps.tensorflow as bps bps.init() # BytePS: pin GPU to be used to process local rank (one GPU per process) gpus = tf.config.experimental.list_physical_devices('GPU') for gpu in gpus: tf.config.experimental.set_memory_growth(gpu, True) if gpus: tf.config.experimental.set_visible_devices(gpus[bps.local_rank()], 'GPU') # Before launching, need to fist download the dataset to ~/.keras/datasets (mnist_images, mnist_labels), _ = \ tf.keras.datasets.mnist.load_data(path='mnist-%d.npz' % bps.rank()) dataset = tf.data.Dataset.from_tensor_slices( (tf.cast(mnist_images[..., tf.newaxis] / 255.0, tf.float32), tf.cast(mnist_labels, tf.int64))) dataset = dataset.repeat().shuffle(10000).batch(128) mnist_model = tf.keras.Sequential([ tf.keras.layers.Conv2D(32, [3, 3], activation='relu'), tf.keras.layers.Conv2D(64, [3, 3], activation='relu'), tf.keras.layers.MaxPooling2D(pool_size=(2, 2)), tf.keras.layers.Dropout(0.25), tf.keras.layers.Flatten(), tf.keras.layers.Dense(128, activation='relu'), tf.keras.layers.Dropout(0.5), tf.keras.layers.Dense(10, activation='softmax') ])
import byteps.tensorflow as bps from byteps.tensorflow.distribute import MirroredStrategy parser = argparse.ArgumentParser( description='TensorFlow Synthetic Benchmark', formatter_class=argparse.ArgumentDefaultsHelpFormatter) parser.add_argument('--rank', default=-1, type=int, help='node rank for distributed training') args = parser.parse_args() tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.INFO) bps.init() args.rank = bps.local_rank() print("my rank ", args.rank) gpus = tf.config.experimental.list_physical_devices('GPU') for gpu in gpus: tf.config.experimental.set_memory_growth(gpu, True) if gpus: tf.config.experimental.set_visible_devices(gpus[bps.local_rank()], 'GPU') def mnist_dataset(batch_size): (x_train, y_train), _ = tf.keras.datasets.mnist.load_data() # The `x` arrays are in uint8 and have values in the range [0, 255]. # We need to convert them to float32 with values in the range [0, 1] x_train = x_train / np.float32(255) y_train = y_train.astype(np.int64)
def create_optimizer(loss, init_lr, num_train_steps, num_warmup_steps, use_tpu): """Creates an optimizer training op.""" global_step = tf.train.get_or_create_global_step() learning_rate = tf.constant(value=init_lr, shape=[], dtype=tf.float32) # Implements linear decay of the learning rate. learning_rate = tf.train.polynomial_decay(learning_rate, global_step, num_train_steps, end_learning_rate=0.0, power=1.0, cycle=False) # Implements linear warmup. I.e., if global_step < num_warmup_steps, the # learning rate will be `global_step/num_warmup_steps * init_lr`. if num_warmup_steps: global_steps_int = tf.cast(global_step, tf.int32) warmup_steps_int = tf.constant(num_warmup_steps, dtype=tf.int32) global_steps_float = tf.cast(global_steps_int, tf.float32) warmup_steps_float = tf.cast(warmup_steps_int, tf.float32) warmup_percent_done = global_steps_float / warmup_steps_float warmup_learning_rate = init_lr * warmup_percent_done is_warmup = tf.cast(global_steps_int < warmup_steps_int, tf.float32) learning_rate = ((1.0 - is_warmup) * learning_rate + is_warmup * warmup_learning_rate) # It is recommended that you use this optimizer for fine tuning, since this # is how the model was trained (note that the Adam m/v variables are NOT # loaded from init_checkpoint.) optimizer = AdamWeightDecayOptimizer( learning_rate=learning_rate, weight_decay_rate=0.01, beta_1=0.9, beta_2=0.999, epsilon=1e-6, exclude_from_weight_decay=["LayerNorm", "layer_norm", "bias"]) if os.environ.get("USE_BYTEPS") and os.environ.get( "USE_BYTEPS").upper() in ["1", "TRUE", "Y"]: print("=================USING DISTRIBUTED OPTIMIZER=================") optimizer = bps.DistributedOptimizer(optimizer) tvars = tf.trainable_variables() grads = optimizer.compute_gradients(loss, tvars) # This is how the model was pre-trained. (grads, _) = tf.clip_by_global_norm(grads, clip_norm=1.0) trace_dir = os.path.join(os.environ.get("BYTEPS_TRACE_DIR", "."), str(bps.local_rank())) dump_computation_graph(trace_dir) train_op = optimizer.apply_gradients(zip(grads, tvars), global_step=global_step) # Normally the global step update is done inside of `apply_gradients`. # However, `AdamWeightDecayOptimizer` doesn't do this. But if you use # a different optimizer, you should probably take this line out. new_global_step = global_step + 1 train_op = tf.group(train_op, [global_step.assign(new_global_step)]) return train_op
def main(_): bps.init() tf.logging.set_verbosity(tf.logging.INFO) bert_config = modeling.BertConfig(256) model_fn = model_fn_builder(bert_config=bert_config, learning_rate=FLAGS.learning_rate, num_train_steps=FLAGS.num_train_steps, num_warmup_steps=FLAGS.num_warmup_steps) max_seq_length = FLAGS.max_seq_length max_predictions_per_seq = FLAGS.max_predictions_per_seq with tf.name_scope("input"): input_ids = tf.placeholder( shape=[FLAGS.train_batch_size, max_seq_length], dtype=tf.int32) input_mask = tf.placeholder( shape=[FLAGS.train_batch_size, max_seq_length], dtype=tf.int32) segment_ids = tf.placeholder( shape=[FLAGS.train_batch_size, max_seq_length], dtype=tf.int32) masked_lm_positions = tf.placeholder( shape=[FLAGS.train_batch_size, max_predictions_per_seq], dtype=tf.int32) masked_lm_ids = tf.placeholder( shape=[FLAGS.train_batch_size, max_predictions_per_seq], dtype=tf.int32) masked_lm_weights = tf.placeholder( shape=[FLAGS.train_batch_size, max_predictions_per_seq], dtype=tf.float32) next_sentence_labels = tf.placeholder( shape=[FLAGS.train_batch_size, 1], dtype=tf.int32) features = { "input_ids": input_ids, "input_mask": input_mask, "segment_ids": segment_ids, "masked_lm_positions": masked_lm_positions, "masked_lm_ids": masked_lm_ids, "masked_lm_weights": masked_lm_weights, "next_sentence_labels": next_sentence_labels } train_op = model_fn(features, None, None, None) infer_shape_ops = add_infer_shape_ops() hooks = [ # Horovod: BroadcastGlobalVariablesHook broadcasts initial variable states # from rank 0 to all other processes. This is necessary to ensure consistent # initialization of all workers when training is started with random weights # or restored from a checkpoint. bps.BroadcastGlobalVariablesHook(0), # Horovod: adjust number of steps based on number of GPUs. tf.train.StopAtStepHook(last_step=205 // bps.size()), ] config = tf.ConfigProto() config.gpu_options.allow_growth = True config.gpu_options.visible_device_list = str(bps.local_rank()) training_batch_generator = train_input_generator(features) with tf.train.MonitoredTrainingSession(hooks=hooks, config=config) as mon_sess: mon_sess = TimelineSession(mon_sess, infer_shape_ops) while not mon_sess.should_stop(): # Run a training step synchronously. feed_dict = next(training_batch_generator) mon_sess.run([train_op], feed_dict=feed_dict)