Reinforcement learning training acceleration system and method based on adaptive tensor parallelism

Abstract

This invention discloses a reinforcement learning training acceleration system and method based on adaptive tensor parallelism, relating to the field of artificial intelligence technology. It includes: an offline measurement module that constructs a system performance database before training begins; an inference execution module that executes the model's autoregressive inference task during the RLHF generation phase and is responsible for collecting the model's inference runtime data in real time; a latency prediction module that, based on the system performance database and runtime data, estimates the remaining generation time cost of the model under different candidate tensor parallelism degrees, as well as the switching time cost required to complete one tensor parallelism switch, in real time during the generation phase; and a switching control module that calculates the expected total latency under different candidate tensor parallelism degrees and, when a trigger condition is met, generates a switching control signal and sends it to the switching execution module. This acceleration system implements a mechanism for dynamic adjustment of tensor parallelism during the generation phase.

Description

Technical Field

[0001] This invention relates to the field of artificial intelligence technology, and in particular to a reinforcement learning training acceleration system and method based on adaptive tensor parallelism. Background Technology

[0002] With the widespread application of Large Language Models (LLMs) in natural language understanding, dialogue generation, and multimodal content generation, further improving the alignment between model outputs and human preferences has become a crucial research direction in the development of artificial intelligence. Reinforcement Learning with Human Feedback (RLHF), a widely adopted post-training paradigm in recent years, can re-optimize model behavior using human preference signals, making the model's generated results more aligned with human intentions and values. Therefore, RLHF has become a core technical approach in the model alignment stage.

[0003] Existing RLHF training workflows typically include three main stages: generation, scoring, and optimization. In the generation stage, the large language model to be optimized generates response text based on input prompts; in the scoring stage, a reward model evaluates the quality of the generated samples; and in the optimization stage, reinforcement learning algorithms (such as Proximal Policy Optimization, PPO) update the policy model parameters based on the reward signal. Theoretically, this workflow can continuously improve the quality and stability of the model's output, but in engineering implementation, the generation stage often becomes the system's performance bottleneck.

[0004] Currently, mainstream RLHF frameworks (such as VeRL and RLHFuse systems) typically employ a static parallelism strategy, using a fixed tensor parallelism and data parallelism configuration throughout the training process. This fixed structure can fully leverage the advantages of distributed parallel computing in the early generation stages when batch sizes are large. However, as the generation process progresses, some input samples terminate prematurely due to short response lengths, leaving only a small number of long-sequence samples for decoding, leading to a gradual decrease in the effective batch size. In this situation, the high-parallelism configuration originally optimized for large batches actually results in higher communication costs, significant idle GPU computing resources in the later stages, and a substantial decrease in overall system throughput.

[0005] Because RLHF employs an autoregressive decoding mechanism in its generation phase, the generation time is highly dependent on the response length distribution of the samples. A few long-sequence samples require several times longer generation times than the average length, thus dominating the overall batch completion time. This typical "long-tailed sample effect" leads to a significant decrease in GPU utilization in the later stages of training. Existing systems typically use batch inference, where all samples share a single synchronous generation step size, requiring the longest response sample to complete before proceeding to the next batch, resulting in underutilization of hardware resources.

[0006] To reduce data transfer and synchronization overhead between different training phases, some RLHF systems have introduced a "phase co-location" strategy, which involves simultaneously executing generation, reward evaluation, and policy optimization tasks on the same set of GPUs. While this strategy improves data utilization and reduces communication latency to some extent, the inefficient operation of the generation phase directly slows down the entire training cycle because the performance bottleneck of the generation phase remains unresolved, leading to a decrease in the overall efficiency of the RLHF training process.

[0007] Furthermore, the existing parallel mechanisms in the RLHF framework are primarily designed for the training phase, focusing on optimizing parameter updates and gradient synchronization, while lacking adaptive scheduling capabilities for load changes within the generation phase. Tensor parallel configurations are fixed at the start of training and cannot be dynamically adjusted based on batch size during the generation phase. As generation enters the tail phase, communication overhead and synchronization costs account for the majority of the overhead, severely limiting system performance and hindering the full utilization of hardware resources.

[0008] In summary, existing RLHF training systems generally suffer from the following problems: the generation phase lacks an adaptive parallelism adjustment mechanism, making it impossible to dynamically adjust resource allocation according to the actual task load; the decoding imbalance caused by long-tailed samples significantly reduces GPU computational efficiency; and the shared-stage architecture makes the generation phase performance bottleneck directly affect the overall training efficiency. These problems collectively contribute to bottlenecks such as low resource utilization, high generation latency, and long iteration cycles in reinforcement learning human feedback training.

[0009] Therefore, there is an urgent need to propose a mechanism that can dynamically adjust the tensor parallelism during the generation stage without changing the existing RLHF algorithm framework. This mechanism should be able to adaptively select the optimal parallel configuration based on real-time inference state and load changes, maximizing GPU utilization and significantly shortening generation time, thereby effectively improving the overall efficiency of reinforcement learning training. Summary of the Invention

[0010] Based on the technical problems existing in the background technology, this invention proposes a reinforcement learning training acceleration system and method based on adaptive tensor parallelism, which realizes a mechanism for dynamic adjustment of tensor parallelism in the generation stage. This mechanism should be able to adaptively select the optimal parallel configuration according to real-time inference state and load changes, maximize GPU utilization and significantly shorten generation time, thereby effectively improving the overall efficiency of reinforcement learning training.

