Package: pipeline-tool
Pipeline toolTable of contents[toc]IntroductionThe field of machine learning is constantly evolving, with increasingly sophisticated models and ever-expanding datasets. Professionals in the field can face significant challenges, especially when it comes to training models that are too large to fit into the memory of a single GPU.In this context, developers have created a tool for distributing a PyTorch machine learning model across multiple GPUs without altering the training process. The tool takes the PyTorch model description as input, interprets each layer independently, and rewrites the model to handle the operations' interdependencies. The result is a new model that can be automatically distributed (creating a pipeline) across multiple GPUs without affecting the quality of the results.In the following, we will present this tool in detail, along with the benefits it can offer to machine learning professionals seeking to train large and complex modelsHow it worksFirst stepTo run the tool, you need to provide some parameters:Number of GPUs: If the user doesn't specify the number of GPUs to use, the tool will automatically detect the available GPUs on the machine running the command. In this case, the model will be trained using all detected GPUs to improve performance.PyTorch Model: The user must provide a PyTorch model that only uses functions and modules from the PyTorch API. In other words, the model should not incorporate custom functions unknown to the PyTorch API. However, you can create custom layers using combinations of functions (always from the PyTorch API).Shapes of the input and output: You will need to provide these to profilememory usage.Data type that will be passed into the model, for example, float.The first step is to add the following imports to your project:fromtorch.distributed.pipeline.syncimportPipefrompipeline_tool.pipeline_toolimportSkippableTracingfrompipeline_tool.pipeline_configimportPipelineConfigNext, you should use the PipelineConfig class, which enables you to prepare the necessary parameters (input and output shape, data type).config=PipelineConfig([X,Y,Z],[X,Y],"dtype")One important thing to know is that the first number given in input/output shape is your batch size.Once you have defined your config and created your model, you can process it as shown in the example below.N_GPUs=2trace=SkippableTracing(N_GPUs,model,config)graph_model=trace.get_modules()Here, we trace the model usingtorch.fxto obtain the GraphModule. This allows us to determine, for each module, its type, parameters, functions (e.g., convolution, activation, multiplication), and their links to other modules.Below is an example of how a simple model is treated:In this basic example, we have a model composed exclusively of PyTorch modules. To describe them accurately, we utilize the trace generated by torch fx.The generated trace appears as follows:OpcodeNameTargetplaceholderxxcall_modulelinear1linear1call_moduleactivationactivationcall_modulelinear2linear2call_modulesoftmaxsoftmaxoutputoutputoutputThis trace allows us to identify each generated layer and provides the following information:Opcode: Indicates the type of operation performed by the layer.Name: Corresponds to the name of the function or operation performed by the layer.Target: Represents the name of the layer as it appears in the description of the PyTorch model.Thus, the trace provides a detailed view of the operations performed by each layer, making it easier to understand and analyze the model.NameModuledeclaration
linear1Linear(in_features=100,out_features=200,bias=True)activationReLU()linear2Linear(in_features=200,out_features=10,bias=True)softmaxSoftmax(dim=None)The retrieval, analysis, and management of all this information enable the generation of a file containing a new model ready for pipelined training on N GPUs.Model splittingNow that we can create a model that can be distributed across multiple GPUs, the question arises: how do we split it intelligently? Currently, the tool proceeds with a somewhat naive approach. We create a dummy dataset to pass through the model and perform a training run. This allows us to measure the memory loads on all GPUs.Initially, the tool divides the layers into two equal parts (in terms of the number of layers) and conducts these memory load measurements.
If the load is not evenly distributed, we rewrite the model (moving layers around) and iterate the dummy run until we achieve uniform distribution on N GPUs.Pipeline Tool ExampleTwo examples are provided inexamples/. They demonstrate how to:Use the pipeline_toolTrain a pipelined modelEvaluating a pipelined modelSave the trained model.Pipeline Tool in Giotto DeepThe Pipeline tool is seamlessly integrated into Giotto-Deep's trainer, requiring no changes to its API.Here's an example of what you need to do:# New importfromgdeep.trainer.trainerimportParallelism,ParallelismType# Create the trainer as beforetrainer=Trainer(wrapped_model,[dl_train,dl_train],loss_function,writer)# Prepare the config of the MHAconfigs=[{'embed_dim':16,'num_heads':8,'dropout':0.1,'batch_first':True},{'embed_dim':16,'num_heads':8,'dropout':0.1,'batch_first':True},{'embed_dim':16,'num_heads':8,'dropout':0.1,'batch_first':True},{'embed_dim':16,'num_heads':8,'dropout':0.1,'batch_first':True},{'embed_dim':16,'num_heads':8,'dropout':0.1,'batch_first':True}]# List of devicedevices=list(range(torch.cuda.device_count()))# Use the Parallelism class created to prepare the trainer for a pipeline runparallel=Parallelism(ParallelismType.PIPELINE,devices,len(devices),pipeline_chunks=2,config_mha=configs)# Call the train function with the new parametertrainer.train(Adam,2,parallel=parallel)ExampleTo experiment with Giotto Deep training using the Pipeline tool in your environment, we have provided two example scripts. Navigate to Giotto's examples folder and run eitherpipeline_basic_image.pyorpipeline_orbit5k.pywith the --pipeline argument to enable the pipeline mode, or without it for regular training.InstallationInstall from sourcesLaunch from the root of the project:pipinstall.This will install pipeline_tool on your Python environment.The necessary imports are:frompipeline_tool.pipeline_configimportPipelineConfigfrompipeline_tool.pipeline_toolimportSkippableTracingBenchmarkingA benchmarking tool is available. This script will test the pipeline_tool on 3 different models:A FFNETA CNNOne VisionTransformerWith these 3 models, we cover the majority of cases that the tool has to deal with. "The CNN and FFNET are two small models and quickly become unable to be split too much, while the VisionTransformer is very large and may not necessarily fit on 1 GPU, and it also contains MultiHeads.This is how we proceed:When you launch the script, set the maximum number of GPUs in the environment (in the example below, 8), then run the first execution with the torch API to create a repository before launching the analyses with the pipeline_tool.The results are as follows:FrameworkModelNumber of GPUsNumber of ChunksTime [s]Alloc [MB]API torchCNN100.53[625]API torchFFNET100.25[520]PipelineCNN121.16[698]PipelineCNN222.03[582, 514]PipelineCNN322.22[582, 0, 514]PipelineCNN423.31[69, 514, 513, 514]PipelineCNN523.97[79, 45, 514, 513, 513]PipelineCNN625.427[79, 51, 0, 514, 513, 513]PipelineCNN725.97[54, 68, 0, 514, 513, 0, 513]PipelineFFNET120.67[668]PipelineFFNET221.4[522, 514]PipelineVisionTransformer222269.89[3979519, 3958031]PipelineVisionTransformer322014.98[2574197, 3021793, 2635734]PipelineVisionTransformer421861.44[2119623, 2228904, 2145561, 2112281]PipelineVisionTransformer521786.39[1706017, 1918617, 1732474, 1925785, 1657517]PipelineVisionTransformer621715.75[1478507, 1691115, 1546192, 1664987, 1546192, 1430453]PipelineVisionTransformer721676.8[1437800, 1277582, 1436969, 1498672, 1277582, 1477999, 1229807]PipelineVisionTransformer821631.67[1210541, 1278218, 1277900, 1312002, 1276935, 1209459, 1209459, 1195574]Result analysisFirstly, we notice that some results are missing; for example, for the FFNET, we only have results on 1 or 2 GPUs, etc. It's simply because when an error occurs, we don't store the result in the benchmark. But there are 4 possible types of error:If the first/last GPU in the chain has only one layer, it cannot be executed.One of the GPUs has 0 layers.Cuda Out of Memory, at least 1 of the GPUs can't handle the amount of layers and data given to it.Finally, an error occurs during training.If none of these errors occur, we store the results.So, based on this, we can see right away that no input is available with the torch API for the VisionTransformer, simply because it doesn't fit on a single GPU. As a result, the pipeline tool allows the model to be separated on multiple GPUs; we can see that with the VisionTransformers that can only be run on more than 1 GPU. Another point to note is that the tool slows down execution time anyway, due to the added communication between GPUs; execution time between the torch API and runs of the Pipeline in 1 GPU show it clearly. So you can't use it routinely and should only use it in really useful cases, i.e., when the model can't fit on a single GPU.Visual explaination of the Pipeline ToolAflowchartexplain all the process made by the tool for generating a complete model.Complex ModelsUnfortunately, a model is never limited to a simple linear sequence of modules taking the output of the previous operation as input... More complex models exist, and it is necessary to handle all possible cases, to trace the model correctly so that it is faithfully reproduced without omitting certain operations.As a result, it is necessary to distinguish PyTorch modules from other elements.We analyze the model received as a parameter and store the elements by their names in a dictionary, which we use to create a correspondence table with the names given by the trace.We can then iterate over the generated trace to differentiate five types of layers:Module: These need to be initialized and thus require their description to be retrieved from the original model.Function: These correspond to simple PyTorch functions executed between tensors or on a tensor (e.g., additions, dimension reductions, etc.).Getitem: These appear in the trace when only a portion of a result in the form of a list needs to be retrieved (e.g., index 1 of a list or the first return value of a function).Getattr: These correspond to retrieving an attribute of a tensor.Propagation: These appear in the trace to propagate tensors to other layers.call_functionLet's explore the concept of call_function with the widely known ResNet model.When we examine the generated trace, we notice a new opcode, distinct from the one we previously discussed in thefirst step:OpcodeNameArgumentsplaceholderx()call_moduleflow_0_conv1(x,)[...]call_moduleflow_0_avgpool(flow_0_layer4_1_relu_1,)# ############################################################################################call_functionflatten(flow_0_avgpool,1)# ############################################################################################call_moduleflow_0_fc(flatten,)call_moduleflow_1(flow_0_fc,)outputoutput(flow_1,)Notice that we also have access to the input arguments of each layer.call_functions are treated differently from call_modules and consequently generate distinct code. Therefore, each call_function is declared as a Torch module that exclusively performs the necessary operation. In the case of the previous trace, let's consider the declaration of the call_functionflatten:classflatten_layer(nn.Module):defforward(self,input):ret=torch.flatten(input,1)returnretFunctions do not necessitate an initialization function. Instead, our tool seeks out the appropriate Torch function based on the name provided in the trace. For instance, when working with the instantiated ResNet18 model, the function "flatten" already exists within the Torch API.The trace allows us to identify the arguments passed to this function. In the case above, the inputs are the output of the previous layer and the integer "1".PropagationAs discussed in the section oncomplex modelsthere are instances where we need to transmit the output of one layer to others that are not inherently connected to it. To facilitate this process, PyTorch provides a useful decorator called "skippable." This decorator introduces two key features:stash: This feature permits us to store a specific value with an associated name, allowing for convenient retrieval later.pop: With this functionality,Let's get a look into an example trace to have a better understanding::OpcodeNameArgumentsplaceholderx()call_moduleflow_0_conv1(x,)[...]call_moduleflow_0_maxpool(flow_0_relu,)call_moduleflow_0_layer1_0_conv1(flow_0_maxpool,)call_moduleflow_0_layer1_0_bn1(flow_0_layer1_0_conv1,)call_moduleflow_0_layer1_0_relu(flow_0_layer1_0_bn1,)call_moduleflow_0_layer1_0_conv2(flow_0_layer1_0_relu,)call_moduleflow_0_layer1_0_bn2(flow_0_layer1_0_conv2,)#############################################################################################call_functionadd(flow_0_layer1_0_bn2,flow_0_maxpool)#############################################################################################call_moduleflow_0_layer1_0_relu_1(add,)[...]call_moduleflow_0_fc(flatten,)call_moduleflow_1(flow_0_fc,)outputoutput(flow_1,)The call_function surrounded have two name in input :flow_0_layer1_0_bn2, which directly stems from the previous layer.flow_0_maxpool, originating from an earlier layer in the model.Our tool is designed to establish connections between layers and retain information about the arguments derived from prior layers.Consequently, when utilizing the skippable decorator in the generated code:[...]@skippable(stash=['flow_0_maxpool_to_add'])classflow_0_maxpool_layer(nn.Module):def__init__(self)->None:super().__init__()self.fc=nn.MaxPool2d(kernel_size=3,stride=2,padding=1,dilation=1,ceil_mode=False)defforward(self,input):ret=self.fc(input)yieldstash('flow_0_maxpool_to_add',ret)returnret[...]classflow_0_layer1_0_bn2_layer(nn.Module):def__init__(self)->None:super().__init__()self.fc=nn.BatchNorm2d(64,eps=1e-05,momentum=0.1,affine=True,track_running_stats=True)defforward(self,input):ret=self.fc(input)returnret@skippable(pop=['flow_0_maxpool_to_add'])classadd_layer(nn.Module):defforward(self,input):flow_0_maxpool=yieldpop('flow_0_maxpool_to_add')ret=torch.add(input,flow_0_maxpool)returnret[...]We ensure that the dependencies between layers are properly preserved.GetitemWithin the trace, certain call_function entries contain the term "getitem" in their names. This indicates that these are not conventional functions but rather indicate the need to access a specific index within a result. Consider the following trace as an example:[...]call_functiongetattr_1(add_3,'shape')call_functiongetitem_4(getattr_1,0)[...]Here, we notice the presence of a getitem operation, which is applied to the result of the previous layer. If we were to translate this trace, it would resemble something like add_3.shape[0] (for an explanation of getattr, please refer to thenext point).The challenge with getitem lies in the limitation of the Torch API, which does not allow the propagation of non-tensor values. Consequently, we must concatenate the getitem operation to the layer from which we require the value, rather than creating an independent layer that cannot effectively transmit its output.GetAttrThere are two distinct types ofgetattroperations:call_function with the Name "getattr": These instances occur when an attribute of modules needs to be accessed..In the provided trace:[...]call_functiongetattr_1(add_3,'shape')[...]call_modulemodel_pooling_layer_scaled_dot_product_attention(expand,add_3,add_3)As previously mentioned, we cannot propagate non-tensor values. The presence of getattr indicates the need to access a specific attribute within a module. In the trace above, the tensor add_3 possesses an attribute "shape" that will be utilized. In such cases, we refrain from creating new modules; instead, we reference the relevant attribute of the tensor when it is passed as a parameter.Here's an illustrative example of generated code to elucidate this approach:[...]@skippable(stash=['add_3_to_expand','add_3_to_model_pooling_layer_scaled_dot_product_attention'],pop=['add_2_to_add_3'])classadd_3_layer(nn.Module):defforward(self,input):add_2=yieldpop('add_2_to_add_3')ret=torch.add(input,add_2)yieldstash('add_3_to_expand',ret)yieldstash('add_3_to_model_pooling_layer_scaled_dot_product_attention',ret)returnret[...]@skippable(pop=['add_3_to_expand'])classexpand_layer(nn.Module):defforward(self,input):add_3=yieldpop('add_3_to_expand')ret=input.expand(add_3.shape[0],-1,-1)returnretget_attr with the Opcode "get_attr": These occurrences arise when a private attribute of a user-created class is requested.In the provided trace:get_attrmodel_pooling_layer_query()We only have the name of the attribute, and it needs to be initialized to propagate or utilize it, we create a module that initializes the attribute based on the provided information. We search for the attribute on the given model and recreate it identically.Here's an example of code to illustrate this process:classmodel_pooling_layer_query_layer(nn.Module):def__init__(self)->None:super().__init__()self.fc=nn.parameter.Parameter(torch.Tensor(1,16),requires_grad=True)defforward(self,input):ret=self.fcreturnretMultiHeadAttention processingUnpredictable management is, however, necessary for MultiHeadAttention. During the module declaration retrieval phase, it is impossible to retrieve those of the MultiHeadAttention. Therefore, the user must provide a dictionary containing the description of all the parameters and their values for the MultiHeadAttention of their model during the tool's initialization.At a minimum, the following parameters must be provided for a MultiHead:embed_dimnum_headsAnd the initialization would be changed to three alternative:Give your hand made dictionnary describing all your MHAmha_config=[{'embed_dim':hidden_val,'num_heads':heads_val,'dropout':0.1,'batch_first':True},{'embed_dim':hidden_val,'num_heads':heads_val,'dropout':0.1,'batch_first':True},{'embed_dim':hidden_val,'num_heads':heads_val,'dropout':0.1,'batch_first':True},{'embed_dim':hidden_val,'num_heads':heads_val,'dropout':0.1,'batch_first':True},{'embed_dim':hidden_val,'num_heads':heads_val,'dropout':0.1,'batch_first':True}]config=PipelineConfig([X,Y,Z],[X,Y],"dtype",mha_config)nb_gpus=2trace=SkippableTracing(nb_gpus,model,config)model_pipe=trace.get_modules()Use MHA dictionnary generator in PipelineConfig class. If you know that all your MHA are identical, you can use this function to create N dictionnary entry identical.config=PipelineConfig([X,Y,Z],[X,Y],"dtype")config.create_mha_conf_equal(embed_dim,num_heads,dropout,batch_first)nb_gpus=2trace=SkippableTracing(nb_gpus,model,config)model_pipe=trace.get_modules()Finaly some default model are setup with classmethod, Persformer, one bigger Persformer and VideoTransformer.# TODO EXPLAINATIONImprovementsIn its current state, the tool works, but it hasn't been designed for performance yet. That's why the following points for improvement are important:Although repartition is currently performed, it is unnecessary when the model fits within a single GPU. The process should automatically avoid splitting when feasible, requiring an initial run on the largest GPU and an error-handling mechanism.Replace the rudimentary repartition method with a more efficient approach, such as employing a dichotomous search.Actually, the tool is searching for the best memory balancing between GPUs. But after some execution time analysis, this solution is not the best concerning execution time. One improvement should be to search for the best execution time instead of the best memory balancing. To put this solution in place:Change the analysis returned by the scriptevaluate_mem.pyto return time and not memory balancing.Find a way to preprocess and create all potential best repartition to avoid testing all possibilities that could be exponential in process time depending on the number of layers.Change the behavior to test all pre-calculated possibilities and not stop, keeping the fastest one.Here is the time analysis made, in italic are the chosen repartition, and in bold, the minimal execution time:ModelNb GPURepartitionEpoch 1Epoch 2Epoch 3Minimal Epoch timeCNN2[7, 7]3.851.841.79CNN2[8, 6]3.811.911.85CNN2[9, 5]4.021.891.841.79CNN3[5, 5, 4]5.082.492.55CNN3[6, 4, 4]5.052.652.66CNN3[7, 3, 4]4.952.622.51CNN3[8, 2, 4]5.132.542.62CNN3[9, 1, 4]4.202.212.212.21CNN4[4, 4, 3, 3]6.663.203.32CNN4[4, 5, 2, 3]6.653.363.40CNN4[5, 4, 2, 3]6.203.323.25CNN4[6, 3, 2, 3]6.143.213.16CNN4[7, 2, 2, 3]6.073.233.28CNN4[8, 1, 2, 3]6.083.313.35CNN4[9, 1, 1, 3]5.392.882.882.88CNN5[3, 3, 3, 3, 2]7.963.963.85CNN5[3, 4, 2, 3, 2]7.813.873.73CNN5[3, 5, 1, 3, 2]7.863.854.05CNN5[3, 6, 1, 2, 2]7.053.613.53CNN5[3, 5, 2, 2, 2]7.873.813.913.53CNN6[3, 3, 2, 2, 2, 2]8.954.984.79CNN6[3, 3, 3, 1, 2, 2]8.104.074.194.07CNN7[2, 2, 2, 2, 2, 2, 2]8.554.604.64CNN7[2, 3, 2, 2, 1, 2, 2]9.595.695.53CNN7[2, 3, 3, 1, 1, 2, 2]9.265.715.69CNN7[2, 3, 4, 1, 1, 1, 2]8.444.564.424.42FFNET2[3, 2]2.391.321.27FFNET2[2, 3]2.411.391.341.27VisionTransformer2[184, 184]471.70470.42470.53VisionTransformer3[123, 123, 122]418.28416.19416.29VisionTransformer4[92, 92, 92, 92]385.19382.62383.39VisionTransformer5[74, 74, 74, 73, 73]370.09367.58367.66VisionTransformer6[62, 62, 61, 61, 61, 61]356.13353.30353.54VisionTransformer6[63, 62, 61, 60, 61, 61]357.41354.58354.84353.30VisionTransformer7[53, 53, 53, 53, 52, 52, 52]347.54345.12345.18VisionTransformer7[54, 53, 52, 53, 52, 52, 52]351.04347.78347.89VisionTransformer7[55, 52, 52, 53, 52, 52, 52]349.60346.24346.29VisionTransformer7[56, 52, 51, 53, 52, 52, 52]349.48346.58346.45VisionTransformer7[57, 52, 50, 53, 52, 52, 52]349.51346.42346.55VisionTransformer7[58, 52, 49, 53, 52, 52, 52]348.30345.28345.35VisionTransformer7[59, 52, 49, 53, 52, 51, 52]348.69345.15345.28345.12VisionTransformer8[46, 46, 46, 46, 46, 46, 46, 46]342.10338.47338.73VisionTransformer8[47, 45, 46, 46, 46, 46, 46, 46]342.17338.51338.44VisionTransformer8[48, 44, 46, 46, 46, 46, 46, 46]339.99336.45336.44336.44Known issueThe pipeline of Persformers in more than 2 GPU have backward process problem for an unknow reason and no error is throw.Actually when a CUDA error OOM is throw we admit that the envrionement don't have enough GPU. In the futur we will implement logic test to see if yes or not the model can be split on the desired config before telling that it is impossible.