A fine-grained, multi-dimensional method for large-scale machine learning performance analysis

Abstract

The present invention discloses a fine-grained, multi-dimensional, large-scale machine learning performance analysis method in the field of computer technology. The method is implemented based on a memory usage performance analysis tool, which includes the Colossal-AI Profiler performance analysis tool and the Colossal-AI Gemini memory policy management tool. The process used by the Colossal-AI Profiler is similar to that of the PyTorch Profiler. The Colossal-AI Profiler is designed as a context manager that can be quickly opened and closed using a "with" statement. It can accurately record Tensor information in different states, thereby obtaining the current specific memory usage situation, helping developers to better optimize memory usage and scheduling strategies during machine learning. It helps developers compare the bandwidth performance of the program with the theoretical upper limit of the data movement bandwidth, conduct objective and accurate performance analysis, obtain the current program's resource utilization, and find the main optimization direction of the program. This method is creative.

Description

Technical Field

[0001] The present invention relates to the field of computer technology, and in particular to a fine-grained, multi-dimensional, large-scale machine learning performance analysis method. Background Art

[0002] In large-scale software projects, performance analysis tools are primarily used to statically analyze various runtime metrics, such as a program's time and space complexity, function call stacks, and program activity overhead. Different performance analysis tools have different applications; for example, many tools can be used to monitor Linux's operating status.

[0003] vmstat – virtual memory statistics;

[0004] iostat – used to report CPU statistics;

[0005] dstat – system monitoring tool;

[0006] iotop - a real-time monitoring tool for Linux processes;

[0007] pidstat - monitors system resources;

[0008] Performance profiling tools are currently widely used in the machine learning field to monitor memory usage, the percentage of time spent on each training component, and hardware information. Among these, PyTorch's performance profiling tool is particularly well-known. Due to the long processing time, complex logic flow, and diverse operator types of deep learning, performance analysis for deep learning requires more data dimensions than single-function performance analysis. To visualize this data, the TensorFlow team launched Tensorboard, a visualization tool for machine learning performance analysis. It has become the de facto standard for visualization tools in the machine learning field. PyTorch, another major deep learning framework, has also developed and integrated its own performance profiling library, libkineto. Kineto, the core of PyTorch's performance profiling tool, is directly integrated into the core PyTorch code and implemented in C++, minimizing the impact of the performance profiling process on the underlying workflow. Based on Kineto, PyTorch has further encapsulated it to provide the torch.profiler class for user convenience. The PyTorch Profiler allows users to specify the type of performance analysis (GPU or CPU) to invoke the corresponding performance analysis code. Furthermore, PyTorch provides a very flexible performance analysis strategy, allowing users to customize the performance analysis workflow at the granularity of iteration steps. After the analysis is completed, the PyTorch Profiler can output the results to TensorBoard for visualization by implementing the PyTorch Tensorboard plugin.

[0009] While mainstream performance analysis tools can accurately record the memory usage of each operation at the current time, they lack sufficient information about the source of the data currently in memory. For machine learning, in-memory data can be broadly categorized as model data and non-model data. Current PyTorch performance analysis tools fail to distinguish these two types of data, hindering developers of machine learning systems from timely tracking memory fragmentation and memory deallocation issues that may arise during development. Furthermore, with the introduction of ZeRO and other memory optimization strategies, Tensors (tensors: high-dimensional arrays, the fundamental data structure for storing neural network parameters) are frequently transferred between main memory and graphics memory, often maintaining different states. For example, in PatrickStar, a Tensor has five different states corresponding to different operations. In this case, simply tracking memory usage alone cannot provide users with a clear understanding of the current model's memory allocation. Therefore, a memory analysis tool that can track these different states is needed.

[0010] Although the PytTorch performance analysis tool can analyze data movement performance, its performance analysis of data movement in different memory spaces has a coarse granularity problem. When analyzing the performance of data movement (communication) between different nodes (the basic unit of computing and communication in machine learning, for example, we call a host a node), its results are biased.

