Novel translation model reasoning method and system based on rwkv

Abstract

The RWKV-based novel translation model inference method and system includes the following steps: 1) Collecting novel translations and extracting parallel corpora, and using dynamic MicroBatch concatenation technology for sequence compression; 2) Introducing a lightweight group query attention mechanism on the basis of the RWKV architecture, directly obtaining KV information from the Embedding layer to build the model; 3) Employing a sublinear complexity hybrid parallel training mode, combined with a global scalar scaling FP16 mixed precision strategy for training; 4) Applying a hierarchical distributed heterogeneous architecture, offloading the optimizer to low-performance devices and performing gradient compression transmission; 5) Outputting the translation using a joint decoder and dynamic batch inference technology. This invention, through the above method and system, effectively reduces the computational complexity and memory usage in the long text translation process, improves model training efficiency and inference throughput, and significantly improves the translation efficiency and contextual coherence of ultra-long texts.

Description

Technical Field

[0001] This invention relates to the fields of artificial intelligence and natural language processing, specifically to a novel translation model reasoning method and system based on RWKV. Background Technology

[0002] China boasts a rich literary heritage and is currently the world's largest producer and consumer of online novels. However, the domestic online literature market is facing challenges such as peak growth and a slowdown in reader growth. The phenomenon of "a surge in works but sluggish reader growth" is hindering the industry's further development. Against this backdrop, actively expanding into overseas markets and promoting the internationalization of domestic online literature has become a realistic and urgent strategic choice. To bring domestic online literature to the world stage, language barriers between different languages ​​must be overcome. Traditional methods of exporting literature mainly rely on human translation, which incurs significant time and economic costs. Especially for online novels, which are produced at a rapid pace, are numerous, and cover a wide range of content, traditional human translation methods are insufficient to meet the requirements of real-time processing and cost control. Statistics show that AI-based machine translation applications can reduce translation costs by nearly 90% and increase speed by more than 100 times, making it an inevitable choice for the large-scale internationalization of online literature.

[0003] In the field of machine translation, Large Language Models (LLMs) have demonstrated unprecedented text understanding and generation capabilities, showcasing entirely new application prospects. Currently, the mainstream approach for LLMs is based on the Transformer architecture, whose core lies in capturing contextual information from the input sequence using a self-attention mechanism. While models based on a pure Transformer architecture perform exceptionally well in various scenarios, they suffer from significant technical bottlenecks: their self-attention mechanism has a quadratic time complexity (O(N)). 2 The time and space complexity of this model is significant. This means that when processing extremely long text sequences, the model incurs unacceptable computation time and memory consumption. Specifically, in novel translation tasks, because novels have very long contexts, and translating the current text often requires in-depth consideration of the logical role of historical texts, novel translation falls under the category of standard ultra-long temporal sequence translation tasks.

[0004] To circumvent the length limitations of the Transformer model, existing techniques often forcibly truncate long texts, breaking them into independent segments for batch translation. This approach forcibly disrupts the contextual relationships between texts, making it difficult for the model to make comprehensive judgments based on distant contexts. Consequently, the translations often fall short in terms of the handling of personal pronouns and the accuracy of referential references, resulting in unnatural contextual coherence. Although architectures such as Transformer-XL, which introduce recurrent memory to improve sequence length, have emerged in recent years, their computational overhead remains enormous.

[0005] On the other hand, linear attention models such as RWKV (Receptance Weighted Key Value) are gaining attention. RWKV combines the advantages of Transformer and Recurrent Neural Networks (RNNs), with its computational and memory complexity increasing linearly with sequence length, thus solving the efficiency problem of Transformer when processing long sequences. However, RWKV, based on the RNN architecture, exponentially decays historical state information in its design. While this reduces memory usage, it also results in a very high compression ratio for its states, making it more difficult to accurately reconstruct detailed historical text information compared to Transformer. In novel translation, this "memory decay" may cause the model to fail to accurately reference specific proper nouns or long-range background information mentioned earlier, thus affecting the rigor of the translation.

[0006] In summary, the current field of machine translation of novels faces a dual challenge: first, traditional Transformer models are extremely inefficient and costly when processing very long novel sequences; second, existing linear attention models, while maintaining high computational efficiency, struggle to accurately memorize and extract extremely long contextual information. Therefore, designing a novel translation model that maintains the advantages of linear complexity while enhancing long-range dependency handling capabilities through improved attention mechanisms is a key technical challenge for improving the efficiency of online literature translation overseas. Summary of the Invention

[0007] The purpose of this invention is to overcome the technical problems of high computational complexity, large memory consumption, insufficient context-dependent modeling ability, and inconsistencies in character titles, proper nouns, and semantic style in long novel translation in existing technologies, and to provide a novel translation model inference method and system based on the fusion of RWKV and attention mechanism.

[0008] The technical solution of this invention is: a novel translation model inference method based on RWKV, the steps of which are as follows: S1. Data collection and preprocessing are performed. Parallel corpora are extracted from different language translations of the novel using unsupervised learning methods. The parallel corpora are then cleaned, segmented, aligned, and dynamically spliced ​​to obtain bilingual text sequences for training.

[0009] The specific method in S1 is as follows: First, collect translations of novels in different languages ​​to construct a bilingual or multilingual parallel corpus; then, perform denoising, cleaning, sentence segmentation, word segmentation, and alignment on the original corpus to obtain standardized text samples that can be used for model training; finally, use dynamic MicroBatch splicing technology to divide the target sequence into text blocks with a fixed width, traverse permutations and combinations with a total length less than or equal to the sequence length, and use a simplex tableau dynamic solver to solve for the optimal number of uses for each splicing scheme, compressing multiple shorter text sequences into a longer training sequence, which is then sent to the training pipeline in S3 to eliminate the invalid computation generated by the traditional padding method, thereby increasing the proportion of effective tokens in all tokens during training and reducing the invalid computation caused by padding tokens.

[0010] Furthermore, the dynamic MicroBatch splicing process in S1, while ensuring that different sub-texts are computed independently, prevents contextual interference between different texts by cutting off the state at the boundary of adjacent text blocks and resetting it to an all-zero matrix.

[0011] S2. Construct a novel translation model based on the fusion of RWKV and attention mechanism. On the basis of RWKV architecture, introduce a lightweight group query attention mechanism and build attention cache by directly obtaining key and value information from the embedding layer.

[0012] The specific method in S2 is as follows: First, a backbone translation network consisting of an Embedding layer, multiple RWKV-Blocks, and an output layer is constructed. The RWKV-Blocks use a stacked structure of ChannelMix and TimeMix to model information in the channel and time dimensions. Then, a lightweight attention mechanism is introduced into the RWKV backbone architecture. This attention mechanism adopts a Cross-Attention architecture and uses GroupQueryAttention, allowing multiple Query attention heads to share a small number of Key and Value attention heads. Furthermore, the Key and Value are directly extracted and cached from the output of the Embedding layer. This cache structure is dynamically invoked during the S5 inference stage. A masking mechanism is used to isolate the cache of different text sequences, thereby supporting accurate reference to long-range backgrounds, reducing the memory consumption caused by multiple layers of repeated caching, and improving the model's ability to call historical text information. This enhances the consistency of proper nouns, character titles, and contextual references in the translation of long novels.

[0013] Furthermore, in the RWKV architecture of S2, TimeMix uses a constant exponential decay method to forget historical information. The Time-Decay parameter is restricted to between 0 and 1 through reparameterization to alleviate the gradient explosion and gradient vanishing problems that may exist in traditional recurrent models under ultra-long text conditions. At the same time, TimeMix defines custom combination operations between Key and Value to expand the state space and enhance the model's ability to preserve long-distance contextual information.

[0014] S3. The novel translation model is trained using a hybrid parallel training mode with sublinear complexity, and a numerically stable FP16 mixed precision training strategy is combined to achieve efficient training.

[0015] The specific method in S3 is as follows: First, the input sequence is divided into multiple training blocks according to the block length. Parallel computation is used within each block, while state is serially passed between blocks, thus forming a sublinear training method between a fully serial mode and a fully parallel mode. Wherein, if the sequence length is N and the block length is M, the number of iterations required for model training is reduced from N to [missing information]. N / M Subsequently, during mixed-precision training, FP16 was used for computationally intensive operations such as matrix multiplication, while FP32 was used for numerically sensitive parts such as LayerNormalization and GroupNormalization. Furthermore, a global scalar scaling γ was introduced in the TimeMix layer and during the RWKV-Block forward propagation. By leveraging the fact that Normalization is insensitive to overall scaling, the absolute value of the tensor was reduced, thereby mitigating the risk of numerical overflow caused by reparameterization and state accumulation, and ensuring the numerical stability of FP16 training.

