def total_train_step(row_offsets, value_tensors, nnz_array, labels): with tf.GradientTape() as tape: # do embedding fprop embedding_results = sparse_model(row_offsets, value_tensors, nnz_array) # convert to PerReplica dense_inputs = tf.split(embedding_results, num_or_size_splits=gpu_count) dense_inputs = PerReplica(dense_inputs) labels = tf.expand_dims(labels, axis=1) labels = tf.split(labels, num_or_size_splits=gpu_count) labels = PerReplica(labels) replica_loss, input_grads = strategy.run(dense_train_step, args=(dense_inputs, labels)) # gather all grads from dense replicas all_grads = tf.concat(input_grads.values, axis=0) # do embedding backward embedding_grads = tape.gradient(embedding_results, sparse_model.trainable_weights, output_gradients=all_grads) sparse_opt.apply_gradients( zip(embedding_grads, sparse_model.trainable_weights)) return strategy.reduce(tf.distribute.ReduceOp.SUM, replica_loss, axis=None)
def total_train_step(indices, values, dense_shape, labels): with tf.GradientTape() as tape: indices, emb_ctx = nvtx_tf.ops.start(indices, message='emb_fprop', domain_name='forward') embedding_result = sparse_model(indices, values, dense_shape, output_shape = [-1, 26, 32] # [batch_size, slot_num, embedding_vec_size] ) embedding_result = nvtx_tf.ops.end(embedding_result, emb_ctx) embedding_result, dense_ctx = nvtx_tf.ops.start(embedding_result, message='dense_fprop', domain_name='forward') labels = tf.expand_dims(labels, axis=1) dense_inputs = tf.split(embedding_result, num_or_size_splits=len(gpus)) dense_labels = tf.split(labels, num_or_size_splits=len(gpus)) dense_inputs_replicas = PerReplica(dense_inputs) dense_labels_replicas = PerReplica(dense_labels) dense_losses, input_grads = distribute_strategy.run(dense_train_step, args=(dense_inputs_replicas, dense_labels_replicas)) all_grads = tf.concat(input_grads.values, axis=0) all_grads = nvtx_tf.ops.end(all_grads, dense_ctx) embedding_grads = tape.gradient(embedding_result, sparse_model.trainable_weights, output_gradients=all_grads) sparse_opt.apply_gradients(zip(embedding_grads, sparse_model.trainable_weights)) return distribute_strategy.reduce(tf.distribute.ReduceOp.SUM, dense_losses, axis=None)
def main(): # create MirroredStrategy with specified GPUs. strategy = tf.distribute.MirroredStrategy( devices=["/GPU:" + str(i) for i in range(gpu_count)]) # create model instance inner the scope of MirroredStrategy with strategy.scope(): model = PluginSparseModel(...) # define optimizer for the variables in DNN model except embedding layer opt = tf.keras.optimizers.SGD() # define loss function for each replica loss_fn = tf.keras.losses.BinaryCrossentropy( from_logits=False, reduction=tf.keras.losses.Reduction.NONE) def _replica_loss(labels, logits): loss_value = loss_fn(labels, logits) return tf.nn.compute_average_loss(loss_value, global_batch_size=batch_size) # define train step for one iteration in a replica @tf.function def _train_step(each_replica, label): with tf.GradientTape() as tape: label = tf.expand_dims(label, axis=1) logit = model(each_replica) replica_loss = _replica_loss(label, logit) replica_grads = tape.gradient(replica_loss, model.trainable_weights) opt.apply_gradients(zip(replica_grads, model.trainable_weights)) return replica_loss # create a tf.data.Dataset to read data dataset = ... # training loop for step, (row_indices, values, labels) in enumerate(dataset): # use this API to broadcast input data to each GPU to_each_replicas = hugectr_tf_ops_v2.broadcast_then_convert_to_csr( model.get_embedding_name, row_indices, values, T=[tf.int32] * gpu_count) to_each_replicas = PerReplica(to_each_replicas) labels = tf.split(labels, num_or_size_splits=gpu_count) labels = PerReplica(labels) # each replica iteration replica_loss = strategy.run(_train_step, args=(to_each_replicas, labels)) # loss reduction in all replicas total_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, replica_loss, axis=None)
def _get_step_inputs(batch, nb_instances): features, labels = batch if isinstance(labels, PerReplica): # need to make a `PerReplica` objects for ``nb_instances`` nb_instances = PerReplica([nb_instances] * len(labels.values)) step_inputs = (features, labels, nb_instances) return step_inputs
def _get_step_inputs(batch): features, labels = batch if isinstance(labels, PerReplica): labels = tf.concat(labels.values, axis=0) nb_instances = tf.reduce_sum(tf.cast(labels != -100, dtype=tf.int32)) if isinstance(labels, PerReplica): # need to make a `PerReplica` objects for ``nb_instances`` nb_instances = PerReplica([nb_instances] * len(labels.values)) step_inputs = (features, labels, nb_instances) return step_inputs
def call(self): slot_num = 26 max_nnz = 3 embedding_vec_size = 16 input_in_each_gpu = [ np.random.normal(size=(self.global_batch_size * slot_num, embedding_vec_size)).astype(np.float32) for _ in range(8) ] input_in_each_gpu_ = PerReplica(input_in_each_gpu) print("[INFO] original inputs: \n", input_in_each_gpu_) target = tf.split(tf.reshape( tf.math.reduce_sum(input_in_each_gpu, axis=0), shape=[self.global_batch_size, slot_num, embedding_vec_size]), num_or_size_splits=8) print("[INFO] target output: \n", target) @tf.function def _step(replica_input): replica_ctx = tf.distribute.get_replica_context() replica_output = sok_unit_test_lib.reduce_scatter_dispatcher( replica_ctx.replica_id_in_sync_group, replica_input, global_batch_size=self.global_batch_size, slot_num=slot_num, max_nnz=max_nnz) return replica_output outputs = self.strategy.run(_step, args=(input_in_each_gpu_, )) print("[INFO] replica output:\n", outputs) for i in range(len(input_in_each_gpu)): tf.debugging.assert_near(target[i], outputs.values[i], message="output %d not meet target.", atol=1e-5, rtol=1e-5)
def _v2_fprop_v1_test(): print("[INFO]: Testing plugin_v2 fprop_experimental vs tf..") if vocabulary_size < slot_num: raise ValueError("vocabulary_size must > slot_num.") # generate initial values init_value, input_keys = generate_embedding_init_value_and_inputs() # -------------------------------- hugectr ops ------------------------------------ # class TestModel(tf.keras.models.Model): def __init__(self, init_value, name_, embedding_type, optimizer_type, max_vocabulary_size_per_gpu, opt_hparams, update_type, atomic_update, scaler, slot_num, max_nnz, max_feature_num, embedding_vec_size, combiner): super(TestModel, self).__init__() self.input_buffer_reset = True if "distributed" == embedding_type else False self.embedding_name = hugectr_tf_ops_v2.create_embedding( init_value=init_value, name_=name_, embedding_type=embedding_type, optimizer_type=optimizer_type, max_vocabulary_size_per_gpu=max_vocabulary_size_per_gpu, opt_hparams=opt_hparams, update_type=update_type, atomic_update=atomic_update, scaler=scaler, slot_num=slot_num, max_nnz=max_nnz, max_feature_num=max_feature_num, embedding_vec_size=embedding_vec_size, combiner=combiner) def build(self, _): self.bp_trigger = self.add_weight(name="bp_trigger", shape=(1, ), dtype=tf.float32, trainable=True) @tf.function def call(self, row_offset, values, nnz, training=True): replica_ctx = tf.distribute.get_replica_context() result = hugectr_tf_ops_v2.fprop_experimental( self.embedding_name, replica_ctx.replica_id_in_sync_group, row_offset, values, nnz, self.bp_trigger, input_buffer_reset=self.input_buffer_reset) return result hugectr_tf_ops_v2.init(visible_gpus=gpus, seed=0, key_type='int64', value_type='float', batch_size=batch_size, batch_size_eval=len(gpus)) strategy = tf.distribute.MirroredStrategy( devices=['/GPU:' + str(i) for i in gpus]) with strategy.scope(): hugectr_model = TestModel( init_value=init_value, name_='test_embedding', embedding_type=embedding_type, optimizer_type='Adam', max_vocabulary_size_per_gpu=(vocabulary_size // len(gpus)) * 2 + 1, opt_hparams=[0.1, 0.9, 0.99, 1e-5], update_type='Global', atomic_update=True, scaler=1.0, slot_num=slot_num, max_nnz=max_nnz, max_feature_num=slot_num * max_nnz, embedding_vec_size=embedding_vec_size, combiner='sum') opt = tf.keras.optimizers.Adam(learning_rate=0.1, beta_1=0.9, beta_2=0.99, epsilon=1e-5) # preprocess inputs dataset_utils = CreateDataset(dataset_names=None, feature_desc=None, batch_size=batch_size, n_epochs=None, slot_num=slot_num, max_nnz=max_nnz, convert_to_csr=None, gpu_count=len(gpus), embedding_type=embedding_type, get_row_indices=None) if "distributed" == embedding_type: row_offsets, value_tensors, nnz_array = dataset_utils._distribute_keys_for_distributed( input_keys) elif "localized" == embedding_type: row_offsets, value_tensors, nnz_array = dataset_utils._distribute_keys_for_localized( input_keys) else: raise ValueError("Not supported embedding_type %s." % embedding_type) # forward function @tf.function def hugectr_train_step(row_offset, values, nnz): with tf.GradientTape() as tape: forward_result = hugectr_model(row_offset, values, nnz) grads = tape.gradient(forward_result, hugectr_model.trainable_weights) opt.apply_gradients(zip(grads, hugectr_model.trainable_weights)) return forward_result # -------------------------------- tf ops ------------------------------------------- # reshape_input_keys = np.reshape(input_keys, [-1, max_nnz]) tf_indices = tf.where(reshape_input_keys != -1) tf_values = tf.gather_nd(reshape_input_keys, tf_indices) sparse_tensor = tf.sparse.SparseTensor(tf_indices, tf_values, reshape_input_keys.shape) tf_embedding_layer = OriginalEmbedding( vocabulary_size=vocabulary_size, embedding_vec_size=embedding_vec_size, initializer=init_value, combiner='sum', gpus=gpus) tf_opt = tf.keras.optimizers.Adam(learning_rate=0.1, beta_1=0.9, beta_2=0.99, epsilon=1e-5) @tf.function def tf_train_step(sparse_tensor): with tf.GradientTape() as tape: tf_forward = tf_embedding_layer( sparse_tensor, output_shape=[batch_size, slot_num, embedding_vec_size]) grads = tape.gradient(tf_forward, tf_embedding_layer.trainable_weights) tf_opt.apply_gradients( zip(grads, tf_embedding_layer.trainable_weights)) return tf_forward # ------------------ comparison ---------------------------------------------------- # for iteration in range(2): replica_row_offsets = PerReplica(row_offsets) replica_values = PerReplica(value_tensors) replica_nnz = PerReplica(nnz_array) hugectr_forward = strategy.run(hugectr_train_step, args=(replica_row_offsets, replica_values, replica_nnz)) if len(gpus) > 1: hugectr_forward = tf.concat(hugectr_forward.values, axis=0) tf_forward = tf_train_step(sparse_tensor) try: tf.debugging.assert_near(hugectr_forward, tf_forward, rtol=1e-4, atol=1e-5) except tf.errors.InvalidArgumentError as error: raise error else: print( "[INFO]: The results from HugeCTR and tf in %d iteration are the same" % (iteration + 1)) # --------------------- release resources -------------------------------------- # hugectr_tf_ops_v2.reset()
def call(self): rows_num_per_sample = 26 max_nnz = 3 all_inputs = np.random.randint( low=1, high=100, size=[self.global_batch_size * rows_num_per_sample, max_nnz]) all_mask = np.random.randint( low=0, high=2, size=[self.global_batch_size * rows_num_per_sample, max_nnz]) all_inputs *= all_mask print("[INFO] original dense all inputs:\n", all_inputs) all_valid_indices = tf.where(all_inputs != 0) all_valid_values = tf.gather_nd(all_inputs, all_valid_indices) all_inputs_sparse_tensor = tf.sparse.SparseTensor( values=all_valid_values, indices=all_valid_indices, dense_shape=all_inputs.shape) print("[INFO] original inputs sparse tensor:\n", all_inputs_sparse_tensor) sparse_tensors = tf.sparse.split(sp_input=all_inputs_sparse_tensor, num_split=8, axis=0) sparse_tensors = PerReplica(sparse_tensors) print("[INFO] to each replica sparse tensors:\n", sparse_tensors) target_values = all_inputs_sparse_tensor.values # target_indices = tf.concat([tf.transpose(sparse_tensor.indices, perm=[1, 0])[0] # for sparse_tensor in sparse_tensors.values], # axis=0) target_indices = tf.transpose(all_inputs_sparse_tensor.indices, perm=[1, 0])[0] target_num_elements = tf.concat([ tf.shape(sparse_tensor.indices, out_type=tf.int64)[0] for sparse_tensor in sparse_tensors.values ], axis=0) target_total_valid_num = tf.size(target_values, out_type=tf.int64) print("[INFO] target_values: \n", target_values) print("[INFO] target_indcies: \n", target_indices) print("[INFO] target_num_elements: \n", target_num_elements) print("[INFO] target_total_valid_num: \n", target_total_valid_num) @tf.function def _step(sparse_tensor): if not isinstance(sparse_tensor, tf.sparse.SparseTensor): raise RuntimeError( "sparse_tensor must be a tf.sparse.SparseTensor") values = sparse_tensor.values # [num_of_valids,] indices = sparse_tensor.indices row_indices = tf.transpose(indices, perm=[1, 0])[0] # [num_of_valids] replica_ctx = tf.distribute.get_replica_context() values_out, indices_out, num_elements, total_valid_num = sok_unit_test_lib.all_gather_dispatcher( replica_ctx.replica_id_in_sync_group, replica_ctx.num_replicas_in_sync, values, row_indices, global_batch_size=self.global_batch_size, rows_num_per_sample=rows_num_per_sample, max_nnz=max_nnz) return values_out, indices_out, num_elements, total_valid_num values_out, indices_out, num_elements, total_valid_num = self.strategy.run( _step, args=(sparse_tensors, )) print("[INFO]: after all gather dispatcher, values = \n", values_out) print("[INFO]: after all gather dispatcher, indices = \n", indices_out) print("[INFO]: after all gather dispatcher, num_elements = \n", num_elements) print("[INFO]: after all gather dispatcher, total_valid_num = \n", total_valid_num) for i in range(len(values_out.values)): tf.debugging.assert_equal( target_values, values_out.values[i][:target_values.shape[0]], message="values %d not meet target." % i) tf.debugging.assert_equal( target_indices, indices_out.values[i][:target_indices.shape[0]], message="indcies %d not meet target." % i) tf.debugging.assert_equal( target_num_elements, num_elements.values[i][:target_num_elements.shape[0]], message="num_elements %d not meet target." % i) tf.debugging.assert_equal( target_total_valid_num, total_valid_num.values[i], message="total_valid_num %d not meet target." % i)
def profile_plugin(embedding_type, gpu_count, vocabulary_size=1737710, slot_num=26, max_nnz=1, batch_size=65536, fprop_version='v1'): file_name = "plugin_v2_" + embedding_type + "_" + str(gpu_count) + "_" + fprop_version dataset = plugin_reader(file_name, gpu_count, fprop_version) # # build model strategy = tf.distribute.MirroredStrategy(devices=["/GPU:" + str(i) for i in range(gpu_count)]) with strategy.scope(): if fprop_version == 'v1': model = PluginSparseModel(batch_size=batch_size, gpus=[i for i in range(gpu_count)], init_value=False, name_='hugectr_embedding', embedding_type=embedding_type, optimizer_type='Adam', max_vocabulary_size_per_gpu=(vocabulary_size // gpu_count) + 1, opt_hparams=[0.1, 0.9, 0.999, 1e-3], update_type='Local', atomic_update=True, scaler=1.0, slot_num=slot_num, max_nnz=max_nnz, max_feature_num=100, embedding_vec_size=32, combiner='sum', num_dense_layers=7, input_buffer_reset=False) elif fprop_version == 'v2': model = PluginSparseModelV2(batch_size=batch_size, gpus=[i for i in range(gpu_count)], init_value=False, name_='hugectr_embedding', embedding_type=embedding_type, optimizer_type='Adam', max_vocabulary_size_per_gpu=(vocabulary_size // gpu_count) + 1, opt_hparams=[0.1, 0.9, 0.999, 1e-3], update_type='Local', atomic_update=True, scaler=1.0, slot_num=slot_num, max_nnz=max_nnz, max_feature_num=100, embedding_vec_size=32, combiner='sum', num_dense_layers=7, input_buffer_reset=False) dense_opt = tf.keras.optimizers.SGD() loss_fn = tf.keras.losses.BinaryCrossentropy(from_logits=False, reduction=tf.keras.losses.Reduction.NONE) def _replica_loss(labels, logits): loss_v = loss_fn(labels, logits) return tf.nn.compute_average_loss(loss_v, global_batch_size=batch_size) if fprop_version == 'v1': @tf.function def _train_step(row_offset, values, nnz, label): with tf.GradientTape() as tape: label = tf.expand_dims(label, axis=1) logit = model(row_offset, values, nnz) replica_loss = _replica_loss(label, logit) replica_grads = tape.gradient(replica_loss, model.trainable_weights) dense_opt.apply_gradients(zip(replica_grads, model.trainable_weights)) return replica_loss elif fprop_version == 'v2': @tf.function def _train_step(each_replica, label): with tf.GradientTape() as tape: label = tf.expand_dims(label, axis=1) logit = model(each_replica) replica_loss = _replica_loss(label, logit) replica_grads = tape.gradient(replica_loss, model.trainable_weights) dense_opt.apply_gradients(zip(replica_grads, model.trainable_weights)) return replica_loss options = tf.data.Options() options.experimental_distribute.auto_shard_policy = tf.data.experimental.AutoShardPolicy.OFF dataset = dataset.with_options(options) input_options = tf.distribute.InputOptions( experimental_prefetch_to_device=True, # TODO: not working.. experimental_replication_mode=tf.distribute.InputReplicationMode.PER_WORKER, experimental_place_dataset_on_device=False ) if fprop_version == "v1": dataset = strategy.experimental_distribute_dataset(dataset, input_options) for step, (row_offsets, value_tensors, nnz_array, labels) in enumerate(dataset): step, step_ctx = nvtx_tf.ops.start(tf.convert_to_tensor(step, dtype=tf.int32), message='Iteration_' + str(step)) replica_loss = strategy.run(_train_step, args=(row_offsets, value_tensors, nnz_array, labels)) total_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, replica_loss, axis=None) tf.print("step:%d, loss:%.5f" %(step, total_loss)) step = nvtx_tf.ops.end(step, step_ctx) elif fprop_version == 'v2': for step, (row_indices, values, labels) in enumerate(dataset): step, step_ctx = nvtx_tf.ops.start(tf.convert_to_tensor(step, dtype=tf.int32), message='Iteration_' + str(step)) to_each_replicas = hugectr_tf_ops_v2.broadcast_then_convert_to_csr( model.get_embedding_name, row_indices, values, T = [tf.int32] * gpu_count) to_each_replicas = PerReplica(to_each_replicas) labels = tf.split(labels, num_or_size_splits=gpu_count) labels = PerReplica(labels) replica_loss = strategy.run(_train_step, args=(to_each_replicas, labels)) total_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, replica_loss, axis=None) tf.print("step:%d, loss:%.5f" %(step, total_loss)) step = nvtx_tf.ops.end(step, step_ctx)