[0011] The reinforcement learning training acceleration system based on adaptive tensor parallelism proposed in this invention includes:

[0012] The offline measurement module acquires the feature distribution of the model inference under different tensor parallelisms and stores it in the system performance database before training begins.

[0013] The inference execution module is used to execute the autoregressive inference task of the model during the RLHF generation phase, and is responsible for collecting the running status data of the model inference in real time.

[0014] The delay prediction module, based on the system performance database and running status data, estimates the remaining generation time cost of the model under the current tensor parallelism and different candidate tensor parallelism in real time during the generation phase, and estimates the switching time cost required to complete a tensor parallelism switch.

[0015] The switching control module calculates the expected total delay under different candidate tensor parallelism based on the remaining generation time overhead and switching time overhead. When the total delay corresponding to a candidate tensor parallelism is greater than the set safety margin, a switching control signal is generated and sent to the switching execution module.

[0016] The switching execution module activates the communication group corresponding to the candidate tensor parallelism based on the switching control signal, and aggregates and re-segments and distributes the model weights within the communication group through set communication operations to match the new tensor parallelism. The new tensor parallelism is then fed back to the inference execution module until all samples are generated.

[0017] Furthermore, the offline measurement module executes multiple sets of test samples by controlling the inference execution module, and records the decoding latency under different tensor parallel configurations and different batch sizes, which is then stored as a feature distribution in the system performance database.

[0018] Furthermore, the inference execution module performs generation operations step by step during the decoding phase, collects running status data at a preset period T, and reports it to the delay prediction module;

[0019] The operational status data includes the number of valid samples in the current batch, the average generated sequence length, and the longest generated sequence length.

[0020] Furthermore, in the delay prediction module, the estimation of the switching time overhead required to complete one tensor parallelism switch is specifically as follows:

[0021] Based on communication bandwidth Model weight data volume KV cache size Calculate the switching time cost from the current tensor parallelism to the candidate tensor parallelism. :

[0022] ;

[0023] in, This represents the processing time of the KV cache. ; This represents the time overhead of deleting the current configuration dependency variables and releasing video memory resources. ,coefficient Characterizes the time taken to release data per unit of storage capacity; Characterizes the time required to reinitialize the target tensor and the runtime environment under the target parallelism. , A latency factor is assigned to each unit of data, representing the average time taken to reallocate a unit of data in video memory. The fixed overhead for environment reset represents the fixed latency of resetting the CUDA runtime environment and switching communication groups; This characterizes the time overhead of resharding model weights across different devices; that is, the communication and copying time required to complete data aggregation and resharding based on the model weight data volume under the given communication bandwidth. .

[0024] Furthermore, the switching control module includes a decision-making unit and a control signal generation unit;

[0025] The decision-making unit calculates the estimated total delay based on the remaining generation time and switching time overhead, using a benefit-cost criterion. :

[0026] ;

[0027] in, The current tensor parallelism The remaining generation time overhead, Parallelism of candidate tensors The estimated remaining generation time cost, To reduce the current tensor parallelism Switch to candidate tensor parallelism Required switching time overhead;

[0028] The control signal generation unit satisfies the trigger condition in one of the candidate tensor parallelisms. At that time, a switching control signal is generated and a switching command is sent to the switching execution module to realize real-time adjustment of the parallel configuration. This is used to set a safety margin, which is intended to prevent the system from switching too frequently.

[0029] Furthermore, the switching execution module includes a communication group management unit, a weight resharding unit, and a KV cache management unit;

[0030] The communication group management unit is used to store and activate the communication group corresponding to the target parallelism.

[0031] The weight resharding unit is used to repartition and distribute model weights based on existing weight copies in video memory after the target communication group is activated, in order to match the new parallel configuration.

[0032] Furthermore, the acceleration system also includes a KV cache management unit for maintaining the consistency of the key-value cache of generated samples during tensor parallelism switching.

[0033] Furthermore, the KV cache management unit is also used for:

[0034] The size of the key-value cache will be migrated as needed during the weighted resharding process. With currently available communication bandwidth The ratio of these values ​​is used as the migration delay;

[0035] When the migration latency is less than the time required for recalculation on the target fragment, and the blocking time caused by the migration is not greater than the set blocking threshold, direct migration is selected.

[0036] Otherwise, recompile the KV cache on demand on the target parallel group and update the KV cache pointer and offset mapping for each sample during migration or recompile.

[0037] Furthermore, after completing weighted resharding and cache synchronization, the system enters the redistribution phase;

[0038] The redistribution phase, based on the new parallelism configuration, redistributes the aggregated samples to GPU devices within the target communication group;

[0039] Each GPU device resumes request execution, and the inference execution module continues autoregressive decoding under the new parallel configuration until all samples are generated.

[0040] An adaptive tensor parallel reinforcement learning training acceleration system includes:

[0041] Before training begins, the feature distribution of model inference under different tensor parallelism is obtained through the offline measurement module and stored in the system performance database;

[0042] The inference execution module performs the autoregressive inference task of the model during the RLHF generation stage and is responsible for collecting the running status data of the model inference in real time.

[0043] Based on the system performance database and runtime status data, the delay prediction module estimates the remaining generation time cost of the model under the current tensor parallelism and different candidate tensor parallelisms in real time during the generation phase, and estimates the switching time cost required to complete a tensor parallelism switch.

