Memory allocation method and device

Abstract

The invention provides a memory allocation method and device.The method comprises the steps that the maximum lexical element handling capacity sent by a host end is received, the maximum lexical element handling capacity is determined based on the length of an input sequence and the number of allocated experts, and the number of the allocated experts is the number of expert models allocated to all lexical elements; applying for a memory matched with the maximum lexical element handling capacity at a device end; dividing a target memory corresponding to each expert model from a memory based on the lexical element quantity processed by each expert model; and allocating each target memory to the corresponding expert model, wherein each target memory is used for storing the lexical elements processed by the corresponding expert model. According to the method, the maximum lexical element processing amount does not need to be synchronized at the host end, so that delay introduced by synchronous operation from the equipment end to the host end can be avoided, communication overhead in the memory allocation process can be reduced, the memory allocation efficiency and the overall processing speed can be improved, and finally the model reasoning performance of the equipment end is improved.

Description

Technical Field

[0001] The present invention relates to the field of artificial intelligence chip technology, and in particular to a memory allocation method and device. Background Art

[0002] In deep learning models, memory allocation refers to the process of allocating storage space on computing devices (e.g., GPUs) for intermediate variables (e.g., tensors) generated during model computation. A reasonable memory allocation strategy is crucial for ensuring efficient model training and inference, especially when working with large-scale models and data.

[0003] Currently, it's typically necessary to first determine the number of tokens processed by each expert model, copy the number of tokens processed by each expert model from the device to the host, and then allocate memory on the host based on the number of tokens processed by each expert model. However, the synchronous operation of copying the number of tokens processed by each expert model from the device to the host introduces significant latency, especially in large-scale models. This latency accumulates as the number of model layers increases, ultimately severely impacting overall computing performance. Summary of the Invention

[0004] The present invention provides a memory allocation method and device to solve the defects existing in the related art.

[0005] The present invention provides a memory allocation method, which is applied to a device and includes the following steps: receiving a maximum word-unit processing capacity sent by the host, wherein the maximum word-unit processing capacity is determined based on the length of the input sequence and the number of assigned experts, where the number of assigned experts is the number of expert models assigned to each word-unit; Applying for memory matching the maximum word unit processing capacity; Based on the amount of tokens processed by each expert model, dividing the target memory corresponding to each expert model from the memory; Each target memory is allocated to a corresponding expert model, and each target memory is used to store word units processed by the corresponding expert model.

[0006] According to a memory allocation method provided by the present invention, the step of determining the maximum word unit processing amount includes: Determining an initial word unit processing amount based on the input sequence length and the number of assigned models; Based on an alignment parameter, an alignment operation is performed on the initial word unit processing amount to obtain the maximum word unit processing amount, and the alignment parameter is determined based on the design structure of the expert model.

[0007] According to a memory allocation method provided by the present invention, performing an alignment operation on the initial word unit processing amount based on an alignment parameter to obtain the maximum word unit processing amount includes: Based on the alignment parameters and the total number of experts, an alignment operation is performed on the initial word-gram processing amount to obtain the maximum word-gram processing amount, where the total number of experts refers to the total number of expert models involved in word-gram processing.

[0008] According to a memory allocation method provided by the present invention, the target memory corresponding to each expert model is divided from the memory based on the amount of tokens processed by each expert model, including: Based on the amount of tokens processed by each expert model and the alignment parameter, a target memory corresponding to each expert model is divided from the memory, and the alignment parameter is determined based on the design structure of the expert model.

[0009] According to a memory allocation method provided by the present invention, the target memory corresponding to each expert model is divided from the memory based on the word unit quantity processed by each expert model and the alignment parameter, including: Based on the amount of tokens processed by each expert model, dividing the initial memory corresponding to each expert model from the memory; Based on the alignment parameters, an offset operation is performed on the initial memory corresponding to each expert model to obtain a target memory corresponding to each expert model.

[0010] According to a memory allocation method provided by the present invention, the initial memory corresponding to each expert model is offset based on the alignment parameter to obtain the target memory corresponding to each expert model, including: Determining a memory offset of the current expert model based on the alignment parameter and an initial memory of the current expert model; The end address of the target memory of the previous expert model is used as the starting address of the initial memory of the current expert model, and based on the memory offset of the current expert model, the initial memory of the current expert model is offset to obtain the target memory of the current expert model.

[0011] The present invention also provides a memory allocation device, which is provided at a device end and includes the following modules: a receiving unit, configured to receive a maximum word unit processing capacity sent by a host end, wherein the maximum word unit processing capacity is determined based on an input sequence length and a number of assigned experts, wherein the number of assigned experts is the number of expert models assigned to each word unit; An application unit, configured to apply for memory matching the maximum word unit processing capacity; a partitioning unit, configured to partition target memories corresponding to the expert models from the memory based on the amount of tokens processed by the expert models; An allocation unit for allocating each target memory to a corresponding expert model, where each target memory is used to store the tokens processed by the corresponding expert model.

[0012] The present invention also provides an electronic device, including a memory, a processor, and a computer program stored on the memory and executable on the processor. When the processor executes the program, the memory allocation method described in any one of the above is implemented.

[0013] The present invention also provides a non-transitory computer-readable storage medium, on which a computer program is stored. When the computer program is executed by a processor, the memory allocation method described in any one of the above is implemented.

