Efficient Multi-Task Large Model Training via Data Heterogeneity-aware Model Management
Yujie Wang Shenhan Zhu Fangcheng Fu Xupeng Miao Jie Zhang Juan Zhu Fan Hong Yong Li Bin Cui

































































 ^1Peking University
 ^2Purdue University
^3Alibaba Group


^1{alfredwang, shenhan.zhu, ccchengff, bin.cui}@pku.edu.cn
^2xupeng@purdue.edu

^3{wanglin.zj, zhujuan.zj, hongfan.hf, jiufeng.ly}@alibaba-inc.com 









§ ABSTRACT 


Recent foundation models are capable of handling multiple machine learning (ML) tasks and multiple data modalities with the unified base model structure and several specialized model components. 
However, the development of such multi-task (MT) multi-modal (MM) models poses significant model management challenges to existing training systems. Due to the sophisticated model architecture and the heterogeneous workloads of different ML tasks and data modalities, training these models usually requires massive GPU resources and suffers from sub-optimal system efficiency.






In this paper, we investigate how to achieve high-performance training of large-scale MT MM models through data heterogeneity-aware model management optimization.
The key idea is to decompose the model execution into stages and address the joint optimization problem sequentially, including both heterogeneity-aware workload parallelization and dependency-driven execution scheduling.

Based on this, we build a prototype system and evaluate it on various large MT MM models. 

Experiments demonstrate the superior performance and efficiency of our system, with speedup ratio up to 71% compared to state-of-the-art training systems.





Recent foundation models are capable of handling multiple machine learning (ML) tasks and multiple data modalities with the unified base model structure and several specialized model components. However, the development of such multi-task (MT) multi-modal (MM) models poses significant model management challenges to existing training systems. Due to the sophisticated model architecture and the heterogeneous workloads of different ML tasks and data modalities, training these models usually requires massive GPU resources and suffers from sub-optimal system efficiency.

In this paper, we investigate how to achieve high-performance training of large-scale MT MM models through data heterogeneity-aware model management optimization. The key idea is to decompose the model execution into stages and address the joint optimization problem sequentially, including both heterogeneity-aware workload parallelization and dependency-driven execution scheduling. Based on this, we build a prototype system and evaluate it on various large MT MM models. Experiments demonstrate the superior performance and efficiency of our system, with speedup ratio up to 71% compared to state-of-the-art training systems.




Efficient Multi-Task Large Model Training via Data Heterogeneity-aware Model Management
    Yujie Wang^1, Shenhan Zhu^1, Fangcheng Fu^1, Xupeng Miao^2, Jie Zhang^3, Juan Zhu^3, Fan Hong^3, Yong Li^3, Bin Cui^1
    September 9, 2024
=========================================================================================================================






PVLDB Reference Format:

. . PVLDB, (): , .

https://doi.org/doi:
[This work is licensed under the Creative Commons BY-NC-ND 4.0 International License. Visit <https://creativecommons.org/licenses/by-nc-nd/4.0/> to view a copy of this license. For any use beyond those covered by this license, obtain permission by emailing mailto:info@vldb.orginfo@vldb.org. Copyright is held by the owner/author(s). Publication rights licensed to the VLDB Endowment. 

Proceedings of the VLDB Endowment, Vol. , No.  
ISSN 2150-8097. 

https://doi.org/doi: 

]footnote-1







PVLDB Artifact Availability:

The source code, data, and/or other artifacts have been made available at <>.








§ INTRODUCTION



























Machine learning (ML) has become an essential tool for understanding and generating knowledge from data and tackling complex tasks for humans. In the past few years, our data management community has put great efforts in developing systems to support the whole ML lifecycle <cit.>, such as data preparation <cit.>, model development <cit.>, model selection <cit.> and model deployment <cit.>.
Recently, as the rapid rise of large-scale foundation models <cit.>, developing these large models is becoming increasingly challenging due to the substantial GPU resource requirements. For example, how to design training systems for those models consisting of billion of parameters over more than dozens of GPUs demands extensive expertise, which has attracted lots of research interests from our community <cit.>. As a result, ML models themselves have become another form of data, and their management techniques are becoming increasingly important <cit.>.

Considering the multi-modal nature of real-world data, ML researchers have shifted their focus to developing model beyond the the language domain (e.g., ChatGPT <cit.>) to many other data modalities (e.g., images <cit.>, speech <cit.>, video <cit.>).
The recent extension further involves composite scenarios <cit.>, where models are capable of processing and interpreting data across several tasks simultaneously.

However, existing large model training systems are mainly designed for a single model with only one input data modality. Despite the extensive research and engineering efforts aimed at optimizing these systems from multiple perspectives, including distributed communication <cit.>, memory management <cit.>, and GPU computation <cit.>, their performance is still limited when it comes to handling the increasingly complex requirements of multi-task (MT) multi-modal (MM) models. We identify two unique obstacles when building training systems for MT MM models.















One is the workload heterogeneity due to the divergent data flows across modalities or tasks. On the one hand, MM models often handle data that vary significantly in structure and size, demanding specialized preprocessing and computational approaches. For example, language models (e.g., GPT-family <cit.>, LLaMA-family <cit.>) are usually equipped with dozens of layers with the same configuration (e.g., hidden size), while vision models may involve uneven layers to compute in various resolutions <cit.>. On the other hand, as depicted in Fig. <ref>, multiple tasks usually leverage distinct data flows and activate individual model components, leading to inter-task workload heterogeneity. 



Due to such heterogeneous modality data flows and sub-models, different modalities and tasks exhibit distinct execution overhead (detailed in Fig. <ref>, <ref>). 
Existing training systems usually overlook such workload heterogeneity and apply sub-optimal training methodologies.


Another is the data flow execution dependency among different model components. Recent MT MM model development usually adopts a sub-model sharing approach <cit.>, where partial model layers containing common knowledge are shared across different modalities and tasks. As shown in Fig. <ref>, each data type also has its own learning component. Within every training iteration, the input data mixed with multiple modalities are simultaneously fed into the sophisticated model, where different model components are intricately activated and updated. 
To avoid redundant resource usage, the shared components are usually responsible for the data flows from multiple sources, resulting in execution barriers and blocking the following model layers.
In addition, the proportion of different data modalities in MT workloads may shift over time due to task addition and completion, introducing further training complexity.

To the best of our knowledge, none of existing training systems can deal with these unforeseen dependency efficiently due to the lack of understanding MT MM model execution.


To tackle these obstacles, this paper introduces , a resource-efficient and high-performance training system for large-scale MT MM models via data heterogeneity-aware model management.
Considering the workload heterogeneity and execution dependency, a naïve solution is to decouple the model structure based on modality and task, replicate the shared components, and deploy them on separate devices. In this way, each sub-model can be optimized by existing systems, but it also brings significant resource wastage and underutilization, as well as additional overheads from replica synchronization.
As an example, Fig. <ref> showcases that such a naïve, decoupled execution suffers from fluctuating device utilization both intra-task and inter-task due to workload heterogeneity, leading to low or even idle GPU utilization for some time slots.
Instead of decoupling, manages to directly train the whole complex model without disjoint sub-model to minimize the resource usage.
A key insight behind 's design is that heterogeneous and dependent sub-models can be decomposed into several sequentially executed and independent stages, each of which contains multiple parallel model modules with similar execution overheads. 




    


    












To achieve resource-efficient and high-performance training of MT MM models, there are three key model management challenges for to address. In the following, we will introduce each challenge and how solves them.

C1: Model Parallelization. First, finding the optimal model parallel configuration for heterogeneous workloads with diverse computational characteristics is a complex combinatorial problem. Existing single-model automatic parallelization approaches (e.g., Alpa <cit.>, Unity <cit.>, Galvatron <cit.>) assume a spatial pipeline stage partition, and each operator (Op) is executed by all devices of the corresponding pipeline stage. Unfortunately, such assumptions only work for homogeneous models, failing to adapt to heterogeneous MT MM models.

Instead of solving the parallel configuration directly, captures the workload heterogeneity at the operator granularity and estimates its execution overheads under different amount of allocated resources and parallel configurations (<ref>). The final configuration decision is left to the later step since it requires to be jointly optimized with considering the execution dependency. also introduces to contract the graph (i.e., fusing continuous identical operators) to avoid redundant estimation overheads and shrink the problem scale (<ref>).

C2: Model Division. Second, breaking down the whole model into sequentially executed stages is straightforward, but it may easily result in inefficiencies.
Determining the optimal division of stages is complicated since the operators differ significantly in their execution overheads and have intricate operator dependencies.


addresses this problem with two steps: 1) 's resource allocator (<ref>) traverses the computation graph following the dependency topology and decides the optimal resource allocation for in each candidate set (i.e., currently executable ). Here we reformulate this issue as a malleable project scheduling problem (MPSP) and subsequently derive the optimal solution. 2) After obtaining the parallel configuration of each , the stage scheduler (<ref>) greedily slices and selects to craft compact stages and minimizes the overall execution time.

C3: Model Mapping. Third, given the stage-based resource allocation and execution schedule plan, how to map them into physical devices is still a problem, 
since different mapping may lead to distinct inter-stage communication overheads and per-device memory consumption.
To further improve the overall system efficiency, carefully considers these trade-offs and the real environment constraints (e.g., inter-device bandwidth, memory capacity) when generating the device placement plan (<ref>).





