In-batch parallel method based on inter-step similarity of diffusion model

Abstract

The invention relates to an in-batch parallel method based on diffusion model inter-step similarity, and provides a core scheme ParaBatch aiming at the problems of high communication overhead, strong equipment dependence, coarse scheduling granularity and the like of an existing diffusion model parallel reasoning scheme: generating subsequent step sample estimation by utilizing distribution similarity of noise prediction of adjacent de-noising steps and multiplexing a preorder noise prediction result; and the noise predictor is combined with a conventional calculation sample into a batch, parallel calculation and result restoration are realized through single forward propagation, and the calling frequency of the noise predictor is remarkably reduced. Meanwhile, a continuous batch processing mechanism with denoising time steps as granularity is provided, intra-batch requests are dynamically adjusted and executed, and queuing and waiting delay is reduced. According to the scheme, efficient parallelism can be realized on single equipment, additional hardware and high-bandwidth communication are not needed, the method is suitable for a continuous batch processing method of a diffusion model, and the capability of processing multiple requests is enhanced.

Description

TECHNICAL FIELD

[0001] The present application belongs to the technical field of diffusion model inference acceleration, and specifically relates to an intra-batch parallel method based on the similarity between adjacent denoising steps of a diffusion model, which is particularly suitable for efficiently parallelizing the diffusion model inference process through a batch processing mechanism in a single device or multi-device environment. BACKGROUND

[0002] With the continuous development of deep learning technology and the continuous improvement of hardware performance, diffusion models have achieved remarkable results in image generation, video generation, and audio generation. However, high inference latency severely limits their application in more extensive scenarios.

[0003] Inference latency mainly comes from two factors: first, the high computational cost of the noise predictor itself, which is usually implemented using architectures such as DiT (Diffusion Transformer); second, the strong sequential nature of the denoising process, which requires the model to call the noise predictor multiple times over dozens or even hundreds of time steps, as shown in Figure 1 It is worth noting that in DiT models that use three-dimensional attention mechanisms, the computational overhead increases quadratically with the increase in spatial resolution and time length.

[0004] To reduce the inference latency of diffusion models, various solutions have been proposed. For example, distillation technology compresses knowledge from large models to small models, effectively reducing computational and storage overhead; other works use post-training quantization to compress full-precision models to 8-bit or 4-bit representations, significantly reducing computational costs without the need for retraining. In addition, some research has explored caching-based strategies to eliminate redundant calculations in the denoising process by caching intermediate results. Although these methods have achieved certain results, they do not fully utilize the potential of distributed computing, which has become the core support of modern deep learning systems.

[0005] To this end, a number of studies have proposed parallelization strategies based on distributed infrastructure to accelerate the inference process of diffusion models. Traditional parallelization strategies, such as data parallelism, pipeline parallelism, and tensor parallelism, are generally not suitable for reducing the inference latency of diffusion models. Data parallelism and pipeline parallelism aim to improve throughput rather than reduce the inference time of a single sample, and thus are almost useless for latency optimization. Although tensor parallelism performs well in accelerating large language models (LLMs), its applicability in diffusion models is limited due to the large-scale activation tensors generated during the denoising process.

[0006] To more fully exploit the computational potential of distributed GPUs, a number of studies have designed specialized acceleration methods for the characteristics of diffusion models. As shown in Figure 2 For example, DistriFusion divides the input latent variables of the noise predictor into p patches and merges them through all-to-all communication before each attention and convolution operation. PipeFusion uses block-level parallelism and reduces communication overhead by reusing stale feature maps. AsyncDiff divides the noise predictor into p stages, each assigned to a different device, and achieves acceleration through pipeline parallelism. xDiT integrates a ring attention mechanism into PipeFusion, further improving inference speed. However, these methods generally have high communication overhead, limiting their application in practical systems, especially in bandwidth-limited commercial devices. These strategies can only achieve ideal performance in data center infrastructure with high bandwidth. In addition, except for Ring Attention, most of the above methods rely on approximate parallelism, which may introduce significant bias in the generated results compared to the original model, affecting the quality of the final output.

[0007] As shown in Figure 1These parallel methods typically rely on operation-wise, layer-wise, or stage-wise communication to exchange key tensors, but this design often comes with significant communication costs. Despite their performance advantages, these methods are difficult to deploy in environments lacking high bandwidth and data center-level network infrastructure.