[0016] S4. The novel translation model is optimized by adopting a hierarchical distributed heterogeneous training strategy, and the optimizer and gradient accumulation copy are offloaded to a low-performance computing device.

[0017] The specific method in S4 is as follows: First, a heterogeneous training environment consisting of high-performance computing devices and low-performance computing devices is constructed. The model's forward propagation, backward propagation, and main tensor operations are deployed on the high-performance computing devices, while the optimizer state and gradient accumulation copies are offloaded to the low-performance computing devices, thereby reducing the memory usage on the high-performance computing devices and increasing the size of the trainable model and the micro-batch size. Subsequently, the gradients transmitted between devices during model training are compressed. When the gradients are propagated back from the high-performance computing devices to the low-performance computing devices, FP16 quantization is used, and on the low-performance computing devices, FP32 is restored for gradient accumulation. Further, based on the sign characteristics of the Sign function in the Lion optimizer, the update amount is encoded and transmitted using a 2-bit compression method, thereby completing the gradient data exchange between the high-performance and low-performance devices during the S3 training process, reducing the communication bandwidth consumption between the high-performance and low-performance computing devices. At the same time, a composite optimization strategy that integrates Adam and Lion is adopted. The Adam optimizer processes the Embedding layer and its corresponding normalization layer, while the Lion optimizer processes the RWKV layer and the output layer, thereby achieving a balance between training stability and training efficiency.

[0018] Furthermore, in the hierarchical distributed training strategy in S4, the low-performance computing device is an inexpensive GPU device with a large video memory or memory capacity, which is used to replace the CPU to perform optimizer unloading and gradient accumulation tasks, thereby avoiding a significant decrease in training speed caused by CPU unloading.

[0019] S5. Decode and infer the trained model using a joint decoder based on TopP, Temperature, and Contrastive, combined with dynamic batch inference technology to achieve efficient novel translation output.

[0020] The specific method in S5 is as follows: First, the trained model is loaded into the inference module, and serial inference is performed on the input text. Then, a joint decoding method of TopP, Temperature, and Contrastive is used to sample the target language sequence from the word probability distribution output by the model to improve the richness and naturalness of the generated text and suppress the generation of duplicate text. Furthermore, dynamic batch inference technology is adopted to maintain multiple inference text sequences simultaneously during the inference process. When a certain inference sequence is missing, the text sequence that arrives first from the task to be processed is selected to fill the gap, and only one token is processed in each round of inference. At the same time, for the attention mechanism, a KV masking mechanism is used to shield the key and value caches from different text sequences to prevent the historical caches between different inference tasks from polluting each other, thereby achieving high throughput and low latency novel translation output.

[0021] A novel translation model inference system based on RWKV includes a data management module, a model training module, a backend inference module, and a user interaction module.

[0022] The data management module is used to extract, clean, segment, align and dynamically concatenate parallel corpora, and generate standardized bilingual text sequences for model training. The model training module is used to construct a translation model that integrates RWKV and attention mechanisms, and to perform sublinear complexity hybrid parallel training, mixed precision training, and heterogeneous distributed training. The backend inference module is used to load the trained novel translation model, execute the translation task based on the joint decoder and dynamic batch inference technology, and ensure cache isolation during the inference process of different text sequences through the KV mask mechanism. The user interaction module provides a visual interface for text input, parameter configuration, task submission, task status query, and translation result display.

[0023] An electronic device includes a memory, a processor, and a computer program stored in the memory and executable on the processor, wherein the processor executes the computer program to implement the aforementioned RWKV-based novel translation model inference method.

[0024] A computer-readable storage medium having a computer program stored thereon, which, when executed by a processor, implements the aforementioned RWKV-based novel translation model reasoning method.

[0025] The beneficial effects of this invention are as follows: By combining the RWKV long sequence modeling structure with a lightweight attention mechanism, this invention significantly reduces memory consumption and computational complexity in long text translation while ensuring the model's ability to continuously model long contexts. Through dynamic MicroBatch concatenation, sublinear complexity training, numerically stable mixed-precision training, and heterogeneous distributed optimization strategies, this invention effectively improves the size of trainable models and training efficiency. By using a joint decoder and dynamic batch inference technology, this invention can output more natural, coherent, and context-consistent novel translations, promoting the transformation of long novel translation from a high-cost, low-efficiency approach to an efficient and deployable intelligent translation method. Attached Figure Description

[0026] Figure 1 This is a schematic diagram of the overall architecture of the RWKV+Attention model provided by the present invention; Figure 2 A detailed diagram of the input and output of the RWKV structure block in loop mode; Figure 3 This is a schematic diagram of the internal logic structure of the ChannelMix layer; Figure 4 Here is a diagram of the TimeMix layer logic structure based on linear attention; Figure 5 A schematic diagram illustrating the process of attention for grouped queries; Figure 6 A comparison diagram of RWKV's grouped computation mode with parallel and loop modes; Figure 7 This diagram illustrates the principle of dynamic MicroBatch data stitching and compression. Figure 8 This is a basic flowchart of TimeMix training in FP16 half-precision; Figure 9 Flowchart for numerically stable training after introducing global scalar scaling γ; Figure 10 Apply a globally scaled forward propagation logic diagram to RWKV-Block; Figure 11 A flowchart illustrating the overall task scheduling and gradient accumulation process for distributed model training; Figure 12 A schematic diagram illustrating the working principle of the Lion optimizer performing gradient compression across heterogeneous devices; Figure 13 State definition diagram for the Sign function under 2-bit quantization compression; Figure 14 This diagram illustrates the division of labor between the Lion and Adam composite optimizers at the model level. Figure 15 This diagram illustrates the key-value masking process for concatenated text during inference. Figure 16 Prototype of the visual control interface for a novel translation system; Figure 17 This is the flowchart for this method. Detailed Implementation

[0027] The present invention will be further described below with reference to the accompanying drawings. It should be noted that the following embodiments are only used to illustrate the technical solution of the present invention and are not intended to limit the scope of protection of the present invention. Equivalent substitutions and conventional adjustments made by those skilled in the art to the number of model layers, hidden dimensions, block length, training accuracy, device deployment methods, and parameter settings without departing from the concept of the present invention should all fall within the scope of protection of the present invention.

[0028] This invention provides a novel translation model based on the fusion of RWKV and Attention. The model uses RWKV as its main architecture and employs a stacked ChannelMix and TimeMix approach to map the input sequence to the output time sequence. ChannelMix is ​​used to perceive data information along the channel dimension, while TimeMix is ​​used to perceive information along both the time and channel dimensions. The overall architecture of the RWKV+Attention model is as follows: Figure 1 As shown, the main body of the model consists of an Embedding layer and multiple RWKV-Blocks. In one instance, the number of RWKV-Blocks is 24.

[0029] In designing the model normalization structure, this invention adopts a Pre-LN approach instead of the Post-LN approach commonly used in Transformers. In the Post-LN structure, the output of the previous layer is added to the result of the current layer's computation before layer normalization. While this can suppress outliers in the outputs of feedforward and self-attention layers, it leads to difficulty in fast convergence in the early stages of training. Furthermore, the gradient is relatively amplified with increasing layers during backpropagation, thus masking the true descent direction and making it difficult for the model to learn effective gradients. To alleviate this problem, traditional Transformers typically use a warm-start approach, using a small learning rate initially and gradually increasing it linearly. However, this approach does not fundamentally solve the slow convergence problem in the early stages of training caused by Post-LN. This invention uses a Pre-LN structure, normalizing the input of each layer instead of normalizing the layer outputs before adding them together. This effectively alleviates the problem of unstable model convergence in the early stages of training and improves training speed.

[0030] Figure 2 The input and output details of the RWKV structure block in loop mode are presented. This structure block introduces multiple types of state information during execution to support continuous modeling over time. This represents the past time step composite information corresponding to RKV in TimeMix; This indicates the input information from the previous time step that was input into TokenShift in TimeMix; This indicates the input information at the previous time step that was input into the TokenShift in ChannelMix; and These represent the Key header and Value header for all past time steps input into TinyAttention, respectively. By introducing these state variables, the RWKV structure block can retain historical context information in recurrent mode without recalculating all historical inputs, thus making it suitable for long text translation scenarios.

