A global cost-aware scheduling method for dynamic shape training of a multi-modal model
By constructing a unified shape specification and cost model, and combining online calibration and adaptive control, the load balancing problem in dynamic shape training of multimodal generative models is solved, achieving efficient global shape sample scheduling and token budget utilization, thereby improving training efficiency and stability.
Patent Information
- Authority / Receiving Office
- CN · China
- Patent Type
- Applications(China)
- Current Assignee / Owner
- UNIV OF SCI & TECH OF CHINA
- Filing Date
- 2026-07-03
- Publication Date
- 2026-08-04
AI Technical Summary
Existing technologies for dynamic shape training of multimodal generative models suffer from problems such as cross-group synchronization waiting, long tail of step time, sampling distribution offset, and low token budget utilization, making it difficult to achieve global load balancing and efficient scheduling.
By constructing a uniform shape specification, combining an analytical prior model and an online calibration model, the training cost of each shape bucket is predicted, and a global shape bucket allocation plan is generated for multiple training groups before each optimization step. A K-step rolling window scheduling and an adaptive controller are used to adjust the batch size and target sampling probability, and the samples are packaged to improve token budget utilization.
It effectively reduces cross-group synchronization waiting time, avoids sampling distribution offset, improves the overall throughput efficiency and token budget utilization of multimodal dynamic shape training, and maintains the stability and efficiency of the training process.
Smart Images
Figure CN122510072A_ABST
Abstract
Description
Technical Field
[0001] This invention relates to the field of multimodal model training technology, and more specifically to a global cost-aware scheduling method for dynamic shape training of multimodal models. Background Technology
[0002] Training multimodal generative models typically requires supporting dynamic input shapes. For example, image generation models need to support images with different resolutions and aspect ratios, video generation models need to support different frame rates and resolutions, and text-to-video generation tasks also need to process text tokens, visual tokens, and audio features simultaneously. Traditional training systems, to simplify data loading and batch construction, often force input samples to be scaled or cropped to a fixed resolution, or manually bin them according to rules such as resolution, aspect ratio, and frame rate. Building upon this, each data set is trained in parallel, with each group independently and randomly sampling shape bins, and batch size or gradient accumulation steps are statically configured for different shape bins.
[0003] In existing technologies, some methods achieve shape consistency within local batches through manual bucketing, ensuring that samples within the same batch have similar spatial dimensions or frame counts. Other methods dynamically adjust batch size based on shape complexity within the same training group to reduce token padding waste. Still others utilize analytical models to estimate the computational cost of different shapes and select shape combinations based on the load of local batches. These methods alleviate the data organization difficulties in dynamic shape training to some extent.
[0004] However, the above methods have the following shortcomings: manual bucketing only solves the shape consistency within a single training group and cannot handle the synchronization time differences caused by shape differences between multiple parallel data groups or sequential parallel groups. The computational cost and memory usage of dynamic shape samples are highly dependent on the number of latent space tokens, the number of spatial tokens, the number of frames, and the parallelism method. Simply allocating by the number of pixels or randomly is difficult to achieve global load balancing. High-resolution or long-frame-count samples are prone to becoming lagging groups, causing other groups to wait, forming a long tail of step time and idle computing resources. If the scheduler only prioritizes low-time-consuming shapes, it will undersample high-cost shapes for a long time, disrupting the target data distribution. In addition, for video or multimodal long sequence samples, it is difficult to efficiently pack samples of different shapes in the same forward computation, resulting in low token budget utilization. Therefore, a method that can perform cost-aware scheduling at the global training group level and maintain the sampling distribution is needed. Summary of the Invention
[0005] The purpose of this invention is to provide a globally cost-aware scheduling method for dynamic shape training of multimodal models, solving the following technical problems:
[0006] In dynamic shape-distributed training for multimodal generative models, how can we achieve globally cost-aware shape sample scheduling among multiple training groups to eliminate cross-group synchronization waits and long tails of step time caused by shape differences, while avoiding long-term sampling distribution offsets and improving token budget utilization?
[0007] The objective of this invention can be achieved through the following technical solutions:
[0008] A globally cost-aware scheduling method for dynamic shape training of multimodal models includes the following steps:
[0009] S1. Obtain the shape metadata of the training samples. The shape metadata includes the spatial resolution, number of time frames, and length of the conditional sequence of the samples. Construct a uniform shape specification based on the shape metadata.
[0010] S2. Divide the samples into multiple shape buckets according to a uniform shape specification. Each shape bucket is associated with the target sampling probability, the upper and lower limits of the batch size, and the parallel mode configuration.
[0011] S3. Construct a cost model to predict the training cost of each barrel shape under the target training configuration. The training cost includes step time and peak memory usage.
[0012] S4. Before each optimization step, generate a shape bucket allocation plan for multiple training groups based on the predicted training cost, the target sampling distribution of each shape bucket, and the device memory constraints.
[0013] S5. Perform training and collect actual operation profiles, which include observation step time, observation peak memory, and idle time of each group.
[0014] S6. Based on the actual operating profile, calibrate the predicted value of the cost model online, or adjust the batch size of the shape bucket and the target sampling probability.
[0015] As a further aspect of the present invention: In step S1, the specific process of constructing a unified shape specification based on shape element information is as follows:
[0016] The uniform shape specification includes at least one of the following: sample type, latent space height, latent space width, latent space frame count, space token count, and latent space token count; for image samples, the frame count is set to 1, and for video samples, both the frame count and the latent space frame count are recorded.
[0017] As a further aspect of the present invention: In step S3, the specific process of predicting the training cost of each shaped bucket under the target training configuration is as follows:
[0018] The training cost also includes at least one of computation time, communication time, and full reduction waiting time; initial predictions are made based on the analytical prior model, and the predicted values of the analytical prior model are scaled and corrected online based on the actual running profile.
[0019] As a further aspect of the present invention: the specific process of online scaling and correction of the predicted values of the analytical prior model based on the actual running contour is as follows:
[0020] The input includes shape features, training configuration features, hardware descriptor features, and the intersection of shape features and hardware descriptor features; the hardware descriptor includes the target hardware peak computing power, computing power at a specified precision, memory bandwidth, intra-node interconnect bandwidth, inter-node interconnect bandwidth, interconnect topology type, and graph compilation overhead.
[0021] As a further aspect of the present invention: the specific process of generating a shape bucket allocation plan for multiple training groups in S4 is as follows:
[0022] The difference in prediction completion time among training groups in the same optimization step is reduced. Candidate shape buckets are constrained or weighted according to the degree of undersampling and oversampling, and the long-term sampling distribution is maintained. A K-step rolling window scheduling is adopted, where K is a set value. The candidate allocation for the next K optimization steps is jointly scored. The score is based on the sum of the maximum prediction step time for each step in the next K steps and the sampling distribution drift penalty term. Only the allocation for the current optimization step is executed.
[0023] As a further aspect of the present invention: after S4, a token budget packaging step is also included:
[0024] Based on the token budget, multiple samples of different shapes are packaged into a batch. For each batch, the token offset, block diagonal attention range, and sample local location encoding of each sample in the packaging sequence are generated. When packaging, the sample combination is selected based on the number of tokens of the candidate samples, the prediction training cost, and the shape affinity. The block diagonal attention range isolates cross-sample attention between different samples. For video samples, the sample local location encoding is generated by time, height, and width triples.
[0025] As a further aspect of the present invention: In step S6, the specific process of adjusting the batch size of the shape bucket and the target sampling probability based on the predicted value of the online calibration cost model of the actual running contour is as follows:
[0026] The average step time, 95th percentile step time, average peak memory usage, and average idle time of each shape barrel are periodically calculated. The batch size, gradient accumulation steps, target sampling probability, or strict cost matching switch of at least one shape barrel are adjusted through an adaptive controller.
[0027] As a further aspect of the present invention: the specific process of adjusting the batch size, gradient accumulation steps, target sampling probability, or strict cost matching switch of at least one shape bucket through the adaptive controller is as follows:
[0028] The adaptive controller employs model predictive control, discretizing candidate control actions into batch size increases, decreases, or remains constant; gradient accumulation steps increase or remain constant; and target sampling probability increases, decreases, or remains constant. It constructs an objective function using throughput, tail step time, and KL divergence of the sampling distribution to select the optimal action from the candidate actions.
[0029] The beneficial effects of this invention are:
[0030] This invention addresses the problems of cross-group synchronization waiting, long-tailed step times, sampling distribution shift, and low token budget utilization in existing dynamic shape training technologies. By constructing a unified shape specification that includes sample type, spatial resolution, number of time frames, and latent space tokens, and combining an analytical prior model with an online calibration model to predict the step time and peak memory usage of each shape bucket, accurate prediction of the training cost for multimodal dynamic shape training is achieved. Before each optimization step, a global shape bucket allocation plan is generated for multiple training groups based on the prediction cost, target sampling distribution, and device memory constraints. This reduces the difference in prediction completion time among groups within the same optimization step. Simultaneously, candidate shape buckets are constrained or weighted according to the degree of undersampling to maintain a long-term sampling distribution, effectively reducing cross-group synchronization waiting time and avoiding distribution drift. Furthermore, a K-step rolling window scheduling is employed to jointly score the allocation of multiple future optimization steps, executing only the current step allocation, thus improving the global optimality of the scheduling. By periodically statistically analyzing the actual running contours and using model predictive control to adaptively adjust the batch size, gradient accumulation steps, and target sampling probability, the system can dynamically adapt to changes during the training process. Furthermore, by using token budgeting, multiple samples of different shapes are packaged into batches with block diagonal attention ranges and local position encodings, isolating attention calculations between different samples and improving token budget utilization while maintaining sample independence. These techniques work together to significantly improve the overall throughput efficiency of distributed dynamic shape training without altering the model structure or sample semantics. Attached Figure Description
[0031] The invention will now be further described with reference to the accompanying drawings.
[0032] Figure 1 This is a flowchart illustrating the present invention. Detailed Implementation
[0033] The technical solutions of the embodiments of the present invention will be clearly and completely described below with reference to the accompanying drawings. Obviously, the described embodiments are only some embodiments of the present invention, and not all embodiments. Based on the embodiments of the present invention, all other embodiments obtained by those skilled in the art without creative effort are within the scope of protection of the present invention.
[0034] Please see Figure 1 As shown, this invention is a globally cost-aware scheduling method for dynamic shape training of multimodal models, comprising:
[0035] Taking the dynamic resolution training scenario of the image generation model as an example, the training dataset contains images with various resolutions and aspect ratios, such as 512×512 pixels, 512×768 pixels, 768×512 pixels, 768×768 pixels, and 1024×1024 pixels. The method of this invention first reads the metadata of each image sample, including the image height, width, and the possible number of conditional text tokens. Based on this metadata, the method constructs a uniform shape specification. The uniform shape specification consists of sample type, height, width, number of frames, number of conditional tokens, latent space height, latent space width, number of latent space frames, number of space tokens, and number of latent space tokens. For image samples, the number of frames is fixed at 1. The latent space height and width are obtained by multiplying the height and width by the latent space downsampling ratio, which is predetermined by the variational autoencoder or similar encoder used. For example, if the downsampling ratio is 8, then the latent space size corresponding to 512×512 pixels is 64×64, and the number of latent space tokens is 64×64=4096. The number of spatial tokens is equal to the product of the latent space height and the latent space width. For video samples, the original frame number and the latent space frame number also need to be recorded. The latent space frame number is obtained by multiplying the original frame number by the temporal downsampling ratio.
[0036] Based on a uniform shape specification, the method divides all samples into multiple shape buckets. Each shape bucket has a unique bucket identifier, a representative shape specification, a target sampling probability, the current batch size, upper and lower bounds of the batch size, gradient accumulation steps, maximum gradient accumulation steps, parallel mode, and sequence parallel segmentation. For example, shape buckets can be divided according to resolution range or aspect ratio range, grouping images with a width of 480 to 512 pixels and a height of 512 to 544 pixels into the same shape bucket, with the representative shape specification being the median resolution of all samples within that bucket.
[0037] After constructing the shape buckets, the method builds a cost model to predict the training cost of each shape bucket under the target training configuration. The training cost includes step time, peak memory usage, computation time, communication time, and full reduction latency. The cost model consists of an analytical prior model and an online calibration model. The analytical prior model estimates the forward computation time, backward computation time, gradient communication time, and parameter update time based on the latent space token count, space token count, frame count, batch size, micro-batch size, gradient accumulation steps, sequence parallel splitting, and parallel mode of the shape bucket. For example, for a shape bucket with 64×64 latent space tokens, a batch size of 8, and training on 8 training groups using data parallel mode, the analytical prior model calculates a base step time prediction. The online calibration model collects the step times observed during actual training, calculates a scaling factor between the observed values and the analytical prior predictions, and uses this scaling factor to scale and correct subsequent predictions. The input to the online calibration model includes not only shape features and training configuration features, but also hardware descriptor features and the interaction terms between shape features and hardware descriptor features. The hardware descriptor records the peak computing power of the target hardware. For example, a certain model of GPU has a peak computing power of 312 TFLOPS at BF16 precision, a memory bandwidth of 2 TB per second, an intra-node interconnect bandwidth of 600 GB per second, an interconnect topology type of full interconnect, and a graph compilation overhead of 50 milliseconds.
[0038] Before each optimization step begins, the scheduler generates a global shape bucket allocation plan. The scheduler first enumerates all candidate shape bucket combinations that satisfy the device's memory limit. The device memory limit is determined by the current remaining memory capacity; for example, if the remaining memory capacity is 40 GB, the predicted peak memory value for each candidate shape bucket cannot exceed this value. For each candidate plan, the scheduler uses a cost model to predict its step time and peak memory. Subsequently, the scheduler generates allocation plans for multiple training groups based on the predicted step time, the target sampling distribution of each shape bucket, the difference between the current cumulative sampling count and the target sampling count, the compensation weight for undersampled shape buckets, and the penalty weight for oversampled shape buckets. The allocation plan reduces the difference in prediction completion time among training groups in the same optimization step. For example, when training on 8 parallel data groups, the scheduler avoids simultaneously allocating multiple high-cost shape buckets to different groups in a single optimization step, and also avoids the same group continuously receiving low-cost shape buckets, leading to long-term sampling distribution shifts. The scheduler uses a K-step rolling window scheduling, where K is set to 5. The scheduler considers not only the allocation for the current optimization step but also simulates the allocation for the next 4 optimization steps. For each candidate allocation scheme in the current step, the scheduler continues to simulate possible allocations for subsequent steps, retaining candidate paths with a bundle width of 3. Finally, the optimal plan is selected based on the objective function score, which includes the sum of the maximum predicted step times for each of the next 5 steps and a penalty term for sample distribution drift after the next 5 steps. The scheduler only executes the allocation corresponding to the current optimization step; the next optimization step involves a re-rolling planning process.
[0039] During training, each training group retrieves samples from its corresponding shape bucket according to the allocation plan output by the scheduler, and performs forward computation, backpropagation, and parameter updates. After execution, the method records the actual running profile. The running profile includes step identifier, bucket identifier, group identifier, shape specification and training configuration, observation step time, observation peak memory usage, observation computation time, observation communication time, group idle time, and full reduction wait time. For example, the observation step time is 320 milliseconds, of which computation time accounts for 220 milliseconds, communication time accounts for 80 milliseconds, and idle time accounts for 20 milliseconds. This running profile data is fed back to the cost model for online calibration of subsequent predictions.
[0040] During training, the method periodically tracks the running status of each shape bucket, including average step time, 95th percentile step time, average peak memory usage, and average idle time. The adaptive controller adjusts the batch size, gradient accumulation steps, target sampling probability, or strict cost-matching switch of the shape buckets based on these statistics. The adaptive controller employs model predictive control. The controller discretizes candidate control actions into three groups: batch size increases by 10%, decreases by 10%, or remains unchanged; gradient accumulation steps increase by 1, or remain unchanged; target sampling probability increases by 0.05, decreases by 0.05, or remains unchanged. The controller simulates the state for the next two optimization steps based on the cost model and constructs an objective function using throughput, 95th percentile step time, and KL divergence of the sampling distribution. It then selects the action from the candidate actions that optimizes the objective function.
[0041] For video generation models or long-sequence multimodal training scenarios, the method further performs a token budget packaging operation. Assume a training batch has a remaining token budget of 16384 tokens. The scheduler selects candidate samples from the queue of samples to be processed. The latent space token count for each video sample is equal to the latent space frame count multiplied by the latent space height multiplied by the latent space width. For example, a 17-frame video sample with a latent space size of 32×32 has a latent space token count of 17×32×32=17408, which exceeds the token budget and therefore cannot be packaged alone. The scheduler selects two short video samples: one with 5 frames and a latent space size of 16×16 (5×16×16=1280 tokens), and the other with 9 frames and a latent space size of 24×24 (9×24×24=5184 tokens). The total token count for the two samples is 6464, which does not exceed the 16384 budget. The scheduler further selects combinations based on predicted training cost and shape affinity to reduce the cost difference between samples within the packaging batch. During packaging, the method generates a token offset for each sample in the packaging sequence for each batch, with the first sample having an offset of 0 and the second sample having an offset of 1280. The method also generates a block-diagonal attention range for each batch, restricting the attention calculation for each sample to its own token range to avoid cross-sample attention leakage. Furthermore, the method generates a sample-local positional encoding for each video sample, generated based on a triplet of time, height, and width, rather than using global sequence positions. For example, for a 5-frame, 16×16 latent space video sample, with time positions ranging from 0 to 4, height positions from 0 to 15, and width positions from 0 to 15, the positional encoding for each token is a combination of indices from these three dimensions.
[0042] In cross-hardware deployment scenarios, the method of this invention achieves cost model migration by constructing hardware descriptors. For example, migrating from an NVIDIA A100 GPU to a Huawei Ascend 910 NPU, the method first reads the descriptor of the target hardware: the Ascend 910 has a BF16 computing power of 320 TFLOPS, a memory bandwidth of 1.6 TB per second, an intra-node interconnect bandwidth of 392 GB per second, an interconnect topology of ring topology, and a graph compilation overhead of 80 milliseconds. The cost model inputs the hardware descriptor, shape features, and training configuration features into the online calibration model. A small number of warm-up runs are performed on the new hardware, such as running 10 optimization steps and collecting actual running profiles. The ratio of observation step time to analytical prior prediction values is calculated to complete the calibration. Afterward, the cost model can generate accurate training cost predictions for this hardware without modifying the scheduling algorithm itself.
[0043] The foregoing has provided a detailed description of one embodiment of the present invention, but this description is merely a preferred embodiment and should not be construed as limiting the scope of the invention. All equivalent variations and modifications made within the scope of the claims of this invention should still fall within the patent coverage of this invention.
Claims
1. A globally cost-aware scheduling method for dynamic shape training of multimodal models, characterized in that, Includes the following steps: S1. Obtain the shape metadata of the training samples. The shape metadata includes the spatial resolution, number of time frames, and length of the conditional sequence of the samples. Construct a uniform shape specification based on the shape metadata. S2. Divide the samples into multiple shape buckets according to a uniform shape specification. Each shape bucket is associated with the target sampling probability, the upper and lower limits of the batch size, and the parallel mode configuration. S3. Construct a cost model to predict the training cost of each barrel shape under the target training configuration. The training cost includes step time and peak memory usage. S4. Before each optimization step, generate a shape bucket allocation plan for multiple training groups based on the predicted training cost, the target sampling distribution of each shape bucket, and the device memory constraints. S5. Perform training and collect actual operation profiles, which include observation step time, observation peak memory, and idle time of each group. S6. Based on the actual operating profile, calibrate the predicted value of the cost model online, or adjust the batch size of the shape bucket and the target sampling probability.
2. The global cost-aware scheduling method for dynamic shape training of a multimodal model according to claim 1, characterized in that, In step S1, the specific process of constructing a unified shape specification based on shape element information is as follows: The uniform shape specification includes at least one of the following: sample type, latent space height, latent space width, latent space frame count, space token count, and latent space token count; for image samples, the frame count is set to 1, and for video samples, both the frame count and the latent space frame count are recorded.
3. The global cost-aware scheduling method for dynamic shape training of a multimodal model according to claim 1, characterized in that, In step S3, the specific process of predicting the training cost of each shaped barrel under the target training configuration is as follows: The training cost also includes at least one of computation time, communication time, and full reduction waiting time; initial predictions are made based on the analytical prior model, and the predicted values of the analytical prior model are scaled and corrected online based on the actual running profile.
4. The global cost-aware scheduling method for dynamic shape training of a multimodal model according to claim 3, characterized in that, The specific process of online scaling and correction of the predicted values of the analytical prior model based on the actual operating profile is as follows: The input includes shape features, training configuration features, hardware descriptor features, and the intersection of shape features and hardware descriptor features; the hardware descriptor includes the target hardware peak computing power, computing power at a specified precision, memory bandwidth, intra-node interconnect bandwidth, inter-node interconnect bandwidth, interconnect topology type, and graph compilation overhead.
5. The global cost-aware scheduling method for dynamic shape training of a multimodal model according to claim 1, characterized in that, In step S4, the specific process of generating a shape bucket allocation plan for multiple training groups is as follows: The difference in prediction completion time among training groups in the same optimization step is reduced. Candidate shape buckets are constrained or weighted according to the degree of undersampling and oversampling, and the long-term sampling distribution is maintained. A K-step rolling window scheduling is adopted, where K is a set value. The candidate allocation for the next K optimization steps is jointly scored. The score is based on the sum of the maximum prediction step time for each step in the next K steps and the sampling distribution drift penalty term. Only the allocation for the current optimization step is executed.
6. The global cost-aware scheduling method for dynamic shape training of a multimodal model according to claim 5, characterized in that, Following S4, a token budget packaging step is also included: Based on the token budget, multiple samples of different shapes are packaged into a batch. For each batch, the token offset, block diagonal attention range, and sample local location encoding of each sample in the packaging sequence are generated. When packaging, the sample combination is selected based on the number of tokens of the candidate samples, the prediction training cost, and the shape affinity. The block diagonal attention range isolates cross-sample attention between different samples. For video samples, the sample local location encoding is generated by time, height, and width triples.
7. The global cost-aware scheduling method for dynamic shape training of a multimodal model according to claim 1, characterized in that, In step S6, the specific process of adjusting the batch size of the shape bucket and the target sampling probability based on the predicted value of the online calibration cost model of the actual running contour is as follows: The average step time, 95th percentile step time, average peak memory usage, and average idle time of each shape barrel are periodically calculated. The batch size, gradient accumulation steps, target sampling probability, or strict cost matching switch of at least one shape barrel are adjusted through an adaptive controller.
8. The global cost-aware scheduling method for dynamic shape training of a multimodal model according to claim 7, characterized in that, The specific process of adjusting the batch size, gradient accumulation steps, target sampling probability, or strict cost matching switch of at least one shape bucket via an adaptive controller is as follows: The adaptive controller employs model predictive control, discretizing candidate control actions into batch size increases, decreases, or remains constant; gradient accumulation steps increase or remain constant; and target sampling probability increases, decreases, or remains constant. It constructs an objective function using throughput, tail step time, and KL divergence of the sampling distribution to select the optimal action from the candidate actions.