Non-volatile memory based deep neural network checkpoint optimization system and method

Abstract

The application provides a non-volatile memory-based deep neural network checkpoint optimization system and method. Before the start of deep neural network training, the corresponding network structure is registered in the non-volatile memory through a client module and a server module, and a data index and a remote direct memory access (RDMA)-based data communication protocol are created. During the neural network training process, the application provides zero-copy, asynchronous, end-to-end neural network data persistence, so that users can do fine-grained checkpointing without affecting the training speed to ensure fault tolerance and data persistence.

Description

Technical Field

[0001] This invention relates to the fields of storage systems and deep learning systems in computer science, and particularly to a deep neural network checkpoint optimization system and method based on non-volatile memory. Background Technology

[0002] As the complexity of neural network model structures and training cluster systems increases, the persistence of training models and data fault tolerance have become important issues in large-scale training systems.

[0003] Existing neural network checkpointing involves multiple hardware and software layers, including GPUs, networks, and storage. The complexity of the software layers prevents high-performance hardware devices from being fully utilized, severely impacting read and write performance during data persistence. Furthermore, the performance limitations of existing solutions result in a lack of fine-grained checkpointing mechanisms during neural network training, leading to low fault tolerance and a lack of rapid data recovery mechanisms. In addition, in the distributed, multi-user training process of modern neural networks, frequent hardware and system failures necessitate high-performance data persistence mechanisms, but such systems are currently lacking. Summary of the Invention

[0004] In view of the shortcomings of the prior art described above, the purpose of this invention is to provide a deep neural network checkpoint optimization system and method based on non-volatile memory to solve the problems of the prior art.

[0005] To achieve the above and other related objectives, this invention provides a deep neural network checkpoint optimization system based on non-volatile memory. The system includes a client module located in a computing node equipped with GPU memory and a server module located in a storage node equipped with non-volatile memory. Before the training of each neural network model begins, the client module sends the network structure obtained by initializing the corresponding neural network model stored in the GPU memory to the server module to construct an index structure of the neural network model in the non-volatile memory, thereby establishing end-to-end communication between the GPU memory and the non-volatile memory. When the server module receives a checkpoint request from the client module during the training of the corresponding neural network model, it directly reads the corresponding model data from the GPU memory into the non-volatile memory based on the index structure of the neural network model.

[0006] In one embodiment of the present invention, the method by which the client module initializes the corresponding neural network model stored in the GPU memory to obtain the network structure includes: collecting GPU memory pointers pointing to each layer of the corresponding neural network model through the neural network framework; registering the GPU address space of each layer of the model as an RDMA memory region using the NVIDIA Peer Memory kernel module based on the GPU memory pointers of each layer of the model, and assigning a unique identifier to each memory region; and aggregating each identifier and the metadata of each layer of the model into a model structure package.

[0007] In one embodiment of the present invention, the client module constructs the index structure of the neural network model in the non-volatile memory in the following manner: after receiving the model structure package, a thread is selected from the thread pool to construct the index structure corresponding to the neural network model in the non-volatile memory based on the model structure package, so as to map each layer of the neural network model to the checkpoint structure one by one.

[0008] In one embodiment of the present invention, the index structure is a three-level index structure, including: a model table at the first level, model metadata at the second level, and model data information at the third level.

[0009] In one embodiment of the present invention, when the server module receives a checkpoint request from the client module during the training of the corresponding neural network model, the method of directly reading the corresponding model data from the GPU memory to non-volatile memory based on the index structure of the neural network model includes: when the client module receives a user checkpoint request during the training of the corresponding neural network model, it obtains the corresponding GPU memory pointer and sends the generated checkpoint request to the server module; based on the checkpoint request from the client module, the server module controls the corresponding thread to directly read the corresponding model data from the GPU memory to non-volatile memory through RDMA read operation based on the corresponding constructed index structure.

