Federal learning acceleration method based on parallel sampling and training of in-batch real-time data of assembly line
By using a pipelined parallel mechanism and weighted training with gradient correction coefficients, the problem of limited computing power on edge devices is solved, improving model accuracy and hardware resource utilization, and achieving accelerated federated learning.
Patent Information
- Application Number
- CN202511793437.X
- Authority / Receiving Office
- CN · China
- Patent Type
- Applications(China)
- Current Assignee / Owner
- Filing Date
- 2025-12-01
- Publication Date
- 2026-03-17
AI Technical Summary
In existing technologies, the limited computing power and massive data volume of edge devices lead to slow training convergence speed, static data sampling causes outdated sample importance and gradient bias, and the utilization rate of hardware resources is low.
A pipelined parallel mechanism is adopted to evaluate the importance of samples in real time and perform weighted training by gradient correction coefficients. Combined with mixed precision inference and a circular buffer, hardware-level concurrency of sampling and training is achieved.
It improved model training speed, enhanced model accuracy, and maximized hardware resource utilization, achieving a training acceleration of 1.5 to 2.5 times.
Smart Images
Figure CN121684100A_ABST
Abstract
Description
TECHNICAL FIELD
[0001] The present application relates to the field of computer technology and artificial intelligence, and particularly relates to a distributed machine learning method and system for accelerating federated learning training and improving model accuracy on resource-constrained edge devices. BACKGROUND
[0002] With the popularity of Internet of Things devices and the generation of a large amount of edge data, federated learning, as a distributed machine learning paradigm that protects privacy, allows edge devices to collaboratively train models without sharing raw data. However, in practical applications, there is a contradiction between the limited computing power (such as limited CPU and GPU performance) of edge devices and the large and redundant amount of local data, which seriously slows down the training convergence speed. Although existing technologies attempt to accelerate training by static data sampling (i.e., selecting a subset based on sample importance before the start of training rounds), this method has significant defects: on the one hand, static sampling causes the sample importance indicator to become outdated with rapid model iterations, failing to accurately reflect the true value of the sample to the current model; on the other hand, importance sampling changes the original data distribution, introducing gradient estimation bias, thereby reducing the final accuracy of the model. In addition, if the sampling and training tasks are simply executed in series at the system level, not only will additional inference overhead be introduced, but also the utilization of hardware resources will be low, and the existing deep learning framework is also inefficient when processing index backpropagation. Therefore, there is an urgent need for a federated learning acceleration solution that can evaluate sample importance in real time, effectively correct gradient bias, and fully utilize end-side heterogeneous computing resources through an efficient parallel mechanism. SUMMARY
[0003] The purpose of the present application is to provide a federated learning acceleration method based on pipeline batch real-time data sampling and training parallelism, aiming to solve the problems of outdated sample importance, gradient bias, and low utilization of computing resources in existing technologies, thereby accelerating training while improving model accuracy.
[0004] Technical solution: A federated learning acceleration method with data sampling and training parallelism, specifically comprising the following steps:
[0005] (1) Construct a pipeline parallel environment: establish a parallel execution sampling task flow and a training task flow on the computing unit of the client device, and allocate shared storage space for data exchange.
[0006] (2) Perform pipeline iteration: in the tth training iteration, concurrently execute the following operations:
[0007] Sampling phase (executed by the sampling task flow): obtain a t-th fixed-size mini-batch, perform forward propagation inference on the mini-batch using the current latest model parameters, and evaluate the importance indicators of each sample; based on the importance indicators, use a preset sampling proposal distribution to screen a micro-batch from the mini-batch, and calculate the gradient correction coefficient of the selected sample; the micro-batch will be used for training in the t+1-th iteration.
[0008] Training phase (executed by the training task flow): obtain the micro-batch generated in the t-1-th iteration and its corresponding gradient correction coefficient; perform back propagation calculation using the micro-batch, and use the gradient correction coefficient to weight and correct the calculated gradient or loss to update the model parameters.
[0009] (3) Use an event synchronization mechanism to coordinate the execution progress of the sampling task flow and the training task flow, and ensure the correctness of the data dependency until the preset local training round is completed.
[0010] Further, the step (2) of the present application comprises the following steps:
[0011] (21) Perform inference on each sample x in the mini-batch and calculate its importance I(x);
[0012] (22) Define a sampling proposal function s(x): if the sample is determined to be important (I(x) = 1), then s(x) = 1 (must be selected); if the sample is determined to be unimportant (I(x) = 0), then s(x) = r, where r is a preset retention rate hyperparameter, r ∈ (0, 1);
[0013] (23) Random sampling based on the probability distribution composed of s(x) to generate a micro-batch
[0014] (24) For any sample x in the micro-batch , its gradient correction coefficient c(x) is equal to the inverse of its sampling probability, i.e. c(x) = 1 / s(x);
[0015] (25) When updating the model, use the weighted loss function to perform back propagation, thereby realizing unbiased estimation of the gradient.
[0016] Further, the step (3) of the present application comprises the following steps:
[0017] (31) Open a ring buffer on the GPU device, containing at least two buffer blocks, respectively used to store the batch data currently being sampled and the batch data currently being trained;
[0018] (32) Use GPU stream technology to realize hardware-level concurrency of sampling and training, wherein the sampling stream uses mixed precision for inference to reduce memory usage and calculation overhead.
[0019] (33) Set up a synchronization control: configured to monitor the state of the first execution stream and the second execution stream, and prevent data races by inserting synchronization barriers or events.
[0020] Beneficial effects: The present application realizes real-time data sampling within the batch, ensures that the evaluation of sample importance is based on the latest model at the "current time", completely solves the problem of outdated importance in static sampling, and enables the model to always focus on the most difficult-to-learn samples. The gradient correction mechanism based on the reciprocal of the sampling probability is introduced, which mathematically guarantees the consistency of the gradient expectation after sampling with the full data gradient expectation, effectively avoiding the performance degradation of the model due to data distribution deviation. The innovative pipeline design overlaps the "sampling" and "training" execution, and uses the concurrency capability of the GPU to mask the additional inference overhead caused by sampling. Combined with mixed precision inference and ring buffer mechanism, the hardware utilization is maximized, and the training speed can be accelerated by 1.5 to 2.5 times under the same hardware conditions. BRIEF DESCRIPTION OF DRAWINGS
[0021] Figure 1 The logical flowchart of the batch real-time data sampling algorithm proposed by the present application;
[0022] Figure 2 The overall architecture schematic diagram of the federated learning system provided by the present application; DETAILED DESCRIPTION
[0023] In order to make the purpose, technical scheme and advantages of the present application clearer and more apparent, the present application will be further described in detail below in combination with the drawings and examples. It should be understood that the specific embodiments described herein are only used to explain the present application, and do not limit the protection scope of the present application.
[0024] Reference Figure 1 The traditional sampling is performed at the epoch level, while the batch real-time data sampling algorithm refines the granularity to the iteration level.
[0025] Step one: data loading. In the tth iteration of local training, the client loads a mini-batch of a fixed size from the local data set
[0026] Step two: Real-time importance evaluation. Utilize the latest local model parameters w t , forward propagate all samples in . Note that only forward process is needed here, no backward propagation. Loss value, gradient norm or uncertainty of classification prediction can be used as importance indicator. The embodiment preferably uses classification loss value. Set a threshold or according to Top-k principle, define indicator function I(x). If loss value of sample x is greater than threshold, I(x) = 1 (considered as difficult example / important sample); otherwise I(x) = 0.
[0027] Step three: Construct micro-batch. Instead of simply discarding non-important samples, probability reservation strategy s(x) is adopted:
[0028]
[0029] Wherein, r is a hyperparameter (for example, 0.3), used to prevent the model from overfitting difficult examples and forgetting learned knowledge. Based on this probability, filter to obtain micro-batch with variable size. Thus, the subsequent calculation amount is reduced.
[0030] Step four: Calculate gradient correction coefficient. In order to eliminate the bias caused by sampling, calculate weight for each sample x of :
[0031]
[0032] Wherein, the weight of important sample is 1; the weight of non-important sample (selected and reserved) is 1 / r. This is equivalent to using a selected simple sample to "represent" multiple discarded simple samples.
[0033] Step five: Weighted training. Only perform backward propagation for micro-batch . When calculating loss function , multiply the corresponding coefficient: This will ensure that the mathematical expectation of the gradient is unbiased.
[0034] Referring to Figure 2 , in order to solve the additional forward propagation overhead introduced by real-time data sampling within batch, the invention designs a pipeline parallel system.
[0035] Step one: Hardware resource allocation. At the host end, start two parallel threads, "sampling thread" and "training thread"; at the device end, create two logically independent CUDA streams, "sampling stream" and "training stream". Since GPU has a large number of CUDA cores, the two streams can be executed concurrently on hardware without blocking each other.
[0036] Step Two: Pipeline Timing Design: The system adopts a "staggered" execution strategy. Assume that at the current moment, the system is performing the logical t-th iteration:
[0037] The training stream works as follows: it currently holds the data generated during the (t-1)th iteration. And its correction coefficients. It is responsible for performing backpropagation, calculating gradients, and updating model parameters from w. t Change to w t+1 .
[0038] How the sampling stream works: It loads new mini-batches in parallel. It uses the current model for forward inference to filter out micro-batches. It is then placed into a buffer for the training stream to use at the next time step t+1.
[0039] Step 3: To prevent the sampling stream from consuming too much GPU computing power and slowing down the training stream, this embodiment enables automatic mixed precision in the sampling stream. During the forward propagation of importance evaluation, model weights and data are converted to half-precision floating-point numbers. This not only reduces memory bandwidth usage but also leverages the acceleration capabilities of GPU Tensor Cores, making the sampling process extremely fast, almost entirely covered by the training time window.
[0040] Step 4: To efficiently transfer data between the two streams, a circular buffer of size 2 is pre-allocated in the GPU memory [2]. While the training stream is reading micro-batches from Buffer [0] for training, the sampling stream is writing new micro-batches to Buffer [1]. Since the data is all in the GPU memory, switching between streams does not require host-to-device or device-to-host memory copying, eliminating the PCIe bandwidth bottleneck.
[0041] Step 5: Since sampling and training are performed in parallel, the execution order must be strictly controlled to prevent data races such as "read-after-write" or "write-after-read". This invention employs lightweight synchronization based on GPU events.
[0042] Finally, it should be noted that the above description is merely a preferred embodiment of the present invention and is not intended to limit the present invention. Although the present invention has been described in detail with reference to the foregoing embodiments, those skilled in the art can still modify the technical solutions described in the foregoing embodiments or make equivalent substitutions for some of the technical features. Any modifications, equivalent substitutions, improvements, etc., made within the spirit and principles of the present invention should be included within the protection scope of the present invention.
Claims
1. A federated learning acceleration method based on parallel real-time data sampling and training within a pipeline batch, characterized in that, Comprise: Batch-level real-time data sampling strategy: re-evaluate the importance of samples based on the latest model for the current mini-batch in each training iteration, and dynamically select a variable size micro-batch for training to avoid the problem of sample importance expiration in traditional static sampling; at the same time, through the gradient correction mechanism, offset the distribution deviation caused by importance sampling, ensure the unbiasedness of gradient expectation, so that the model can reduce the amount of calculation while maintaining faster convergence and higher accuracy. Parallel mechanism of sampling and training: parallelizing sample importance evaluation and model training, using the micro-batch sampled in the t-1th iteration in the tth iteration training the model while in another execution flow a new mini-batch is processed sampling, thereby constructing a cross-iteration pipeline execution, avoiding the additional forward overhead and GPU idling caused by sequential execution, improving CPU-GPU resource utilization, realizing overlapping execution of sampling and training, and significantly reducing training time.
2. The method of claim 1, wherein, The sampling stage specifically comprises the following sub-steps: performing forward inference on each sample x in the mini-batch , calculating its importance I(x); defining a sampling proposal function s(x): if the sample is determined to be important, set the sampling probability s(x)=1; if the sample is determined to be unimportant, set the sampling probability s(x)=r, where r is a pre-set retention rate hyperparameter and 0 3. The method of claim 1, wherein, The calculation of the gradient correction coefficient of the selected sample is specifically: for any sample x in the micro-batch The gradient correction coefficient of the sample x is equal to the reciprocal of the sampling probability s(x), that is, c(x)=1 / s(x); during the model update in the training stage, the weighted loss function is used for back propagation, and the weighted loss function is represented as wherein is the original loss value for the sample.
4. The method of claim 1, wherein, The method for determining whether a sample is important or unimportant is: setting a threshold or defining an indicator function according to the Top-k principle; if the classification loss value of sample x is greater than the threshold or is within the Top-k range, it is determined to be an important sample, otherwise it is determined to be an unimportant sample.
5. The method of claim 1, wherein, In step three, the event synchronization mechanism is used to coordinate the execution progress, which specifically includes: opening a ring buffer on the GPU device, which contains at least two buffer blocks for storing the current sampling batch data and the current training batch data; using GPU stream technology to realize the hardware-level concurrency of the sampling task stream and the training task stream; configure the synchronization control, monitor the task flow state and insert the synchronization barrier or event to prevent read-write data competition.
6. The method of claim 1, wherein, In the sampling task stream, mixed precision is used for forward propagation inference, and model weights and data are converted to half-precision floating-point numbers for calculation.
7. The method of claim 1, wherein, The pipeline iteration adopts a staggered execution strategy: at the t-th iteration moment in logic, the training task stream utilizes the micro-batch generated by the t-th iteration to perform back propagation and parameter update; meanwhile, the sampling task stream loads a new mini-batch in parallel and generates a micro-batch for the t+1-th iteration using the current model parameters and writes it into the free block of the ring buffer.
8. The method of claim 1, wherein, A pipeline construction module is configured to establish a sampling task flow and a training task flow executed in parallel on the client device and allocate a shared storage space; a sampling processing module is configured to run the sampling task flow, acquire a mini-batch, perform forward inference using current model parameters to evaluate sample importance, and screen a micro-batch and calculate a gradient correction coefficient; a training processing module, configured to run the training task flow, and utilize the micro-batch perform back propagation, and weight the gradient or loss according to the gradient correction coefficient to update the model; and a synchronization control module, configured to monitor and coordinate the execution order of the sampling task flow and the training task flow by utilizing an event synchronization mechanism.
9. A federated learning training framework, the modules of which comprise: Batch-level real-time data sampling module, responsible for real-time evaluation and screening of the most valuable data in each iteration, while ensuring mathematical unbiasedness. Sampling training parallel module, responsible for coordinating the concurrent execution of sampling and training tasks on heterogeneous hardware to maximize hardware utilization.