[0044] The switching control module calculates the expected total delay under different candidate tensor parallelism based on the remaining generation time overhead and switching time overhead. When the total delay corresponding to a candidate tensor parallelism is greater than the set safety margin, a switching control signal is generated and sent to the switching execution module.

[0045] The switching execution module activates the communication group corresponding to the candidate tensor parallelism based on the switching control signal, and aggregates and re-segments and distributes the model weights within the communication group through set communication operations to match the new tensor parallelism. The new tensor parallelism is then fed back to the inference execution module until all samples are generated.

[0046] The advantages of the adaptive tensor parallelism-based reinforcement learning training acceleration system and method provided by this invention are as follows: By introducing an adaptive tensor parallelism switching mechanism in the RLHF generation stage, dynamic acceleration of the inference process is achieved, which can significantly improve the overall efficiency of reinforcement learning training while maintaining model accuracy and training consistency. The switching control module calculates the expected total latency under different tensor parallelism levels based on the prediction results of the latency prediction module and evaluates the difference between the switching benefit and the switching cost. The switching operation is triggered only when the benefit (i.e., the expected total latency) exceeds a set safety margin, thereby avoiding performance fluctuations caused by frequent switching. This strategy can effectively improve GPU utilization in the long-tail sample stage. Through pre-caching communication groups, dynamic selection of KV cache, and in-memory weight resharding mechanism, switching can be completed without restarting the inference engine, and the switching latency can be controlled within a few seconds. Attached Figure Description

[0047] Figure 1 This is a schematic diagram of the structure of the present invention;

[0048] Figure 2 A schematic diagram illustrating the tensor parallelism execution provided for existing SIE methods;

[0049] Figure 3 This is a schematic diagram of tensor parallelism switching provided in this embodiment;

[0050] Figure 4 The overall system flowchart provided in this embodiment. Detailed Implementation

[0051] The technical solution of the present invention will now be described in detail through specific embodiments. Many specific details are set forth in the following description to provide a thorough understanding of the invention. However, the present invention can be implemented in many other ways different from those described herein, and those skilled in the art can make similar modifications without departing from the spirit of the invention. Therefore, the present invention is not limited to the specific embodiments disclosed below.

[0052] like Figures 1 to 4 As shown, the reinforcement learning training acceleration system based on adaptive tensor parallelism proposed in this invention includes:

[0053] The offline measurement module acquires the feature distribution of the model inference under different tensor parallelisms and stores it in the system performance database before training begins.

[0054] The inference execution module is used to execute the autoregressive inference task of the model during the RLHF generation phase, and is responsible for collecting the running status data of the model inference in real time.

[0055] The delay prediction module, based on the system performance database and running status data, estimates the remaining generation time cost of the model under the current tensor parallelism and different candidate tensor parallelism in real time during the generation phase, and estimates the switching time cost required to complete a tensor parallelism switch.

[0056] The switching control module calculates the expected total delay under different candidate tensor parallelism based on the remaining generation time overhead and switching time overhead. When the total delay corresponding to a candidate tensor parallelism is greater than the set safety margin, a switching control signal is generated and sent to the switching execution module.

[0057] The switching execution module activates the communication group corresponding to the candidate tensor parallelism based on the switching control signal, and aggregates and re-segments and distributes the model weights within the communication group through set communication operations to match the new tensor parallelism. The new tensor parallelism is then fed back to the inference execution module until all samples are generated.

[0058] This embodiment addresses the technical problems in existing RLHF training processes, such as low GPU resource utilization, high generation latency, and limited overall iteration efficiency due to long-tailed samples in the generation phase. It proposes a solution that dynamically adjusts the parallelism configuration of the large language model generation phase at runtime. This solution can dynamically adjust the parallelism configuration of the large language model generation phase without changing the existing reinforcement learning framework, thus resolving the technical problems of idle GPU resources, high generation latency, and low overall iteration efficiency caused by long-tailed samples in existing RLHF training. A mechanism for dynamically adjusting the tensor parallelism in the generation phase is implemented. This mechanism should adaptively select the optimal parallel configuration based on real-time inference state and load changes, maximizing GPU utilization and significantly shortening generation time, thereby effectively improving the overall efficiency of reinforcement learning training.

[0059] The acceleration system in this embodiment includes an offline measurement module, an inference execution module, a latency prediction module, a switching control module, and a switching execution module. Communication channels based on message queues or shared memory mechanisms are established between these modules.

[0060] (1) Offline measurement module;

[0061] The offline measurement module performs multiple rounds of performance testing during the initialization phase to obtain the feature distribution of the model under different tensor parallelisms. This establishes a performance benchmark system for specific models and hardware environments, providing data support for subsequent latency prediction modules and switching decisions. Through the offline modeling process, performance profiling can be completed without affecting the formal training task, avoiding additional interference from online measurements to the inference task and ensuring the objectivity and repeatability of performance evaluation results.

[0062] (2) Delay prediction module;

[0063] The latency prediction module is used in the generation phase to estimate the remaining generation time cost of the model in real time under the current tensor parallelism. This module comprehensively utilizes runtime status data fed back from the inference execution module. Simultaneously, it combines this data with the system performance database (i.e., performance profiling data) established by the offline measurement module to dynamically estimate the inference latency under different tensor parallelisms. Furthermore, the latency prediction module can also adjust the latency based on the current system's communication bandwidth. Model weight data volume and KV cache size Estimate the switching time cost required to complete one tensor parallelism switch. This provides information support for subsequent switching decisions.