[0010] In one embodiment of the present invention, when the server module receives a data recovery request from the client module, it actively writes the corresponding model data from the non-volatile memory to the GPU video memory based on the index structure of the neural network model.

[0011] In one embodiment of the present invention, the client module communicates with the server module via the IPoIB protocol or the TCP protocol.

[0012] In one embodiment of the present invention, the neural network framework is the PyTorch software library.

[0013] To achieve the above and other related objectives, this invention provides a deep neural network checkpoint optimization method based on non-volatile memory, used in a deep neural network checkpoint optimization system based on non-volatile memory. The method includes: a client module located within a computing node equipped with GPU memory and a server module located within a storage node equipped with non-volatile memory. The method comprises: before the start of training for each neural network model, the client module sends the network structure obtained by initializing the corresponding neural network model stored in the GPU memory to the server module to construct an index structure of the neural network model in the non-volatile memory, thereby establishing end-to-end communication between the GPU memory and the non-volatile memory; when the server module receives a checkpoint request from the client module during the training of the corresponding neural network model, it directly reads the corresponding model data from the GPU memory into the non-volatile memory based on the index structure of the neural network model.

[0014] In one embodiment of the present invention, the method further includes: when the server module receives a data recovery request from the client module, actively writing the corresponding model data from the non-volatile memory to the GPU memory based on the index structure of the neural network model.

[0015] As described above, this invention is a deep neural network checkpoint optimization system and method based on non-volatile memory, which has the following beneficial effects: Before the deep neural network training begins, this invention registers the corresponding network structure in non-volatile memory through client and server modules, and creates a data index and a data communication protocol based on Remote Direct Memory Access (RDMA); furthermore, during neural network training, this invention provides zero-copy, asynchronous, end-to-end neural network data persistence, allowing users to perform fine-grained checkpointing without affecting training speed to ensure fault tolerance and data persistence. Attached Figure Description

[0016] Figure 1 The diagram shown is a schematic representation of a deep neural network checkpoint optimization system based on non-volatile memory, according to an embodiment of the present invention.

[0017] Figure 2 The diagram shown is a schematic representation of a deep neural network checkpoint optimization system based on non-volatile memory, according to an embodiment of the present invention.

[0018] Figure 3 The diagram shown is a schematic representation of a deep neural network checkpoint optimization system based on non-volatile memory, according to an embodiment of the present invention.

[0019] Figure 4 The diagram shown is a flowchart of a deep neural network checkpoint optimization method based on non-volatile memory according to an embodiment of the present invention.

[0020] Figure 5 The diagram shown is a schematic representation of the data read / write process in one embodiment of the present invention. Detailed Implementation

[0021] The following specific examples illustrate the implementation of the present invention. Those skilled in the art can easily understand other advantages and effects of the present invention from the content disclosed in this specification. The present invention can also be implemented or applied through other different specific embodiments, and various details in this specification can also be modified or changed based on different viewpoints and applications without departing from the spirit of the present invention. It should be noted that, unless otherwise specified, the following embodiments and features described therein can be combined with each other.

[0022] It should be noted that in the following description, reference is made to the accompanying drawings, which illustrate several embodiments of the invention. It should be understood that other embodiments may also be used, and changes in mechanical composition, structure, electrical system, and operation may be made without departing from the spirit and scope of the invention. The following detailed description should not be considered limiting, and the scope of the embodiments of the invention is defined only by the claims of the published patents. The terminology used herein is for the purpose of describing particular embodiments only and is not intended to limit the invention. Spatially related terms, such as “upper,” “lower,” “left,” “right,” “below,” “below,” “lower part,” “above,” “upper part,” etc., may be used herein to illustrate the relationship between one element or feature shown in the figures and another element or feature.

[0023] Throughout this specification, when it is said that a part is "connected" to another part, this includes not only "direct connection" but also "indirect connection" by placing other elements in between. Furthermore, when it is said that a part "includes" a certain constituent element, unless otherwise stated otherwise, this does not exclude other constituent elements, but rather means that other constituent elements may also be included.