[0014] The present invention also provides a computer program product, including a computer program. When the computer program is executed by a processor, the memory allocation method described in any one of the above is implemented.

[0015] The memory allocation method and device provided by the present invention enable the device side to receive the maximum token processing amount sent by the host side. Then, before each expert model processes the corresponding tokens, the corresponding target memory can be allocated for each expert model on the memory of the device side in advance, without sending the above maximum token processing amount to the host side before each expert model processes the corresponding tokens, that is, without synchronizing the maximum token processing amount on the host side. Thus, the delay introduced by the synchronization operation from the device side to the host side can be avoided, and further the communication overhead during the memory allocation process can be reduced, the memory allocation efficiency and the overall processing speed can be improved, and finally the model inference performance of the device side can be enhanced. BRIEF DESCRIPTION OF THE DRAWINGS

[0016] In order to more clearly illustrate the technical solutions in the present invention or related technologies, the following will briefly introduce the drawings required for use in the embodiments or related technology descriptions. Obviously, the drawings in the following description are some embodiments of the present invention. For those of ordinary skill in the art, other drawings can be obtained based on these drawings without creative efforts.

[0017] Figure 1 It is a flowchart of the memory allocation method provided by the present invention.

[0018] Figure 2 It is a schematic diagram of the target memory of each expert model provided by the present invention.

[0019] Figure 3 It is a schematic diagram of the structure of the memory allocation device provided by the present invention.

[0020] Figure 4 It is a schematic diagram of the structure of the electronic device provided by the present invention. DETAILED DESCRIPTION OF THE EMBODIMENTS

[0021] To make the objectives, technical solutions, and advantages of the present invention clearer, the following will clearly and completely describe the technical solutions in the present invention in conjunction with the accompanying drawings in the present invention. Apparently, the described embodiments are some, but not all, of the embodiments of the present invention. All other embodiments obtained by those of ordinary skill in the art based on the embodiments in the present invention without creative efforts fall within the scope of protection of the present invention.

[0022] Currently, in the Mixture of Experts (MoE) model, the Permutation operator is responsible for distributing the tokens in the input tensor to different expert models according to the output of the Router. The specific process is as follows: First, according to the expert selection results (hit expert indices) output by the Router, a new index tensor indices_trans with a shape of [expert_num, seq_len] is generated. This index tensor records the absolute addresses (i.e., indices) of each token in the original input tensor in the corresponding expert models. At the same time, the number of tokens selected for each expert model is counted on the device side (the device side usually corresponds to hardware such as GPUs, GPGPUs, TPUs, etc.).

[0023] Next, the number of tokens selected for each expert model is synchronized to the host side (the host side usually corresponds to hardware such as CPUs) to allocate corresponding memory spaces for each expert model on the host side, and then the memory spaces allocated for each expert model are synchronized to the device side, so that the device side can move the tokens responsible for each expert model to the corresponding memory spaces, completing the distribution of tokens to the expert models.

[0024] However, a major bottleneck of the above solution is that the number of tokens that each expert model needs to process must be determined before moving the tokens responsible for each expert model, and these numbers of tokens need to be copied from the device side to the host side. Then, the host side allocates memory for each expert model according to these numbers of tokens. This synchronization operation from the device side to the host side introduces significant latency. Especially in large-scale models, due to the increase in the number of model layers, this latency will accumulate linearly, significantly affecting the overall computing performance.

[0025] In response to this, the present invention provides a memory allocation method. The execution entity of the present invention is the device side (device end), and the device side can be an artificial intelligence chip such as a GPU, GPGPU, TPU, etc. In the present invention, the device side receives the maximum number of token processing amounts sent by the host side (host end), and executes the memory method in the present invention. Among them, the host side can be a CPU. In the present invention, the host side and the device side can perform data transmission and communication through a bus such as PCIe, so that the host side can send the maximum number of token processing amounts to the device side through the bus. The heterogeneous architecture of the host side and the device side is a mode of division of labor and cooperation. The host side is responsible for control and management, and the device side is responsible for computationally intensive tasks. Connected by a high-speed bus, they jointly complete complex tasks, so as to make full use of the advantages of different hardware and improve the performance of the overall system.

[0026] Among them, Figure 1 is a schematic flowchart of the memory allocation method provided by the present invention. As Figure 1 shown, this method is applied to the device side (device end), and includes step 110, step 120, step 130, and step 140.

[0027] Step 210: Receive the maximum number of token processing amounts sent by the host side. The maximum number of token processing amounts is determined based on the input sequence length and the number of allocation experts. The number of allocation experts is the number of expert models allocated to each token.

[0028] Here, the input sequence length (seq_len) refers to the number of tokens contained in the text sequence of the input model. A token is the basic unit after the text is tokenized. For example, in natural language processing, a sentence can be split into multiple tokens. Among them, the input sequence length determines the total number of tokens that the model needs to process.

[0029] In the model architecture, the expert model refers to a sub-model used to process specific tokens. The number of allocation experts (topk) refers to the number of expert models allocated to each token. Different tokens may be allocated to different numbers of expert models for processing. Among them, the number of allocation experts can be determined based on the complexity of the token, the importance of the task, or the load balancing requirements of the model, or can be set according to actual needs. For example, more expert models are allocated to tokens of key information to improve processing accuracy, and fewer expert models are allocated to tokens of ordinary information to save resources.