[0031] Figure 3 The internal logical structure of the ChannelMix layer is presented. ChannelMix introduces the TokenShift mechanism. TokenShift is a time-series shifting mechanism that enables the feedforward network (FFN) in the traditional Transformer to have temporal awareness. This mechanism fuses information from the current time step and the previous time step using a learnable fixed ratio. Although TokenShift differs from the attention mechanism's approach of allowing arbitrary interaction between positions and the entire sequence, allowing only temporal interaction between adjacent positions, this interaction is sufficient to effectively improve the performance of the FFN. For text sequence prediction tasks, the current and previous positions typically contribute more to the current prediction, while the contribution of more distant positions gradually weakens; therefore, this mechanism aligns with the actual characteristics of text sequences.

[0032] In one example, the computation process of ChannelMix is ​​as follows: First, the input from the previous time step when position 0 is corrected; then, the input from the previous time step and the input from the current time step are mixed separately to obtain intermediate results for the Key branch and the gate branch; next, the shrinkage rate required for the final output is calculated, and the Key head is calculated; then, the SquaredReLU activation function is introduced to enhance the nonlinear fitting ability; finally, the scaled gate result is multiplied by the value mapping result corresponding to the Key to obtain the output at the current position. Through this structure, ChannelMix essentially forms a mini RNN with a time step length of 2, giving it a certain ability to process information across time steps without relying primarily on attention layers to complete information interaction in the time dimension, as is the case with traditional Transformers.

[0033] Furthermore, within ChannelMix, the gating branch and the key branch use different learnable parameters to fuse the inputs from historical time steps and the current time step. The gating branch ultimately passes through a sigmoid activation function, controlling its value range between [0,1]. This gating result is then scaled with subsequent calculations to obtain the final output. Simultaneously, ChannelMix employs the SquaredReLU activation function to introduce stronger non-linear expressive power into the model, improving its fitting ability. However, while SquaredReLU enhances the expression of non-linear features, it also more easily introduces numerical accuracy issues, which are further addressed in the subsequent mixed-precision training and numerical stabilization sections.

[0034] While ChannelMix already possesses the ability to perceive information between the current position and the previous position, it still lacks the ability to perceive information in the global time dimension. Therefore, this invention further introduces a TimeMix layer in RWKV to replace the self-attention layer in Transformer, which has quadratic time and space complexity. TimeMix can achieve sequence perception capabilities similar to self-attention while maintaining linear time complexity during inference and sublinear time and space complexity during training.

[0035] Figure 4 The logical structure of the TimeMix layer based on linear attention is presented. Parallelograms represent TokenShift operations, linear transformation modules represent matrix multiplication between the input and predefined model parameters, and red rectangles represent parameterless operations, including activation functions and matrix addition. TimeMix uses a WKV header with constant exponential decay to achieve time-dimensional information awareness. This decay method enables TimeMix to implicitly learn the positional relationships between different tokens. For information at a past position n... In terms of its relation to the current time step t Contribution satisfy ,in It is used to characterize the attenuation contribution of past location information to the current location. Since the distance from different locations to the current time step is different, their corresponding attenuation weights are also different. Therefore, TimeMix can indirectly learn location information without explicitly introducing traditional location encoding.

[0036] In TimeMix, input the previous time step. As a state It is initialized with all zeros. Information from the previous time step and the current time step is fused by addition. This method allows both the information from the previous time step and the information from the current time step to be fully connected to the fully connected layer without splitting the channels, thus avoiding affecting the expressive power of the fully connected layer and enabling the model to more freely choose between the information from the previous time step and the information from the current time step.

[0037] Compared to ChannelMix, TimeMix can be seen as a novel linear attention structure. This structure retains the constant memory usage characteristic of RNNs during inference, while possessing learning capabilities and text representation abilities close to those of Transformers. Therefore, RWKV is highly competitive in terminal deployment scenarios and ultra-long text processing scenarios, especially suitable for tasks with strong long-range context dependencies, such as novel translation.

[0038] Building upon the RWKV backbone structure, this invention further introduces an attention mechanism to enhance the model's ability to utilize historical text information and improve novel translation performance. The Transformer employs a multi-head attention mechanism to handle information interaction along the time step dimension and remembers historical information through a KV-Cache. This caching method saves the calculated Key and Value header information for each token. Compared to the state space complexity of O(n log n), this significantly improves performance. The RNN structure has a state space complexity of O(n). Multi-head attention mechanisms have a stronger advantage in remembering and retrieving historical information. This is because multi-head attention structures do not compress information in the time dimension, resulting in a lower compression ratio for the original information in the KV cache, making it easier to restore historical text content from the cache. This structure has a clear advantage in tasks where the model needs to copy previous content or strictly refer to historical text to complete the translation. In contrast, RWKV, based on the RNN architecture, processes historical state information through exponential decay. This approach allows RWKV to use less GPU memory with the same number of parameters, but it also means that RWKV is more difficult to accurately restore text information compared to Transformer.

[0039] For translation tasks, one of the most basic approaches is literal translation, which involves reorganizing the target language text to make it fluent and natural after establishing word-to-word mappings. This word-to-word translation method aligns well with the caching design of Transformers, and the multi-head attention mechanism enables the model to more easily learn the word mapping relationships between the source and target languages, thereby improving translation performance. Based on this consideration, this invention introduces an attention mechanism into RWKV to enhance the model's ability to process and utilize historical text information, thereby improving the overall translation performance in novel translation tasks.

[0040] The attention mechanism introduced in this invention differs from the Transformer's approach, which heavily relies on self-attention to process temporal interaction information; instead, it employs a simplified design. Unlike the Self-Attention architecture used in common GPT models, this invention uses a Cross-Attention architecture. This design allows RWKV to extract information from historical text as needed. Furthermore, unlike common Transformers that compute the Key and Value information of the attention mechanism on the output of the previous layer, this invention directly obtains the Key and Value from the model's Embedding layer output. In this way, the structure of the attention mechanism is significantly simplified, and the KV cache only needs to be stored once by caching the Embedding output. Compared to traditional attention mechanisms that cache Key and Value at each layer, this approach, while introducing only a small amount of recomputation cost, reduces the GPU memory consumption caused by caching the Key and Value of the attention mechanism to a fraction of the original approach. ,in This represents the number of model layers.

[0041] Furthermore, the Embedding layer output contains the most original and richest input information of the model. This information has not undergone information distortion introduced by subsequent calculations, and can preserve the original input content relatively completely. Based on this characteristic, this invention enables RWKV to more accurately reference historical text and more easily perform operations such as text copying. At the same time, the caching mechanism of KV Cache can offset some of the recalculation problems caused by only caching the Embedding.

[0042] Considering that traditional MultiHeadAttention consumes a lot of GPU memory, which limits the maximum sequence length that the model can infer under fixed GPU memory conditions and affects the decoding speed during inference, this invention further adopts the GroupQueryAttention structure to compress the number of attention heads in KV Cache and reduce the GPU memory usage of the attention mechanism. Figure 5 The working process of GroupQueryAttention is described. GroupQueryAttention can be viewed as a milder form of MultiQueryAttention, allowing multiple query attention heads to share a key attention head and a value attention head. In the specific implementation, the number of key attention heads must be the same as the number of value attention heads, denoted as . And ensure that the number of Query attention heads is an integer multiple of the number of KV attention heads, denoted as This grouping method avoids caching the same number of Key and Value attention heads as the number of queries, thereby reducing GPU memory usage, decreasing the number of GPU memory accesses to the attention layer, and improving model decoding speed, while having a relatively small impact on the final model performance.

[0043] Figure 5 The left half illustrates how GroupQueryAttention achieves a one-to-one pairing with the Query by doubling the Key and Value; the right half shows the calculation process of GroupQueryAttention, which is the same as the standard multi-head attention mechanism. Unlike traditional Transformers that rely on explicit positional encoding, the attention mechanism in this invention does not require additional explicit positional encoding. This is because in RWKV, TimeMix decays past information exponentially at a certain ratio using TimeDecay. During this process, the model implicitly learns positional information, which is then further passed to the query parameters of the attention layer. Simultaneously, each Query in the attention layer corresponds to a different number of key attention heads, forming a diagonal matrix. The diagonal matrix naturally possesses asymmetry, which allows the attention layer to implicitly learn the positional relationships between key attention heads. Therefore, the attention mechanism in this invention can complete positional relationship modeling within the existing model architecture without introducing an additional explicit positional encoding method based on sine and cosine functions.