[0024] The terms "first," "second," and "third," etc., used herein are for the purpose of describing various parts, components, regions, layers, and / or segments, but are not limiting. These terms are used only to distinguish one part, component, region, layer, or segment from others. Therefore, the "first part," "component," "region," "layer," or "segment" described below may refer to a "second part," "component," "region," "layer," or "segment" without departing from the scope of this invention.

[0025] Furthermore, as used herein, the singular forms “a,” “an,” and “the” are intended to include the plural forms as well, unless the context indicates otherwise. It should be further understood that the terms “comprising,” “including,” indicate the presence of the stated feature, operation, element, component, item, kind, and / or group, but do not preclude the presence, occurrence, or addition of one or more other features, operations, elements, components, items, kinds, and / or groups. The terms “or” and “and / or” as used herein are interpreted as inclusive, or mean any one or any combination thereof. Thus, “A, B, or C” or “A, B, and / or C” means “any one of: A; B; C; A and B; A and C; B and C; A, B, and C.” Exceptions to this definition arise only when combinations of elements, functions, or operations are inherently mutually exclusive in some manner.

[0026] This invention discloses a deep neural network checkpoint optimization system based on non-volatile memory. Before the deep neural network training begins, the corresponding network structure is registered in non-volatile memory through client and server modules, and a data index and a data communication protocol based on Remote Direct Memory Access (RDMA) are created. During the neural network training process, this invention provides zero-copy, asynchronous, end-to-end neural network data persistence, allowing users to perform fine-grained checkpointing without affecting training speed to ensure fault tolerance and data persistence.

[0027] The present invention will now be described in detail with reference to the accompanying drawings, so that those skilled in the art can readily implement it. The present invention can be embodied in many different forms and is not limited to the embodiments described herein.

[0028] like Figure 1 This illustration shows a schematic diagram of a deep neural network checkpoint optimization system based on non-volatile memory, according to an embodiment of the present invention.

[0029] The system includes:

[0030] The client module 101 is located within the computing node equipped with GPU memory 2. It serves as an extension of the neural network framework, is called by the user's neural network training program, and is connected to GPU memory 2. GPU memory 2 stores the corresponding neural network model data.

[0031] The server module 102 is located in the storage node equipped with non-volatile memory 3. It runs as an independent application on the storage node equipped with non-volatile memory and provides data indexing, communication and non-volatile memory management services.

[0032] The client module 101 is communicatively connected to the server module 102;

[0033] Before the training of each neural network model begins, the client module 101 sends the network structure obtained by initializing the corresponding neural network model stored in the GPU memory 2 to the client module 102 to construct the index structure of the neural network model on the non-volatile memory, so as to establish end-to-end communication between the GPU memory 2 and the non-volatile memory 3.

[0034] When the server module 101 receives a checkpoint request related to the neural network model from the client module 102 during the training of the corresponding neural network model, it directly reads the corresponding model data (checkpoint file) from the GPU memory 2 into the non-volatile memory 3 based on the index structure of the neural network model.

[0035] Therefore, during the checkpoint writing process, this invention uses a data transmission flow that is the opposite of the traditional data flow to reduce user-kernel mode switching overhead. The checkpoint file is not written by the client module 101, but rather the server module 102 actively reads the checkpoint file from GPU memory 2 into non-volatile memory 3.

[0036] In another embodiment, during data recovery, the server module 102 also actively writes the data into the client GPU memory 2 in an end-to-end manner; that is, when the server module 102 receives a data recovery request from the client module 101, it actively writes the corresponding model data from the non-volatile memory 3 into the GPU memory 2 based on the index structure of the neural network model.

[0037] Combining the above-mentioned checkpointing and data recovery methods, this data flow separates client-side training from server-side data persistence, achieving an asynchronous checkpointing mechanism that is imperceptible to the client.