[0008] Therefore, there is an urgent need to design a parallel method with lower communication overhead to achieve efficient inference in commercial hardware environments. Furthermore, traditional parallel methods require additional computing equipment to support the degree of parallelism. However, current computing equipment is very expensive, thus the high cost hinders the further application of these parallel methods. Summary of the Invention

[0009] This invention aims to address the problems of high communication overhead, strong device dependence, coarse scheduling granularity, and degraded generation quality in existing parallel inference methods for diffusion models. It provides an intra-batch parallel method based on the similarity between adjacent denoising steps of the diffusion model, enabling efficient, low-latency, and low-cost diffusion model inference in single-device or multi-device environments.

[0010] To achieve the above objectives, the technical solution of the present invention is as follows: This invention proposes ParaBatch, an intra-batch parallelism method that leverages the similarity between adjacent steps in a diffusion model. The core idea is that the noise predicted by adjacent steps in a diffusion model is very similar, thus allowing the reuse of noise from a certain number of steps to obtain a batch of predicted samples. These predicted samples are then repaired within a batch using a noise predictor, thereby achieving intra-batch parallelism. To address the step-by-step execution characteristic of the diffusion model, we propose continuous batch processing, enabling dynamic shrinking of the current batch. We use ParaBatch as the backend of our diffusion model service system. The detailed technical method includes the following steps: Step 1: The user sends a request to the front end of the diffusion model service system. After receiving the request, the front end puts it into the request queue to wait for scheduling.

[0011] Step two: When the scheduler schedules a request to the backend at a step-by-step granularity, the backend dynamically adds the request to the current batch.

[0012] Step four: The backend of the diffusion model service system supports continuous batch processing, calling the model at a step-by-step granularity. Furthermore, this backend supports the ParaBatch intra-batch parallel method proposed in this invention.

[0013] Step five, the design of ParaBatch needs T-step denoising in the inference of diffusion model, where T is a value specified by the user. For the t-step denoising, the input of the diffusion model (also known as a noise predictor) is the t-step sample, and the output is the t-step noise. Then the planner obtains the next sample by subtracting the t-step noise from the t-step sample, and takes it as the input of the next diffusion model. After T steps of iteration, the final output is obtained. In this process, the planner only contains a small amount of linear calculation, and almost no calculation overhead. While the diffusion model needs to perform thousands of heavy matrix multiplication, which accounts for almost all the calculation overhead. ParaBatch takes advantage of the inter-step similarity of the noise predicted by the diffusion model, and obtains a batch of samples by reusing the noise of a certain step, and then performs batch calculation within a batch. ParaBatch takes advantage of the batching effect of the diffusion model, especially the audio diffusion model, to achieve speedup without additional computational resources. By combining with continuous batching, the diffusion model service can be efficiently performed.

[0014] The beneficial effects of the present application are: (1) The intra-batch parallel method ParaBatch is proposed, which takes advantage of the batching effect of the audio diffusion model to achieve inference acceleration without additional device overhead.

[0015] (2) A continuous batching method suitable for diffusion models is proposed, which performs calculation scheduling at a step granularity to achieve dynamically shrinking batching. The queuing time of requests is reduced, and the space of batching is increased.

[0016] (3) Experimentally, the present application is tested on AudioLDM2, and the experiment shows that the present application can achieve 6.56x speedup with almost no performance loss. BRIEF DESCRIPTION OF DRAWINGS

[0017] Figure 1 is the inference process of the existing diffusion model Figure 2 is a diffusion model service system deployment diagram of the present application Figure 3 is the existing diffusion model service system that cannot dynamically shrink the batch Figure 4 is a schematic diagram of the continuous batching method suitable for diffusion models Figure 5 is a schematic diagram of the intra-batch parallel method ParaBatch of the present application DETAILED DESCRIPTION The technical solutions of the present application will be described in detail below with reference to the drawings and examples, but the protection scope of the present application should not be limited thereby.

[0018] The implementation of the method can rely on existing deep learning frameworks (such as PyTorch, TensorFlow) and standard GPU hardware environments. The following is an example of a typical audio diffusion model (such as AudioLDM2) service system.