[0044] During model training, this invention employs cross-entropy as the primary loss function. Cross-entropy loss uses entropy to measure the uncertainty between the predicted and true labels, and by reducing entropy, it allows the model to learn the mapping relationship between input data and target labels. Its form is as follows:

[0045] in, N Indicates the size of the vocabulary. The model predicts the first... i The possibilities of each category, This indicates that the real label corresponds to the first i The possibilities of each category. In this invention, the true label is represented using one-hot encoding, meaning that only the position corresponding to the target word in the true sequence is 1, and all other positions are 0. This form is equivalent to approximating the conditional probability distribution that the model needs to learn as a fixed distribution that only takes the value 1 at the target position. One-hot encoding can effectively simulate the true probability distribution of the model output, but this method itself is a relatively hard approximation, which will limit the richness of the model output expression to some extent. Therefore, this invention subsequently enhances the richness of the output text through various decoding methods.

[0046] Cross-entropy is a commonly used and efficient classification loss function with good optimization properties, which helps the model learn the mapping relationship between input data and target labels more stably during gradient descent. To further reduce numerical accuracy issues during model training and improve the stability of the training process, an auxiliary loss function is introduced, which takes the following form:

[0047] As can be seen, the auxiliary loss function is actually a variant of the softmax denominator, equivalent to imposing an additional constraint on the numerical range of the denominator. The softmax function can be expressed as:

[0048] In practical training, softmax is often used in conjunction with the cross-entropy loss function to reduce redundant computation and mitigate numerical overflow and vanishing issues caused by exp and log operations. The fused cross-entropy form is as follows:

[0049] Considering that the label uses one-hot encoding, where only one position is 1 and the rest are 0, it can be further simplified to:

[0050] Unlike the common approach of uniformly calculating loss at each position, this invention introduces an exemption mechanism that significantly reduces the loss value at positions that meet specific conditions. For novel translation tasks, readers are typically more concerned with the fluency of the translation and the natural flow of the context than with perfectly reproducing the details of the original text. Therefore, by reducing loss at certain positions, the model can generate richer expressions, resulting in more natural and diverse output text.

[0051] Furthermore, unlike methods that calculate loss only for predicted positions, this invention also calculates loss for the input text sequence. The input sequence loss is multiplied by a position-dependent scaling function and then concatenated with the output loss. This scaling function is designed as follows:

[0052] in, i Represents the first in the input sequence i The model uses a set number of tokens, where α represents the scaling factor reaching 1 after a specified number of tokens, and β represents that the loss is not calculated for the first few tokens of the input sequence. This design allows the model to apply loss constraints with different weights to different positions in the input sequence, thus providing more flexible control over the training process.

[0053] Furthermore, to adapt to the demands of novel translation for diverse generation and fluency, and to further reduce the impact of overfitting, this invention, referencing the Focal Loss concept, introduces an exemption mechanism to the loss function: if the top-K high-probability candidate words output at the current position already contain the word corresponding to the true label, the loss value at that position is reduced by a certain proportion. In this way, the model can achieve greater expressive freedom while maintaining basic semantic correctness, thereby generating more natural and diverse translations.

[0054] To achieve efficient training of RWKV under long sequence conditions, this invention employs a sublinear implementation method that falls between the cyclic and parallel modes. A comparison demonstration of its group computation mode with parallel and cyclic modes is shown below. Figure 6 As shown. RWKV defaults to providing the RNN form, although the RNN form has a mathematical complexity of O(n log n). However, in practical engineering implementations, it is difficult to create an efficient GPU-based implementation by combining existing operators without customizing them. This is because RNNs have relatively more loop steps, and RWKV is more complex than LSTM, leading to frequent activation of the CUDA Kernel during training, resulting in higher communication overhead and significant additional computational costs. Therefore, although the RWKV RNN form theoretically has superior computational complexity, it is not the optimal training mode in practice.

[0055] On the other hand, unlike traditional RNNs such as LSTM, RWKV does not use a non-linear activation function between consecutive states. Therefore, it is easier to write parallel modes where each time step is computed independently, such as... Figure 6 As shown. In parallel mode, each position is computed independently, and the computational complexity is O(n log n). The space complexity is also In this case, the model complexity and space complexity of RWKV are the same as those of Transformer. To achieve a balance between a fully serial RNN with linear computational complexity and a parallel model with quadratic complexity, this invention provides a sublinear implementation with a complexity of N×M, where N is the sequence length and M is the block length. The block length M determines the upper limit of parallel processing of the model, and the speedup effect brought by M is limited by the computing power of the computing device.

[0056] Figure 6This illustrates a comparison of text grouping methods under different modes. In the recurrent mode, the entire text sequence is processed sequentially in chronological order; in the parallel mode, all positions participate in the computation simultaneously; in the grouped computation mode, when GroupSize is 4, the text is divided into several groups, with positions within the same group being computed in parallel, while different groups are processed sequentially in chronological order. Compared to the RNN mode, the number of loops required for the model to complete the prediction of the entire sequence can be reduced from N to [missing information]. Therefore, this combination of parallel and serial methods can significantly reduce the number of loops in the training process and improve training efficiency.

[0057] In one example, to maximize training efficiency, Tensorflow 2.14 can be used for training. First, the code is debugged using dynamic graph mode, then switched to static graph mode for formal training. Since switching from dynamic to static graph requires compilation, and the compilation time is quadratic with the number of loops N, excessively long loops not only reduce model training efficiency but also increase the time required to compile from dynamic to static graph. The efficiency of different training modes is shown in Table 1: Table 1: Efficiency Comparison under Different Training Modes

[0058] As shown in Table 1, under long sequence conditions, the parallel mode cannot complete training due to memory overflow, while the serial mode, although theoretically less complex, has long compilation time and low training speed. The hybrid mode proposed in this invention can complete compilation within an acceptable time while maintaining high training efficiency. Therefore, this sublinear complexity implementation achieves a balance between training efficiency and memory consumption, making it suitable for training long-text novel translation models.

[0059] To further illustrate the sublinear complexity implementation, RWKV replaces the self-attention mechanism with TimeMix. TimeMix is ​​a linear attention layer that uses constant exponential decay to forget historical information. U The Time-First parameter is used to mark the input information at the current location, in order to distinguish information from the current location from information from the historical state; W The Time-Decay parameter is used to forget past information. To ensure that the value of Time-Decay is within the range of [0,1], the following reparameterization method is used:

[0060] In this way, regardless of the original parameters How are values ​​obtained after reparameterization? The gradient is always within the range [0,1], which allows the model to only attempt to forget historical information without amplifying it unbounded. This mechanism helps alleviate the gradient explosion and vanishing problems that may exist in traditional RNNs or LSTMs under extremely long texts, laying the foundation for RWKV to work stably in extremely long contexts.

[0061] In TimeMix, the symbol " "" represents a custom operator used to define how the key and value are combined. This calculation method enables RWKV to generate a larger state space, allowing the model to remember more historical information with a lower signal-to-noise ratio, reducing forgetting of previous content in very long text sequences, and improving the model's stability in different task scenarios. Its calculation method is as follows: The input is a shape of... Tensors A and B, where For sequence length, The width of the hidden layer and the size of the header are [value]. Firstly, according to and The relationship determines the number of attention heads, and the Key tensor is rearranged as follows: The form rearranges the Value tensor into Form; then element-wise multiplication is performed on the rearranged tensor to obtain the tensor. This serves as the basis for subsequent state calculations. The core calculation of TimeMix can be represented as:

[0062] in, It indicates that WKV is in the Output at any moment For TimeFirst, For the first The rate of contraction at time t, For the first The key to attention at any moment For the first The attention head for the Value at time step. The above parallel form can be further derived into a loop form:

[0063] Therefore, TimeMix can be expressed as a decaying accumulation of historical terms from a parallel perspective, or it can be rewritten as a process using states. The recursive loop format balances training efficiency with long text reasoning ability.

[0064] To further achieve block-based parallel training, without considering the shrinkage vector... Under the premise of matrix multiplication, the output at the current time step It can be written as:

[0065] in, Let be the position attenuation function. To obtain the position attenuation mask, the position attenuation parameters are first reparameterized, resulting in:

[0066] Then calculate the position. Attenuation coefficient:

[0067] Based on this, a list of attenuation weights determined by the distance from the last position is constructed, and a Time-First parameter is introduced at the current position. After further processing such as copying, flattening, cropping, and reshaping, the final position attenuation mask is obtained. The final form of the mask is:

[0068] By using the positional attenuation mask described above, we can obtain The values ​​obtained under various location conditions are then substituted into the rkv calculation formula to obtain the rkv value within each group and the corresponding output.

[0069] If the sequence is grouped into fixed lengths and computed independently within each group, the transmission of state and gradient information between groups will be interrupted. To address this issue, this invention further introduces a loop module to transmit state information from the previous group between adjacent groups, thereby connecting adjacent groups from left to right. The following explanation uses a special case with an input time series length of 6 and a block length of 3.

[0070] In this example, the output at each time step in parallel mode can be written as:

[0071] in, to Corresponding to the output of the first loop block, to This corresponds to the output of the second loop block. By comparing the expansion formulas of the first and second loop blocks in fully parallel mode, it can be found that the second loop block contains the following shared part:

[0072] and With the above There is a multiple relationship. The shared parts are respectively of Times. Therefore, an obvious computation time step for The contribution method is to simply calculate And according to Scale and add as bias The corresponding position. When only the time step is considered. ,for When making contributions, by putting , Substitute And it was calculated using parallel mode. and for By adding up the contributions, we can obtain the result. Actual output value.

[0073] Furthermore, the state of the second loop block It can be written as:

[0074] By comparison, it can be found that, without considering position decay... and The middle of the time step The components contributing to the output have the same structure, and there is a fixed multiple relationship between them. Therefore, it is possible to further write based on... of The recursive form is:

[0075] This establishes the state. With state The connection between them. Meanwhile, when ignoring... China regarding When it is part of the time, and The computational structure remains consistent, differing only in the indices. This means that the state of any loop block can be obtained through a unified state calculation function combined with a scaling and addition method. Therefore, block-based parallel training is no longer just about independent computation within a group. Instead, it uses a loop module to pass state information from the previous group between adjacent groups, connecting them from left to right. This ensures that state and gradient information continue to be passed between groups.

[0076] Thus, the most complex RKV calculation in TimeMix is ​​completed. After completing the above calculation, TimeMix further uses GroupNormalization to standardize the data. This operation scales the data to the linear space corresponding to the standard normal distribution, thereby improving gradient stability during training and providing a foundation for subsequent mixed-precision training and numerical stabilization mechanisms.

[0077] To accommodate token sequences of varying lengths in the training set, this invention introduces a data concatenation scheme during the data preprocessing stage. The principles of data concatenation and compression are as follows: Figure 7 As shown. The length of natural text is usually limited by its specific content, and there are often significant length differences between different samples. To adapt to this characteristic, existing training data processing methods typically revolve around truncation and padding. Padding refers to using invalid tokens to fill in shorter text; in the model, such words are usually represented by... <pad>representations. The tokenizer-decoded-out <pad>The word segmentation is not retained in the final output text sequence, and it is included in the label sequence during training. <pad>The location of the marker is not included in the calculation of the model loss.

[0078] A common truncation method is to cut very long texts to the same length as shorter text sequences. While this method can transform text sequences of varying lengths into multiple text sequences of uniform length without introducing unnecessary computation, it reduces the data richness of the input text, preventing some longer texts from being fully trained. Another common method is to truncate the shorter texts... <pad>Padding is performed to make the length equal to the target sequence length, and then a masking mechanism is used to exclude the padded portion from the loss calculation, effectively removing it from the gradient descent process. While this method is commonly used, it slows down model training because it introduces a large number of padding tokens that participate in computation but not training. A compromise is to truncate longer texts and pad shorter texts to unify their length to the target sequence length. This approach allows most text to participate in training, but extremely long texts still cannot be fully trained, and it still introduces a large amount of padding computation, resulting in wasted computational resources.

[0079] To measure the effective training amount of a model during training, this invention defines the training amount. The effective number of tokens used in the model's training throughout the entire training cycle is expressed as:

[0080] in, This indicates the number of texts used in the training process. Indicates the length of the sequence used for training. This represents the proportion of tokens that participate in gradient updates during training out of all tokens. From this definition, we can see that one way to improve model performance is to increase the number of training texts. However, the amount of training is directly related to economic and time costs. Under cost constraints, simply increasing the amount of training is not enough. It is often unrealistic. Therefore, within a given training budget, improving... Increasing the proportion of effective computation in all computations is a feasible performance improvement solution.

[0081] Based on the above considerations, this invention proposes a dynamic MicroBatch concatenation technique, where the concatenation process is completed during the dataset preprocessing stage. This concatenation method can significantly improve the real token throughput per second during training, and introduces almost no additional computation. Its time complexity depends only on the length of the training text and is essentially independent of the overall dataset size, with a time complexity of O(n log n). ,in, Indicates the length of the text sequence. This indicates the size of the text block during the training process.

[0082] In one example, this invention divides the target sequence length into a series of adjacent blocks with a constant width, and names them text blocks. The text block is the most basic scheduling unit in the splicing process. Text sequences whose length is not an integer multiple of the text block length are processed by... <pad>Padding is done upwards to an integer multiple of the nearest text block length. For example, when the text sequence length is 11 and the text block length is 4, the padded text sequence length is 12. Introducing text blocks not only reduces the computational complexity of the concatenation process but also aligns with the loop block length in the RWKV training process, allowing the RWKV model to trim gradients as needed during training, ensuring the training process proceeds normally.

[0083] Figure 7 The paper presents the original text, a text sequence processed using conventional padding methods, a text sequence obtained using the method of this invention, and a text sequence obtained using a truncation method. It can be seen that the "empty" positions in the text sequence are significantly reduced after applying the method of this invention. By flexibly combining text blocks, this invention can compress multiple shorter text sequences into a longer text sequence, thereby reducing gaps within the same MicroBatch and significantly improving computational efficiency.

[0084] Specifically, the data splicing scheme of the present invention includes two steps: the first step is to calculate all possible splicing schemes; the second step is to use a simplex tableau dynamic solver to solve for the number of uses of each splicing scheme.

[0085] In the first step, all permutations and combinations with a total length less than or equal to the target sequence length are iterated. Taking the case where the target text sequence contains 3 text blocks as an example, all possible concatenation schemes are shown in Table 2. Each row in the table represents a concatenation and combination scheme, and each column represents the number of text blocks of the corresponding length that the scheme can accommodate.

[0086] Table 2: All possible splicing schemes when the number of text blocks is 3

[0087] As can be seen from Table 2, there is significant redundancy in the solutions generated through permutations and combinations. For example, solutions X5 and X6 are clearly less effective than solution X1. However, these redundant solutions do not need to be manually removed beforehand, as they will be automatically identified and discarded during the subsequent solution stages.

[0088] Generalizing the above special cases, we can obtain a more general set of splicing schemes, as shown in Table 3. In the table, the first dimension represents different splicing schemes, and the second dimension represents the number of text blocks of various lengths that can be used under each splicing scheme.

[0089] Table 3: Splicing schemes obtained after promoting the entire R set

[0090] In the second step, the scheme representations in Table 3 are transformed into a linear programming optimal solution problem, which is then solved using a simplex tableau dynamic solver. Specifically, the objective function, which minimizes the sum of the usage quantities of all schemes, can be expressed as:

[0091] At the same time, constraints are established to ensure that the sum of text blocks of corresponding lengths in all solutions is not less than the number of text blocks of each type in the dataset, that is:

[0092] in, These represent the number of times each splicing scheme was selected. This represents the required number of text blocks of different lengths in the dataset. This is achieved by solving... This allows us to obtain the optimal number of characters to use for each splicing scheme. Then, we combine texts of different lengths according to the corresponding scheme to obtain the compressed training sequence.

[0093] Although this invention compresses multiple texts into a single longer text using the aforementioned method, each subsequence should still be computed independently. Since the preprocessing stage has already ensured that the length of each text is an integer multiple of the text block length through padding, the boundary between two adjacent texts also corresponds to the boundary between text blocks. During training, this invention employs a sublinear RWKV implementation. Within a training block, parallel mode is used with independent computation at each position, while between two training blocks, a serial mode is used, connecting adjacent training blocks through states to pass information. Therefore, simply interrupting the states at the boundary between corresponding text blocks of two texts and resetting the states obtained from the previous training block to an all-zero matrix is ​​sufficient to achieve context separation between different texts. Furthermore, this invention introduces additional input to record whether each state needs to be interrupted.