Our contributions are summarized as follows:


    
  * We present , a high-performance and resource-efficient training system for large-scale MT MM models via data heterogeneity-aware model management optimization.
    
  * We propose a jointly optimization framework to achieve heterogeneity-aware workload parallelization and dependency-driven execution scheduling.
    
    
    
    
  * We build a general runtime engine to perform the stage-based schedule, automatically resolving execution dependencies at the stage boundaries.
    
  * We evaluate on various large MT MM models, and the results demonstrate the superior performance and efficiency of compared with the state-of-the-art baselines, with the speedup ratio up to 71%.






















For example, speech recognition and image captioning tasks shall share the text encoder but feed the visual- and audio-inputs into different encoders <cit.>.
Many empirical results have also shown that such a joint multi-task training paradigm achieves better multi-modal capabilities for generalist models than performing single-task training separately <cit.>.










§ PRELIMINARY





 §.§ Multi-Task Multi-Modal Models










  §.§.§ Multi-Modal Application of Foundation Models

The advent of foundation models <cit.> has revolutionized deep learning (DL).
Beginning with the birth of BERT <cit.> and GPT <cit.> based on Transformer <cit.> structure, followers such as the GPT series <cit.>,  T5 <cit.>, OPT <cit.>, and the LLaMA series <cit.> have set new benchmarks across a range of language tasks.

Foundation models have also been successfully adapted for tasks of other data modalities, including image processing <cit.>, audio processing <cit.>, video analysis <cit.>.
Multi-modal models <cit.> leverage these foundation models to integrate information from multiple data modalities. They can be primarily categorized into two types.
The first category fuses modality information via contrastive learning objectives <cit.>, with CLIP <cit.> being a notable example, and ImageBind <cit.> further extending CLIP to six modalities. 
These models typically have a multi-tower structure, where each modality has its own encoder. They take paired modality data (e.g., image-text pairs for CLIP), extract features via modality encoders, and perform cross-modal alignment using contrastive objectives.
The second category merges modality features through a language model's generative loss <cit.>. These models usually consist of multi-tower modality encoders and a cross-modal module. Modality encoders extract features from each modality, which are then fed into the cross-modal module for feature fusion. 
Recently, with the success of open-sourced large language models (LLMs) <cit.>, researchers have started to enhance multi-modal models with powerful pretrained LLMs <cit.>. 
These multi-modal LLMs, also falling into the second category.




  §.§.§ Multi-Task Multi-Modal Models

Recently, researchers have begun to construct multi-task multi-modal (MT MM) models <cit.>, enabling the support for diverse multi-modal tasks within a unified model.
This is because each modality encompasses various tasks, and each task often involves multiple modalities as well.
The general model structure and the training flow is illustrated in the upper side of Fig. <ref>.
MT MM models reflects researchers' aspiration towards a general-purpose AI.
Flamingo <cit.> is among the first to handle multiple vision-language tasks.


OFASys <cit.> proposes a general MT MM learning paradigm,
as shown in Fig. <ref>, 
designing distinct modality encoders and cross-modal modules for different tasks and modalities, allowing the activation of different components as required by the task and modality at hand. 
For example, speech recognition and image captioning tasks shall activate and share the text encoder but feed the visual- and audio-inputs into different encoders.

Many empirical results have also shown that such a joint multi-task training paradigm achieves better multi-modal capabilities for MT MM models than performing single-task training separately <cit.>.




 §.§ Parallelisms in Distributed Training


As model sizes and training data volumes grow, modern DL systems commonly utilize clusters of multiple GPUs for distributed training, thereby enhancing efficiency. 
Various parallelisms are employed to manage model parameters or training data in a distributed manner. 
Data parallelism (DP) <cit.> splits the input data, with each device handling a portion of the data storage and computation, and synchronizing model gradients across devices. 
Model parallelism <cit.> partitions model parameters, with each device responsible for storing and computing a segment of the model. Model parallelisms can be categorized into two popular types: tensor parallelism (TP) partitions the model vertically <cit.>, while pipeline parallelism (PP) <cit.> splits the model horizontally, organizing model computations into a pipeline.


Contemporary distributed training systems, such as Megatron-LM <cit.> and DeepSpeed <cit.>, leverage multiple parallelisms and implement a hybrid parallelism approach for model training. For example, Megatron-LM introduces 3D parallelism, which concurrently utilizes DP, TP, and PP, enhancing training efficiency. 


Researchers have also developed advanced automatic parallelism techniques facilitate the tuning of optimal parallelism combinations.
These automatic parallelism <cit.> approaches integrate multiple parallelism dimensions, employ sophisticated optimization workflows, and automatically determine the most efficient hybrid parallelism strategy, significantly reducing the reliance on human effort.
However, these existing training system are mainly deigned for single task and single model training, with limited performance on the complex scenario of training MT MM models.











  §.§.§ Foundation Models



The advent of foundation models <cit.> has revolutionized deep learning.
Beginning with the birth of BERT <cit.> and GPT <cit.> based on Transformer structure <cit.>, subsequent iterations such as the GPT series <cit.>, RoBERTa <cit.>, T5 <cit.>, OPT <cit.>, and the LLAMA series <cit.> have set new benchmarks across a range of language tasks.

Foundation models have also been successfully adapted for a variety of applications, including computer vision <cit.>, audio processing <cit.>, video analysis <cit.>.










  §.§.§ Multi-Modal Learning





Transformer models have also been applied to multi-modal learning.
Researchers leverage the robust understanding capabilities of pretrained language models to learn information from multiple modalities simultaneously <cit.>. 
Based on the modality feature fusion approach and model architecture, multi-modal models can be categorized into two types.

The first category fuses modality information via contrastive learning objectives <cit.>, with CLIP <cit.> being a notable example.
These models typically have a multi-tower structure, where each modality has its own encoder, such as ViT <cit.> for vision. They take paired modality data (e.g., image-text pairs for CLIP), extract features via modality-specific encoders, and perform cross-modal alignment using contrastive objectives, enabling zero-shot capabilities. 
Extensions like <cit.> broaden CLIP's application to more modalities.
ImageBind <cit.>, introducing a model capable of handling six modalities, aligns various modalities with vision using image-paired data and contrastive objectives. 


The second category merges modality features through a language model's generative loss <cit.>. These models usually consist of multi-tower modality encoders and a cross-modal encoder. Modality encoders extract features from each modality, which are then concatenated and fed into the cross-modal encoder for feature fusion. 

OFA <cit.> and others <cit.> use lightweight encoders for all modalities and a single unified Transformer-based cross-modal encoder. 

ALBEF <cit.>, BLIP <cit.>, OFASys <cit.> and others <cit.> retain parameterized modality-specific encoders (e.g., ViT) and a cross-modal encoder, enabling better learning of intra-modal and inter-modal interactions.


Recently, with the success of open-sourced large language models (LLMs) <cit.>, researchers have started to enhance multi-modal models with powerful pretrained LLMs <cit.>. 
These multi-modal LLMs, also falling into the second category, utilize multi-tower pretrained modality encoders to extract modality features, mapping them into the LLM feature space through modality adaptors, and then fusing modal information with pretrained LLMs. 







  §.§.§ Multi-Task Multi-Modal Learning


Multi-task learning is a natural aspect of multi-modal learning, where a single modality encompasses various tasks. For instance, the visual modality includes tasks such as image captioning, image classification, and visual question answering. 
A single task often involves multiple modalities as well; for example, the image captioning task involves both visual and textual modalities.
Consequently, researchers have begun constructing generalist models through multi-task multi-modal learning approaches <cit.>, enabling these models to handle diverse multi-modal tasks within a single framework. 
This endeavor reflects researchers' aspiration towards a general-purpose AI. 
Flamingo <cit.> was among the first to utilize interleaved vision-language data, adopting an in-context few-shot learning approach for mastering multiple vision-language tasks. 
Gato <cit.>, OFA <cit.>, and Unified-IO <cit.> employ lightweight modality encoders to transform multi-modal and multi-task data into sequences, learning multi-task abilities using a unified Transformer encoder or an encoder-decoder architecture. 
OFASys <cit.> proposes a general multi-task multi-modal learning paradigm, designing distinct modality encoders and cross-modal encoders for different tasks and modalities, allowing the activation of different components as required by the task and modality at hand. Such approach aligns with the design of Google's Pathways <cit.>, where different tasks have their own data flow and activate distinct parts of the model. 

In multi-task multi-modal learning, the data flow and model computations for different modalities and tasks vary, resulting in heterogeneous workloads. Moreover, the dynamic nature of tasks and modalities during training presents challenges for efficient multi-task multi-modal learning.







 §.§ Parallelisms in Distributed Training







  §.§.§ Data & Model Parallelsim

As model sizes and training data volumes grow, modern DL systems commonly utilize clusters of multiple GPUs for distributed training, thereby enhancing efficiency. 
Various parallelisms are employed to manage model parameters or training data in a distributed manner. 
Data parallelism (DP) <cit.> splits the input data, with each device handling a portion of the data storage and computation, and synchronizing model gradients across devices. 
Model parallelism <cit.> partitions model parameters, with each device responsible for storing and computing a segment of the model. Model parallelisms can be categorized into two popular types: tensor parallelism (TP) partitions the model vertically <cit.>, while pipeline parallelism (PP) <cit.> splits the model horizontally, organizing model computations into a pipeline.




















  §.§.§ Hybrid Parallelism and Automatic Parallelism