[0038] In one embodiment, the client module 101 communicates with the server module 102 via the TCP protocol.

[0039] In one embodiment, the neural network framework is the PyTorch software library; the client module 101 is located in the PyTorch software library.

[0040] In one embodiment, the neural network framework is connected to the GPU memory via a PCIe bus.

[0041] In one embodiment, the client module 101 initializes the corresponding neural network model stored in the GPU memory 2 to obtain the network structure in the following ways:

[0042] Collect GPU memory pointers pointing to each layer of the corresponding neural network model using a neural network framework;

[0043] The NVIDIA Peer Memory kernel module is used to register the GPU address space of each layer of the model as an RDMA memory region based on the GPU memory pointer of each layer of the model, and to assign a unique identifier to each memory region.

[0044] Each identifier is aggregated into a model structure package, corresponding one-to-one with the metadata of each layer of the model, to characterize the structure of this neural network model; after initialization, the client module 101 sends the structure package to the server module 102.

[0045] In one embodiment, such as Figure 2 As shown, the way the client module 102 constructs the index structure of the neural network model on the non-volatile memory 3 includes: after receiving the model structure package from the client module 101, selecting an available thread 401 from the thread pool 4 and assigning the package to the thread, the thread constructs an index structure corresponding to the neural network model in the non-volatile memory 3 based on the model structure package, so as to map each layer of the neural network model to the checkpoint structure one by one.

[0046] In one embodiment, the index structure is a three-level index structure, including: a model table at the first level, model metadata at the second level, and model data information at the third level.

[0047] Specifically, the thread maintains a three-level index structure in non-volatile memory based on the information in the structure package: model table - model metadata - model data. The way this three-level index structure is constructed includes: the thread first registers the model name in the model table, and then creates a new model metadata block in non-volatile memory to record the model structure, so that each neural network layer can be mapped to the checkpoint structure. Then, the thread allocates a non-volatile memory region for each neural network layer. Finally, the server module establishes end-to-end communication between the GPU memory and the non-volatile memory.

[0048] In one embodiment, when the server module 102 receives a checkpoint request from the client module during the training of the corresponding neural network model, the method of directly reading the corresponding model data from the GPU memory 2 to the non-volatile memory 3 based on the index structure of the neural network model includes:

[0049] When the client module 101 receives a user checkpoint request during the training of the corresponding neural network model, it obtains the corresponding GPU memory pointer and sends the generated checkpoint request to the server module 102.

[0050] Based on the checkpoint request from the client module, the server module 102 controls the thread corresponding to the model to obtain the RDMA memory area of ​​the model data in the CPU video memory and read the model data into the storage area of ​​non-volatile memory based on the corresponding constructed index structure. Then, through the RDMA read operation, the corresponding model data is directly read from the GPU video memory 2 into the non-volatile memory 3.

[0051] To better illustrate the above-mentioned deep neural network checkpoint optimization system based on non-volatile memory, the present invention provides the following specific embodiments.

[0052] Example 1: A deep neural network checkpoint optimization system based on non-volatile memory. For example... Figure 3 The diagram shown is a schematic of the structure of the deep neural network checkpoint optimization system based on non-volatile memory in this embodiment.

[0053] The system includes:

[0054] The computing node includes: CPU video memory connected via a PCIe bus and main memory containing the PyTorch software library; the client module described in this application is located in the PyTorch software library of the computing node;

[0055] The storage node includes: locally connected persistent memory and a memory containing the server module described above;