[0030] In addition, the maximum number of tokens processed refers to the maximum number of tokens that the model can process under the current input sequence length and the current number of assigned experts. Among them, the input sequence length directly determines the total number of tokens to be processed. If the input sequence length is long, even if the number of assigned experts for each token is small, the total number of tokens processed will be large. In addition, the more the number of assigned experts for each token, even if the input sequence length is short, the total number of tokens to be processed will also be large, because each token will be processed by multiple expert models, increasing the complexity and resource requirements of the processing.

[0031] That is to say, both the input sequence length and the number of assigned experts will affect the maximum number of tokens processed, that is, the two jointly determine the total number of tokens that the model needs to process. Optionally, the maximum number of tokens processed = input sequence length × number of assigned experts.

[0032] For example, the input text is: "The weather is really nice today". After word segmentation, the input sequence obtained is ["today", "weather", "really", "nice"], and the input sequence includes 4 tokens, that is, the input sequence length is 4. If each token is assigned to 2 expert models for processing, that is, the number of assigned experts is 2, then the maximum number of tokens processed = input sequence length × number of assigned experts = 4 × 2 = 8.

[0033] It should be noted that the host side determines the maximum number of tokens processed in advance before the expert model performs model operations, and copies the maximum number of tokens processed to the device side in advance, so as to avoid the delay introduced by the synchronization operation from the device side to the host side in the traditional method.

[0034] In addition, the device side does not need to recalculate the maximum number of tokens processed, and can directly use the maximum number of tokens processed for subsequent calculations. For example, calculations such as Feed-Forward Network (FFN) calculations and Unpermutation calculations can directly use the maximum number of tokens processed for calculations without having to recalculate the maximum number of tokens processed, saving computing resources.

[0035] Step 120: Apply for memory matching the maximum number of tokens processed.

[0036] Specifically, the device side can be understood as the hardware device where the model actually runs and performs calculations, such as a GPU or a specific accelerator. After receiving the maximum number of tokens to be processed, the size of the memory space required for the model to run can be determined based on the maximum number of tokens to be processed. Furthermore, according to the required memory space size, memory matching the maximum number of tokens to be processed can be applied for on the device side. This memory can be understood as the storage space for storing intermediate results, model parameters, and final outputs during the model calculation process. Among them, this memory can be Usharp video memory, which refers to a memory area on the GPU for storing data required during the operation of the UdonSharp script.

[0037] As an alternative embodiment, after determining the maximum number of tokens to be processed, corresponding memory can be directly applied for on the device side according to the size of the maximum number of tokens to be processed. For example, a memory application function can be directly called according to the maximum number of tokens to be processed to allocate a memory space of the corresponding size.

[0038] Step 130: Based on the number of tokens processed by each expert model, divide the target memory corresponding to each expert model from the memory.

[0039] Specifically, the number of tokens processed by each expert model refers to the number of tokens allocated to each expert model for processing. For example, for the input sequence "The cat sat on the mat", the tokens processed by expert model 1 include "The", "cat", "on", "mat", and the tokens processed by expert model 2 include "sat", "the". That is, the number of tokens processed by expert model 1 is 4, and the number of tokens processed by expert model 2 is 2.

[0040] Since different tokens are allocated to different numbers of expert models, the number of tokens processed by each expert model may be different. For example, for an expert model that processes complex or critical information, the corresponding number of tokens processed is relatively large; for an expert model that processes simple or non-critical information, the corresponding number of tokens processed is relatively small.

[0041] Among them, the tokens processed by each expert model can be determined based on the following steps: According to the expert selection result output by the Router, determine the expert model to which each token is allocated, and count the number of tokens allocated to each expert model, which is the number of tokens processed by the expert model.

[0042] Since this memory matches the maximum number of tokens to be processed, this memory is sufficient to accommodate the space required for all expert models to process tokens. In addition, the number of tokens processed by different expert models is different, so the storage space required during the calculation of the expert models is different, and thus the target memory required by different expert models is different. This target memory can be understood as the dedicated memory required for an expert model to process the allocated tokens. For an expert model that processes a larger number of tokens, the memory space it requires is larger; for an expert model that processes a smaller number of tokens, the memory space it requires is smaller.

[0043] On this basis, in the embodiments of the present invention, based on the fact that the number of tokens processed by each expert model is different, the target memory of each expert model is partitioned from the memory. For example, if expert model 1 needs to process a larger number of tokens and requires a larger memory space, a larger continuous memory area can be partitioned from the memory as the target memory of expert model 1. If expert model 2 needs to process a smaller number of tokens and requires a smaller memory space, a smaller continuous memory area can be partitioned from the memory as the target memory of expert model 2.

[0044] As an alternative embodiment, when partitioning the target memory of each expert model from the memory, the memory can be partitioned for each expert model in ascending order of the expert model number. The expert model number can be determined based on a preset value during model design, or based on the importance or computational complexity of the expert model (for example, an expert model with a high computational complexity is assigned a smaller number, so that the target memory can be preferentially allocated to this expert model to ensure that it has sufficient memory resources for calculation). The embodiments of the present invention do not make specific limitations on this. For example, since the number of expert model 1 is the smallest and a certain number of tokens have been allocated to it, a continuous memory segment starting from the starting address in the memory can be partitioned to expert model 1; since the number of expert model 2 is the second smallest and a certain number of tokens have been allocated to it, the end address of the target memory corresponding to expert model 1 can be used as the starting address of the target memory of expert model 2, and based on the number of tokens processed by expert model 2, the required space of expert model 2 is determined, and the end address of the target memory of expert model 2 is determined according to this space.