Contemporary distributed training systems, such as Megatron-LM <cit.> and DeepSpeed <cit.>, leverage multiple parallelisms and implement a hybrid parallelism approach for model training. For example, Megatron-LM introduces 3D parallelism, which concurrently utilizes DP, TP, and PP, enhancing training efficiency. However, these systems, resembling toolboxes, necessitate manual tuning of parallelism combinations to achieve optimal efficiency, which is both challenging and labor-intensive. Therefore, researchers have developed advanced automatic parallelism techniques to minimize the need for manual intervention. These automatic parallelism <cit.> approaches integrate multiple parallelism dimensions, employ sophisticated optimization workflows, and automatically determine the most efficient hybrid parallelism strategy, significantly reducing the reliance on human effort.







 §.§ Cluster Scheduling for DL Jobs







Multi-task multi-modal learning prompts consideration of a related field: cluster scheduling for DL (Deep Learning) jobs. This subsection outlines job scheduling research and contrasts it with multi-task multi-modal learning, aiming for clarity in explaining their differences.

GPU clusters typically provide training services for multiple DL jobs from numerous users. To enhance cluster utilization and the training efficiency, GPU clusters often design cluster schedulers to coordinate resource allocation and the execution order among multiple jobs. 
Some cluster schedulers <cit.> allocate resources to jobs based directly on the requirements users specify. 
There are also cluster schedulers <cit.> automatically allocating resources to each job based on the job scalability to the computing resource. Many of these schedulers aims to minimize job completion time (JCT). For instance, Optimus <cit.> leverages a performance model to predict the job training speed as a function of allocated resources, and introduces the concept of marginal gain to guide resource allocation, aiming to minimize job completion time.


Here, we highlight that the multi-task multi-modal learning scenario targeted by and cluster scheduling represent two distinct scopes of problems: 
(1) Despite the superficial similarity, multi-tasks and multi-jobs are quite different concepts. In cluster scheduling, there is no inherent connection or dependency between jobs, no need for model sharing, nor any data dependency. In contrast, multi-task multi-modal learning involves dependencies among multiple tasks, where tasks may share and simultaneous update the same part of a model within a single iteration. An example is when a image captioning task and an image classification task share and both update the same vision encoder.
This necessitates a unique approach to scheduling multi-tasks, focusing on iteration-level management.
(2) Multi-task multi-modal learning is characterized by a diverse range of tasks and modalities, each activating different parts of the model, leading to heterogeneous workload distributions both across and within tasks. On the other hand, cluster scheduling typically deals with resource allocation at a broader job level. This coarse-grained method of allocating resources, when applied to the more complex scenario of multi-task multi-modal learning, may result in sub-optimal efficiency. 
Nonetheless, even with these differences, we can still draw inspiration from some of the coarse-grained resource allocation strategies in cluster scheduling as guidance.







§ MOTIVATION
 
To be written...Maybe motivation and challenge can be illustrated in Introduction?



















 §.§ System Challenge of Multi-task Learning



  * Heterogeneity: Workload heterogeneity among tasks and inside task. Need some profiling results.

  * Dynamicity: Tasks may change dynamically along training.

  * Dependency: Model sharing dependency. Op dependency, operation topological order.


Why current system cannot handle these challenges?

What's the difference and relationship between MTP and other lines of work: 

Multi-task multi-modal training system is an unexplored dimension.

Mainstream training system: Megatron-LM Deepspeed

1) Parallelism optimization works, Megatron-LM Deepspeed, MTP multi-task, resource allocation; 

2) Job scheduling works, job are independent, doesn't require model synchronization. These works usually schedule the jobs/tasks at task level, ignoring the heterogeneous workload between the computation Ops inside each task.
























§ SYSTEM OVERVIEW







is a highly efficient and scalable training framework designed for MT MM models. 
It facilitates effective operator-level resource allocation and sophisticated distributed training scheduling for MT MM training. 
(check this summary later)

Fig. <ref> depicts the system architecture of , comprising the execution planner and the training framework. Given user-defined multi-modal tasks and the GPU cluster, the execution planner initiates a graph contraction process (<ref>), contracting the original computation graph into a MetaGraph composed of (Fig. <ref>), with each characterizing a unique workload. Subsequently, the scalability estimator (<ref>) estimates the execution time and resource scalability for each , producing scaling curves (Fig. <ref>). Following this, the resource allocator (<ref>) allocates appropriate amount of resources for each , and the resulting allocation plan (Fig. <ref>) is processed by the stage scheduler (<ref>) to develop a stage-based execution schedule. Device placement (<ref>) strategies are then employed to assign to appropriate devices. This results in the execution plan (Fig. <ref>). Finally, the runtime engine (<ref>) utilizes this plan to instantiate the model on each device and facilitate an efficient multi-task multi-modal training process.





















































































































§ SYSTEM DESIGN
 



is a highly efficient and scalable training framework designed for MT MM models. 




Fig. <ref> depicts its system architecture, comprising the execution planner and the training framework. 
Given the diverse user-defined training tasks and the GPU cluster, the goal of is to devise the most efficient execution plan to facilitate effective MT MM training.



  §.§.§ Problem Formulation

We formalize the optimization problem of as follows.
Firstly, interprets the input tasks as a unified directed acyclic computation graph 𝒢 = (𝒱, ℰ), where each node i ∈𝒱 represents a computational operator and each edge ⟨ i,j ⟩∈ℰ denotes the data flow from operator i to j. 
Each task activates specific operators and parameters with unique data flows. 
For instance, a vision-related task activates a vision Transformer layer as an operator, with image features serving as the data flow.
The left side of Fig. <ref> displays an example of a computation graph.
Then, given the computation graph 𝒢 and the GPU cluster with N devices, aims to minimize the maximal operator completion time C.
Specifically, we need to find an execution plan P, which assigns each operator i ∈𝒱 with an AS-tuple ⟨ n_i, s_i ⟩∈𝒰, such that the operator i is Allocated n_i devices and is Scheduled to execute from time s_i.
Here the set 𝒰 = {⟨ n, s ⟩ | n ∈ℕ, s ≥ 0} is formed by all valid AS-tuples. We further denote the execution time of operator i when allocated n_i devices as t_i=T_i(n_i).
Then, the optimization problem is formulated as follows:











    _P={i →⟨ n_i,s_i ⟩ |
     i ∈𝒱, ⟨ n_i,s_i ⟩∈𝒰} C max_i ∈𝒱{s_i + t_i}
    s.t.   ∑_t ∈ (s_i, s_i+t_i), i ∈𝒱 n_i ≤ N    for ∀ t ∈ℝ^+ 
        s_i+t_i ≤ s_j    for ∀⟨ i,j ⟩∈ℰ

Here (<ref>) is the allocation capacity constraint for any time t, and (<ref>) is the operator dependency constraint.



  §.§.§ Sketch of Solution

Before stepping into the solution of , we would like to first present an overview for better readability.
First, initiates a graph contraction process (<ref>), contracting the original graph 𝒢 into a MetaGraph 𝒢_M composed of  (Fig. <ref>), where each characterizes a unique workload. 
This process further decouples into different , ensuring that there are no dependencies among within the same .
Second, the scalability estimator (<ref>) estimates the execution time and resource scalability for each , producing scaling curves (Fig. <ref>). 
Following this, the resource allocator (<ref>) deduces the allocation plan for each individually (Fig. <ref>). 
Given the allocation plan, the stage scheduler (<ref>) slices the and organizes them into , and produces the -based schedule for execution. 
Subsequently, device placement (<ref>) strategies are then employed to assign to appropriate devices, resulting in the execution plan (Fig. <ref>).
Finally, the runtime engine (<ref>) utilizes this plan to instantiate the model on each device and facilitate an efficient MT MM training process.






 §.§ Graph Contraction
 
























  §.§.§ Depicting Workload Heterogeneity with 

is designed to minimize the execution time by optimizing resource allocation and scheduling for each operator within 𝒢. 
This optimization process necessitates an understanding of the workload characteristics for each operator i ∈𝒱, which can be reflected by its execution time function t_i = T_i(n_i), which varies with the device allocation amount n_i. 
Given that 𝒢 typically includes a large number of operators while many of them share similar workload characteristics (such as stacked Transformer layers), initiates a graph contraction process to streamline the complicated graph. It categorizes operators based on their computational workload characteristics, as illustrated in Fig. <ref>.
In this process, operators are contracted into a if they meet the following criteria:



  * There is a data flow between operator i and j, i.e.,  ⟨ i,j ⟩∈ℰ, and both the out-degree of operator i and the in-degree of operator j are 1, ensuring that they are direct predecessors and successors to each other.

  * Operator i and j share the same computational operator type, parameter size, and input data size, confirming identical computational workloads.


During the graph contraction procedure, we traverse the original graph 𝒢 in topological order, contracting operators based on the specified criteria until no further pairs of operators meeting these conditions remain. This results in a contracted MetaGraph 𝒢_M = (𝒱_M, ℰ_M), with each node m ∈𝒱_M representing a that consists of L_m consecutive operators in 𝒢. 
Since operators in the same share the same workload, we slightly abuse the notation and denote the execution time function for each operator in m as T_m(n).






  §.§.§ Disentangling Dependency with 