[0011] For performance analysis of data movement in different memory spaces, such as Figure 2 In the PyTorch performance analysis tool, only the execution time ( Figure 1 The PyTorch performance analysis tool lacks bandwidth analysis, which makes it difficult for us to objectively evaluate program performance. In addition, due to differences in physical hardware, data is transmitted between GPUs (graphics processing units, graphics cards) using NVLINK (a bus and its communication protocol developed and launched by NVIDIA), while data is transmitted between the CPU and GPU using the PCIE (PCI-Express (peripheral component interconnect express) is a high-speed serial computer expansion bus standard) bus. The transmission bandwidth of the former is much higher than that of the latter. The PyTorch performance analysis tool does not distinguish between the two, so we cannot accurately evaluate the performance of the program.

[0012] For data movement (communication) performance analysis between different nodes, such as Figure 3 As shown, it lacks performance analysis for bandwidth. Furthermore, its data volume metric is incorrectly calculated, and using it to calculate bandwidth can produce overstated results, potentially misleading our evaluation of program performance. Based on this, the present invention designs a fine-grained, multi-dimensional, large-scale machine learning performance analysis method to address the aforementioned issues. Summary of the Invention

[0013] The purpose of the present invention is to provide a fine-grained, multi-dimensional, large-scale machine learning performance analysis method to solve the problems raised in the above background technology.

[0014] To achieve the above objectives, the present invention provides the following technical solution: a fine-grained, multi-dimensional, large-scale machine learning performance analysis method, implemented based on a memory usage performance analysis tool. The memory usage performance analysis tool includes the Colossal-AI Profiler performance analysis tool and the Colossal-AI Gemini memory policy management tool. The process used by the Colossal-AI Profiler is similar to that of the PyTorch Profiler. The Colossal-AI Profiler is designed as a context manager that can be quickly opened and closed using the "with" statement, including the following steps:

[0015] S1: Import the corresponding package and then use the with statement to initialize the entire performance analysis process. During the initialization process, you can customize the options for performing this performance analysis.

[0016] S2: If you choose to track Tensors in different states, StatefulTensorMemoryProfilerExtention will be enabled;

[0017] S3: Colossal-AI.gemini.GeminiMemoryManager is the manager of all token vectors, and the memory state tracker uses it to collect data;

[0018] S4: After collecting the data, use the memory usage performance analysis tool to output the collected data in JSON format to the specified path.

[0019] Preferably, the options for customizing the performance analysis during the initialization process in step S1 include: a. the device, memory, and video memory to be analyzed; b. the strategy adopted for this performance analysis. Users can customize the performance analysis strategy to eliminate the impact of the performance analysis process or model initialization events on memory usage during training; c. the output of auxiliary information. Users can choose to output the shape of the Tensor, the source code location where the corresponding operation occurs, and the FLOPs information obtained by the test; d. whether to track Tensors in different states. If this option is enabled, the Colossal-AI.engine.Engine instance needs to be passed in, and the various states of the Tensor need to be tracked through the hook mechanism.

[0020] Preferably, the main process of the extension work in step S2 is: 1: During the initialization process, a new memory state tracker instance and an engine memory state tracking instance are created; 2: When performance analysis is turned on, a hook is added to the engine and activated; 3: This hook uses the memory state tracker for sampling during each operation, including forward propagation or backward propagation.

[0021] Preferably, the global singleton of Colossal-AI.gemini.GeminiMemoryManager in Colossal-AI in step S3 is GST_MGR, which can manage all labeled Tensors; wherein, the Tensor states are divided into five types, namely FREE, HOLD, HOLD_AFTER_FWD, HOLD_AFTER_FWD and COMPUTE, and the Tensor switches between different states according to certain rules; the memory state tracker directly calls GST_MGR to obtain the number of Tensors currently in each state.

[0022] Compared with the existing technology, the beneficial effects of the present invention are: the present invention can accurately record the information of Tensors in different states, so as to obtain the specific situation of current memory usage, help developers better optimize memory usage and scheduling strategies in the machine learning process, help developers compare the bandwidth performance of the program with the theoretical upper limit of data movement bandwidth, conduct objective and accurate performance analysis, obtain the current program's resource utilization, and find the main optimization direction of the program, which is creative. BRIEF DESCRIPTION OF THE DRAWINGS

[0023] In order to more clearly illustrate the technical solutions of the embodiments of the present invention, the following briefly introduces the drawings required for describing the embodiments. Obviously, the drawings described below are only some embodiments of the present invention. For ordinary technicians in this field, other drawings can be obtained based on these drawings without creative work.

[0024] Figure 1 Memory information graph collected by the existing technology PyTorchProfiler;

[0025] Figure 2 A schematic diagram summarizing the execution time of existing technologies;

[0026] Figure 3 It is a statistical diagram of communication operations in the prior art;

[0027] Figure 4 This is a sampling diagram of the present invention;

[0028] Figure 5 This is a schematic diagram of GeminiMemoryManager of the present invention;

[0029] Figure 6 This is a schematic diagram of the Tensor state transition of the present invention. DETAILED DESCRIPTION

[0030] The following will clearly and completely describe the technical solutions in the embodiments of the present invention in conjunction with the accompanying drawings. Obviously, the described embodiments are only part of the embodiments of the present invention, not all of the embodiments. Based on the embodiments of the present invention, all other embodiments obtained by ordinary technicians in this field without making any creative efforts shall fall within the scope of protection of the present invention.

[0031] In the description of the present invention, it should be understood that the terms "upper", "lower", "front", "back", "left", "right", "top", "bottom", "inside", "outside", etc., indicating directions or positional relationships, are based on the directions or positional relationships shown in the accompanying drawings, and are only for the convenience of describing the present invention and simplifying the description, rather than indicating or implying that the device or element referred to must have a specific direction, be constructed and operated in a specific direction, and therefore should not be understood as limiting the present invention.

[0032] The present invention provides a technical solution: a fine-grained, multi-dimensional, large-scale machine learning performance analysis method, implemented based on a memory usage performance analysis tool. The memory usage performance analysis tool includes the Colossal-AI Profiler performance analysis tool and the Colossal-AI Gemini memory policy management tool. The process used by the Colossal-AI Profiler is similar to that of the PyTorch Profiler. The Colossal-AI Profiler is designed as a context manager that can be quickly opened and closed using the "with" statement, including the following steps:

[0033] S1: Import the corresponding package and then use the with statement to initialize the entire performance analysis process. During the initialization process, you can customize the options for performing this performance analysis.

[0034] S2: If you choose to track Tensors in different states, StatefulTensorMemoryProfilerExtention will be enabled;

[0035] S3: Colossal-AI.gemini.GeminiMemoryManager is the manager of all token vectors, and the memory state tracker uses it to collect data;

[0036] S4: After collecting the data, use the memory usage performance analysis tool to output the collected data in JSON format to the specified path.

[0037] Among them, the options for customizing the performance analysis during the initialization process in step S1 include: a. The device, memory, and video memory to be analyzed; b. The strategy adopted for this performance analysis. Users can customize the performance analysis strategy to eliminate the impact of the performance analysis process or model initialization events on memory usage during training; c. The output of auxiliary information. Users can choose to output the shape of the Tensor, the source code location where the corresponding operation occurs, and the FLOPs information obtained by the test; d. Whether to track Tensors in different states. If this option is enabled, you need to pass in the Colossal-AI.engine.Engine instance and use the hook mechanism to track the various states of the Tensor.

[0038] The main process of the extension work in step S2 is as follows: 1: During the initialization process, create a new memory state tracker instance and an engine memory state tracker instance; 2: When performance analysis is turned on, add a hook to the engine and activate this hook; 3: This hook uses the memory state tracker to sample each time it performs an operation, including forward propagation or backward propagation, such as Figure 4 shown.

[0039] In step S3, the global singleton of Colossal-AI.gemini.GeminiMemoryManager in Colossal-AI is GST_MGR, which can manage all labeled Tensors. The states of Tensors are divided into five categories, such as Figure 5 As shown, they are FREE, HOLD, HOLD_AFTER_FWD, HOLD_AFTER_FWD and COMPUTE, and they are tracked separately. Tensor switches between different states according to certain rules, such as Figure 6 As shown; the memory state tracker directly calls GST_MGR to obtain the number of Tensors currently in each state. When the performance analysis tool is finished running, we will read the saved data (json file) for a more fine-grained analysis. For the performance analysis of memory usage, we use the collected OP-level model / non-model memory data to draw a graph of the changes in various types of memory usage over time. For the performance analysis of data movement, we can divide data movement in machine learning into two categories: 1: Data movement within a node; data movement in different memory spaces within a node; 2: Data movement between nodes: data movement completed through network communication.

[0040] Data movement performance analysis within a node

[0041] For memory movement performance analysis, based on the PyTorch performance analysis tool, we divide the data movement within the node into three categories according to the type of kernel operation:

[0042] 1. From Host to Device

[0043] 2. From Device to Host

[0044] 3. From device to device, where Host: CPU and its memory; Device: GPU and its memory;

[0045] Existing PyTorch performance analysis tools do not distinguish between these three types of memory movement behaviors. However, because these three types of memory movement behaviors physically use different hardware, their transmission bandwidth limits vary. Data is transmitted between GPUs using NVLINK, while data is transmitted between the CPU and GPU using the PCIe bus. The transmission bandwidth of the former is much higher than that of the latter. Without this distinction, it is difficult to determine the bandwidth utilization of memory movement behaviors. Therefore, we distinguish these three types of memory movement behaviors based on the type of kernel operation. The transmission bandwidth of each type of behavior is calculated to facilitate comparison with the theoretical transmission bandwidth limit.

[0046] Data movement performance analysis between nodes

[0047] Data movement between nodes is completed through network communication, so when we analyze the performance of data movement between nodes, we are actually analyzing the performance of network communication.

[0048] Existing PyTorch performance analysis tools exhibit biased results when analyzing communication performance. They generally assume that the total amount of data is the communication volume of all communication operations. However, different types of communication operations often employ different communication modes, so their communication volumes cannot be generalized. To address this issue, we categorize communication operations by type and apply a specific bias correction to the communication volume of each type of operation, thereby obtaining more realistic communication volume and bandwidth data.

[0049] When using the fine-grained, multi-dimensional, large-scale machine learning performance analysis method of the present invention, taking NCCL (NVIDIA Collective Communication Library) as an example, the deviation correction of the communication volume of various communication operations is described, where N is the number of nodes and D is the amount of data each node participates in the communication:

[0050] 1.All-gather:

[0051] 2.All-reduce:

[0052] 3. All-to-all: (N-1)D

[0053] 4.Broadcast: D

[0054] 5.Gather:(N-1)D

[0055] 6.Reduce: D

[0056] 7.Reduce-scatter: (N-1)D

[0057] 8.Scatter:(N-1)D

[0058] Only after the above formula is corrected can the actual communication volume be obtained, thereby calculating the actual bandwidth performance, facilitating comparison with the theoretical bandwidth upper limit and calculating the bandwidth utilization.

[0059] Throughout this specification, references to terms such as "one embodiment," "example," or "specific example" indicate that the specific features, structures, materials, or characteristics described in conjunction with that embodiment or example are included in at least one embodiment or example of the present invention. In this specification, schematic representations of these terms do not necessarily refer to the same embodiment or example. Furthermore, the specific features, structures, materials, or characteristics described may be combined in any suitable manner in any one or more embodiments or examples.

[0060] The preferred embodiments of the present invention disclosed above are intended only to help illustrate the present invention. These preferred embodiments do not exhaustively describe all details, nor do they limit the present invention to the specific embodiments described. Obviously, many modifications and variations are possible based on the content of this specification. These embodiments are selected and described in detail in this specification to better explain the principles and practical applications of the present invention, thereby enabling those skilled in the art to better understand and utilize the present invention. The present invention is limited only by the claims and their full scope and equivalents.

Claims

1. A fine-grained, multi-dimensional, large-scale machine learning performance analysis method based on memory usage performance analysis tools. The memory usage performance analysis tools include the Colossal-AIProfiler performance analysis tool and the Colossal-AIGemini memory policy management tool. The process used by Colossal-AIProfiler is similar to that of PyTorch Profiler. Colossal-AIProfiler is designed as a context manager that can be quickly opened and closed using the "with" statement. It is characterized by: The steps include: S1: Import the corresponding package and then use the with statement to initialize the entire performance analysis process. During the initialization process, you can customize the options for performing this performance analysis. S2: If you choose to track Tensors in different states, StatefulTensorMemoryProfilerExtention will be enabled; S3: Colossal-AI.gemini.GeminiMemoryManager is the manager of all token vectors, and the memory state tracker uses it to collect data; S4: After collecting the data, use the memory usage performance analysis tool to output the collected data in JSON format to the specified path.

2. The fine-grained, multi-dimensional, large-scale machine learning performance analysis method according to claim 1, characterized in that: The options for customizing the performance analysis during the initialization process in step S1 include: a. the device, memory, and video memory to be analyzed; b. the strategy used in the performance analysis. Users can customize the performance analysis strategy to eliminate the impact of the performance analysis process or model initialization events on memory usage during training; c. Output of auxiliary information. Users can choose to output the shape of the Tensor, the source code location where the corresponding operation occurs, and the FLOPs information obtained by the test. d. Whether to track Tensors in different states. If this option is enabled, you need to pass in a Colossal-AI.engine.Engine instance and use the hook mechanism to track the various states of the Tensor.

3. The fine-grained, multi-dimensional, large-scale machine learning performance analysis method according to claim 1, characterized in that: The main process of the extension work in step S2 is as follows: 1: During the initialization process, create a new memory state tracker instance and an engine memory state tracking instance; 2: After turning on performance analysis, add a hook to the engine and activate it; 3: This hook samples the memory state tracker at each operation, including forward propagation or backward propagation.

4. The fine-grained, multi-dimensional, large-scale machine learning performance analysis method according to claim 1, characterized in that: In step S3, the global singleton of Colossal-AI.gemini.GeminiMemoryManager in Colossal-AI is GST_MGR, which can manage all labeled Tensors. The Tensor states are divided into five types: FREE, HOLD, HOLD_AFTER_FWD, HOLD_AFTER_FWD, and COMPUTE. Tensors transition between different states according to certain rules. The memory state tracker directly calls GST_MGR to obtain the number of Tensors currently in each state.

Images