def multi_output_partition_info(ref_model, rank): if rank == p_0_rank: partition_info = pinfo.PartitionInfo('0') in_0 = pinfo.EndpointInfo(0, ref_model.inputs[0].shape, tf.float32) partition_info.real_input_infos = {0: in_0} out_edge_0 = pinfo.EndpointInfo( 0, ref_model.get_layer('ten_classes').output.shape, tf.float32) out_edge_1 = pinfo.EndpointInfo( 1, ref_model.get_layer('two_classes').output.shape, tf.float32) partition_info.edge_output_infos = {0: out_edge_0, 1: out_edge_1} elif rank == p_1_rank: partition_info = pinfo.PartitionInfo('1') in_edge_0 = pinfo.EndpointInfo( 0, ref_model.get_layer('ten_classes').output.shape, tf.float32) in_edge_1 = pinfo.EndpointInfo( 1, ref_model.get_layer('two_classes').output.shape, tf.float32) partition_info.edge_input_infos = {0: in_edge_0, 1: in_edge_1} out_0 = pinfo.EndpointInfo(0, ref_model.outputs[0].shape, tf.float32) out_1 = pinfo.EndpointInfo(1, ref_model.outputs[1].shape, tf.float32) partition_info.real_output_infos = {0: out_0, 1: out_1} return partition_info
def alexnet_partition_info(ref_model, rank): if rank == p_0_rank: partition_info = pinfo.PartitionInfo('0') in_0 = pinfo.EndpointInfo(0, ref_model.inputs[0].shape, tf.float32) partition_info.real_input_infos = {0: in_0} out_edge_0 = pinfo.EndpointInfo( 0, ref_model.get_layer('split_layer0').output.shape, tf.float32) partition_info.edge_output_infos = {0: out_edge_0} elif rank == p_1_rank: partition_info = pinfo.PartitionInfo('1') in_edge_0 = pinfo.EndpointInfo( 0, ref_model.get_layer('split_layer0').output.shape, tf.float32) partition_info.edge_input_infos = {0: in_edge_0} out_edge_0 = pinfo.EndpointInfo( 1, ref_model.get_layer('split_layer1').output.shape, tf.float32) partition_info.edge_output_infos = {0: out_edge_0} elif rank == p_2_rank: partition_info = pinfo.PartitionInfo('2') in_edge_0 = pinfo.EndpointInfo( 1, ref_model.get_layer('split_layer1').output.shape, tf.float32) partition_info.edge_input_infos = {0: in_edge_0} out_0 = pinfo.EndpointInfo(0, ref_model.outputs[0].shape, tf.float32) partition_info.real_output_infos = {0: out_0} return partition_info
def to_microbatched(model, micro_batch_size, num_micro_batches, num_batches, num_test_batches): rank = tnt.get_rank() partition_generator = pgen.GraphPartitionGenerator(model) rank_mapper = rmapper.RankMapper(num_ranks = tnt.get_size(), pipeline_graph = partition_generator.get_pipeline_graph()) partition_id = rank_mapper.get_partition_for_rank(rank) partition_graph = partition_generator.get_partition_graph(partition_id) partition_info = pinfo.PartitionInfo(partition_id = partition_id, partition_graph = partition_graph) core_model_builder = cm_builder.CoreModelBuilder(model, partition_id, partition_graph) core_model = core_model_builder.get_model() connection_table = rank_mapper.get_connections_for_rank(rank) pipeline_communicator = tnt.PipelineCommunicator(connection_table, num_micro_batches) shared_model_builder = shared.SharedModelBuilder(partition_info, core_model, pipeline_communicator, micro_batch_size) shared_model = shared_model_builder.get_model() microbatched_model_builder = microbatched.MicrobatchedModelBuilder(partition_info, shared_model, micro_batch_size, num_micro_batches) ds = load_microbatched_datasets(micro_batch_size, num_micro_batches, num_batches, num_test_batches, partition_info) pipeline_communicator.setup_infrastructure(micro_batch_size) return microbatched_model_builder, ds
def simple_partition_info(ref_model, rank): partition_info = pinfo.PartitionInfo('0') in_0 = pinfo.EndpointInfo(0, ref_model.inputs[0].shape, tf.float32) partition_info.real_input_infos = {0: in_0} out_0 = pinfo.EndpointInfo(0, ref_model.outputs[0].shape, tf.float32) partition_info.real_output_infos = {0: out_0} return partition_info
def test_partition_info(self, model_and_partitions): model, partition_gen, expected_num_partitions, expected_partition_gen, _ = model_and_partitions rank_mapper = rmapper.RankMapper( num_ranks=expected_num_partitions, pipeline_graph=partition_gen.get_pipeline_graph()) for rank in range(expected_num_partitions): partition_id = rank_mapper.get_partition_for_rank(rank) partition_info = pinfo.PartitionInfo( partition_id=partition_id, partition_graph=partition_gen.get_partition_graph( partition_id)) assert partition_info == expected_partition_gen(model, rank)
def __init__(self, model, group, partition_generator, rank_mapper, num_pipeline_stages = None): super().__init__(model = model, group = group) self._model_name = model.name self.built = False self.compile_properties = None self.num_pipeline_stages = num_pipeline_stages connection_table = rank_mapper.get_connections_for_rank(self.rank) self.pipeline_communicator = tnt.PipelineCommunicator(connection_table, self.num_pipeline_stages) self.initialized = False partition_id = rank_mapper.get_partition_for_rank(self.rank) partition_graph = partition_generator.get_partition_graph(partition_id) self.partition_info = pinfo.PartitionInfo(partition_id, partition_graph) core_model_builder = cm_builder.CoreModelBuilder(model, partition_id, partition_graph) self.model = core_model_builder.get_model() self.nano_batch_size = None self.built = False