[0045] Step 140: Allocate each target memory to the corresponding expert model, and each target memory is used to store the tokens processed by the corresponding expert model.

[0046] Specifically, after determining the target memory for each expert model, allocate each target memory to the corresponding expert model, and store the tokens processed by each expert model in the corresponding target memory. Further, after storing the tokens processed by each expert model in the corresponding target memory, each expert model performs computational processing on the processed tokens, and the computational results can also be stored in the target memory for subsequent calculations or outputs. Optionally, methods such as pointer passing or memory mapping can be used to allocate each target memory to the corresponding expert model.

[0047] In the memory allocation method provided by the embodiments of the present invention, the device side receives the maximum token processing amount sent by the host side. Then, before each expert model processes the corresponding tokens, the corresponding target memory can be allocated for each expert model on the memory of the device side in advance, without sending the above maximum token processing amount to the host side before each expert model processes the corresponding tokens, that is, without synchronizing the maximum token processing amount on the host side. Thus, the delay introduced by the synchronization operation from the device side to the host side can be avoided, and further the communication overhead during the memory allocation process can be reduced, the memory allocation efficiency and the overall processing speed can be improved, and finally the model inference performance of the device side can be enhanced.

[0048] Among them, the above expert model can be applied to text processing. At this time, the input sequence can be any parameter used, input, or generated in tasks such as text classification, sentiment analysis, and text generation, such as the semantic feature vector of the text. After performing memory allocation based on the method of the above embodiments, each expert model performs computational processing on the processed tokens, and the computational results can also be stored in the target memory. Finally, the computational results of all expert models are fused, and inference is performed based on the computational results to obtain the corresponding output text.

[0049] Based on the above embodiments, determining the maximum token processing amount based on the input sequence length and the number of allocated models includes: Determining the initial token processing amount based on the input sequence length and the number of allocated models; Performing an alignment operation on the initial token processing amount based on the alignment parameter to obtain the maximum token processing amount, where the alignment parameter is determined based on the design structure of the expert model.

[0050] Specifically, the initial token processing amount is determined based on the input sequence length and the number of allocated models, without considering the actual design structure of the expert model. However, the design structure of the expert model usually includes one or more FFN layers, and the FFN layer consists of two linear transformations and an activation function. If the alignment operation is not performed on the initial token processing amount, it may cause the dimensions of the input and output not to match during the matrix operation of the FFN (Feed-Forward Network) layer, thereby causing computational errors or performance degradation.

[0051] The alignment operation refers to adjusting the initial token processing amount to an integer multiple of the alignment parameter. The alignment parameter is determined based on the design structure of the expert model. For the expert model, the alignment parameter can be the hidden layer dimension size of the FFN layer in the expert model, or a value that can divide the hidden layer dimension size evenly.

[0052] For example, assume that the hidden layer dimension size of the FFN layer is 256. The alignment parameter can be 256 or an integer multiple of 256. If the initial token processing amount is 127, after the alignment operation, the maximum token processing amount can be adjusted to 256 (or its multiple, such as 512) to ensure the correctness of the FFN layer calculation. If not aligned, the input data of the FFN layer may not be loaded effectively, resulting in the input data not being correctly fed into the FFN layer for calculation.

[0053] In this regard, in the embodiments of the present invention, after determining the initial token processing amount, an alignment operation is performed on the initial token processing amount based on the alignment parameter to ensure that the obtained maximum token processing amount matches the design structure of the expert model, avoiding the problem of dimension mismatch during the FFN layer calculation.

[0054] Based on any of the above embodiments, an alignment operation is performed on the initial token processing amount based on the alignment parameter to obtain the maximum token processing amount, including: An alignment operation is performed on the initial token processing amount based on the alignment parameter and the total number of experts. The total number of experts refers to the total number of expert models participating in token processing.

[0055] Specifically, the total number of experts refers to the total number of expert models participating in token processing. Considering that there is an FFN layer in the expert model, the calculation of the FFN layer needs to ensure that the input dimension matches the hidden layer dimension (i.e., the input dimension is the same as the hidden layer dimension, or the input dimension is an integer multiple of the hidden layer dimension). Among them, the input of the FFN layer can be understood as the input of the expert model (i.e., the tokens processed by the expert model), that is, the input dimension of the FFN layer is the number of tokens processed by the expert model.

[0056] If only an alignment operation is performed on the initial token processing amount based on the alignment parameter, it may lead to the problem of dimension mismatch during FFN calculation due to the change in the number of expert models, thereby affecting the overall performance of the model. For example, assume that the initial token processing amount does not consider the number of expert models. When the number of expert models increases, the input dimension of the FFN may not be aligned with the hidden layer dimension (such as the number of tokens is not an integer multiple of the FFN hidden layer dimension), resulting in calculation errors.