[0064] Among them, the switching time overhead The calculation formula is:

[0065] ;

[0066] in, This represents the processing time of the KV cache. ; This represents the time overhead of deleting the current configuration dependency variables and releasing video memory resources. ,coefficient Characterizes the time taken to release data per unit of storage capacity. Characterizes the time required to reinitialize the target tensor and the runtime environment under the target parallelism. This characterizes the time overhead of resharding model weights across different devices; that is, the communication and copying time required to complete data aggregation and resharding based on the model weight data volume under the given communication bandwidth. .

[0067] It is understandable; It is linearly related to memory release operations, in the formula In the middle, coefficient The performance of the graphics memory management system is negatively correlated with the performance of the computing device; that is, the stronger the device performance, the higher the coefficient. The smaller the coefficient It can be obtained through a single measurement run; The time represented specifically covers the time spent resetting the CUDA runtime environment, reallocating KV cache space, and changing the communication group. , A latency factor is assigned to each unit of data, representing the average time taken to reallocate a unit of data in video memory. Fixed overhead for environment reset, representing the fixed latency of resetting the CUDA runtime environment and switching communication groups.

[0068] Key-value caching is a caching mechanism used to accelerate inference in large language models. It improves performance by storing intermediate results (keys and values) from attention mechanisms to reduce redundant computations. CUDA Graph is a technology introduced by NVIDIA to optimize GPU computing efficiency. It significantly reduces communication overhead between the CPU and GPU by packaging a series of GPU operations into a static graph structure and submitting them for execution all at once.

[0069] Specifically, the latency prediction module combines current operational status data with an offline measured system performance database, and fits the latency using a linear interpolation algorithm. This model can accurately predict the estimated decoding time under different candidate tensor parallelism levels, ensuring the correctness of switching decisions.

[0070] (3) Switch control modules;

[0071] The switching control module is used to calculate the expected latency under different candidate tensor parallelisms based on the remaining generation time overhead (i.e., the expected decoding time) and switching time overhead provided by the latency prediction module, and to comprehensively evaluate the performance gains that may be obtained before and after the switching.

[0072] The switching control module uses quantitative analysis to determine whether to perform a tensor parallelism switching operation, so as to select the optimal parallel strategy during the generation process, thereby achieving efficient utilization of system resources and improving overall operating efficiency.

[0073] Specifically, the switching control module includes a decision-making unit and a control signal generation unit. The decision-making unit calculates the expected total delay based on the remaining generation time and switching time overhead, using a benefit-cost criterion. This allows for the selection of the optimal parallel configuration to guide the adjustment of the system's parallelism. The control signal generation unit, upon meeting the triggering conditions... At that time, a switching control signal is generated and a switching command is sent to the switching execution module to realize real-time adjustment of parallel configuration and optimization of system operating efficiency. This is the safety margin set.

[0074] Among them, total delay The calculation formula is as follows:

[0075] ;

[0076] in, The current tensor parallelism The remaining generation time overhead, Parallelism of candidate tensors The estimated remaining generation time cost, To reduce the current tensor parallelism Switch to candidate tensor parallelism Required switching time overhead.

[0077] That is, when the decision-making unit determines that the predicted current remaining time is less than the sum of the switching time overhead and the remaining time after the switching, the system maintains the current parallel configuration and does not perform a switching operation. When the decision-making unit determines that the predicted current remaining time is greater than the sum of the switching time overhead and the remaining time after the switching, the control signal generation unit generates a switching control signal and sends it to the switching execution module.

[0078] By coordinating the decision-making unit and the control signal generation unit, a joint assessment of the benefits and costs of switching actions can be achieved, ensuring that parallelism switching only occurs when performance improvement is substantial. This collaborative mechanism enhances the stability and controllability of adaptive tensor parallel scheduling, avoiding resource waste caused by frequent switching.

[0079] (4) Inference execution module;

[0080] The inference execution module performs autoregressive inference tasks on the large language model during the RLHF generation phase. This module is responsible for collecting real-time execution status information of the model inference, including the number of valid samples in the current batch, the average generated sequence length, and the longest generated sequence length, and reporting this information to the switching control module at set intervals. By continuously monitoring the execution status, the inference execution module provides precise data support to the switching control module, thereby assisting the system in dynamically selecting the optimal parallel strategy and achieving efficient scheduling and resource management of the generation process.

[0081] The inference execution module employs a distributed communication library (such as NCCL or GLOO) to achieve cross-device data synchronization and collaborative computing. The module supports parallel generation across multiple GPUs through operations such as broadcasting, reduction, and full connection, ensuring data consistency and computational efficiency under different parallel configurations.

[0082] When the system performs a tensor parallelism switching operation, the inference execution module pauses the allocation of new requests and merges currently incomplete requests into new batches to ensure the consistency of task status and generation context. After the switching operation is completed, the system resumes request execution, achieving a smooth adjustment of parallelism while ensuring the continuity of the generation process and the correctness of computation.

[0083] (5) Switch execution modules;

[0084] After receiving the switching signal from the switching control module, the switching execution module is responsible for performing the reallocation of the tensor parallel communication group, the re-sharding of model weights, and the migration of cached data, thereby realizing online switching of tensor parallelism.

[0085] The switching execution module internally includes a communication group management unit and a weight resharding unit. The communication group management unit is responsible for storing and activating the communication group corresponding to the target parallelism, ensuring smooth data communication across devices. The weight resharding unit is used to repartition and distribute the model weights based on the existing weight copies in the video memory after activating the target communication group, in order to match the new parallel configuration.