To facilitate operator-level resource allocation and scheduling, we further introduce an abstraction called , which signifies the level of dependency. at the same level are independent to each other.
The level of each can be derived by a Breadth-First-Search (BFS), with the level assigned based on the search depth, which inherently ensures no dependency among the of same level.
By doing so, the problem (<ref>) can be dissected into several simplified sub-problems for different . 
Next, we introduce how derives the allocation and scheduling for each individually, and merges them into the final plan. 









 §.§ Scalability Estimator
 






As  differ in operator types and/or input data sizes, they characterize heterogeneous workloads and thus necessitate different amount of resources. 
Furthermore, there's no doubt that these  have distinct resource scalability (i.e., how its execution time varies w.r.t. the amount of allocated resources).
For instance, the left side of Fig. <ref> shows the execution time of different , T_m(n), in Multitask-CLIP (an multi-task extension of CLIP <cit.>, refer to <ref> for details). 
Some show almost linear decreases in execution time as resources increase (e.g., Task2-Vision), while others decrease much more slowly (e.g., Task1-Text).
The right side of Fig. <ref> further shows the value of ς_m(n)=T_m(1)/T_m(n), which measures how much the operator accelerates when using more GPUs, and a value of ς_m(n) closer to n signifies better resource scalability.
As can be seen, different not only have varying execution time, but also exhibit different resource scalability, posing a significant challenge for resource allocation.


In response to this issue, employs a scalability estimator to accurately capture the execution time and the resource scalability of each . 
Previous works <cit.> have designed effective estimation methods for distributed training, commonly utilizing the αβ modelling <cit.>.
However, although this may work well for homogeneous workloads (e.g., large language models with homogeneous layers), we find that it does not fit the workload heterogeneous nature of MT MM models. 
This is because different have distinct workload and resource scalability, and the invoked kernels may vary across different per-device workloads, therefore causing distinct performance.
In a nutshell, our scalability estimator adopts the piecewise αβ modelling for more accurate estimation of heterogeneous MT MM workloads.
Given the target MT MM model, it profiles several discrete data points (n_i, T_m(n_i)) for each under different parallel configurations, and then fits the curve of piecewise αβ function. 
To estimate the execution time T_m(n), it locates the range that n falls into, and returns the estimated time according to the corresponding piecewise function. 
In practice, the profiling and estimating process for each MT MM model takes within 5 minutes, which is negligible compared to the massive training time.
In Fig. <ref>, the scatter points represent empirical measurements, while the curves depict the function estimated by our scalability estimator, which we denote as scaling curves. 
As can be seen, our scalability estimator effectively and accurately estimates the execution time T_m(n) for each . 

More details are illustrated in Appendix  <ref>.



[ffc 8.31 version]
As differ in operator types and/or input data sizes, they characterize heterogeneous workloads and thus necessitate different amount of resources. 
To estimate the execution time of each operator in m over n devices, i.e., T_m(n), prior works generally utilize the αβ modelling <cit.>. 
The rationale behind is that the total workload is evenly distributed across n devices.
However, although this may work well for homogeneous workloads (e.g., large language models with homogeneous layers), we find that it does not fit the workload heterogeneous nature of MT MM models. 
This is because different have distinct resource scalability (i.e., how its execution time varies w.r.t. the amount of allocated resources).
To elaborate, Fig. <ref> shows the execution time and scalability of different in Multitask-CLIP (an multi-task extension of CLIP <cit.>, refer to <ref> for details). 
As can be seen, different not only have varying execution time, but also exhibit different resource scalability. Some show almost linear decreases in execution time as resources increase (e.g., Task2-Vision), while others decrease much more slowly (e.g., Task1-Text).
Undoubtedly, this poses a significant challenge for resource allocation of MT MM models.

In response to this issue, employs a scalability estimator to accurately capture the execution time of each . 
In a nutshell, scalability estimator adopts the piecewise αβ modelling.
Given the target MT MM model, it profiles several discrete data points (n_i, T_m(n_i)) for each under different parallel configurations, and then fits the curve of αβ function between any two neighboring points. 
To estimate the execution time T_m(n), it locates the range that n falls into, and returns the estimated time according to the corresponding piecewise function. 


In practice, the profiling and estimating process for each MT MM model takes within 5 minutes, which is negligible compared to the massive training time.

In Fig. <ref>, the scatter points represent empirical measurements, while the curves depict the function estimated by our scalability estimator, which we denote as scaling curves. 
As can be seen, our scalability estimator effectively and accurately estimates the execution time T_m(n) for each . 













[wyj 8.31 version]
As  differ in operator types and/or input data sizes, they characterize heterogeneous workloads and thus necessitate different amount of resources. 
Furthermore, there's no doubt that these  have distinct resource scalability (i.e., how its execution time varies w.r.t. the amount of allocated resources), posing a significant challenge for resource allocation.
For instance, the left side of Fig. <ref> shows the execution time of different , T_m(n), in Multitask-CLIP (an multi-task extension of CLIP <cit.>, refer to <ref> for details). 
Here T_m(n) denotes the execution time of each operator in m over n devices.

Some show almost linear decreases in execution time as resources increase (e.g., Task2-Vision), while others decrease much more slowly (e.g., Task1-Text).

The right side of Fig. <ref> further shows the value of ς_m(n)=T_m(1)/T_m(n), which measures how much the operator accelerates when using more GPUs, and a value of ς_m(n) closer to n signifies better resource scalability.
As can be seen, different not only have varying execution time, but also exhibit different resource scalability.






















In response to this issue, employs a scalability estimator to accurately capture the execution time and the resource scalability of each .
Previous works <cit.> have provided effective estimation methods for distributed training. 
For instance, communication cost can be estimated via bandwidth and the communication volume specified by parallel configurations , and the αβ modeling <cit.> is commonly used for estimation. 



Ideally speaking, the total workload is evenly distributed across n devices. 
However, our empirical observations indicate that the invoked kernels may vary across different per-device workloads.
Therefore, we further apply a piecewise αβ function for more accurate modeling. Refer to Appendix  <ref> for details.


Ideally speaking, the total workload is evenly distributed across n devices. 
However, our empirical observations indicate that the execution time T_m(n) only scales linearly within a range. 
Thus, we characterize T_m(n) by a generalized piecewise αβ function:



    T_m(n) = α_m,i + β_m,i×c_m + β_m,i'×w_m/n,∀ n ∈ [n_i-1, n_i], i = 1,…,k


where k is the number of pieces, α_m,i represents the coefficient of fixed overheads (e.g., kernel launch costs), β_m,i and β_m,i' represent the reciprocal of execution efficiency (e.g., GPU computation speed and network bandwidth), w_m/n denotes the distributed workload of m across n devices (e.g., computational workload), and c_m denotes the workload that doesn't scale with n (e.g., communication volume of data parallelism).
Such piecewise function indicates that under varying resource scales, due to changes in the per-device workload, coefficients such as 
α, β and β' might differ, as the invoked kernels may vary across different workloads.









Our scalability estimator adopts a profiling-based method to model this function. It begins by profiling several discrete data points (n_i, T_m(n_i)) for each under different parallel configurations.

Subsequently, it fits the curve of αβ function to model the execution time T_m(n) as introduced above. 



In practice, the profiling and estimating process for each MT MM model takes within 5 minutes, which is negligible compared to the massive training time.

In Fig. <ref>, the scatter points represent empirical measurements, while the curves depict the function estimated by our scalability estimator, which we denote as scaling curves. 
As can be seen, our scalability estimator effectively and accurately estimates the execution time T_m(n) and resource scalability ς_m(n) for each . 




















 §.§ Resource Allocator
 


In this subsection, we introduce the resource allocator of , which allocates appropriate computational resources to each . 

We begin by transitioning the problem (<ref>) into the sub-problem on one . 
We then detail our allocation strategies, which first relax constraints and optimize the continuous problem, and then discretize the optimal solution to obtain practical allocation plans.





  §.§.§ Problem Formulation on 

We first re-formulate the problem (<ref>) on one with a set of denoted by 𝒱_M.

In this formulation, we split each into different execution part, by assigning it with several ASL-tuples ⟨ n,s,l ⟩∈𝒰_M, such that l consecutive operators of this are scheduled to execute from time s with n devices.
Here 𝒰_M = {⟨ n, s, l ⟩ | n,l ∈ℕ, s ≥ 0} is formed by all valid ASL-tuples.
For each m ∈𝒱_M, its execution plan is a set of ASL-tuples P_m. 



For a , the execution plan P consists of P_m for all m ∈𝒱_M, i.e., P={m → P_m}.
Given m ∈𝒱_M and one ASL-tuple p=⟨ n_m^(p),s_m^(p),l_m^(p)⟩∈ P_m, we denote the execution time span, end time, and time interval by 
t_m^(p) = T_m(n_m^(p)) · l_m^(p), 
e_m^(p) = s_m^(p)+t_m^(p), 
and I_m^(p)=(s_m^(p),e_m^(p)), 
respectively.

The problem can be re-written as:

    _P={m → P_m | m ∈𝒱_M, P_m ⊂ 2^𝒰_M}Cmax_m ∈𝒱_M, p ∈ P_m{e_m^(p)}
    s.t.    ∑_t∈ I_m^(p), m ∈𝒱_M, p ∈ P_m n_m^(p)≤ N    for ∀ t ∈ℝ^+ 
        I_m^(p_1)∩ I_m^(p_2) = ∅   for ∀ m∈𝒱_M, p_1,p_2∈ P_m 
       ∑_p ∈ P_m l_m^(p) = L_m    for ∀ m ∈𝒱_M