[0057] To this end, an embodiment of the present invention introduces the total number of experts, which is used to ensure dimensional matching of FFN calculations. The alignment parameter is usually the FFN hidden layer dimension or an integer multiple thereof, which is used to adjust the initial word unit processing capacity to a dimension that is compatible with the FFN structure. When the initial word unit processing capacity is aligned with the total number of experts, the maximum word unit processing capacity can be dynamically adjusted according to the total number of experts, avoiding the dimension mismatch problem during FFN calculations, thereby ensuring the normal operation of the FFN structure in the expert model and improving the stability and performance of the overall model.

[0058] Optionally, the maximum word processing capacity can be determined using the following formula: Maximum word processing capacity = initial word processing capacity + alignment parameter × total number of experts Furthermore, considering that if the initial word unit processing capacity is an integer multiple of the alignment parameter, if the maximum word unit processing capacity is determined based on the above formula, there may be over-alignment, which may lead to a waste of resources.

[0059] For example, assuming the hidden layer dimension of the FFN is 64 (i.e., the alignment parameter is 64), the calculated initial word processing capacity is also 64, indicating that the initial word processing capacity is aligned with the hidden layer dimension. If there are three expert models, that is, the total number of experts is 3, when using the formula "maximum word processing capacity = initial word processing capacity + alignment parameter × total number of experts", the maximum word processing capacity is 64 + 64 × 3 = 256.

[0060] Although in this case the maximum word unit processing capacity is four times the hidden layer dimension, that is, the dimensions are aligned, in fact, there is a large surplus in the calculation of the maximum word unit processing capacity, which exceeds the actual needs of the model. As a result, the memory requested on the device side based on the maximum word unit processing capacity may also be redundant and wasted.

[0061] To address this issue, embodiments of the present invention preferably use the formula "maximum word processing capacity = initial word processing capacity + alignment parameter × (total number of experts - 1)" to calculate the maximum word processing capacity: maximum word processing capacity = 64 + 64 × (3 - 1) = 192. In this case, the maximum word processing capacity is three times the hidden layer dimension. While the dimensions are aligned, the calculated maximum word processing capacity is smaller, reducing unnecessary resource usage and enabling more efficient device memory allocation, thus avoiding memory waste.

[0062] Based on any of the above embodiments, the target memory corresponding to each expert model is divided from the memory based on the amount of tokens processed by each expert model, including: Based on the number of tokens processed by each expert model and the alignment parameters, the target memory corresponding to each expert model is divided from the memory.

[0063] Specifically, considering that the number of tokens processed by each expert model may not be an integer multiple of the hidden layer dimension, in this case, the input dimension of each expert model is not aligned with the hidden layer dimension, and it may not be possible to effectively load the input of each expert model during calculation, which may lead to errors in FFN calculation.

[0064] In this case, it is necessary to align the input dimension of each expert model with the hidden layer dimension, that is, adjust the input dimension of each expert model to an integer multiple of the hidden layer dimension. Since the size of the input dimension of each expert model is the same as the size of the number of tokens processed by each expert model, adjusting the size of the number of tokens processed by each expert model can achieve the adjustment of the input dimension of each expert model. On this basis, based on the alignment parameter, a zero-padding operation is performed on the number of tokens processed by each expert model (such as adding a number of zero-valued tokens at the end of the expert model token sequence so that the total number of tokens reaches an integer multiple of the alignment parameter), so that the size of the number of tokens processed by the expert model after the zero-padding operation is aligned with the hidden layer dimension, thereby achieving the alignment of the input dimension and the hidden layer dimension of each expert model.

[0065] Next, based on the size of the number of tokens processed by each expert model after the zero-padding operation, the target memory corresponding to each expert model is partitioned from the memory, so that the partitioned target memory is sufficient to store the tokens processed by each expert model after the zero-padding operation.

[0066] For example, the number of tokens processed by expert model 1 is 1254, and the alignment parameter is 256. Obviously, 1254 is not an integer multiple of 256. At this time, a zero-padding operation can be performed on the tokens processed by expert model 1, that is, 26 zero-valued tokens are added after 1254 tokens, so that the total number of tokens reaches 1280. Thus, the number of tokens processed by expert model 1 after the zero-padding operation is 1280, and 1280 is an integer multiple of 256, meeting the alignment requirement. Then, the target memory corresponding to expert model 1 can be partitioned from the memory based on "1280".

[0067] It can be seen that based on the number of tokens processed by each expert model and the alignment parameter, the target memory corresponding to each expert model is partitioned from the memory in the embodiments of the present invention, which can achieve the alignment of the input dimension and the hidden layer dimension, avoid calculation errors caused by dimension mismatch, ensure the normal operation of the model, and improve the calculation efficiency and model performance.

[0068] Based on any of the above embodiments, partitioning the target memory corresponding to each expert model from the memory based on the number of tokens processed by each expert model and the alignment parameter includes: Based on the number of tokens processed by each expert model, partitioning the initial memory corresponding to each expert model from the memory; Based on the alignment parameter, performing an offset operation on the initial memory corresponding to each expert model to obtain the target memory corresponding to each expert model.

[0069] Specifically, the initial memory refers to the memory space allocated according to the number of tokens actually processed by the expert model without alignment adjustment. For example, if the number of tokens processed by the expert model is 1254, the initial memory corresponding to the expert model is divided from the memory based on "1254", and this initial memory can store 1254 tokens.