[0094] To verify the actual performance of the compression algorithm described above, this invention was tested on the "Chinese and English bilingual" subset of the WikiMatrix multilingual public dataset. In the test, the training sequence length was set to 384 tokens, and the training block size was set to 48 tokens. The actual test results are shown in Table 4: Table 4: Compression Ratio Table for the WikiMatrix Dataset

[0095] To improve training efficiency and reduce memory usage while maintaining model performance, this invention employs mixed-precision training during model training. Mixed-precision training combines the advantages of half-precision and single-precision floating-point formats. During model training, data that is not sensitive to numerical values ​​is stored and calculated using half-precision, while data requiring high numerical stability is calculated using single-precision. This accelerates training while maintaining model performance as close as possible to that of single-precision training.

[0096] In one example, the half-precision data format uses FP16. FP16 uses 2 bytes to store a value, while single-precision float32 uses 4 bytes, thus FP16 can halve the memory usage. FP16's 16 bits include 1 sign bit, 5 exponent bits, and 10 mantissa bits. However, due to its smaller dynamic range and limited precision, direct use for model training can easily lead to underflow or overflow issues. Therefore, during training, FP16 is used for calculations in parts less sensitive to numerical stability, such as fully connected layers and activation function intermediate values; while single-precision floating-point calculations are used for parts more sensitive to numerical precision, such as LayerNormalization and GroupNormalization. The basic workflow of TimeMix under FP16 half-precision training is as follows: Figure 8 As shown.

[0097] In most cases, training with mixed precision and loss amplification directly achieves relatively stable and efficient training. However, training with FP16 on RWKV presents significant challenges, especially as model size increases. For example, using... Figure 8 When training a 90M-scale model using the method shown, the model is relatively stable in the early stages of training, but its stability gradually decreases as the model size increases and the training process deepens. When the model size increases to 0.4B, it becomes increasingly difficult to continue using this method. Figure 8 Stable training is achieved as shown. Therefore, to address the numerical problems encountered by RWKV during FP16 training, this invention further proposes a corresponding numerical stabilization scheme.

[0098] In RWKV, reparameterization is one of the areas prone to numerical overflow. Time-Decay relies on reparameterization to construct parameters with a range between [0,1], in the form:

[0099] in, The range of values ​​for covers the entire value range, while It is restricted to the range [0,1]. If the process is calculated directly using FP16, then the first step is calculated... At that time, When any value of is greater than 11.1, It will overflow because it exceeds the representable range of FP16. .although It is still possible to calculate and obtain a constant value of 0, but if such a value appears in Time-Decay, it will cause the corresponding parameter to be essentially "dead". Therefore, in this invention, the reparameterized part of Time-Decay is calculated using single precision to avoid numerical overflow caused by reparameterization.

[0100] Furthermore, during model training, as the model's capabilities improve, the Time-Decay value gradually approaches 1, thus reducing the degree of forgetting historical information. In RWKV, each time step accumulates the input from the previous time step after scaling by Time-Decay. If the accumulated value exceeds the upper limit represented by FP16, numerical overflow will occur and the result will become... This can lead to training crashes. To mitigate the numerical overflow caused by this accumulation, this invention employs a method of scaling all elements proportionally. Specifically, one input in an addition or matrix multiplication operation that may cause numerical overflow is divided by a constant scalar. If the entire calculation process involves only linear computation, then it can be guaranteed that the final result is only relative to the original result. The scaling relationship is an integer multiple, without changing the relative proportions between corresponding elements.

[0101] Specifically in RWKV, such as Figure 9 As shown, this invention performs scalar scaling on the KV (Key-Value) array. Since all operations before GroupNormalization are linear, it can be guaranteed that the corresponding elements in the output after scalar scaling always have the same scaling factor as the output obtained without scaling. This conclusion is significant because GroupNormalization is insensitive to overall scaling; mathematically, scaling all elements uniformly results in equivalent normalized outputs. Based on this characteristic, by applying uniform scaling to intermediate tensors, the absolute value of intermediate computational tensors can be reduced without altering the normalized result, thereby mitigating the risk of numerical overflow.

[0102] To illustrate the standardization process, this invention employs the following standardization implementation. For the input tensor... If it is a one-dimensional or multi-dimensional tensor, then only the last dimension is included in the standardization calculation. First, calculate the mean of the tensor on the last dimension:

[0103] Then calculate the standard deviation of the tensor in the last dimension:

[0104] Further calculate the normalized value of each element:

[0105] in, This is a very small constant used to prevent undefined behavior caused by a denominator of 0. Based on the above standardization method, this invention further employs GroupNormalization to group and standardize the data. Specifically, for data with a width of... D Input tensor X By group size C Rearrange it as:

[0106] Then, each group is standardized separately:

[0107] Finally, flatten each group back to its original shape to match the input, resulting in the output tensor. :

[0108] Based on the above standardization formula and group standardization formula, it can be seen that, without considering... Under the premise of input X Performing overall scaling does not affect the normalized output. In practice, since the input tensor has already been scaled, the default... The value is relatively large and can no longer be considered a negligible minimum. To ensure the stability of numerical calculations, this invention will... Adjusted to Furthermore, the FP32 computational layer normalization and group normalization processes were continued. After these adjustments, the numerical accuracy issues caused by GroupNormalization were effectively alleviated.

[0109] Building upon the aforementioned local scaling, this invention further proposes a global numerical scaling method, which applies the forward propagation logic after global scaling as follows: Figure 10 As shown. Although the aforementioned method can solve the numerical overflow problem caused by FP16 in most cases, numerical overflow at random locations may still occur in a few cases as training progresses. To further reduce such problems, this invention applies a constant numerical scaling to the output of the Embedding layer and makes this scaling consistent throughout the entire forward propagation process, thereby reducing the absolute numerical size of the tensor and stabilizing the training process.

[0110] Specifically, thanks to the insensitivity of Normalization to overall scaling, the output obtained after Normalization of the scaled input is numerically identical to the output obtained after Normalization of the original input. However, the same scaling still needs to be applied to the outputs of the TimeMix layer, ChannelMix layer, and TinyAttention layer to ensure that the scaling ratio of each branch output remains consistent with that of the main branch. Through this method, the present invention can apply global scaling to the entire forward propagation main link of the RWKV-Block without changing the normalized calculation results, thereby further reducing the risk of random numerical overflow during FP16 training and improving the stability of the training process.

[0111] To improve model training speed and expand the scale of trainable models, this invention jointly designs key strategies in the training and inference processes. Figure 11 This diagram illustrates the process of a complete gradient accumulation iteration during training. The blue boxes represent the computational flow within a specific high-performance computing device for a particular batch, while the red boxes represent the information exchange flow between different high-performance computing devices within the same batch. The gradients obtained within the different red dashed boxes are sent to a low-performance computing device for gradient summation, where partial gradient updates are performed. The remaining gradients are then returned to the high-performance computing device to continue the subsequent update process.

[0112] Regarding distributed training strategies, this invention employs a hierarchical distributed architecture. Traditional distributed training typically assumes that each computing device has identical computational performance and characteristics, an assumption common in traditional data parallelism and model parallelism schemes. Model tensor parallelism evenly slices the model and distributes it across multiple computing devices. During backpropagation, data involving cross-device communication must be transferred to the current device before computation can continue. Therefore, compared to data parallelism that only synchronizes gradients, model tensor parallelism introduces a significant amount of inter-device communication. Generally, model tensor parallelism is more suitable for scenarios where the model size is too large and a single device cannot fully store the model.

[0113] The proposed solution is based on a data parallel architecture and further employs a hierarchical distributed approach that combines high-performance computing devices with inexpensive low-end computing devices that have large video memory or system memory capacity. Specifically, components that consume a large amount of video memory but participate in computations less frequently or with relatively small computational loads are offloaded from the high-performance devices to the next lower level of computing devices. Furthermore, this distributed solution allows one low-performance computing device to work collaboratively with multiple high-performance computing devices, and deduplicates variables that were originally mirrored across the high-performance devices on the low-performance devices, retaining only one copy. This significantly saves video memory usage on the high-performance devices, and this memory saving effect becomes more pronounced as the number of high-performance computing devices increases.