Compared with the original problem (<ref>), the sub-problem (<ref>) on gets rid of the dependency constraint, while the  constraint (<ref>) enforces the execution intervals of ASL-tuples in P_m to be pairwise disjoint, because operators within the same cannot execute simultaneously, and (<ref>) ensures all operators are executed for each .







  §.§.§ Optimum of the Continuous Problem



If we relax the constraints, allowing GPU resources and operators to be continuously divisible (i.e., n and l in ASL-tuples are not limited to integers), the problem is transformed into a well-established problem, malleable project scheduling problem (MPSP), with malleable projects and continuously divisible resources <cit.>.

We denote the optimal solution of this relaxed problem by .
Prior works <cit.> have given the following theorem.

    
    If the execution time functions T_m(n), n∈ℝ^+, are positive and non-increasing for every m∈𝒱_M, then = {m → P_m} satisfies that P_m = {⟨ n_m^∗, 0, L_m⟩}, ∀ m ∈𝒱_M, where the optimum objective C^∗ and allocations n_m^∗ can be found from
    
    
    T_m(n_m^∗) · L_m = C^∗ for ∀ m ∈𝒱_M
             and ∑_m ∈𝒱_M n_m^∗ = N.

    

From Theorem <ref>, it follows that in the optimal situation, all start simultaneously, execute all their operators, and finish together.
They share an identical end time e_m = C^∗, which is exactly the minimized operator completion time.







To achieve , our allocator utilizes the scaling curves from <ref> to acquire an estimation of T_m(n), and performs a bisection search procedure over C^∗ with the following equation. The details are illustrated in Appendix <ref>.


    ∑_m∈𝒱_M T_m^(C^∗/L_m) = N.














  §.§.§ Bi-point Discretized Allocation

From the continuous problem, we've determined the optimal time C^∗, as well as the optimal allocations for each , n_m^∗, which is a real number.
To reinstate n's as integers, our allocator  computes each 's proper discrete allocations individually.
For every m, it uses two discrete ASL-tuples ⟨n_m, ·, l_m⟩, ⟨n_m, ·, l_m⟩ to linearly represent the continuous, optimal solution ⟨ n_m^∗, 0, L_m ⟩ in . To preserve the optimum property of , we require the discretized allocation plan to satisfy the following two conditions:
























0.35

    L_m = l_m + l_m


0.05






0.58

    C^∗ = T_m(n_m) ·l_m + T_m(n_m) ·l_m

1em



Cond. (<ref>) ensures these two discrete ASL-tuples complete the workload of m, and Cond. (<ref>) ensures their total execution time is exactly equal to the minimum operator completion time C^∗ in , thus perserving the optimum property.
Here we first select n_m, n_m as the closest valid integer numbers such that n_m^∗∈ [n_m, n_m], and l_m,l_m∈ℝ^+ are derived naturally.
For instance, as shown in Fig. <ref>, 2 with n_2^∗=1.5, L_2=12 in  is discretized as n_2=2, n_2=1 and l_2=8.4, l_2=3.6 in this step.
Here we impose the valid constraint on the allocation n for m for practical reasons.
For instance, if an is applied data parallelism, its allocation n is supposed to divide its global batch size B_m to avoid resource under-utilization due to uneven partition of samples.
For another example, if an is applied tensor parallelism or sequence parallelism with degree 2, its allocation n is supposed to be divisible by this degree, thus n=3, 5, 7 as invalid.
Such valid constraint ensures the allocation plan for each is practical.
Specially, allocation with n_m=0 is treated as a dummy allocation (e.g., 3 in Fig. <ref>), which preserves the optimum property of Cond. (<ref>) but will then be ignored.


Then, we reinstates l's as integers by rounding l_m, l_m to the nearest integers.
If the rounded l equals 0, this ASL-tuple will be ignored.

This rounding procedure preserves the integrity of Cond. (<ref>) and introduces only minor bias to Cond. (<ref>).
Finally, the discretized ASL-tuples of all form the allocation plan.
Note that the allocation plan only ensures the longest execution time among all is approximately C^∗, yet it does not specify the start time for each ASL-tuple, which is determined by stage scheduler in <ref>.





































 §.§ Stage Scheduler
 
In this subsection, we describe how schedules the execution of guided by the allocation plan generated by the resource allocator.
We first introduce the concept of , which is a scheduling unit of .
Then we introduce our -greedy scheduling algorithm, which schedules the execution of greedily for each . 
Finally, the operator dependencies among are reinstated by merging the -based schedules together.






















  §.§.§ Definition of 

It is worthy to note that, although Theorem <ref> implies that all share the same start and end time in the continuous form, this property does not hold after the discretization process. 
The reason is that the execution time of ASL-tuples may vary, or the resources are insufficient to execute all tuples concurrently. 
To cope with this problem, we devise a fine-grained scheduler that slices the and selects a few of them to execute concurrently, where the slicing and selection aim to achieve that (1) the devices are occupied as many as possible, and (2) their execution time are as close as possible. 
To ease the description, we define as the scheduling unit, which corresponds to one concurrent execution as aforementioned. 
Next, we introduce our greedy algorithm that crafts the to form the scheduling plan.










  §.§.§ -greedy Scheduling













As outlined in Alg. <ref>, the scheduler iteratively crafts in a greedy manner. 
Below we introduce how one is crafted with Fig. <ref> as an example.


  * First, the scheduler greedily proposes ASL-tuples to form a candidate set, aiming to utilize as many devices as possible (line 3). For instance with Fig. <ref>, the scheduler proposes the first ASL-tuple of 1 to craft 1 since it occupies all devices. Similarly, for 2, it proposes the ASL-tuples of 1, 2, and 4, which correspond to 4, 2, 2 devices, respectively, in order to make full use of all devices.

  * If the candidate set fails to occupy all devices, the cluster resources will be underutilized. To address this issue, we extend the allocated resources in specific tuples to ensure all devices are utilized (line 4). 
For instance, in 4 of Fig. <ref>, the allocation of 4 is extended from 1 device to 2 devices.
Resource extension is prioritized for with larger remaining execution time, with the hope of balancing the remaining workload among the .

  * In most cases, the proposed ASL-tuples differ in execution time. If we directly craft a with them, it would be inefficient since there must be idle devices. 
Fortunately, this can be avoided by dissecting the ASL-tuples to align their time span (i.e., only a few number of operators in the are scheduled in this ). 
For instance, in 2 of Fig. <ref>, the proposed ASL-tuples for 1, 2, and 4 correspond to 9, 14, and 3 operators, respectively. To align the execution time, the ASL-tuples for 1 and 2 are dissected, with only 1 and 2 operators of them being scheduled, while the remaining 8 and 13 operators left to be scheduled in subsequent .
Our scheduler simply aligns the time span w.r.t. the ASL-tuple with shortest execution time (e.g., the one for 4 in the previous example), and computes the aligned time span as the duration of current stage (line 5).

  * After the time span alignment, the scheduler concludes the current (lines 6-7), including specifying the start time for operators that are scheduled in this , and removing them from the remaining set.







The scheduler aims to schedule the execution start time s for each ASL-tuple of in the allocation plan generated in <ref>.
performs scheduling stage by stage, as outlined in Alg. <ref>. 
We first initialize a candidate set cand with all ASL-tuples in the allocation plan. 
For each stage, we select a subset of tuples from cand to create a working set work (lines 3-4), aiming to utilize all N devices as fully as possible, and schedule the corresponding for execution in current stage.
However, as the execution time of these ASL-tuples vary, when some complete, the resource allocation must be dynamically adjusted to prevent idle devices, thereby avoiding resource wastage.
Therefore, we define the duration of the current stage T_stage by the shortest execution time among the ASL-tuples in work; the stage concludes once any tuple finishes (line 5).
All the tuples in work are split along operators to form a scheduled tuple, which forms sched and aligns the execution time to T_stage, and the unscheduld tuples, which form unexe, is returned to cand for scheduling in subsequent stages (line 10).
For instance, in Fig. <ref>, when scheduling stage 2, tuple ⟨ 2, ·, 3⟩ for 4 determines the stage duration, and the tuple ⟨ 2, ·, 14⟩ for 2 is split to ⟨ 2, ·, 2⟩, which aligns the current stage duration, and ⟨ 2, ·, 12⟩, which will be scheduled later.
This mechanism ensures that all the scheduled ASL-tuples in each stage stops simultaneously as much as possible to maintain workload balance. 


Notably, if the working set fails to fully occupy all N devices, the cluster resources will be underutilized. Therefore, we apply a greedy approach when selecting work, prioritizing tuples with larger allocation n (line 3), which allows more efficient use of device fragments left by larger tuples by those with smaller n.
Furthermore, we also perform a resource extension procedure to deal with the device fragments, which extends the allocation amount for certain ASL-tuples to fully utilize all devices (line 6-7). 
For example, in Fig. <ref> stage 4, the allocation of tuple of 4 is extended from 1 GPU to 2 GPUs.
Specifically, we greedily select the tuples in sched with larger normalized remaining time (line 6), NormRT_m = T_m(n_m^∗) · l_remain, for resource extension, which keeps the balance of the remaining workload of each to avoid resource wastage. 
Here n_m^∗ is the optimal allocation for m in , and l_remain is the number of its remaining operators.
Refer to Appendix <ref> for details.
During the extension of tuple p, if necessary, additional operators can be transferred from the same ’s ASL-tuples in unexe to align its execution time with T_stage.
After resource extension, the schedule k in current stage is finalized (line 8-9), and the scheduler updates the candidate set and the start time for the next stage (line 10-11).
When the cand is empty, the scheduler returns the stage-based schedule as output, consisting of multiple consecutive stages.



































































  §.§.§ Merging 