[0086] The communication group management unit pre-builds all possible tensor parallel communication groups during system initialization and caches them in memory. When switching tensor parallelism, the switching execution module can directly call the cached communication group information without rebuilding the communication groups, thereby significantly reducing switching latency and improving the efficiency and response speed of the system's online dynamic adjustment of parallel configuration.

[0087] This embodiment employs an in-memory weight resharding mechanism (i.e., weight resharding unit). During the switching of tensor parallelism, the weight copies in the GPU memory are directly used for repartitioning and synchronous distribution, without the need to reload the weight file from external storage, thereby significantly reducing I / O overhead and synchronization waiting time.

[0088] The weight resharding operation is implemented through collective communication. Each GPU device performs block transmission according to the new tensor parallelism partitioning scheme, thereby redistributing the model weights among the devices.

[0089] In one embodiment, the acceleration system further includes a KV cache management unit. This unit maintains the consistency of the generated sample KV cache during tensor parallelism switching. The KV cache is an intermediate state generated and continuously maintained during the generation phase, used to reuse historical attention calculation results, thereby significantly reducing the computational overhead of subsequent generation processes. The KV cache management unit manages the KV cache during the switching process. Furthermore, the KV cache management unit can maintain the KV cache within the original KV cache area, unifying the KV cache before and after the switch within the same KV cache area. Therefore, it is not necessary to create a new KV cache area, eliminating the problem of memory fragmentation. The KV cache management unit dynamically selects a migration strategy or a recalculation strategy based on the system network bandwidth and cache capacity to ensure the integrity of cached data and computational continuity during the switching process, while optimizing data transmission overhead and improving the overall system operating efficiency.

[0090] When network bandwidth is sufficient, the KV cache management unit adopts a direct migration method to completely migrate the key-value cache of the generated samples to the target device; when network bandwidth is limited, the unit selects a recalculation strategy to ensure a balance between computational continuity and system efficiency during the switching process.

[0091] The acceleration system in this embodiment can be deployed in distributed GPU clusters or single-machine multi-GPU environments. The system automatically identifies the available tensor parallelism range through hardware topology detection and performs dynamic switching operations only within the allowed range, thereby ensuring efficient utilization of communication bandwidth and optimizing the allocation of computing resources to achieve flexible and efficient parallel generation or training.

[0092] Furthermore, the acceleration system in this embodiment is based on the SGLang inference engine (i.e., a large model inference engine), which is compatible with existing RLHF frameworks (such as VeRL). The system adopts a plug-in interface design, which can realize functional expansion and module integration without modifying the original training process, thereby providing flexible and scalable inference and scheduling capabilities.

[0093] To ensure efficient data flow, this embodiment adopts a single inference engine mode. The single inference engine mode achieves data continuity and computational consistency during the generation phase by uniformly managing generation tasks, model parameters, key-value caches, and communication status within a single inference process. Unlike multi-engine distributed deployment, the single inference engine mode eliminates the need for data synchronization between engines, significantly reducing communication synchronization and data copying overhead.

[0094] In single inference engine mode, the inference execution module is responsible for centrally scheduling the execution of all generation requests. During the initialization phase, the system builds a unified communication group management module to maintain consistent communication channels and topology information across different tensor parallel configurations. The weights of each GPU device are re-sharded and remapped within the GPU memory. The inference execution module uses collective communication instructions to synchronously update the sharded weights, avoiding I / O latency caused by reloading model files from external storage.

[0095] The aggregation process of generated results in the single inference engine mode is more efficient. Since the system uses only a single inference engine for centralized management, there is no need to perform aggregate communication between multiple inference execution modules to summarize the generated results, thus significantly reducing communication latency and synchronization overhead.

[0096] Example 1;

[0097] Methods for accelerating reinforcement learning training based on adaptive tensor parallelism include:

[0098] Step 1: Initialize the RLHF training framework: This step establishes the basic runtime environment for the RLHF training process. The system loads the pre-trained large language model parameters, reward model, and policy optimization module. During initialization, the framework registers the scheduling policies for each stage of the task, including the generation stage, reward evaluation stage, and policy update stage. The RLHF framework can be built on top of existing VeRL systems to maintain compatibility with existing reinforcement learning interfaces.

[0099] Step 2: Initialize the Training Engine: This step configures the distributed training environment, establishing data parallelism, pipeline parallelism, and tensor parallelism. Model parameters, optimizer states, and training batch data are mapped to the GPU, and communication dependencies are automatically calculated to reduce AllReduce and memory copying overhead. Simultaneously, the optimizer and learning rate scheduler are initialized for subsequent updates to the target model.

[0100] Step 3: This step establishes the inference environment for the large language model generation stage. The system first loads the inference execution module (e.g., SGLang) and synchronizes the model parameters loaded in the training engine to the inference execution module. Subsequently, the model parameters are sharded and allocated according to the preset tensor parallelism. Simultaneously, the system captures the CUDA computation graph and initializes the KV cache to accelerate the generation process. During the inference engine initialization phase, a system performance database is also built to provide reference information for the latency prediction module in the generation stage, used to predict generation latency in advance. The CUDA computation graph is the core model in the NVIDIA CUDA architecture used to optimize parallel computing. Its core idea is to decompose computational tasks into reusable graph structures to achieve efficient execution. CUDA (Compute Unified Device Architecture) is a parallel computing architecture and programming model introduced by NVIDIA, designed to leverage the powerful parallel computing capabilities of GPUs (Graphics Processing Units) to accelerate general-purpose computing tasks, not just graphics processing.