[0070] Furthermore, considering that the input dimension of each expert model needs to be aligned with the hidden layer dimension, that is, when the input dimension of the expert model is not an integer multiple of the hidden layer dimension, it is necessary to perform a zero-padding operation on the tokens processed by the expert model based on the alignment parameter (such as padding 0 in the middle), which means that the memory space required by the expert model is increased, that is, the initial memory is not enough to store the tokens processed by the expert model after the zero-padding operation.

[0071] On this basis, based on the alignment parameter, an offset operation is performed on the initial memory corresponding to each expert model, that is, the additional memory space required is calculated according to the alignment parameter, and an offset is made on the basis of the initial memory, so as to expand the memory space, so that the target memory corresponding to each expert model obtained can store the token data after the zero-padding operation and alignment.

[0072] For example, if the number of tokens processed by the expert model is 1254, after calculation, 26 zero-valued tokens need to be added after the 1254 tokens to make the total number of tokens reach 1280, which is aligned with the hidden layer dimension (256). The initial memory can store 1254 tokens, but cannot store 1280 tokens. At this time, an offset operation needs to be performed on the initial memory, that is, on the basis of the initial memory address, the memory space occupied by 26 tokens is offset backward, so as to expand the memory space to a size that can accommodate 1280 tokens, so as to provide enough storage space for the token data after the zero-padding operation and ensure that the calculation of the FFN layer can be carried out correctly.

[0073] Based on any of the above embodiments, an offset operation is performed on the initial memory corresponding to each expert model based on the alignment parameter to obtain the target memory corresponding to each expert model, including: Determine the memory offset of the current expert model based on the alignment parameter and the initial memory of the current expert model; Use the end address of the target memory of the previous expert model as the start address of the initial memory of the current expert model, and based on the memory offset of the current expert model, perform an offset operation on the initial memory of the current expert model to obtain the target memory of the current expert model.

[0074] Specifically, the memory offset refers to the additional memory space that the expert model needs to allocate to accommodate the zero-padded tokens mentioned above. The current expert model refers to the expert model that is currently performing memory allocation and offset operations, and the previous expert model refers to the expert model that completed memory allocation before the current expert model. Among them, the starting address of the target memory of the first expert model is the starting address of the memory corresponding to the maximum token processing capacity.

[0075] Figure 2 is a schematic diagram of the target memory of each expert model provided by the present invention. As Figure 2 shown, the number of tokens processed by expert model 1 is 1254, the initial memory of expert model 1 is the memory corresponding to ①, the alignment parameter is 256. Since 1254 is not an integer multiple of 256, the initial memory at ① needs to be offset backward by the memory space occupied by 26 tokens (i.e., the memory corresponding to ②, which is also the memory offset of expert model 1). The offset memory space can be filled with zeros for out-of-bounds (OOB) to expand the memory space to a size that can accommodate 1280 tokens. That is, the memory composed of ① and ② is the target memory. Similarly, the number of tokens processed by expert model 2 is 2307, the initial memory of expert model 2 is the memory corresponding to ③, the alignment parameter is 256. Since 2307 is not an integer multiple of 256, the initial memory at ③ needs to be offset backward by the memory space occupied by 253 tokens (i.e., the memory corresponding to ④, which is also the memory offset of expert model 2), so as to expand the memory space to a size that can accommodate 2560 tokens. That is, the memory composed of ③ and ④ is the target memory. Among them, the starting address of the target memory of expert model 2 is the ending address of the target memory of expert model 1. The number of tokens processed by expert model 3 is 1468, the initial memory of expert model 3 is the memory corresponding to ⑤, the alignment parameter is 256. Since 1468 is not an integer multiple of 256, the initial memory at ⑤ needs to be offset backward by the memory space occupied by 68 tokens (i.e., the memory corresponding to ⑥, which is also the memory offset of expert model 3), so as to expand the memory space to a size that can accommodate 1536 tokens. That is, the memory composed of ⑤ and ⑥ is the target memory. Among them, the starting address of the target memory of expert model 3 is the ending address of the target memory of expert model 2.

[0076] Based on any of the above embodiments, after allocating each target memory to the corresponding expert model, the following steps are further included: Based on each expert model, use the tokens stored in each target memory for inference to obtain the local inference results output by each expert model; Fuse the local inference results output by each expert model to obtain the inference result of the input sequence.

[0077] Specifically, the local inference result refers to the output obtained by each expert model after reasoning based on the tokens in the target memory, which represents the local understanding of the input sequence by each expert model. The inference result of the input sequence refers to the final output obtained by the model after reasoning on the entire input sequence. It is the overall understanding of the input sequence by the model and is a comprehensive manifestation of the local understandings of all expert models.

[0078] As an alternative embodiment, a weight can be assigned to the local inference result of each expert model, and the local inference results of each expert model are added based on this weight to obtain the inference result of the input sequence. Among them, the weights of the local inference results corresponding to each expert model can be determined based on factors such as the performance of the expert model, the confidence of the expert model, and the importance of the tokens processed by the expert model. For example, the weight of an expert model can be determined according to its accuracy on the validation set. The higher the accuracy of the expert model, the higher the weight. Or, the confidence of an expert model can be determined according to the entropy value of the probability distribution output by the expert model. The lower the entropy value of the expert model, the higher the confidence and the higher the weight. Or, the importance of a token can be measured by its TF-IDF value in the sentence. The more important the token processed by an expert model, the higher the weight of this expert model.