As stated in <ref>, are decoupled into to disentangle operator dependencies. invokes the aforementioned allocation and scheduling for each individually, and merges their -based schedules together as the final execution schedule.






 §.§ Device Placement
 










Given the -based schedule, which consists of the allocation amount and the execution time of each , we now discuss how determines the specific devices to allocate to each , known as device placement. 

Device placement affects the inter-communication overhead, as well as the memory consumption of each device.
employs several guidelines based on empirical insights or observations to optimize device placement for , as detailed below.




  §.§.§ Intra-Device-Island Placement

Placement within a device island is always preferred for each and each data flow between . 
A device island consists of a group of devices connected by high-bandwidth interconnects (e.g., NVLink, PCIe), typically comprising adjacent devices, such as adjacent GPUs within one node. 
For , prioritizing placement within the device island reduces the potential intra-communication costs. 
For example, on a cluster with two GPUs per node, it's more efficient to place a to a contiguous device group like GPU 0 and 1 within one node rather than a scattered group across two nodes.
For data flow between across , intra-island placement reduces transmission costs leveraging the high intra-island bandwidth or even faster intra-device copying.

For example, if data flow exists between m and m^', and they are assigned 1 and 2 GPUs, respectively, strives to place m on GPU 1 and m^' on GPU 0 and 1.
This arrangement allows data flow via intra-island communication (i.e., 1 to 0) or intra-device copying (i.e., 1 to 1), avoiding the inter-island communication costs that would occur if m were on device 1 and m^' on device 2 and 3 on the other node.


























  §.§.§ Prioritizing High Communication Workloads

When the ideal scenarios outlined above are not achievable — that is, when it's infeasible to place all within the device island nor to align all data flows on the same device group — and data flows with higher communication volumes should be prioritized. 
estimates the communication volume of each and each data flow to prioritize placing those with higher volumes within a device island and aligning high-volume data flows on the same device group. 
For instance, in Fig. <ref>, the data flow volumes between red and blue are significantly higher than that between yellow ones. Therefore, prefer to place the data flow between red and blue ones within the device island, while place the data flow between yellow ones across the island. 
This guideline ensures that the most communication-intensive components receive the most efficient hardware configuration to minimize communication overhead.








  §.§.§ Device Memory Balance

As each device holds heterogeneous , the memory overhead varies across devices.
Placing too many memory-intensive on a single device may cause out-of-memory errors. 
Therefore, actively strives to balance the memory load across all devices during device placement.
Specifically, estimates the memory consumption of each , record the available memory capacity of each device during placement, and prioritizes placing to the device with the highest available memory capacity.
Besides, for sharing the same parameters, we prioritize placing them on the same device to minimize redundant storage.

Based on these guidelines, performs device placement by greedily, prioritizing the minimization of communication overhead, such as inter-transmission, while simultaneously maintaining device memory balance. 
When out-of-memory occurs due to imbalanced placement, will consider alternative placements with sub-optimal communication costs and better memory balance. If necessary, backtracking is employed to adjust the placements from earlier to effectively address the out-of-memory issues.




































































































































































































































































































































































































 §.§ Runtime Engine
 














The runtime engine is responsible for running the execution plan to facilitate efficient multi-task multi-modal training. 

This process is more complex than conventional single-task training, as each device handles heterogeneous and local computation graphs.
The runtime engine operates in four main steps:




  * Localization. 

Initially, localizes the execution plan to each device. 
Specifically, each device instantiates the corresponding of each locally, and initializes the required model components and parameters.




  * Intra-task Data Dependency. 

Secondly, inserts transmission operators to connect the across to handle the data flow dependencies, including activations from the forward pass and gradients from the backward pass. 

According to the device groups of and data format requirements, operations such as copy, shard, concat, send, and receive are used to transmit data flows with minimal overhead. 
For example, a simple copy is sufficient for that share the same device group. 
However, for complicated cases, more complex send and receive operations are necessary to transmit the data appropriately.


This step not only correctly handles data flow dependencies between but also links the on each device into a complete local computation graph ready for execution.


  * Inter-task Model Dependency. 

Then, manages parameter device groups for synchronization among various tasks by maintaining a global parameter device group pool. 


Specifically, during each iteration, for each parameter W_j, all tasks or modalities that activate it on different devices contribute to its gradient computation.
These gradients need to be accumulated and synchronized to facilitate parameter sharing. 
Therefore, before the training process, scans all devices to determine the device group D_i for each parameter W_j, which represents W_j is shared and should be synchronized within group D_i.
For efficiency, manages parameters with the same device group collectively and maintains a global parameter device group pool {D_i→{W_j}}, where each device group D_i corresponds to a set of parameters {W_j}.


  * Training Step. 
After the first three steps, the training process is ready to begin. In each iteration of , each device executes the forward and backward propagation of the local computation graph in a -by-manner, which is comprised of the interleaved execution of and transmission of data flow.
Following the forward and backward phases, performs group-wise parameter synchronization to maintain the parameter consistency. 
Specifically, each parameter set {W_j} is synchronized within its corresponding device group D_i in the parameter device group pool.






































































§ IMPLEMENTATION






is an efficient and scalable MT MM training system built on PyTorch with 10K Loc in Python: 2.1K LoC for the execution planner and 7.9K LoC for the runtime framework.
We implement the data flow transmission with NCCL batched P2P primitives and the parameter device groups with NCCL communication groups.
provides the users with simple, user-friendly and flexible API for defining MT MM training workloads.
Specifically, training tasks in are represented as Task, and users can define various multi-modal tasks by customizing PyTorch modules and connecting them flexibly through the add_flow API in . 
For example, a user can create a vision task by linking a vision encoder with a language model, or an audio task by linking an audio encoder with a language model.
Alternatively, user can also define different computational logic for various tasks implicitly within a single unified model. 
can automatically split the modules and construct Tasks via PyTorch FX Tracer, streamlining the process of task definition.
After the definition of multi-modal tasks, conducts the optimization workflow automatically, as illustrated in Fig. <ref>, and the runtime engine provides efficient and scalable model training process.


















§ EXPERIMENTS




 §.§ Experimental Setups
 









































































  §.§.§ Competitors.

We evaluate the efficiency of by comparing it with state-of-the-art distributed training systems, Megatron-LM <cit.> and DeepSpeed <cit.>. 
As discussed in <ref>, these systems are primarily developed for single-task training and do not cater specifically to the complexities of multi-task multi-modal training scenarios. 
To further explore the advantages of 's flexible resource allocation and scheduling capabilities, we introduce several baselines implemented on that represent typical strategies for multi-task training. 
The features of these competitors are summarized in Table <ref>.







(1) Megatron-LM:
Megatron-LM <cit.> is a widely used system that employs hybrid parallelism schemes, such as 3D parallelism, tailored for single-task training. 
The out-of-the-box approach to multi-task multi-modal training on Megatron-LM involves executing each task sequentially in each iteration, with each task utilizing the entire cluster.
Need to rewrite as decoupled training of multiple tasks



(2) DeepSpeed:
DeepSpeed <cit.> is another popular distributed training system leveraging sharded data parallelism (ZeRO <cit.>) and, similar to Megatron-LM, is typically optimized for single-task training. 
Its out-of-the-box paradigm for multi-task multi-modal training is also sequential task execution, with each task consuming all cluster resources.



(1)&(2) Megatron-LM & DeepSpeed:

Megatron-LM <cit.> and DeepSpeed <cit.> are widely used state-of-the-art training systems tailored for single-task training. 
The naïve approach to train MT MM models on these systems is to decouple all sub-models on separate devices (<ref>), which requires plenty of resources and is impractical. 
Therefore, we decouple sub-models on temporal dimension within each iteration, where each sub-model takes up the whole cluster within a short time period, and is dependently and sequentially executed.



(3) -Seq:

This baseline on allocates all available devices to each task and execute tasks sequentially within each iteration, similar to Megatron-LM and DeepSpeed. 

It reflects the performance of our system without specific optimizations for MT MM workloads. 


(4) -Uniform:

This baseline demonstrates a basic, workload-unaware task-level resource allocation strategy for multi-task multi-modal training.
It allocates available devices uniformly to each task, and executes each task in parallel within each iteration. 





(5) -Optimus:

This baseline represents a workload-aware task-level resource allocation strategy, which adapts allocations according to the workload at the task level granularity.
It's inspired by Optimus <cit.>, an effective cluster job scheduling system which proposes a greedy resource allocation scheme and iteratively assigns devices to the job that has the largest marginal gain.