[0056] At the start of a new neural network model training task, the client module first collects GPU memory pointers to each layer in the neural network model using PyTorch. Then, the client module uses the NVIDIA Peer Memory kernel module to register the GPU address space of each layer as an RDMA memory region, assigning each memory region a unique identifier. After registration, the client module aggregates these identifiers and the metadata (including size, data type, etc.) of each layer into a model structure packet to characterize the neural network model structure. Finally, the client module sends this structure packet to the server module via a TCP connection. Upon receiving the structure packet, the server module selects an available thread from the thread pool using the allocation mechanism and assigns the packet to that thread. This thread maintains a three-level index structure in non-volatile memory based on the information in the structure packet: model table - model metadata - model data. The thread first registers the model name in the model table and creates a new model metadata block in non-volatile memory to record the model structure, mapping each neural network layer to a checkpoint structure. Then, the thread allocates a non-volatile memory region for each neural network layer. Finally, the server module establishes end-to-end communication between the client's GPU memory and the server's non-volatile memory.

[0057] The client module described in this application can intercept the user's Checkpoint request, obtain the model data pointer in the GPU memory, and interact with the server module described in this application via TCP and RDMA connections. The server module described in this application is located in the storage node. The server module processes the client module's requests, performs data management and indexing on persistent memory, and writes the model data to persistent memory via RDMA.

[0058] In this embodiment, the client is a plugin or add-on module of the PyTorch deep learning library, while the server is implemented as a standalone executable file. Combining the client and server achieves the goal of improving Checkpoint performance without affecting the user's computation process.

[0059] Similar to the principles of the above embodiments, the present invention provides a deep neural network checkpoint optimization method based on non-volatile memory.

[0060] The following specific embodiments are provided in conjunction with the accompanying drawings:

[0061] like Figure 4 This diagram illustrates a flow chart of a deep neural network checkpoint optimization method based on non-volatile memory, as described in an embodiment of the present invention.

[0062] The system is applied to a deep neural network checkpoint optimization system based on non-volatile memory, comprising: a client module located in a computing node equipped with GPU memory and a server module located in a storage node equipped with non-volatile memory; the system can implement the functions of the deep neural network checkpoint optimization system based on non-volatile memory in the above embodiments.

[0063] The method includes:

[0064] Step S1: Before the training of each neural network model begins, the client module sends the network structure obtained by initializing the corresponding neural network model stored in the GPU memory to the client module to build the index structure of the neural network model on the non-volatile memory, so as to establish end-to-end communication between the GPU memory and the non-volatile memory.

[0065] In one embodiment, the client module initializes the corresponding neural network model stored in the GPU memory to obtain the network structure by: collecting GPU memory pointers pointing to each layer of the corresponding neural network model through the neural network framework; using the NVIDIA Peer Memory kernel module to register the GPU address space of each layer of the model as an RDMA memory region based on the GPU memory pointers of each layer of the model, and assigning a unique identifier to each memory region; and aggregating each identifier and the metadata of each layer of the model into a model structure package.

[0066] In one embodiment, the client module constructs the index structure of the neural network model in the non-volatile memory by: upon receiving the model structure package, selecting a thread from the thread pool to construct the index structure corresponding to the neural network model in the non-volatile memory based on the model structure package, so as to map each layer of the neural network model to the checkpoint structure one by one.

[0067] In one embodiment, the index structure is a three-level index structure, including: a model table at the first level, model metadata at the second level, and model data information at the third level.

[0068] In one embodiment, the client module communicates with the server module via the IPoIB protocol or the TCP protocol.

[0069] In one embodiment, the neural network framework is the PyTorch software library.

[0070] Step S2: When the server module receives a checkpoint request from the client module during the training of the corresponding neural network model, it directly reads the corresponding model data from the GPU memory to the non-volatile memory based on the index structure of the neural network model.

[0071] In one embodiment, step S2 includes: when the client module receives a user checkpoint request during the training of the corresponding neural network model, it obtains the corresponding GPU memory pointer and sends the generated checkpoint request to the server module; based on the checkpoint request from the client module, the server module controls the corresponding thread to directly read the corresponding model data from the GPU memory to non-volatile memory through RDMA read operation based on the corresponding constructed index structure.