[0079] The memory allocation device provided by the present invention will be described below. The memory allocation device described below can be correspondingly referred to the memory allocation method described above.

[0080] Based on any of the above embodiments, Figure 3 is a schematic structural diagram of the memory allocation device provided by the present invention. As shown in Figure 3 shown, this device is set at the device end and includes: A receiving unit 310, configured to receive the maximum token processing amount sent by the host end. The maximum token processing amount is determined based on the input sequence length and the number of allocated experts, and the number of allocated experts is the number of expert models allocated to each token; An application unit 320, configured to apply for memory matching the maximum token processing amount; A partitioning unit 330, configured to partition the target memory corresponding to each expert model from the memory based on the token amount processed by each expert model; An allocation unit 340, configured to allocate each target memory to the corresponding expert model, and each target memory is used to store the tokens processed by the corresponding expert model.

[0081] Based on any of the above embodiments, the determination steps of the maximum token processing amount include: Based on the input sequence length and the number of allocated models, determine the initial token processing amount; Based on the alignment parameter, perform an alignment operation on the initial token processing amount to obtain the maximum token processing amount, where the alignment parameter is determined based on the design structure of the expert model.

[0082] Based on any of the above embodiments, perform an alignment operation on the initial token processing amount based on the alignment parameter to obtain the maximum token processing amount, including: Based on the alignment parameter and the total number of experts, perform an alignment operation on the initial token processing amount to obtain the maximum token processing amount, where the total number of experts refers to the total number of expert models participating in token processing.

[0083] Based on any of the above embodiments, divide the target memory corresponding to each expert model from the memory based on the token amount processed by each expert model, including: Divide the target memory corresponding to each expert model from the memory based on the token amount processed by each expert model and the alignment parameter, where the alignment parameter is determined based on the design structure of the expert model.

[0084] Based on any of the above embodiments, divide the target memory corresponding to each expert model from the memory based on the token amount processed by each expert model and the alignment parameter, including: Divide the initial memory corresponding to each expert model from the memory based on the token amount processed by each expert model; Based on the alignment parameter, perform an offset operation on the initial memory corresponding to each expert model to obtain the target memory corresponding to each expert model.

[0085] Based on any of the above embodiments, perform an offset operation on the initial memory corresponding to each expert model based on the alignment parameter to obtain the target memory corresponding to each expert model, including: Determine the memory offset amount of the current expert model based on the alignment parameter and the initial memory of the current expert model; Use the end address of the target memory of the previous expert model as the start address of the initial memory of the current expert model, and perform an offset operation on the initial memory of the current expert model based on the memory offset amount of the current expert model to obtain the target memory of the current expert model.

[0086] Based on any of the above embodiments, allocate each target memory to the corresponding expert model, and then further include: Based on each expert model, perform inference using the tokens stored in each target memory to obtain the local inference results output by each expert model; Fuse the local inference results output by each expert model to obtain the inference result of the input sequence.

[0087] Figure 4 It is a schematic structural diagram of the electronic device provided by the present invention, as Figure 4As shown in the figure, the electronic device may include: a processor 410, a communications interface 420, a memory 430, and a communication bus 440. Among them, the processor 410, the communications interface 420, and the memory 430 communicate with each other through the communication bus 440. The processor 410 can call the logical instructions in the memory 430 to execute a memory allocation method, which includes: receiving the maximum token processing amount sent by the host side, where the maximum token processing amount is determined based on the input sequence length and the number of allocated experts, and the number of allocated experts is the number of expert models allocated to each token; applying for memory matching the maximum token processing amount; partitioning, from the memory, the target memory corresponding to each expert model based on the token amount processed by each expert model; and allocating each target memory to the corresponding expert model, where each target memory is used to store the tokens processed by the corresponding expert model.

[0088] In addition, when the logical instructions in the above-mentioned memory 430 are implemented in the form of software functional units and sold or used as independent products, they can be stored in a computer-readable storage medium. Based on such an understanding, the technical solution of the present invention, in essence, or the part that contributes to the related technology, or a part of this technical solution, can be embodied in the form of a software product. This computer software product is stored in a storage medium and includes several instructions for causing a computer device (which may be a personal computer, a server, or a network device, etc.) to execute all or part of the steps of the methods described in various embodiments of the present invention. The foregoing storage medium includes: various media such as USB flash drives, mobile hard disks, read-only memories (ROMs), random access memories (RAMs), magnetic disks, or optical discs that can store program codes.

[0089] On the other hand, the present invention also provides a computer program product. The computer program product includes a computer program that can be stored on a non-transitory computer-readable storage medium. When the computer program is executed by a processor, the computer can execute the memory allocation method provided by the above-mentioned various methods. The method includes: receiving the maximum token processing amount sent by the host side, where the maximum token processing amount is determined based on the input sequence length and the number of allocated experts, and the number of allocated experts is the number of expert models allocated to each token; applying for memory matching the maximum token processing amount; partitioning, from the memory, the target memory corresponding to each expert model based on the token amount processed by each expert model; and allocating each target memory to the corresponding expert model, where each target memory is used to store the tokens processed by the corresponding expert model.