[0101] Step 4: Adaptive Tensor Parallel Generation: This step is the core of this embodiment. It involves online evaluation and dynamic adjustment of the tensor parallelism of the large language model during the generation phase of RLHF training. This reduces the dominant effect of long-tailed samples on batch completion time and improves GPU resource utilization. This step is collaboratively completed by the offline measurement module, inference execution module, latency prediction module, switching control module, switching execution module, communication group cache unit, and KV cache management unit.

[0102] Before the reinforcement learning human feedback training begins, the offline measurement module samples the performance of the large language model under different tensor parallel configurations. The offline measurement module controls the inference execution module to execute multiple sets of test samples, recording the decoding latency under different degrees of parallelism and batch sizes. The obtained performance data is stored in the system performance database as reference information for the online prediction phase.

[0103] like Figure 4 As shown, when the generation phase begins, the inference execution module performs a one-time key-value calculation on each input context before generation begins and establishes a KV cache in local video memory or addressable cache. The inference execution module records the initial context pointer and the initial value of the number of generation steps for each sample, for use in subsequent decoding and possible migration.

[0104] like Figure 4 As shown, the generation phase follows an autoregressive generation pattern. The inference execution module performs generation operations step by step during the decoding phase, according to a preset cycle. Collect operational status data, that is, be able to generate each After 1 character, collect status data ( (Typically set to 1). The collected operational status data includes the number of valid samples in the current batch, the number of steps generated for each sample, the average generated sequence length, and the longest generated sequence length. The inference execution module reports this operational status data to the delay prediction module and the switching control module.

[0105] like Figure 1 As shown, the delay prediction module takes the running status data and calls the system performance database of the offline measurement module as input data, and performs parallelism calculations on multiple candidate tensors. Parallel computing estimates remaining generation time overhead The delay prediction module outputs the estimated remaining generation time cost for the parallelism of each candidate tensor. The results are then provided to the switching control module.

[0106] The latency prediction module estimates the current tensor parallelism based on communication bandwidth, model weight data volume, KV cache size, and the fixed overhead of switching initialization. Switch to candidate tensor parallelism The required switching time overhead .

[0107] The switching control module is used to receive the output of the delay prediction module. and The switching control module calculates the expected total delay based on the benefit-cost criterion. , The switching control module controls the parallelism of one of the candidate tensors. Meet the triggering conditions When switching is enabled, a switching control signal is generated and a switching command is sent to the switching execution module to achieve real-time adjustment of the parallel configuration; otherwise, the current parallel configuration is maintained and execution continues. This is the safety margin set.

[0108] like Figure 1 As shown, the switching control module balances the switching time overhead. With remaining generation time overhead The system determines whether a tensor parallel switching mechanism is triggered. If the determination is successful, the switching control module sends a switching signal to the switching execution module on each GPU, thereby driving subsequent tensor parallel state transitions.

[0109] like Figure 2 and 3 The diagram shown illustrates the switching between the existing SIE method and the tensor parallelism of this embodiment. Figure 2The existing SIE method's generation phase execution flow often relies on mature inference engines, such as vLLM (large language model inference engine) or SGLang (source: SGLang: Efficient Execution of Structured Language Model Programs), and these engines lack phase awareness capabilities. Therefore, in this process, the parallel group uses the same resources and statically identical tensor parallel configuration in both the alignment and tail phases. However, the adaptive tensor parallelism of this embodiment can modify the resources used by the parallel group by switching the tensor parallelism method between the alignment and tail phases. Specifically, in... Figure 3 In the process, after receiving the switching information, the switching execution module loads the communication topology information corresponding to the target parallelism from the communication group cache unit and activates the corresponding communication group. The communication group cache unit pre-builds and stores multiple available candidate tensor parallelisms during the initialization phase. The values ​​correspond to communication mappings to avoid synchronization blocking caused by rebuilding communication groups at runtime. In addition, during the switch, incomplete data from different parallel groups is merged into the new parallel group and continues to be executed. This method can speed up the execution time of the tail stage and improve throughput.

[0110] The weight resharding unit is used to repartition the model weights based on the existing weight copies in GPU memory after activating the target communication domain. The repartitioning can adopt a block-based or row / column-based sharding strategy and complete the cross-device data redistribution through collective communication operations (such as All-to-All, Reduce-Scatter, and All-Gather). All-to-All is a many-to-many collective communication operation; Reduce-Scatter is a collective communication operation for distributed deep learning models, whose core function is to split the data and distribute it to multiple nodes for processing, and then return the aggregation result to all nodes; All-Gather is a collective communication operation used to collect data from each process into all processes.

[0111] The weighted resharding unit maintains the source-target sharding index mapping and confirms the resharding is complete after all devices have completed the index consistency check.

[0112] The KV cache management unit is used to maintain the consistency of the overall KV cache. Specifically, it will migrate the KV cache size as needed during weighted resharding. With currently available communication bandwidth The ratio of the two values ​​is used as the migration delay. ,Right now When migration delay If the time required for recompiling on the target slice is less than the time required for recompiling, and the blocking time caused by migration is not greater than the set blocking threshold, then direct migration is selected; otherwise, the KV cache is recomputed on demand on the target parallel group, and the KV cache pointer and offset mapping of each sample are updated during migration or recompiling.

[0113] In addition to the switching operations mentioned above, in Figure 3 The sample merging process shown also includes a sample recovery step. This step is executed in conjunction with the KV cache migration process to reconstruct the decoding state of the target samples, thereby avoiding the additional overhead required to rebuild the KV cache. After the recovery is completed, the inference engine can continue to perform autoregressive decoding in the updated parallel configuration until the generation process of all samples is finished.

[0114] Step 5: Train the model based on the generated results: After the adaptive tensor parallel generation stage is completed, reward calculation and policy optimization are performed on the generated samples, thereby achieving closed-loop update of RLHF training. This stage establishes an efficient data interaction channel between the generation and training engines, ensuring that the generated results can be quickly transmitted back and used in model updates.

[0115] In this embodiment, after the generation phase is completed, the inference execution module submits all generated response samples and their corresponding context information to the reward evaluation module. The reward evaluation module uses a pre-trained reward model to score the quality of the generated samples. The input to the reward model includes the prompt text, the generated text, and the reference sample, and the output is a numerical reward signal reflecting the degree of human preference. The reward signal is used to measure the quality of the generated samples and serves as the optimization objective in the reinforcement learning phase.

[0116] After the reward evaluation is completed, the system will generate samples, reward signals, and corresponding policy probabilities, package them into training data batches, and send them back to the policy optimization module of the training engine. This module calculates the policy gradient and updates the model parameters based on the Proximal Policy Optimization (PPO) algorithm or an equivalent reinforcement learning algorithm.

[0117] After the model parameters are updated, the system synchronizes the updated weights back to the inference engine to ensure that the next generation phase uses the same parameter state as the training phase. This synchronization process is completed through collective communication or broadcasting.

[0118] For example, to verify the effectiveness of the adaptive tensor parallel mechanism in this embodiment, performance tests were conducted on a reinforcement learning human feedback training task in a multi-GPU cluster environment. The test platform used 8×A40 GPU servers, with 8 billion model parameters, and employed a mixed-precision training mode.

[0119] The training dataset is derived from a long text writing task, containing prompts and answers. It is used to construct the reward model for the reinforcement learning phase and is divided into training and validation sets at a ratio of 80% to 20%. To verify the effectiveness of this invention, different maximum generation lengths were selected for testing in this embodiment, set to 2048, 4096, and 6144 respectively, to evaluate the performance of this invention at different generation scales.

[0120] The testing process compares and evaluates a fixed parallel configuration with the adaptive tensor parallel configuration proposed in this invention. To accelerate the generation phase, existing methods typically use the smallest possible tensor parallelism based on model size and GPU memory capacity. Therefore, in this test, the existing SIE method selects the minimum tensor parallelism that can support the model, i.e., the tensor parallelism degree. Solution; the adaptive configuration in this embodiment is in Dynamic switching within the range. Evaluation metrics include average generation throughput (number of words generated per second).

[0121] Table 1 shows the evaluation results of the adaptive tensor parallelism and existing SIE methods in this embodiment at different lengths. The evaluation metric is end-to-end throughput. The results show that the adaptive tensor parallelism mechanism proposed in this embodiment can improve the overall training throughput under the same hardware conditions.

[0122] Table 1 Evaluation Results

[0123]

[0124] This embodiment has the following significant advantages:

[0125] This embodiment enables efficient allocation of computing resources and optimization of inference performance during reinforcement learning human feedback training (RLHF) by dynamically adjusting the parallel configuration of the generation phase through an adaptive tensor parallel mechanism. This mechanism solves the problems of idle GPU resources caused by long-tailed samples, excessive latency in the generation phase, and limited training iteration speed in existing technologies.

[0126] In this embodiment, the switching control module calculates the expected total latency under different tensor parallelisms based on the prediction results of the latency prediction module, and evaluates the difference between the switching benefits and switching costs. Switching is triggered only when the benefits (i.e., the expected total latency) exceed a set safety margin, thereby avoiding performance fluctuations caused by frequent switching. This strategy can effectively improve GPU utilization during the long-tail sample stage, increasing the overall training throughput by approximately 58.8%.

[0127] This embodiment achieves switching without restarting the inference engine by using a pre-cached communication group, dynamic selection of KV cache, and a weight resharding mechanism within video memory, and the switching latency can be controlled within a few seconds.

[0128] Based on the data in Table 1, this embodiment can reduce the average generation latency by 42.7% in typical RLHF long-tail tasks and achieve an end-to-end training throughput improvement of approximately 1.6 times under the same hardware conditions.

[0129] The technical solution of this embodiment has high compatibility and scalability, and can be directly integrated into mainstream RLHF frameworks (such as VeRL), making it suitable for cloud clusters, local multi-GPU servers, and hybrid computing environments.

[0130] This embodiment introduces an adaptive tensor parallel switching mechanism in the RLHF generation stage, which dynamically accelerates the inference process and can significantly improve the overall efficiency of reinforcement learning training while maintaining model accuracy and training consistency.

[0131] The above description is only a preferred embodiment of the present invention, but the scope of protection of the present invention is not limited thereto. Any equivalent substitutions or modifications made by those skilled in the art within the scope of the technology disclosed in the present invention, based on the technical solution and inventive concept of the present invention, should be covered within the scope of protection of the present invention.

Claims

1. A reinforcement learning training acceleration system based on adaptive tensor parallelism, characterized in that, include: The offline measurement module acquires the feature distribution of the model inference under different tensor parallelisms and stores it in the system performance database before training begins. The inference execution module is used to execute the autoregressive inference task of the model during the RLHF generation phase, and is responsible for collecting the running status data of the model inference in real time. The delay prediction module, based on the system performance database and runtime status data, estimates the remaining generation time cost of the model in real time during the generation phase under the current tensor parallelism and different candidate tensor parallelisms. It also estimates the switching time cost required to complete one tensor parallelism switch. Specifically: Based on communication bandwidth Model weight data volume KV cache size Calculate the switching time cost from the current tensor parallelism to the candidate tensor parallelism. : ； in, This represents the processing time of the KV cache. ; This represents the time overhead of deleting the current configuration dependency variables and releasing video memory resources. ; Characterizes the time required to reinitialize the target tensor and the runtime environment under the target parallelism. ; Represents the time cost of resplitting model weights across different devices. ,coefficient Characterizes the time taken to release data per unit of storage capacity. A latency factor is assigned to each unit of data, representing the average time taken to reallocate a unit of data in video memory. The fixed overhead for environment reset represents the fixed latency of resetting the CUDA runtime environment and switching communication groups; The switching control module calculates the expected total delay under different candidate tensor parallelism based on the remaining generation time overhead and switching time overhead. When the total delay corresponding to a candidate tensor parallelism is greater than the set safety margin, a switching control signal is generated and sent to the switching execution module. The switching execution module activates the communication group corresponding to the candidate tensor parallelism based on the switching control signal, and aggregates and re-segments and distributes the model weights within the communication group through set communication operations to match the new tensor parallelism. The new tensor parallelism is then fed back to the inference execution module until all samples are generated.

2. The acceleration system according to claim 1, characterized in that, The offline measurement module executes multiple sets of test samples by controlling the inference execution module, and records the decoding latency under different tensor parallel configurations and different batch sizes, which is then stored as a feature distribution in the system performance database.

3. The acceleration system according to claim 1, characterized in that, The inference execution module performs generation operations step by step during the decoding phase, collects running status data at a preset period T and reports it to the delay prediction module. The operational status data includes the number of valid samples in the current batch, the average generated sequence length, and the longest generated sequence length.

4. The acceleration system according to claim 1, characterized in that, The switching control module includes a decision-making unit and a control signal generation unit; The decision-making unit calculates the estimated total delay based on the remaining generation time and switching time overhead, using a benefit-cost criterion. : ； in, The current tensor parallelism The remaining generation time overhead, Parallelism of candidate tensors The estimated remaining generation time cost, To reduce the current tensor parallelism Switch to candidate tensor parallelism Required switching time overhead; The control signal generation unit satisfies the trigger condition in one of the candidate tensor parallelisms. At that time, a switching control signal is generated and a switching command is sent to the switching execution module to realize real-time adjustment of the parallel configuration. To set a safety margin.

5. The acceleration system according to claim 1, characterized in that, The switching execution module includes a communication group management unit, a weighted resharding unit, and a KV cache management unit; The communication group management unit is used to store and activate the communication group corresponding to the target parallelism. The weight resharding unit is used to repartition and distribute model weights based on existing weight copies in video memory after the target communication group is activated, in order to match the new parallel configuration.

6. The acceleration system according to claim 1, characterized in that, The acceleration system also includes a KV cache management unit, which is used to maintain the consistency of the key-value cache of generated samples during the tensor parallelism switching process.

7. The acceleration system according to claim 6, characterized in that, The KV cache management unit is also used for: The size of the key-value cache will be migrated as needed during the weighted resharding process. With currently available communication bandwidth The ratio of these values ​​is used as the migration delay; When the migration latency is less than the time required for recalculation on the target fragment, and the blocking time caused by the migration is not greater than the set blocking threshold, direct migration is selected. Otherwise, recompile the KV cache on demand on the target parallel group and update the KV cache pointer and offset mapping for each sample during migration or recompile.

8. The acceleration system according to claim 1, characterized in that, After completing weighted resharding and cache synchronization, the system enters the redistribution phase; The redistribution phase, based on the new parallelism configuration, redistributes the aggregated samples to GPU devices within the target communication group; Each GPU device resumes request execution, and the inference execution module continues autoregressive decoding under the new parallel configuration until all samples are generated.

9. The acceleration method of the acceleration system according to claim 1, characterized in that, include: Before training begins, the feature distribution of model inference under different tensor parallelism is obtained through the offline measurement module and stored in the system performance database; The inference execution module performs the autoregressive inference task of the model during the RLHF generation stage and is responsible for collecting the running status data of the model inference in real time. Based on the system performance database and runtime status data, the delay prediction module estimates the remaining generation time cost of the model under the current tensor parallelism and different candidate tensor parallelisms in real time during the generation phase, and estimates the switching time cost required to complete a tensor parallelism switch. The switching control module calculates the expected total delay under different candidate tensor parallelism based on the remaining generation time overhead and switching time overhead. When the total delay corresponding to a candidate tensor parallelism is greater than the set safety margin, a switching control signal is generated and sent to the switching execution module. The switching execution module activates the communication group corresponding to the candidate tensor parallelism based on the switching control signal, and aggregates and re-segments and distributes the model weights within the communication group through set communication operations to match the new tensor parallelism. The new tensor parallelism is then fed back to the inference execution module until all samples are generated.

Images