[0114] In one instance, this invention offloads the model optimizer and gradient accumulation copies required to a low-cost computing device. In traditional data parallelism, optimizer and gradient accumulation copies are typically stored as mirrors on each computing device. After offloading, only one copy of the optimizer and one copy of the gradient accumulation are retained on the low-performance computing device, saving the GPU memory overhead corresponding to one copy of the optimizer and one copy of the gradient accumulation on each high-performance device. Since the number of model parameters often dominates during training, and the number of optimizer parameters is usually proportional to the number of model parameters (one to two times the number of model parameters for different optimizers), coupled with gradients and gradient accumulation vectors of equal size to the number of model parameters, the GPU memory occupied by the optimizer and gradient accumulation is not negligible. By transferring the optimizer and gradient accumulation tasks to low-performance computing devices, in one instance, the maximum trainable model size is increased from 0.1B to 0.4B.

[0115] In existing distributed training frameworks, DeepSpeed ​​is a commonly used approach. When using offloadstage, DeepSpeed ​​offloads the optimizer from the GPU to the CPU, thereby significantly reducing GPU memory consumption during training. However, offloading the model to the CPU can significantly slow down training. To address this, this invention, based on TensorFlow's MirrorStrategy, uses inexpensive GPUs to perform the offloading task instead of the CPU, thus conserving GPU memory while minimizing a significant drop in training speed.

[0116] While offloading from low-performance computing devices significantly frees up memory on high-performance devices and allows for further increases in micro-batch size, the heterogeneous architecture combining high-performance and low-performance computing devices necessitates cross-device communication. During communication, high-performance devices cannot perform other computational tasks, thus prolonged communication reduces their utilization. To reduce the time required for model gradients to be propagated back from high-performance devices to low-performance computing devices, this invention uses the FP16 data type to quantize gradients. After the gradients are propagated back to the low-performance computing device, they are converted to FP32 for gradient accumulation. Simultaneously, to minimize precision loss due to the lower representation precision of FP16 compared to FP32, gradient scaling on the high-performance device is no longer performed by the global batch size; instead, the gradients are directly summed, increasing their value through accumulation. This reduces the precision loss and the possibility of values ​​returning to zero caused by the conversion from FP32 to FP16.

[0117] To further reduce the time required for the optimized parameters to be sent back to various high-performance computing devices after the optimizer is updated, this invention breaks down the optimization process according to the characteristics of the optimizer. Figure 12 This illustrates the working relationship of the Lion optimizer in gradient compression across heterogeneous devices. Although this invention ultimately employs a composite optimizer combining Lion and Adam, since most parameters in the model are updated by Lion, communication compression of Lion can still significantly improve the overall communication efficiency of the composite optimizer.

[0118] This scheme leverages the characteristics of the Sign function in the Lion optimizer, whose output depends only on the sign of the input value and whether the input value is zero. Initially, it was attempted to ignore the case where the Sign function outputs 0, since the Sign function can only output 0 when the input is strictly equal to 0. However, during actual training, because RWKV uses SquareReLU as an intermediate activation function in ChannelMix, the outputs of some neurons are truncated to 0 due to ReLU activation, resulting in the gradient value of the corresponding neuron being 0; simultaneously, the Lion's momentum parameter... Since zero initialization is used, gradients will appear in the early stages of training. With momentum The situation is also 0. When the Sign function outputs 0, if 1-bit compression is used, this state cannot be correctly represented. This would result in erroneously quantizing 0 as -1 or 1, leading to an incorrect gradient direction. The model cannot receive the correct gradient and ultimately fails to converge. Actual training results show that using 1-bit compression causes the RWKV model loss to continuously increase, and the model shows no signs of convergence.

[0119] To accurately represent the value range of the Sign function, this invention uses 2 bits to compress the update amount and selects three states as the valid values ​​for Sign. Figure 13 This illustrates the state definition of the Sign function under 2-bit quantization compression. In practical implementation, two Boolean tensors can be used to represent a parameter state of the model, and bitwise operations can be used to achieve efficient encoding and decoding of the parameters. The compressed update, after being transmitted to a high-performance device, can be processed based on the model parameters on the high-performance device. The updated model parameters are calculated based on the learning rate, and the parameter update is completed. This Lion splitting method can significantly optimize the communication speed between different devices and improve training efficiency without introducing additional resident memory.

[0120] Regarding optimizer selection, this invention further employs a training optimization approach that integrates Adam and Lion. Lion is a compact and memory-efficient gradient optimizer. Compared to adaptive optimizers like AdamW, which require storing both first and second moments, Lion only needs to store momentum, thus reducing resident memory consumption by about half. When training microbatches are small, model parameters and optimizers typically dominate memory usage; therefore, reducing optimizer memory usage can significantly increase the trainable model size or microbatch size, thereby improving model training performance. On the other hand, adaptive optimizers with first and second moments introduce more computation compared to Lion, which only retains momentum, thus reducing training speed. Existing research has shown that Lion can achieve an average training speedup of 2% to 15% compared to AdamW and Adafactor across different tasks.

[0121] However, in the embedding layer, symbol-update optimizers such as Lion or Tiger can cause special phenomena. Previous studies have observed that in the embedding layer optimized by Tiger or Lion, some tokens are gradually optimized into pairs of fixed values ​​with the same absolute value, making them difficult to distinguish and thus affecting the model's performance in certain scenarios. Further observation shows that tokens containing outliers are highly correlated with their frequency of occurrence in the training set; tokens with lower frequency are more likely to exhibit outliers during gradient updates, while the problem is relatively less pronounced for tokens with higher frequency. This tendency for low-frequency word embeddings to become identical significantly impacts the model's performance when handling low-frequency words.

[0122] Based on Lion's update characteristics, without considering the effects of learning rate and gradient decay, the Lion optimizer updates model parameters with a fixed step size of [1, 0, -1]. The reason for the Embedding outlier is due to errors in the gradient update formula. In the middle, due to the existence of the sign function sign, only when momentum The parameter gradient only changes when the sign of the token changes. Based on the low-frequency token, the gradient changes in most training steps. Assuming the value is approximately 0, we can obtain:

[0123]

[0124] when t When taking the maximum value, it can be approximated as From this, we can further derive:

[0125] Figure 14 The division of labor between the Lion and Adam composite optimizers at the model level is illustrated. Based on the above analysis, this invention identifies the cause of outliers in the embedding and adopts a composite optimization scheme accordingly: the Adam optimizer is used to process the embedding layer and its corresponding LayerNorm normalization layer, while the RWKV layer and output layer continue to use the Lion optimizer. This scheme corresponds to the second approach in existing solutions, which bypasses Lion and uses only other optimizers to process the embedding module, leaving the remaining model parameters to be optimized by Lion.

[0126] In one example, this invention uses a cosine decay function with warm-start as the learning rate function and applies different learning rate scaling factors to the Adam and Lion optimizers. Because Lion incorporates the unique structure of the Sign function, the learning rate used by the Lion optimizer is one-tenth that of the Adam optimizer under the same number of iterations. To maintain the same weight decay effect for each weight in a practical sense, the corresponding weight decay coefficients are adjusted proportionally. Although Adam and Lion use different initial learning rates, they still use the same weight-based proportional weight decay. Through this composite optimizer design, this invention achieves a good balance between performance and stability.

[0127] After the model completes training, to correctly sample individual words from the word segmentation probability distribution at each position, this invention further employs a joint model decoder. The model outputs the probability distribution of each candidate word at each position; therefore, the design of the decoder directly affects the richness, fluency, and overall performance of the final generated text. A straightforward decoding approach is to use a greedy strategy, that is, to select the word with the highest probability as the current output each time. However, this approach only selects a locally optimal solution at the current time step, without considering the overall impact of future time steps on the current and past time steps, thus making it difficult to obtain a globally better output result.

[0128] To address this, this invention employs a joint decoder based on TopP, Temperature, and Contrastive. This joint decoding approach does not rely on a single deterministic strategy but instead uses a random sampling approach under a specified distribution, providing better decoded text results without introducing additional inference iterations. Compared to greedy decoding, this joint decoder can provide richer text output during the generation process and suppress the repeated generation of the same text content. Furthermore, by adjusting the threshold of TopP, the range of the candidate token set during model decoding can be adjusted, thereby affecting the diversity and degrees of freedom of the model's generated text.