[0090] In another aspect, the present invention also provides a non-transitory computer-readable storage medium, on which a computer program is stored. When the computer program is executed by a processor, it implements a memory allocation method provided by the above-mentioned various methods. The method includes: receiving a maximum token processing amount sent by a host side, where the maximum token processing amount is determined based on an input sequence length and the number of allocation experts, and the number of allocation experts is the number of expert models allocated to each token; applying for memory matching the maximum token processing amount; partitioning, based on the token amounts processed by each expert model, target memory corresponding to each expert model from the memory; and allocating each target memory to the corresponding expert model, where each target memory is used to store the tokens processed by the corresponding expert model.

[0091] The device embodiments described above are merely illustrative. The units described as separate components may or may not be physically separated, and the components shown as units may or may not be physical units, that is, they may be located in one place or distributed to multiple network units. Some or all of the modules can be selected according to actual needs to achieve the purpose of the solution of this embodiment. A person of ordinary skill in the art can understand and implement it without creative efforts.

[0092] Through the description of the above embodiments, those skilled in the art can clearly understand that each embodiment can be implemented by means of software plus a necessary general hardware platform, and of course, it can also be implemented by hardware. Based on this understanding, the essence of the above technical solution, or the part that contributes to the related technology, can be embodied in the form of a software product. The computer software product can be stored in a computer-readable storage medium, such as ROM / RAM, magnetic disk, optical disk, etc., and includes several instructions for causing a computer device (which can be a personal computer, a server, or a network device, etc.) to execute the methods described in each embodiment or some parts of the embodiments.

[0093] Finally, it should be noted that the above embodiments are only used to illustrate the technical solutions of the present invention, and are not intended to limit them; although the present invention has been described in detail with reference to the foregoing embodiments, those of ordinary skill in the art should understand that they can still modify the technical solutions described in the foregoing embodiments, or perform equivalent replacements for some of the technical features; and these modifications or replacements do not make the essence of the corresponding technical solutions deviate from the spirit and scope of the technical solutions of the embodiments of the present invention.

Claims

1. A memory allocation method, characterized in that, Applied to the device side, including: Receiving the maximum number of token processes sent by the host side, where the maximum number of token processes is determined based on the input sequence length and the number of assigned experts, and the number of assigned experts is the number of expert models assigned to each token; Applying for memory matching the maximum number of token processes; Based on the number of tokens processed by each expert model, partitioning the target memory corresponding to each expert model from the memory; Allocating each target memory to the corresponding expert model, and each target memory is used to store the tokens processed by the corresponding expert model.

2. The memory allocation method according to claim 1, wherein The steps for determining the maximum number of token processes include: Based on the input sequence length and the number of assigned models, determining the initial number of token processes; Based on the alignment parameter, performing an alignment operation on the initial number of token processes to obtain the maximum number of token processes, where the alignment parameter is determined based on the design structure of the expert model.

3. The memory allocation method according to claim 2, characterized in that, The performing an alignment operation on the initial number of token processes based on the alignment parameter to obtain the maximum number of token processes includes: Based on the alignment parameter and the total number of experts, performing an alignment operation on the initial number of token processes to obtain the maximum number of token processes, where the total number of experts refers to the total number of expert models participating in token processing.

4. The memory allocation method according to any one of claims 1 to 3, characterized in that, The partitioning the target memory corresponding to each expert model from the memory based on the number of tokens processed by each expert model includes: Based on the number of tokens processed by each expert model and the alignment parameter, partitioning the target memory corresponding to each expert model from the memory, where the alignment parameter is determined based on the design structure of the expert model.

5. The memory allocation method according to claim 4, wherein The partitioning the target memory corresponding to each expert model from the memory based on the number of tokens processed by each expert model and the alignment parameter includes: Based on the number of tokens processed by each expert model, partitioning the initial memory corresponding to each expert model from the memory; Based on the alignment parameter, performing an offset operation on the initial memory corresponding to each expert model to obtain the target memory corresponding to each expert model.

6. The memory allocation method according to claim 5, wherein, The performing an offset operation on the initial memory corresponding to each expert model based on the alignment parameter to obtain the target memory corresponding to each expert model includes: Based on the alignment parameter and the initial memory of the current expert model, determining the memory offset of the current expert model; Using the end address of the target memory of the previous expert model as the start address of the initial memory of the current expert model, and based on the memory offset of the current expert model, performing an offset operation on the initial memory of the current expert model to obtain the target memory of the current expert model.

7. A memory allocation device, characterized in that, Set on the device side, including: A receiving unit for receiving the maximum number of token processes sent by the host side, where the maximum number of token processes is determined based on the input sequence length and the number of assigned experts, and the number of assigned experts is the number of expert models assigned to each token; An application unit for applying for memory matching the maximum number of token processes; A partitioning unit for partitioning the target memory corresponding to each expert model from the memory based on the number of tokens processed by each expert model; An allocation unit for allocating each target memory to the corresponding expert model, and each target memory is used to store the tokens processed by the corresponding expert model.

8. An electronic device, comprising a memory, a processor, and a computer program stored on the memory and running on the processor, characterized in that, When the processor executes the computer program, the memory allocation method described in any one of claims 1 to 6 is implemented.

9. A non-transitory computer-readable storage medium having a computer program stored thereon, characterized in that, When the computer program is executed by the processor, the memory allocation method described in any one of claims 1 to 6 is implemented.

10. A computer program product, comprising a computer program, characterized in that, When the computer program is executed by the processor, the memory allocation method described in any one of claims 1 to 6 is implemented.

Images