[0072] In one embodiment, the method further includes: when the server module receives a data recovery request from the client module, it actively writes the corresponding model data from the non-volatile memory to the GPU memory based on the index structure of the neural network model.

[0073] Since the implementation principle of the deep neural network checkpoint optimization method based on non-volatile memory has been described in the previous embodiments, it will not be repeated here.

[0074] To better illustrate the above-mentioned deep neural network checkpoint optimization method based on non-volatile memory, the present invention provides the following specific embodiments.

[0075] Example 1: A method for writing checkpoints to a deep neural network based on non-volatile memory. Figure 5 This is a schematic diagram of the data read / write process in this embodiment.

[0076] The method includes:

[0077] When performing a data checkpoint operation:

[0078] 1. When the client initializes the model, it sends the model structure to the server.

[0079] 2. The server constructs an index structure in non-volatile memory using the received model structure, and sends a confirmation message to the client after successful construction.

[0080] 3. After the client receives the confirmation message, the model initialization ends and the calculation process begins.

[0081] 4. When the computation begins, the client sends a Checkpoint request to the server. The server then uses RDMA read operations to copy the required model data from the client's GPU memory to the server's non-volatile memory according to the index structure.

[0082] 5. The client has already completed the RDMA operation before updating the weights, so the client can perform the data calculation process without interruption.

[0083] It should be noted that during model training, a checkpoint write is required after each round of forward propagation, backward propagation, and weight update.

[0084] This embodiment provides a pure user-space, non-volatile memory-based fast checkpoint data path and management method through the aforementioned modules. Significant performance improvements are demonstrated in both benchmark tests and real-world application scenarios. For neural network models used in practical applications such as computer vision and natural language processing, this embodiment improves data persistence performance by up to 9.2 times and data recovery performance by 7.0 times compared to existing checkpointing mechanisms. Furthermore, this embodiment implements a zero-copy, end-to-end data transmission mechanism and an asynchronous data transmission mechanism with zero client wait time.

[0085] In summary, the deep neural network checkpoint optimization system and method based on non-volatile memory of this invention registers the corresponding network structure in non-volatile memory before the deep neural network training begins through client and server modules, and creates data indexes and a data communication protocol based on Remote Direct Memory Access (RDMA). Furthermore, during neural network training, this invention provides zero-copy, asynchronous, end-to-end neural network data persistence, allowing users to perform fine-grained checkpointing without affecting training speed to ensure fault tolerance and data persistence. Therefore, this invention effectively overcomes the various shortcomings of existing technologies and has high industrial application value.

[0086] The above embodiments are merely illustrative of the principles and effects of the present invention and are not intended to limit the invention. Any person skilled in the art can modify or alter the above embodiments without departing from the spirit and scope of the present invention. Therefore, all equivalent modifications or alterations made by those skilled in the art without departing from the spirit and technical concept disclosed in the present invention should still be covered by the claims of the present invention.

Claims

1. A deep neural network checkpoint optimization system based on non-volatile memory, characterized in that, The system includes: Client modules located in compute nodes equipped with GPU memory and server modules located in storage nodes equipped with non-volatile memory; Before training each neural network model begins, the client module initializes the network structure of the corresponding neural network model stored in the GPU memory and sends it to the client module to build the index structure of the neural network model on the non-volatile memory, so as to establish end-to-end communication between the GPU memory and the non-volatile memory. When the server module receives a checkpoint request from the client module during the training of the corresponding neural network model, it directly reads the corresponding model data from the GPU memory into non-volatile memory based on the index structure of the neural network model. Furthermore, the client module initializes the corresponding neural network model stored in the GPU memory to obtain the network structure in the following ways: Collect GPU memory pointers pointing to each layer of the corresponding neural network model using a neural network framework; The NVIDIA Peer Memory kernel module is used to register the GPU address space of each layer of the model as an RDMA memory region based on the GPU memory pointer of each layer of the model, and to assign a unique identifier to each memory region. Each identifier is aggregated into a model structure package, corresponding one-to-one with the metadata of each layer of the model. The client module constructs the index structure of the neural network model on the non-volatile memory in the following ways: Upon receiving the model structure package, a thread is selected from the thread pool to construct an index structure corresponding to the neural network model in non-volatile memory based on the model structure package, so as to map each layer of the neural network model to the checkpoint structure. When the server module receives a checkpoint request from the client module during the training of the corresponding neural network model, the method of directly reading the corresponding model data from the GPU memory to non-volatile memory based on the index structure of the neural network model includes: When the client module receives a user checkpoint request during the training of the corresponding neural network model, it obtains the corresponding GPU memory pointer and sends the generated checkpoint request to the server module. Based on the checkpoint request from the client module, the server module controls the corresponding thread to directly read the corresponding model data from the GPU memory to non-volatile memory through RDMA read operations based on the constructed index structure.