[0129] To improve the actual throughput during the inference phase, this invention further employs dynamic batch inference technology. This technology significantly improves inference speed by increasing the parallelism during the model inference process. In one example, the model uses a serial inference mode. This serial inference method provides a high degree of freedom while maintaining high inference performance.

[0130] During inference, the model retains at most [number] records at a time. There are inference text sequences, each inference sequence denoted as . Each inference sequence may contain multiple sub-time series, each corresponding to a time series prediction task. The first sub-sequence in the first sequence can be denoted as... Let the nth subsequence in the mth sequence be denoted as When a gap exists in a certain inference sequence, the first arriving sequence task from the text sequence to be processed is selected to fill the gap in the inference sequence.

[0131] In the actual inference process, each inference sequence processes only one token per round, and decides whether to retain the token based on the current state of the sub-time series. When the model decodes a termination token in a sub-sequence during inference, that sub-sequence is removed from the task queue, and a new task is selected and added to the current task sequence. In this way, tasks can be continuously loaded and recycled within the limited inference sequence slots, thereby improving the overall inference throughput.

[0132] Regarding the attention mechanism, this invention employs a masking mechanism to shield the key-value caches from past time series, preventing cached information from different text sequences in the past from interfering with the current text sequence. The KV masking process during the concatenation is as follows: Figure 15 As shown, by isolating the caches between different text sequences during the inference process, dynamic batch inference can improve throughput without compromising the independence and contextual correctness of each text sequence.

[0133] Based on the above training and reasoning mechanisms, this invention further constructs a novel translation implementation system, the prototype of which is shown in the figure below. Figure 16 As shown in the example, in one instance, the system uses console requests to the Flask backend to achieve front-end and back-end interaction. The backend uses the Flask framework and defines a `request_task` interface for the backend inference module to submit task requests. This interface is mounted under the ` / api / interface / translation / ` route and is used to receive POST requests with JSON structures and return the task ID and estimated completion time. The estimated completion time can be estimated based on the assumption that the number of output tokens is equal to the number of input tokens, i.e., the total sequence length is twice the length of the input sequence, combined with the current inference speed.

[0134] Furthermore, this invention defines a `get_info` interface for obtaining the current task execution status. This interface is also mounted under the ` / api / interface / translation / ` route, receives a POST request with the task ID, and returns whether the decoding task is complete and the currently generated decoding sequence. Unlike a single-phase interface that directly returns the final result with a single request, this invention uses a two-phase approach to complete task submission and result retrieval separately. This approach reduces programming complexity and mitigates the risk of communication interruption due to network fluctuations or long HTTP connection establishment times.

[0135] In terms of front-end implementation, this invention uses Dash to draw the front-end page and utilizes the dash_bootstrap_components component to beautify the user interface. Through the above-mentioned front-end and back-end collaboration, a complete novel translation control system can be formed, realizing the submission of model inference tasks, query of execution status, and display of translation results. In one example, the system can receive Chinese novel excerpts as input and output the corresponding English translations, thereby completing the novel translation task.< / pad> < / pad> < / pad> < / pad> < / pad>

Claims

1. A novel translation model inference method based on RWKV, characterized by the following steps: include: Step 1. Data preprocessing: Parallel corpora are extracted from the novel translation using unsupervised learning methods. The short sequences are compressed into standardized text sequences using dynamic MicroBatch splicing technology, and the state is reset at the boundary of adjacent text blocks. The resulting sequences are used as input data for model training in Step 3. Step 2. Model Architecture Design: A lightweight grouped query attention mechanism is introduced on the basis of the RWKV architecture, and the Key and Value information directly output from the Embedding layer is configured as a KV cache. This architecture serves as the main training component in Step 3 and provides the basis for long-range background calls during inference in Step 5. Step 3. Parallel Model Training: Input the sequence obtained in Step 1 into the architecture constructed in Step 2, adopt a sublinear complexity block mode, combine a global scalar scaling FP16 mixed precision strategy, and introduce a loss waiver mechanism to iteratively optimize the model parameters. The training process is supported by distributed support in Step 4. Step 4. Distributed heterogeneous optimization: A composite optimization strategy that integrates Adam and Lion is adopted to offload the optimizer and gradient replicas to low-performance computing devices. Based on the characteristics of the Sign function, the update amount of the Lion optimizer is compressed and transmitted in 2 bits to realize parameter update interaction between heterogeneous devices. Step 5. Model Inference Output: Load the parameters obtained from training in Step 3 and the KV cache configured in Step 2, and use the joint decoder, dynamic batch inference technology and asynchronous task interaction mechanism to output the translation results.

2. The inference method for novel translation models based on RWKV according to claim 1, characterized in that: In step 1, the dynamic MicroBatch splicing technology divides the target sequence into text blocks with a fixed width, traverses all permutations and combinations with a total length less than or equal to the sequence length, and uses a simplex tableau dynamic solver to solve for the optimal number of tokens to use for each scheme, thereby increasing the proportion of effective tokens calculated during training.

3. The inference method for novel translation models based on RWKV according to claim 2, characterized in that: The specific method in step 1 is as follows: Step 1.1) Collect data on translations of the novel in different languages ​​and construct a bilingual or multilingual parallel corpus; Step 1.2) The original corpus is denoised, cleaned, segmented into sentences, segmented into words, and aligned to obtain standardized text samples that can be used for model training; Step 1.3) Employ dynamic MicroBatch stitching technology to divide the target sequence into text blocks with a fixed width, traverse permutations and combinations with a total length less than or equal to the sequence length, and use a simplex tableau dynamic solver to solve for the optimal number of uses for each stitching scheme, compressing multiple shorter text sequences into a longer training sequence, increasing the proportion of effective tokens in all tokens during training, and reducing the invalid computation caused by filling tokens.

4. The inference method for novel translation models based on RWKV according to claim 1, characterized in that: In step 2, the attention mechanism adopts a Cross-Attention architecture. By caching the original information of the Embedding layer as a KVCache, the memory consumption caused by the attention mechanism is reduced to 1 / 2N of the original, where N is the number of model layers.

5. The inference method for novel translation models based on RWKV according to claim 1, characterized in that: The specific method in step 3 is as follows: the input sequence is divided into blocks of length M, with parallel computation within each block and serial state transfer between blocks, thereby reducing the number of iterations from the sequence length N. N / M Meanwhile, a global scalar scaling γ is introduced in the TimeMix layer to process tensor values ​​by taking advantage of the scaling insensitivity of Normalization.

6. The inference method for novel translation models based on RWKV according to claim 1, characterized in that: The specific solution in step 4 is as follows: split the Lion optimizer onto inexpensive computing devices, and use the characteristics of the Sign function to compress gradient information into 2 bits for transmission, thereby reducing bandwidth consumption between high-performance devices and low-performance computing devices.

7. The inference method for novel translation models based on RWKV according to claim 1, characterized in that: The specific method in step 5 is as follows: After the model outputs the probability distribution of the lexical at each position, a joint decoder based on TopP, Temperature and Contrastive is used to sample the candidate lexical and generate the target language sequence; during the inference process, a maximum of n inference text sequences are retained, and each inference sequence is serially decoded token by token in round. When any subsequence decodes a termination token, the subsequence is removed and added to a new task to be processed; at the same time, a mask is applied to the key and value caches of different text sequences to shield the influence of historical caches from other text sequences.

8. A reasoning method system for implementing the RWKV-based novel translation model reasoning method according to any one of claims 1-7, characterized in that: It includes a data management module, a model training module, a backend inference module, and a user interaction module. The data management module is used for the extraction, cleaning, and dynamic splicing of parallel corpora, and generates standardized bilingual text sequences for model training. The model training module is used to build a translation model that integrates RWKV and attention mechanism, and to perform sublinear complexity hybrid parallel training, mixed precision training and heterogeneous distributed training. The backend inference module is used to load the trained novel translation model, execute the translation task based on the joint decoder and dynamic batch inference technology, and ensure cache isolation during the inference process of different text sequences through the KV mask mechanism. The user interaction module provides a visual interface for text input, parameter configuration, task submission, task status query, and translation result display.

9. An electronic device comprising a memory, a processor, and a computer program stored in the memory and executable on the processor, wherein the processor, when executing the computer program, implements the RWKV-based novel translation model reasoning method according to any one of claims 1-7.

Images