[0019] 1.1 System deployment and environment configuration First, deploy a complete diffusion model service system, whose architecture is shown in Figure 2 The system includes: Front-end connection point: receives user-generated requests (such as text prompts and generation parameters) and adds them to the request queue.

[0020] Request queue: uses a priority queue or FIFO queue to manage pending requests.

[0021] Scheduler: a core component that implements the step granularity scheduling strategy described in the present invention.

[0022] Execution backend: loads pre-trained diffusion models (such as AudioLDM2 noise predictors) and planners (such as DDIM or PNDM Scheduler), and implements continuous batch processing and ParaBatch parallel mechanisms.

[0023] The workflow is as follows: First, the user sends a request to the front-end of the diffusion model service system. After the system receives the request, it is placed in the request queue waiting for the scheduler to schedule. When the execution backend finishes executing a whole batch, the result is returned to the scheduler, which returns it to the user. Then, the scheduler packages the requests in the request queue into a batch and sends it to the execution backend.

[0024] The traditional diffusion model service system needs to wait for the current batch to complete execution before loading new requests from the request queue and sending them to the backend for execution. This leads to three problems: 1. Increases the waiting time of the request. As shown in Figure 3 When the current batch is not executed, the requests in the request queue must wait, even if the backend can accommodate more requests, it cannot be scheduled; 2. Increases the execution time of the request. As shown in Figure 3 When a request in the batch is executed, it must wait for all requests in the batch to be executed before it can be returned; 3. Affects the batch processing space of the backend. When there are already processed requests in the batch that have not been returned, it will affect the batch processing space of the backend.

[0025] High inference latency severely limits the application of diffusion models in more extensive scenarios, such as service systems. In service systems, we expect the execution backend to return the execution result in time. However, diffusion models have high latency by nature, and inference latency mainly comes from two factors: first, the high computational cost of the noise predictor itself, which is usually implemented with DiT (Diffusion Transformer) architecture; second, the strong sequence nature of the denoising process itself, which requires the model to call the noise predictor multiple times over dozens or even hundreds of time steps. Reducing the inference latency of diffusion models can increase the response capability of service systems and improve user experience.

[0026] Referring to Figure 4 , the specific steps of continuous batch processing are as follows: After system initialization, the execution backend maintains a current batch, which is initially empty.

[0027] Dynamic batch-out: At the end of each denoising step, check the number of steps each request in the current batch has executed. If a request has reached its total inference steps (such as 50 steps), immediately remove it from the current batch, assemble the generated results and return them to the user, and release the occupied GPU memory.

[0028] Dynamic batch-in: The scheduler runs a lightweight integer programming solver (such as using the OR-Tools library) at the end of each time step based on the latest system state (such as the free slots of the current batch, GPU utilization, and the waiting time of requests in the queue). The solver aims to maximize throughput or meet target latency, and outputs which new requests should be taken from the request queue to join the current batch.

[0029] Through the above dynamic in-out mechanism, the batch size is adjusted in real time, avoiding the idle resources and request blocking caused by waiting for the entire batch to complete in traditional systems (such as the problem shown in Figure 3 ).

[0030] The present invention schedules at the step granularity. Traditional diffusion model service systems schedule at the batch granularity, and after executing all requests in the batch, new requests can be executed. The present invention changes the scheduling granularity to steps, that is, each step can reorganize the current batch being executed, dynamically kick out the executed requests, and dynamically pull in the requests in the request queue.

[0031] 3. Implementation of ParaBatch batch-parallel method The application utilizes a key property of diffusion models: the noise predicted between adjacent time steps is very close in diffusion model denoising process. Therefore we can reuse the noise obtained in previous time steps to skip some calculations. At time step T, the noise predictor uses the current input to generate a noise estimate , and then the scheduler computes the next state , and time step based on . At time step , the noise predictor is called again to estimate and generate . In contrast, a reuse strategy can be used to reduce inference latency: the model first predicts based on to generate ; then, at time step , instead of re-calling the predictor, the model directly reuses to estimate the next noise sample . We refer to this strategy as the direct reuse process.