Despite differences between job scheduling and multi-task training (<ref>), we apply a similar principle and devise the marginal gain as (T^_m(n) - T^_m(n')/(n'-n), i.e., the task completion time reduction scaled by the allocation increment from n to n'.
Here n' is the next valid allocation number larger than n.
This baseline is aware of inter-task heterogeneity, whereas unaware that of intra-task.


















(6) -STMM:

This baseline represents a naïve multi-task (MT) extension of single-task (ST) multi-modal (MM) model training systems.
It decouples multi-tasks, and for each single MM task it allocates resources to different modality encoders, akin to DistMM <cit.>, a recent system designed for ST MM models. 
Then it executes tasks sequentially.

Contrary to -Optimus, -STMM is aware of intra-task workload heterogeneity, whereas unaware that of inter-task.




















  §.§.§ Experimental Workloads


We conduct experiments on three different workloads of MT MM models, namely Multitask-CLIP <cit.>, OFASys <cit.>, QWen-VAL <cit.>.


The configuration of these models are summarized in Table <ref>.





(1) Multitask-CLIP:

Multitask-CLIP is a generalized version of CLIP <cit.>, which extends CLIP to 6 modalities and multiple contrastive learning tasks of paired data modalities.

We utilize the same model structure and configuration of ImageBind <cit.>.






We select 10 different contrastive learning tasks for evaluation, each with distinct workloads.



(2) OFASys:




OFASys<cit.> is a more general MT MM training workload, allowing modalities and tasks to activate the model components flexibly as needed. 
OFASys utilizes modality-specific adaptors for different modalities, e.g., ViT for vision data, and adopts a unified encoder-decoder LM with generative loss.
We select 7 different multi-modal tasks for evaluation.







(3) QWen-VAL:

QWen-VAL is a larger-scale MT MM model with up to 9.25 billion parameters, supporting three modalities, including text, vision, and audio. 
It adopts the same structure and configuration of the popular open-sourced multi-modal LLMs, QWen-VL <cit.> and QWen-Audio <cit.>. 

It has modality encoders for audio and vision, and the extracted modality-specific features are combined with text tokens and together fed into the unified LLM, QWen <cit.>.
We select three tasks for evaluation, i.e., vision-language (VL) task, audio-language (AL) task, and vision-audio-language (VAL) task, representing different combinations of modalities. 


































  §.§.§ Protocols

We conduct all the experiments on an 8-node GPU cluster. Each node consists of 8 NVIDIA A800 80 GB GPUs equipped with NVLink, and the nodes are interconnected by 400 Gbps InfiniBand network. 

Since the baseline systems do not support automatic planning given a targeted MT MM model training workload, to achieve a fair comparison, we manually tune their parallel configurations and memory optimization techniques (e.g., data parallelism degree, tensor parallelism degree, ZeRO stage, activation checkpointing, and etc.) to achieve the best performance.



For each system on each workload, we evaluate the system performance on different cluster sizes and report the iteration time averaged over 100 iterations. 
































 §.§ End-to-End Performance

















    







    

    






















Fig. <ref> displays end-to-end comparisons between and baseline systems across various model workloads, multi-modal task configurations, and cluster sizes. 





  §.§.§ Comparison with SOTA systems.

In general, compared to state-of-the-art (SOTA) training systems, i.e., Megatron-LM and DeepSpeed, achieves speedup ratios of up to 67% and 71%, respectively.
Below we delve into the performance advantages of . 





To begin with, consistently outperforms the competitors across different task configurations and numbers of tasks. Notably, excels when handling a larger number of tasks.

When comparing with Megatron-LM and DeepSpeed on the 10-task Multitask-CLIP workloads, achieves speedup ratios ranging from 31% to 63% and 33% to 71% compared to Megatron-LM and DeepSpeed, respectively.
Similar results are shown for the 7-task OFASys workloads, with the speedup ratios ranging from 31% to 67% and 33% to 71%, respectively.
This underscores 's excellent scalability with increasing task counts.




In addition, consistently achieves optimal performance across various cluster sizes.
For instance, on Multitask-CLIP, compared to SOTA systems, achieves the highest speedup ratios of 37%, 33%, and 71% on 8, 16, and 32 GPUs, respectively. Similarly, on OFASys, achieves acceleration ratios up to 71%, 46%, and 51% on 8, 16, and 32 GPUs, respectively. These results highlight 's excellent scalability w.r.t. cluster size. Notably, maintains high efficiency even when the scalability of SOTA systems begins to diminish — that is, when the increase in resources does not correspond to significant speed improvements. For example, in the 4-task Multitask-CLIP scenario, expanding the cluster size from 16 to 32 GPUs results in only modest speedup of 1.21× and 1.17× for Megatron-LM and DeepSpeed, respectively, whereas still achieves a 1.45× speedup. This efficiency stems from ’s carefully designed resource allocation and scheduling mechanisms.
Unlike existing systems that naively allocate all resources across all operators and tasks, ensures that each operator is allocated suitable resources when the cluster size increases, to maintain high computational efficiency.




More importantly, also exhibits excellent scalability w.r.t. model size. 
On larger models QWen-VAL with 9.25 billion parameters, achieves a maximum speedup of 1.16× on 32 GPUs and 1.63× on 64 GPUs, compared to SOTA systems. 
Notably, when training the QWen-VAL over 64 GPUs, shows remarkable scalability: it achieves a 1.78× speedup when scaling from 32 to 64 GPUs, whereas Megatron-LM and DeepSpeed only achieve 1.27× and 1.26× speedups, respectively. 
This is unsurprising since allocates cluster resources across different operators more flexibly, thereby avoiding the unsatisfactory scalability of with light workloads, as discussed in <ref>.














  §.§.§ Comparison with other baselines


Next, we discuss the performance of the variants of . 

Since -Seq has a comparable performance against Megatron-LM and DeepSpeed in most cases — which is reasonable as all three counterparts execute tasks sequentially — we focus on the comparison with task-level resource allocation strategies, i.e., -Uniform and -Optimus, as well as the single-task strategy, i.e., -STMM.







We find that the workload-unaware uniform allocation of -Uniform performs well only in limited scenarios, achieving a maximum speedup ratio of 28% over DeepSpeed. 
This suggests that resource allocation can enhance computational efficiency to some extent. 
However, it generally underperforms compared to SOTA systems due to its tendency to distribute resources evenly, leading to unbalanced workloads across tasks and system performance being constrained by the most resource-intensive tasks. 


In contrast, -Optimus, which allocates resources based on task workloads, shows better performance, especially in larger-scale cluster scenarios, with the speedup ratio up to 44% compared to DeepSpeed. 
However, there are still many scenarios where -Optimus underperforms, sometimes even falling behind DeepSpeed. 
This is because -Optimus’s task-level resource allocation overlooks the workload heterogeneity within tasks, thereby limiting training efficiency. 
Moreover, its coarse granularity of task-level allocation can sometimes fail to achieve ideal load balancing among tasks, often resulting in performance being constrained by the slowest task. 




In comparison, the operator-level strategy employed by enables finer-grained resource allocation and load balancing, consistently achieving higher efficiency compared to task-level strategies. 
We find that even in scenarios where -Optimus has already surpassed the performance of SOTA systems, still manages to achieve a speedup ratio of up to 45% over -Optimus (4-task Multitask-CLIP on 8GPUs), verifying its superiority.


As for -STMM, we find it perform better than SOTA systems in most cases, with the speedup ratio up to 20%, benefiting from its intra-task workload awareness and resource allocation.
However, it's designed for single-task (ST) multi-modal (MM) models, which decouples tasks and optimizes each task separately, and such single-task strategy is not the global optimum for multi-task cases. 
The lack of awareness of inter-task heterogeneity limits its performance, causing it to underperform compared to the task-level strategy -Optimus in many cases.
For OFASys, -STMM shows almost similar performance to SOTA systems. 
This is because -STMM gains acceleration by parallelizing sub-models of the multi-tower structure.

In contrast, OFASys utilizes a lightweight text adaptor, so most tasks that pair a modality with text are dominated by the other modality, making the intra-task parallelization of sub-models ineffective.
Compared to -STMM, jointly optimizes the allocation and scheduling of all tasks and operators, taking into account both intra-task and inter-task workload heterogeneity. This enables to consistently outperform the single-task strategy of -STMM, achieving a speedup ratio of up to 59%.















 §.§ Case Study








To better understand the advantages and performance gain of over the other competitors, we further conduct an in-depth case study of Multitask-CLIP (4 tasks, 16 GPUs). Fig. <ref> presents system performance considering three key metrics: cluster average utilization over time, average utilization per device, and computational utilization of each .

Firstly, -Seq, which executes the tasks sequentially with all resources, experiences fluctuating utilization due to the workload heterogeneity, leading to generally low overall utilization. 
-Uniform, which allocates resources uniformly at the task level, improves cluster utilization to some extent at the iteration beginning, 

but as tasks with light workloads finish, more devices become idle, declining overall utilization. 
-Optimus partially mitigates this imbalance with workload-aware allocation, though it still suffers from utilization drops due to its coarse granularity of task-level allocation. 

-STMM manages to enhance utilization via intra-task resource allocation for each task compared to -Seq, but the ignorance of inter-task heterogeneity limits its utilization.



In contrast, maintains consistently high utilization and the shortest iteration times thanks to its joint optimization of the unified computation graph of all tasks and operators, which addresses the heterogeneity both within and among tasks.






Furthermore, significantly elevates the utilization of all devices and all , showcasing its superior handling of workload balance through operator-level strategies. 
In contrast, -Seq shows lower utilization across all devices and . 

Although task-level strategies can enhance the computational efficiency of certain devices, the coarse granularity of allocation inevitably leads to workload imbalances, leaving many devices underutilized, sometimes even worse than -Seq, resulting in poor average cluster utilization.

-STMM also improves the utilization of certain devices and , but the results are still unsatisfactory as it fails to capture the workload differences among tasks, and fails to reach the global optimal allocation and scheduling plan for multi-tasks.




Overall, 's unified optimization of MT MM models captures both intra-task and inter-task heterogeneity, and effectively balances workloads. 

Thus, it consistently enhances utilization across all operators and devices, and maintains high computational efficiency across the cluster.








 §.§ Time Breakdown






Fig. <ref> shows the runtime breakdown for and -Seq across various workloads, primarily consisting of forward and backward propagation, parameter synchronization, and inter-send and receive. We've isolated parameter synchronization from the backward phase for individual analysis.
In MT MM training, we find that forward and backward propagation dominate the runtime, typically accounting for 80%-95% due to the large number of tasks and computational demands. focuses on reducing this significant time component through flexible resource allocation and scheduling. 
Parameter synchronization usually consumes a small fraction of the time, about 5%-15%, since it only occurs after accumulating gradients from multiple tasks.

Notably, consistently achieves equal or lower synchronization cost compared to -Seq. For instance, on 32 GPUs with QWen-VAL, cuts synchronization time to just 55% of that of -Seq. 
Although not the primary optimization focus, 's operator-level design inherently reduces synchronization overhead. This is achieved by synchronizing each parameter only within the device group that activates it and leveraging 's device island placement to convert potentially inter-island synchronizations into more local communications within device islands. 
Furthermore, we find that while introduces extra overhead for inter-send and receive, this overhead remains minimal, typically not exceeding 6%, thanks to the device placement mechanism that avoids unnecessary communications. 
Detailed ablation study of device placement is in <ref>.







 §.§ Component Analysis








  §.§.§ Optimality Analysis of Execution Planner


We analyze the optimality of execution planner in Fig. <ref>. 
Specifically, we compare the iteration time of to the theoretical optimal time C^∗ derived from Theorem <ref> in <ref>. 
As discussed in <ref>, when relaxing the constraints of the optimization problem (<ref>) and allowing the continuous divisibility of GPU resources n and operator number l, Theorem <ref> offers the theoretical optimum  and corresponding optimal time C^∗. 
Such a solution is unachievable due to these relaxed constraints, but serves as a theoretical upper bound of performance. 
The execution planner preserves most of the optimum property when finding the practical solution (e.g., Cond. (<ref>) and Cond. (<ref>) in <ref>), but still introducing minor biases (e.g., reinstating l's to integers in <ref>, resource extension in in <ref>). 
In Fig. <ref>, we calculate and estimate the theoretical optimum C^∗ according to Theorem <ref>, and compare it with 's iteration time. 
We find that across various task configurations and cluster sizes, the deviation between and theoretical optimum is consistently low, below 7%. 
This observation underscores the effectiveness of in offering a practical and near-optimal execution plan for MT MM models.
Besides, efficiently generates the execution plans within 3 seconds across all experiments, which is negligible compared to model training time.





  §.§.§ Ablation on Device Placement
 

We conduct an ablation study on the device placement strategy in <ref>, focusing on its impact on inter-communication overhead, which is the extra overhead introduced by our system. 

Specifically, we compare 's device placement strategy with a sequential placement strategy, which naïvely assigns each with consecutive devices. 
Our results indicate that the inter-communication overhead of the sequential placement strategy is approximately 3-6 times greater than that of , taking up to 27% of the end-to-end training time, which is considerably high.
However, with 's placement strategies, this overhead only takes up to 6%.
This demonstrates the effectiveness of our locality-aware placement, which significantly reduces the extra communication overhead.



















































 §.§ Dynamicity Performance






We evaluate the performance of various systems during dynamic changes of the multi-task workloads, a common occurrence in MT MM training. For instance, tasks with fewer training data may exit early, and new tasks may join partway through training. We simulate these dynamic changes by altering the training task set.

When the multi-task workloads change, the current model is first saved, and the new set of tasks and the saved model is loaded to continue training.
Fig. <ref> illustrates the performance of each system under such conditions. consistently achieves optimal training efficiency and the shortest overall training time. This advantage is due to 's adaptability to dynamically changing workloads, enabling it to adopt an appropriate execution plan for the efficient training of MT MM models.




§ RELATED WORKS
 


  §.§.§ Cluster Scheduling for DL Jobs









GPU clusters often design cluster schedulers to coordinate resource allocation and the execution order among multiple DL jobs. 
Some cluster schedulers <cit.> allocate resources to jobs based directly on user-specified requirements. 
Others <cit.> automatically allocate resources to each job based on the job scalability to the computing resource. Many of these schedulers aim to minimize job completion time (JCT). 

For instance, Optimus <cit.> introduces the concept of marginal gain to guide resource allocation, aiming to minimize job completion time.
Here, we highlight the difference of these works and MT MM model training. 
Unlike the independence among jobs in cluster scheduling, MT MM training involves execution dependencies among tasks.
Furthermore, while traditional scheduling focuses on job-level allocation, MT MM training requires finer-grained strategies to address intra-task workload heterogeneity.







  §.§.§ Training Optimization on Heterogeneous Cluster

This line of research focuses on optimizing the distributed training efficiency of DL models on heterogeneous GPU clusters <cit.>. 
While these works primarily concentrate on optimizing single model training and address hardware heterogeneity, mainly focus on more complex MT MM models, and addresses the challenges posed by the workload heterogeneity of MT MM models.





  §.§.§ Data and Model Management Optimization

The data management community has developed effective systems for managing data and models in specific domains, such as graph-structured data and models <cit.>, recommendation system data <cit.>, tabular data <cit.>, and video data <cit.>, etc. However, no existing work focuses on optimizing multi-task (MT) multi-modal (MM) data and model management, which is the key problem that addresses.




































































§ CONCLUSION















Efficient training of MT MM models faces significant system challenges due to the workload heterogeneity and complex execution dependency. 
In this paper, we propose to facilitate efficient training of MT MM models via data heterogeneity-aware model management optimization, which jointly optimizes heterogeneity-aware workload parallelization and dependency-driven execution scheduling.
Extensive experiments demonstrate the consistent superior performance of , outperforming existing state-of-the-art training systems with speedup ratio up to 71%.











ACM-Reference-Format












§ DETAILS OF SCALABILITY ESTIMATOR
 
characterizes the execution time of m over n devices, T_m(n), by a generalized piecewise αβ function:



    T_m(n) = α_m,i + β_m,i×c_m + β_m,i'×w_m/n,∀ n ∈ [n_i-1, n_i], i = 1,…,k


where k is the number of pieces, α_m,i represents the coefficient of fixed overheads (e.g., kernel launch costs), β_m,i and β_m,i' represent the reciprocal of execution efficiency (e.g., GPU computation speed and network bandwidth), w_m/n denotes the distributed workload of m across n devices (e.g., computational workload), and c_m denotes the workload that doesn't scale with n (e.g., communication volume of data parallelism).
Such piecewise function indicates that under varying resource scales, due to changes in the per-device workload, coefficients such as 
α, β and β' might differ, as the invoked kernels may vary across different workloads.





§ DETAILS OF BISECTION SEARCH FOR OPTIMUM OF CONTINUOUS PROBLEM

Alg. <ref> illustrates our bisection search algorithm to solve the optimum of malleable project scheduling problem, MPSP.
The function  finds the value of T_m^ (C).
It first finds the closest valid allocations of m, denoted as n_m and n_m, such that C∈ [T_m(n_m),T_m(n_m)].
It then returns

    n_m = (C - T_m(n_m)) ·n_m +  (T_m(n_m) - C) ·n_m/T_m(n_m) - T_m(n_m),

which is the linear combination of n_m, n_m
such that T_m(n_m) = C.






§ COMPARISON ON SINGLE-TASK MULTI-MODAL WORKLOAD




We also compare with baseline systems on single-task (ST) multi-modal (MT) scenario, which is a special case of MT MM training, as shown in Fig. <ref>. We are pleased to observe that even in the single-task scenario, outperformed SOTA systems by up to 48%. This is attributed to 's fine-grained, operator-level resource allocation and scheduling, which recognize not only the inter-task workload heterogeneity but also intra-task operator workload variations — a capability beyond the reach of task-level strategies as well as SOTA systems.




It's worth noting that -STMM has similar performance to on ST MM scenario, which is reasonable.







§ MEMORY CONSUMPTION






We also conduct a comparative analysis of memory consumption between and the other competitors. Fig. <ref> depicts the memory usage for each device in the scenario of Multitask-CLIP (4 tasks, 16 GPUs). Our findings indicate that generally exhibits lower memory consumption than SOTA systems such as Megatron-LM and DeepSpeed. This efficiency stems from 's operator-level strategy and selective parameter storage feature, where only devices that activate a specific operator need to maintain its corresponding parameters, thereby minimizing redundant storage. Additionally, we've observed that task-level strategies, e.g., -Optimus, experience significant memory imbalances. This issue is even more pronounced with -Uniform, which we have omitted in Figure <ref> for clarity. 
In contrast, maintains an excellent balance of memory consumption across devices, a success that is attributed to our device placement strategies.






