2. The deep neural network checkpoint optimization system based on non-volatile memory as described in claim 1, characterized in that, The index structure is a three-level index structure, including: the model table at the first level, the model metadata at the second level, and the model data information at the third level.

3. The deep neural network checkpoint optimization system based on non-volatile memory as described in claim 1, characterized in that, When the server module receives a data recovery request from the client module, it actively writes the corresponding model data from the non-volatile memory to the GPU memory based on the index structure of the neural network model.

4. The deep neural network checkpoint optimization system based on non-volatile memory as described in claim 1, characterized in that, The client module communicates with the server module via the TCP protocol.

5. The deep neural network checkpoint optimization system based on non-volatile memory as described in claim 1, characterized in that, The neural network framework is the PyTorch software library.

6. A checkpoint optimization method based on non-volatile memory in deep neural networks, characterized in that, A checkpoint optimization system for deep neural networks based on non-volatile memory includes: a client module located in a computing node equipped with GPU memory and a server module located in a storage node equipped with non-volatile memory. The method includes: Before the training of each neural network model begins, the client module will send the network structure obtained by initializing the corresponding neural network model stored in the GPU memory to the client module to build the index structure of the neural network model on the non-volatile memory, so as to establish end-to-end communication between the GPU memory and the non-volatile memory. When the server module receives a checkpoint request from the client module during the training of the corresponding neural network model, it directly reads the corresponding model data from the GPU memory into non-volatile memory based on the index structure of the neural network model. Furthermore, the client module initializes the corresponding neural network model stored in the GPU memory to obtain the network structure in the following ways: Collect GPU memory pointers pointing to each layer of the corresponding neural network model using a neural network framework; The NVIDIA Peer Memory kernel module is used to register the GPU address space of each layer of the model as an RDMA memory region based on the GPU memory pointer of each layer of the model, and to assign a unique identifier to each memory region. Each identifier is aggregated into a model structure package, corresponding one-to-one with the metadata of each layer of the model. The client module constructs the index structure of the neural network model on the non-volatile memory in the following ways: Upon receiving the model structure package, a thread is selected from the thread pool to construct an index structure corresponding to the neural network model in non-volatile memory based on the model structure package, so as to map each layer of the neural network model to the checkpoint structure. When the server module receives a checkpoint request from the client module during the training of the corresponding neural network model, the method of directly reading the corresponding model data from the GPU memory to non-volatile memory based on the index structure of the neural network model includes: When the client module receives a user checkpoint request during the training of the corresponding neural network model, it obtains the corresponding GPU memory pointer and sends the generated checkpoint request to the server module. Based on the checkpoint request from the client module, the server module controls the corresponding thread to directly read the corresponding model data from the GPU memory to non-volatile memory through RDMA read operations based on the constructed index structure.

7. The deep neural network checkpoint optimization method based on non-volatile memory as described in claim 6, characterized in that, The method further includes: When the server module receives a data recovery request from the client module, it actively writes the corresponding model data from the non-volatile memory to the GPU memory based on the index structure of the neural network model.

Images