[0032] If the noise predicted at each time step can be reused for the next time steps, the number of calls to the noise predictor will be reduced to , and correspondingly, the model inference latency will also be reduced to . However, this aggressive reuse approach can cause a distributional mismatch with the original denoising trajectory, resulting in a significant drop in model performance. To alleviate the above problem, we propose a "reuse-then-refine" mechanism that combines the reuse strategy with a prediction refinement process, effectively reducing the quality loss caused by reuse while maintaining computational efficiency. Compared to directly reusing the noise predicted in the previous step, this mechanism first uses the reused noise to generate the noise sample at the current time step, and then inputs this sample into the noise predictor to predict the refined noise at the next time step, thereby obtaining the repaired predicted noise. Based on this mechanism, we propose a new distributed sampling method, ParaBatch. This method parallelizes the noise prediction tasks between adjacent time steps and then converts them into parallel within the same batch, thereby achieving batch-wise parallelization of the diffusion model inference process.

[0033] As Figure 1As shown, in the initial "warm-up steps", the diffusion model performs the denoising process in the standard order. After the warm-up, the sample and the predicted noise are obtained. Then the predicted noise from the previous step is reused, and the scheduler predicts the next state accordingly. The computation overhead of the scheduler can be ignored. Then ParaBatch packs the sample and into a batch, and schedules the model to execute the batch, obtaining the corresponding noise values and . Then the planner is called to calculate the sample and for the next time step using these noises. This loop is repeated in subsequent time steps, thereby achieving efficient parallelization of adjacent time step noise prediction across devices. Under the setting of a parallelism degree of , ParaBatch executes a complete parallel loop every steps. Since the scheduler operation has very small computation, the entire system can be regarded as completing a completely parallel noise prediction process. Therefore, the total inference delay of the entire denoising process will be shortened to of the original, the theoretical speedup ratio is , and no additional devices are required. This speedup is guaranteed by the batching effect of the audio diffusion model.

Claims

1. A batch-in parallel method based on similarity between diffusion model neighboring denoising steps, characterized in that, The method comprises the following steps: receiving a user inference request and adding it to a request queue; a scheduler generates a scheduling strategy dynamically based on the backend computing load, the target delay of the request, and the total inference step number, with the denoising time step as the minimum scheduling granularity, through an optimization algorithm; a continuous batch processing mechanism is implemented in the execution backend, and the batch composition is dynamically adjusted at the end of each denoising time step: the requests that have reached the specified inference step number are removed from the current batch, and new requests waiting for scheduling in the request queue are added to the current batch; by utilizing the similarity of noise prediction results between adjacent denoising time steps of the diffusion model, noise prediction output of the t-th step is reused to directly generate noise sample estimates of the t-k-th step, where k is the reuse span; samples obtained through the normal denoising path and sample estimates generated by reusing noise are merged into a batch, which are jointly input into the noise predictor for parallel forward calculation, and accurate noise prediction values of all samples in the batch are obtained at one time; the planner calculates the state of each sample at the next time step based on the accurate noise prediction values; the above steps of reuse, batch construction, parallel prediction, and state update are iteratively performed to realize batch-parallel acceleration of the inference process of the diffusion model.

2. The batch-in parallel method based on similarity between diffusion model based neighboring denoising steps of claim 1, wherein, The minimum scheduling granularity is the denoising time step, which specifically includes: the scheduler monitors the computing load state of the execution backend and the waiting condition of the request queue in real time; an integer linear programming algorithm is used to solve the optimal request scheduling sequence, with the optimization goal of minimizing the average waiting time, maximizing the throughput, or meeting the target delay; at the boundary of each denoising time step, the composition of the execution batch is dynamically updated according to the scheduling sequence.

3. The batch-in parallel method based on similarity between diffusion model based neighboring denoising steps of claim 1, wherein, The reuse of noise prediction output of the t-th step to directly generate noise sample estimates of the t-k-th step is specifically: skip the noise predictor calculation of the t-1-th to t-k+1-th steps; directly use the planner and the noise prediction output of the t-th step to calculate the noise sample estimates of the t-k-th step through deterministic or stochastic difference equations.

4. The batch-in parallel method based on similarity between diffusion model based neighboring denoising steps of claim 3, wherein, The reuse span k is a tunable hyperparameter, or a dynamically determined numerical value based on model characteristics, noise levels, or empirical data, which is used to balance the acceleration ratio and generation quality.

Images