raise accuracy. Emerging non-technical business needs, such as privacy and regulatory compliance, may also necessitate learning over groups. For instance, online advertising platforms build disaggregated partner-specific models, with each groups' training data organized as a separate pipeline [15].

During ML model building, model selection is typically inevitable to control underfitting vs overfitting, e.g., via hyperparameter tuning [33,57]. Practitioners often compare tens to hundreds of models [16,48]. Learning over groups only amplifies this load many times, since model selection is needed for each group's model. For instance, if one compares 30 models on a group and has 50 groups, this whole process results in the training of 1500 models. At this scale, it is impractical to be building models one by one. High-throughput parallel ML systems are needed to train models en masse.

Example. Consider a data scientist at a Web advertising team modeling click-through rate. She tries a logistic regression model on the whole population. She then has a hunch that separate models per country and mobile platform can raise accuracy, as Figure 1(A) shows. She materializes each group's data subset on a distributed platform such as HDFS and then runs model selection for each group. After all model configurations of all groups finish, she picks the best model per group for further analyses.

As the example shows, learning over groups proceeds in two steps: (1) ETL to create and load groups' data subsets and (2) schedule ML model selection for all groups. Given the high volume of models to be trained, parallelism on a cluster is critical. There are two dominant existing approaches to execute this workload in parallel: task parallelism (TP) or data parallelism (DP). We now briefly explain both of them, contrast their tradeoffs to explain why we need new approaches, and then present our novel approach. Existing Approaches and Drawbacks. ETL for TP needs each group's data subset in its entirety. Then, randomly assign each worker the full dataset of a group. Repeat for each group. TP must balance data sizes assigned to workers to avoid imbalances. But in full generality, this becomes an NP-Hard multi-way partitioning problem [55]. Note that TP can raise the storage footprint substantially, since the dataset is fully copied across workers. In contrast, ETL for DP is more straightforward: split each group's data subset evenly across all workers. DP does not raise the storage footprint.

Scheduling model selection works as follows. TP spawns training for each group independently on a worker; workers do not talk to each other. But due to imbalanced task assignments and training times on workers, TP often results in idle times. This is especially problematic when group sizes are disproportionately distributed. In contrast, DP trains each model of each group using the power of all workers simultaneously. But this can incur enormous communication cost for distributed training on a cluster, especially for ML workloads based on stochastic gradient descent (SGD) [38,56]. DP is also overkill for small groups. In general, real-world datasets typically have a mix of both small and large groups.

Recent work in the DB literature proposed a hybrid of TP and DP named Model Hopper Parallelism (MOP) for SGD workloads [34,48]. Naturally, one may wonder if MOP can resolve the above TP-DP dichotomy. Alas, we find that MOP does not directly suit the setting of learning over groups for two reasons. First, MOP is akin to DP in sharding datasets across all workers. This is still overkill for small groups. Second, MOP is tied to SGD's access patterns. But we seek to study learning over groups for other key ML access patterns too.

Desiderata. We seek to optimize learning over groups with the following desiderata. (1) Generality: Benefit multiple kinds of key ML access patterns. (2) Scalability: Scale along multiple axes, including number of groups, group sizes, cluster sizes, and model selection search space sizes. (3) Efficiency: Avoid the issues of TP, DP, and MOP, while retaining all their benefits. (4) Non-disruptive integration: Ideally, achieve all the above without needing to change the internal code of existing popular ML systems.

Our Approach. We introduce the analogue of GROUP BY for ML at scale, which allows bulk specifications of model selection for groups at once. For the underlying physical layer, we perform an in-depth analysis of the access patterns of learning over groups for 3 popular classes of ML: generalized linear models (GLMs), deep learning (DL), and gradient boosted decision trees (GBDT).

We devise analytical cost models to account for multiple aspectscomputation, network, and memory-to compare the efficiency of alternative approaches. Using our analytical models, we explored the tradeoff space thoroughly and gleaned insights for designing our new approach, which we explain next.

We call our approach grouped learning. It has three parts: (1) a simple and highly general epoch-level scheduling template, (2) a non-uniform data partitioning strategy for ETL, and (3) a novel form of parallel ML execution we call Gradient Accumulation Parallelism (GAP). The first part is inspired by Cerebro [48] and helps us meet the desiderata of generality and non-disruptive integration. It lets us unify all 3 of those classes of ML and support multiple forms of model selection. We support grid/random searches for now.

For the second and third parts, we formalize our optimization problem from first principles as an MILP. It is NP-Hard. So, we decompose ETL and scheduling. For ETL, we adopt and adapt an algorithm called wrap-around from the operations research literature [31]. It lets us cut communication costs substantially, while avoiding replication. Scheduling is tied to data placement. Most small groups are trained in a task-parallel manner. Large groups get sharded-typically non-uniformly-across workers. For such groups, we use GAP to reduce communication costs further. GAP is a new form of "bulk asynchronous" parallelism, a sibling of MOP. We prototype all of our ideas into a system we call Kingpin on top of the state-of-the-art distributed computation engine, Ray [46]. Figure 2(C) illustrates our system architecture. We evaluate Kingpin empirically on two large ML benchmark datasets: Criteo [14] and Cityscapes [13]. Kingpin matches or outperforms existing approaches to learn over groups, with speedups up to 14x for GLMs and GBDT and 4x for DL. Deeper analysis of resource utilization logs also validates that Kingpin's gains come from avoiding unnecessary network communications and disk I/O and reducing idle times on workers, thus affirming the benefits of our new techniques. In summary, this paper makes the following contributions:

• To the best of our knowledge, this is the first work to enable and optimize the analogue of GROUP BY for ML at scale, a process we call learning over groups. • We perform an in-depth analytical comparison of existing approaches to learn over groups and their tradeoffs for 3 main classes of ML: GLMs, DL, and GBDT.

• Based on our analyses, we devise a novel approach, grouped learning, that mitigates the issues of existing approaches, while still being easy to implement. • As part of our approach, we adopt and adapt the wrap-around algorithm for shard placement and devise a novel form of bulk-asynchronous parallelism, GAP. • We implement our ideas in a scalable ML system Kingpin, built on top of Ray and existing ML tools. A thorough empirical evaluation shows that Kingpin matches or surpasses strong existing baselines.

Gradient-based optimization is a highly popular mathematical substrate in ML. Many popular ML models are defined as minimization problems over model parameters and training data [26]. Given training data 𝐷 = {(𝑥 𝑖 , 𝑦 𝑖 )} 𝑛 𝑖=1 , 𝑥 𝑖 , 𝑦 𝑖 denote features and targets; 𝑛 is the number of examples. Many ML models compute argmin 𝑤 𝐿(𝑤), where 𝐿 is their loss function and 𝑤 are model parameters, using optimization procedures such as gradient descent. We now explain them briefly; for more mathematical details, please see [18,50].

The simplest optimization procedure is batch gradient descent (BGD). It does a full pass over s𝐷 to compute the gradient of 𝐿 at an initial 𝑤 (0) , given by: ∇𝐿(𝑤) =  𝑛 𝑖=1 ∇𝑙 (𝑤 𝑇 𝑥 𝑖 , 𝑦 𝑖 ). The update step follows: 𝑤 (1) ← 𝑤 (1) -𝛼∇𝐿(𝑤 = 𝑤 (0) ); this is the "descent." BGD repeats this process multiple times, each called an epoch, until convergence to an optimal. Second-order batch methods are more popular than BGD for convex losses, like in GLMs [39]. They compute the second derivative of 𝐿, the Hessian, too. A popular secondorder method is limited-memory BFGS (LBFGS). It approximates the Hessian with a fixed-size history of the last 𝑘 gradients.

Stochastic Gradient Descent (SGD) is more popular than BGD for non-convex losses, like in DL, especially over unstructured data. One takes a random sample of 𝐷, called a mini-batch, to estimate ∇𝐿. Sampling is typically done without replacement, which can be done at scale using a random shuffle of 𝐷 followed by a sequential pass. SGD typically converges much faster than BGD on large-scale data because it performs many updates to 𝑤 within a single epoch. Optionally, the dataset is reshuffled between epochs [18].

Finally, GBDTs are a popular form of ML over structured data that also use gradients. They build an ensemble of weak models (typically short decision trees) to minimize 𝐿. A difference to BGD/SGD is that while GD descends along the gradient to update model parameters, gradient-boosting in GBDT performs descent against the gradient by adding new models. Thus, GBDT needs a full pass over 𝐷 to run inference using all existing models.

Looking across the 3 common forms of optimization in ML, we can organize them into 3 main kinds of data access patterns based on amenability to distributed execution: 1) Algebraic: BGD and LBFGS compute algebraic aggregates, akin to SQL SUM [23]. Thus, an epoch is easily parallelized on a cluster with independent shards and workers; partial gradients are added by an aggregator in the end. This is logically equivalent to singleworker execution.

2) Sequential: An SGD epoch needs sequential access to the shuffled dataset. During a pass, each mini-batch gradient update depends on the previous, making SGD inherently hard to parallelize at the full-dataset scale. We call this pattern sequential-parallelizable. Techniques such as Parameter Server [38] and Horovod [56] are sometimes used for data-parallel SGD.

3) Sampling: In this pattern, gradients used for updates are downsampled from the full dataset. This is typical in distributed GBDT tools such as LightGBM [30]. The histogram-based GBDT performs parallel gradient computations as with LBFGS and constructs multiple histograms from local gradients [19]. Then it finds the best splits on the combined global histogram.

Unless otherwise mentioned, we assume that 𝐷 fits in the cluster's total memory. We do allow data spills to disk; this is treated as an extension, explained later. We assume the ETL step materializes group data subsets and stores them in cheap (possibly ephemeral) networked storage. We fetch data during training by reading over the network; we cache data on workers if possible and necessary. For model selection, we focus on grid/random searches with a fixed number of hyperparameter configurations (configs). Grid/random searchers are the most popular forms of models selection in practice [11] but as we will explain later, Kingpin's design is amenable to easily plugging in AutoML procedures in future work.

Table 1 lists our notation. For simplicity sake, we assume all groups reuse the same 𝑠; this is also standard in ML practice. For simplicity of exposition in the next section, ℎ denotes number of training epochs with one variable but Kingpin supports convergencebased criteria as well and does not need a fixed number of epochs up front. Likewise, the single 𝑚 is for simplicity of exposition; different configs/groups can have different model sizes.

The problem of learning over groups is the following: orchestrate the given model selection workloads for all given groups on a given provisioned cluster. The goal is to minimize completion time, also called makespan. We want to satisfy all 4 desiderata listed in Section 1. This is challenging from a formal standpoint because we need to kill two birds with one stone: ETL for non-uniform partitioning and scheduling of model selection. It is challenging from the systems standpoint because we seek to build a unified system for 3 key access patterns in ML across GLMs, DL, and GBDT.

We first explain the existing approaches for the 3 classes of ML in terms of access patterns. We distill them into analytical cost models to offer a more in-depth understanding. Our goal is not to build some sort of cost-based optimizer but to study and explore the tradeoff space analytically to derive insights that help us devise our new approach. Section 4 will present our approach, which dominates these existing approaches due to formally grounded reasons.

In the ML world, the most common way to train models of groups in task parallelism (TP). A group's dataset is copied to each worker. A scheduler gets each worker to train some groups and/or configs. Each worker trains its assigned config(s) until convergence.

We identify two flavors: group task parallelism (GTP) and model task parallelism (MTP), based on the granularity of tasks. The former defines tasks around groups; the latter, around configs. GTP places entire model selection of a given group on one worker to finish end to end. MTP breaks apart the model selection of a group to place one config on a worker at a time. MTP is more prevalent when compute resources are abundant.

Both flavors load a group's data over the network once. But MTP caches data on local disk so that it is not read again for a different config of the same group. Note that if there are more configs than workers (𝑠 > 𝑝), a worker may train more than one config per group even in MTP.

Compute Cost. Gradient computations would account for most of the compute time, while some cycles would also be used to update models. Therefore, we model the total compute cost as:

where 𝑓 grad (1) is the time to compute the gradient for one example, 𝑓 𝑢𝑝𝑑𝑎𝑡𝑒 is the time to perform update to models, and 𝜂 is the number of updates to models per iteration. 𝜂 = 𝑛 𝑏 for 𝑠𝑒𝑞𝑢𝑒𝑛𝑡𝑖𝑎𝑙 workloads, and 𝜂 = 1 for 𝑎𝑙𝑔𝑒𝑏𝑟𝑎𝑖𝑐 and 𝑠𝑎𝑚𝑝𝑙𝑖𝑛𝑔 workloads.

I/O and Memory Cost. Iterative ML computations result in rounds of scans over the whole dataset. Training on one group requires a full scan of the dataset's corresponding subset. For GTP, we cache each group's data in memory for training all configs. For MTP, if 𝑠 > 𝑝, we cache data on disk and load them 𝑠 -𝑝 times to train additional configurations. Thus, I/O and memory cost is:

where 𝑓 mem is the cost of reading one example in memory, 𝑓 disk is the cost of loading one example from disk and 𝑓 disk gives non-zero only in MTP.

Training models in TP requires little to no network communication except fetching data from storage before training.

So the network cost is a lump-sum cost of loading data to workers. For MTP, we copy the full dataset 𝐷 to each worker.

where 𝑓 network is the cost of fetching one example over network.

Idle Cost. Workers may go idle in GTP when too few groups and configs (or too many workers) are present, or groups are highly skewed. For MTP, skew is less of a concern if we have enough tasks to parallelize (𝑠 • 𝑔 > 𝑝). Random placement of tasks on workers will amortize such idle times. The max idle time MTP can have is when training the largest group for one config on one worker.

where 𝑛 max is the number of examples in the largest group 𝐺 max .

Total Runtime. We cannot merely add the three costs directly because the actual runtime depends on the underlying hardware. For example, in a memory-optimized cluster, the I/O cost will contribute less to the total runtime. Thus, we model total runtime with some cost parameters:

We tune the cost parameters 𝛼, 𝛽, 𝛾 and 𝛿 using offline calibration runs, akin to RDBMSs. In particular, 𝛽 may dominate when 𝐷 does not fit in memory. For the rest of this section, we reuse this template of analytical cost models and cost parameters.

This approach shards the dataset of each group on each worker. Thus, it improves upon TP by exploiting all workers for each config, which can reduce worker idling for contiguous periods. But this also means DP needs communication among workers to aggregate partial gradients across workers. So, DP's cost model differs from TP primarily in the network and idle costs.

Compute and I/O costs of DP are similar to GTP. Due to space constraints, we provide the details in the technical report [40].

Why are TP and DP not "good enough" for learning over groups? In short, each bakes in some tradeoffs on key scalability axes that make it substantially suboptimal in many realistic scenarios. Using our analytical cost models, we contrast them to expose such tradeoffs in Figure 3. Note that compute cost does not differ much across them; I/O costs do differ but can be optimized by faster storage and network. The key differentiating factors that make this tradeoff space non-trivial are inherent differences in communication cost/complexity and varying chances of worker idle times. The plots use cost parameters calibrated based on real empirical runs on Criteo dataset (Section 5); due to space constraints we provide their details in the technical report [40].

The overall takeaway is neither TP nor DP dominates the other on all axes. MTP is worse than both GTP and DP mainly due to its Country has 18 groups; ratio of largest group, 𝑛 max /𝑛, is 0.26. We set 𝑝 = 4 workers. For each plot here, we vary one variable, while fixing all other workload properties. repeated calls to reload data. GTP and DP are often comparable but have many crossovers. When the dataset size goes down or number of groups goes up, GTP dominates; but with more workers or more skew in group sizes, DP dominates. As a teaser for comparison, our approach in Kingpin (Section 4) is shown too-it matches or dominates all these alternatives on all these scalability axes. We summarize key systems design issues with both TP and DP:

(1) TP Suffers from Imbalances. GTP faces idle times due to imbalances in groups across workers. Its optimal scheduling problem is NP-Hard via a reduction from multi-way number partitioning [55]. But even with a near-optimal scheduling heuristic, group size skews will still cause idle times due to the fundamental indivisibility of a task. As an extreme example, if a dataset has one very large group, the time to train that group on a worker will utterly bottleneck the cluster. Thus, TP as practiced today has inherent efficiency limits.

(2) TP Wastes Storage/Network. MTP mitigates GTP's issues with imbalances by placing different configs of a group on different workers. But this requires a full copy of that group's data to each worker. While storage is relatively cheaper, it is still a concern at scale, e.g., 1TB blows up to 10TB on a 10-node cluster! Of course, reading from remote storage (e.g., S3) can avoid such blowups. But due to the iterative nature of gradient-based ML this ends up being highly wasteful of the network instead! Recent work showed that this approach can have even 100x higher network costs [48]. Thus, TP as practiced today has inherent scalability limits as well.

(3) DP has Inherent Communication Costs. DP avoids the above issues of TP by splitting each config across all workers. Alas, this leads to new bottleneck: inevitable communication rounds and synchronization to aggregate gradients. Our analytical cost models show that this cost grows linearly with many key factors: number of configurations 𝑠, number of epochs ℎ, number of batches (for SGD) and number of groups 𝑔. In fact, regardless of the class of ML, DP gets significantly slower as 𝑔 increases. Thus, DP as practiced today also has inherent efficiency limits.

We now dive into our approach, grouped learning (GL), a novel hybrid of TP and DP. We start with some intuition and an overview, then formalize our problem, present our algorithms, and finally present its analytical cost model and a key extension.

Our intuition is as follows: DP is too fine-grained, while TP is too coarse-grained. This leads to their fundamental issues laid out in Section 3.3. Philosophically, with GL we "take things apart" to go down from TP's level to avoid its issues (imbalances, resource bloat) and "put things back better" to go up from DP's level to avoid its issues (high communication costs). Concretely, GL reduces idle times from imbalances (vs TP), avoids data copying bloat (vs TP), and avoid needless communication/synchronization (vs DP). We achieve all this by assembling three things:

(1) A simple and general two-level scheduling template, inspired in part by Cerebro [48]. By decoupling per-epoch and acrossepoch scheduling, we can support many kinds of iterative gradientbased ML (GLMs, DL, and GDBT) and many model selection heuristics in a unified way, including with varying numbers of epochs per config. The template itself is not novel but our application of it to the problem of learning over groups is novel.

(2) A non-uniform data partitioning strategy for ETL based on a suitable algorithm from the operations research world: wraparound. It enables GL to holistically optimize data and computation placement across groups and configs at every epoch boundary. The wrap-around algorithm itself is not novel but our application of it to this large-scale ML systems setting is novel.

(3) A novel form of parallel ML execution per epoch: Gradient Accumulation Parallelism (GAP). To the best of our knowledge, GAP is only the second-known form of "bulk asynchronous" parallelism, inspired by its (complementary) sibling MOP [48]. GAP works for all 3 major ML access patterns laid out in Section 2.2.

The precise workflow of GL is as follows, given the cluster, full dataset, groups, and model selection workloads. First, partition the dataset to favor TP using wrap-around. A group's data is not split across workers unless really needed. If a group's data is "too" large, it is sharded across workers. Second, place the epochs of the current set of configs on the workers and shards based on the scheduler's decision. Third, execute the training of configs in a hybrid-parallel manner, using TP for those that can and GAP for those that cannot. Next, we explain the mechanics of GAP first and then formalize our scheduling and partitioning problems.

GAP does the following: given a large set of configs and potentially sharded data per group, execute gradient-based ML using a hybrid of TP and DP. The novelty of GAP lies in its operation on sharded data of one group. Basically, it accumulates gradients-related artifacts and ships them across workers asynchronously and in parallel. The degree of parallelism is determined by the number of configs and shards. We dive into how GAP works for each class of ML. Figure 4(C-D) illustrate the workflow on Algebraic and Sequential. Due to space constraints, we present the illustration of Sampling (GBDT) in the technical report [40].

(1) Algebraic (LBFGS): Split gradient computation algebraically as usual; multiple workers simultaneously train the same config to obtain local partial gradients. DP would aggregate partial gradients eagerly in one go, but GAP has this twist: aggregate partial gradients lazily and asynchronously over time, as dictated by our scheduler.

(2) Sequential (SGD): Parallelize at the granularity of a group; two workers can not simultaneously train the same config. This ensures logical equivalence to sequential SGD, helping accuracy. When a worker finishes training a group's config, it moves on to another config of that group while asynchronously shipping the previous config to another worker. In this sense, GAP is inspired by MOP.

(3) Sampling (GBDT): Parallelize the sampled statistics on workers as usual. We repeat swapping of parallelizable artifacts asynchronously among workers until GAP finishes one iteration of all configs of all groups, akin to shipping partial gradients in Algebraic.

Overall, GAP has the following invariants based on the ML access patterns, with the last 2 being potentially relaxed for Algebraic and Sampling workloads: (1) Data sharding: 𝐷 is sharded across workers, not fully copied. (2) Worker exclusivity: A worker handles only one config at a time. (3) Config isolation: A config is handled by at most one worker at a time. (4) Group isolation: Configs of a group do not run concurrently across workers.

Our scheduler tackles the two-fold problem of ETL/partitioning and scheduling of configs. We now state it formally as a mixed-integer linear program (MILP). Table 2 summarizes the extra notation. The objective is to minimize makespan 𝐶 to train all configs of all groups for one epoch. We do a pilot run to get the runtime of a config on a group's dataset to set 𝑉 .

𝐶,𝑁 ,𝑋 ,𝑂,𝑃,𝑄

(1)

𝑋 𝑖,𝑗 ≥ 0 (1b)

The variables to optimize over are the start times 𝑋 , group isolation assignments 𝑂, worker exclusivity assignments 𝑃, config isolation assignments 𝑄, and data partitioning 𝑁 . The constraints enforce the four invariants of GAP as listed above: (1c) ensures data sharding; (1d) and (1e) ensure group isolation; (1f) and (1g) ensure worker exclusivity; (1h) and (1i) ensure config isolation; (1a) and (1b) define makespan and sanity of start times.

One might now wonder if we can use an MILP solver such as Gurobi. But our problem turns out to be a fusion of two classic NP-Hard problems: multi-way number partitioning [55] and open-shop scheduling [22]. Since the total number of configs across groups can even be in the thousands, MILP solvers may be too slow.

Due to the above, we adopt the following two-step heuristic approach that is both efficient and offers near-ideal makespans in practice, certainly significantly faster than the prior art approaches (TP and DP). First, partition a set of numbers (in our case, data for groups) into a collection of subsets so that each collection's sums are as equal as possible. Second, find an optimal schedule to execute these collections/groups to minimize the makespan.

If we relax 𝑁 in the MILP to real space, the ETL part becomes more manageable. In the operations research literature such a relaxation is called "pliable shop" scheduling [31]. That prior work proposed a linear-time algorithm that is a good fit for our (relaxed) scenario: wrap-around algorithm. The idea is to compute the expected capacity/makespan and keep adding data for groups onto a worker. If adding a group would exceed the current worker's optimal capacity, we shard the data and continue adding the overflow part to next worker. The order groups are put onto workers itself entails a schedule to manage executions. Based on all of our notation so far, the optimal makespan will now be:

Prior work [31] also showed that the wrap-around algorithm is provably optimal; we refer the interested reader to [44] for their proof. Note that max 𝑖 ∈ [1,...,𝑔] (𝑛 𝑖 ) comes from the rationale that we cannot execute a job in parallel but in a strictly sequential manner. Overall, this algorithm is helpful to ensure sequential execution order on sharded data. If we end up needing to shard in between an example (non-integral splits), we round it by rolling over the extra example to the next worker. Just one example off from the optimal has virtually no impact on the makespan. Thus, we modify it slightly and present the pseudocode above. Figure 4(B) illustrates the algorithm. 𝑐𝑢𝑟 _𝑓 𝑖𝑙𝑙𝑒𝑑 += 𝑛 𝑖 , 𝑛 𝑖 = 0 19: Return 𝐴 data and hyperparameter search space. Workers read their assigned data over the network. GAP is used for training all configs for that epoch. We accumulate training artifacts and ship them around workers asynchronously. When a worker finishes training a config on its (shard of a) group, it propagates gradient-related results to another worker with the same group data. The last worker will finally add up all gradients, run the update step for the optimization procedure, and then broadcast updated models back to relevant workers. Thus, all workers are kept busy almost always during training instead of waiting for instructions from a centralized manager.

See Figure 4 again for the end-to-end illustration of how wraparound and GAP for all 3 classes of ML workloads. Note that synchronization over the network is only needed when a group's data is not entirely on one worker. All in all, grouped learning scales seamlessly on all axes explained in Section 3: dataset size, cluster size, number of groups, and group sizes.

We now present our approach's analytical cost models that we used for Figure 3 in Section 3. The compute and I/O costs are the same as that of DP; we skip them for brevity.

Network Cost. This is a key advantage of GAP: it does not incur much network cost, especially for SGD, even if a group's dataset is sharded across workers. The actual cost per epoch is as follows. Note it is is linear in the number of shards of a group, which in the worst case is only the number of workers 𝑝 (say, for a super large group). This is in contrast to DP for SGD (e.g., Horovod or PyTorch DDP), which are orders of magnitude higher.

Idle Cost. A worker may still be idle when waiting for its peers to finish their parts. But this synchronization happens only once per epoch, no matter how many groups or configurations we have.

So far, our approach can be easily layered onto existing ML popular frameworks for LBFGS and SGD-based ML (e.g., PyTorch) without modifying their internal code. But for GBDT, an implementation nuance is that this is not possible with just the pure wrap-around algorithm. We will explain the implementation part of this nuance later in Section 4.8. But for now we present an algorithmic extension that will let us support GBDT systems (e.g., LightGBM) seamlessly as well, preserving the high generality of our approach.

Basically, we decouple gradient computation and its subsequent use for the ML/optimization procedures. We then modify the MILP to account for this. There are only so many ways to partition a group's dataset: 1, 2, or upto 𝑝 shards. Formally, this is akin to setting a constraint on 𝑁 𝑖,𝑗 as follows:

∀𝑐 ∈ [1, ..., 𝑝], 𝑁 𝑖,𝑗 = 0 𝑜𝑟 𝑁 𝑖,𝑗 = 𝑛 𝑖 𝑛 • 𝑐 It is also equivalent to having the same upper and lower bound for each type of split. Even if we relax it to linear space, this bound setting makes the problem a variant of the pliable job scheduling problem, which is NP-Hard [31]. Thus, we propose a new heuristic. for each split do Q.push((𝑐𝑢𝑟 _𝑓 𝑖𝑙𝑙𝑒𝑑 + 𝑛 𝑖 / best_num_splits, 𝑤𝑜𝑟𝑘𝑒𝑟 ))

13: Return 𝐴 Constrained Wrap-around Algorithm. Intuitively, we want to split a group into as few shards as possible because the more shards there are, the higher the network cost. Ideally we do not shard a group at all, resembling TP, albeit with no copying. As a heuristic, we always split a group to its nearest ratio with the optimal makespan: 𝑟𝑜𝑢𝑛𝑑 (𝑛 𝑖 • 𝑝/𝑛). We call this heuristic the constrained wrap-around algorithm. Algorithm 2 presents the pseudocode. Figure 5 presents an illustration. The key difference now is that we first sort the groups in 𝐺 in decreasing order of size to decide whether to split large groups first. Smaller groups usually have that ratio < 1; so, we need not split them at all. Empirically, we find that this heuristic achieves near-optimal makespans for realistic workloads; due to space constraints, we present more details of this comparison in the technical report [40].

Kingpin is a new ML system that implements our ideas of grouped learning, including GAP and the wrap-around algorithm. As a strategic decision, we chose to prototype Kingpin on top of Ray, a recent highly scalable runtime engine for compute-intensive ML/AI workloads [46]. We weighed a prototype on Spark too but found Ray better for 2 reasons. First, Ray has new systems capabilities that avoid extra I/Os for intermediates relative to Spark, both for baselines (TP and DP) and for Kingpin. Second, Spark's JVM-based data caching leads to extra overheads for copying and shipping data to ML tools (we use Pytorch and LightGBM). Due to space constraints, we leave the detailed specifics in the technical report [40].

4.8.1 Overview. We adopt an extensible architecture that can talk to multiple ML tools, as in Figure 2(C). Ray helps parallelize workloads on physical machines. We still use Spark, albeit only for materializing group data subsets if the full dataset comes in raw.

A user specifies the following in Kingpin APIs: group metadata (with syntactic sugar), model_creator, optimizer_creator, hyperparameter search space, and cluster configuration (number of CPUs, GPUs, etc.). Our scheduler then runs the wrap-around algorithm and spawns workers to load corresponding data. All workers are started as native Ray actors connected by PyTorch's distributed primitives [41]. Intermediate results (e.g., gradient artifacts) are saved in Ray's in-memory object store, with the option to spill to disk for fault tolerance. In the end, users will get back serialized copies of all groups' models, including selected best models.

For ML training, Kingpin essentially acts as an abstraction enveloping popular ML tools without modifying their internal code, e.g., for tensor compilation or hardware optimizations. Currently, we support the APIs for PyTorch (for LGBFS and SGD) and Light-GBM (for GBDT). It is relatively straightforward to add support for their rival tools if needed (e.g., TensorFlow or XGBoost).

As a detail about the LightGBM integration, it ships histograms during distributed training. But these histogram data structures are baked deep into its C++ code and are hard to decouple. LightGBM provides developers an extensibility option to compute gradients using a custom loss function. But to use this, we need to rewrite existing loss functions written in C++ (e.g., cross-entropy) in Python, which may become a backward compatibility issue for our software. We decided to avoid this deeper dependency with a different approach: extend the wrap-around algorithm itself in our scheduler-as explained in Section 4.7-so that we can restrict ourselves to using only the user-facing APIs of LightGBM.

We now present an extensive empirical evaluation seeking to answer three questions. (1) How does Kingpin compare with the existing approaches of MTP, GTP, and DP on runtimes, accuracy, and resource utilization? (2) Did our analytical cost-model based discussion correctly predict the trends? (3) Does Kingpin offer good scalability on key axes?

Datasets. We use two large ML benchmark datasets with different groups: Criteo Sponsored Search Data [14] and Cityscapes [13].

Criteo is a product ad click binary classification task with numeric and categorical features. We pick Country and Partner as meaningful groups for 2 datasets. As standard ML practice, we use random hashing to convert categorical features into feature vectors. It took about 90 min to featurize each dataset. Cityscapes is a popular image dataset for semantic understanding of urban street scenes in German and Swiss cities. Table 4 lists the dataset and group statistics. Figure 6 provides the group size distributions. We hold out a random 10% of data in each group as its validation set. We also partition all data into small files (default: 25MB) to enable the wrap-around algorithm to virtually split data at "any" point.

Workloads. We run end-to-end test for all three ML access patterns.

(1) Algebraic: Logistic regression with LBFGS (history size 10). (2) Sequential: A UNet [53] neural network with 8 filters for semantic segmentation, Adam as the SGD method, and 𝑏 = 32. (3) Sampling: GBDT. For all 3, we run standard grid search for hyperparameter tuning. Table 3 Baselines. GTP, MTP, DP, and Switching Parallelism (SP) implemented on top of Ray's native ML infrastructure are our baselines. Ray offers native and robust TP scheduling, while RaySGD runs DP using PyTorch's DDP [41]. For GBDT, DP runs LightGBM in data-parallel mode on sockets [4]. We make all baselines as strong as can be with our best-effort engineering. All baselines use the same data loader we implement for Kingpin. Also, we scale DP to run more processes when a workload alone cannot saturate compute resources on one worker; the same applies to GTP. We spawn multiple configs to train in parallel on one worker using shared memory store in Ray [1]. For MTP, we do not perform this optimization because it will lead to training groups simultaneously on one worker. Loading all their corresponding data to memory at once is infeasible. SP is a middle ground that switches between DP and GTP. It runs DP on large groups and GTP on small groups. Large groups in our narrow definition are those in the 50th percentile.

We load all data shards to workers' local memory first before training for all experiments. Different approaches will trigger this process multiple times-thus, we must include data loading time from remote storage in the end-to-end runtimes. For GTP and MTP, we use a FIFO queue to schedule out groups and configs. For DP, we train one group at a time on all worker nodes. For GBDT, Kingpin uses the general extension in Section 4.7.

Total Runtimes. Table 5 shows the results. We see Kingpin is either significantly faster than the baselines in most cases cases, Overall, Kingpin dominates all the existing approaches. On Algebraic over Criteo-Partner (a much larger group), Kingpin's runtime is comparable to what it saw for Country, thus showing its ability to scale well with large numbers of groups. In contrast, DP becomes dramatically slower due to its inherently high communication costs. GTP is also comparable to what it saw for Country; MTP, slightly slower than before. SP is surprisingly faster than before as well as both GTP and DP.

On Sampling (GBDT) over Criteo for both grouping attributes, Kingpin is comparable to GTP. DP is marginally slower on Country but sees 4x slowdown on Partner due its larger number of groups. MTP's is still very slow, due to its repeated data loading overheads. SP's performance over Partner is between GTP's and DP's, while the result over Country is worse than both.

Finally, on Sequential (UNet) over Cityscapes, Kingpin is comparable to GTP; their GPU utilizations are also comparable. MTP is only 1.4x slower than Kingpin; this gap is lower now because deep learning is compute-bound and MTP's data loading disadvantage is less significant here. DP is 4.3x slower than Kingpin due to its inherent communication costs even though PyTorch DDP is a stateof-the-art system for data-parallel SGD. SP's runtime sits exactly in the middle of GTP's and DP's.

Learning curves. For Algebraic and Sampling, all approaches (GTP, MTP, DP, and Kingpin) have virtually indistinguishable convergence behaviors. So, we elide them for brevity. For Sequential (SGD), different approaches go through a group's dataset in different orders, likely leading to different learning curves. For all of them, we shuffled each group's dataset once up front, which often suffices for SGD [18]. Figure 7 shows the learning curves on the largest group of in Cityscapes (Wurzburg), as well as a CDF of the validation accuracies achieved by the best models for all groups after 10 epochs. The learning curves on the largest group largely overlap, showing that all approaches are suitable from the accuracy standpoint-their runtimes and resource efficiency are what differentiate them. Recall that Kingpin, just like GTP and MTP, ensures sequential-equivalence for SGD. In the CDF, we see that DP leads to somewhat worse models than the others; this is just an artifact of the batch size being fixed, since PyTorch DDP raises the "effective" batch size [48]. Tuning 𝑏 too will mitigate that.

Summary. Kingpin is the most resource-efficient approach for learning over groups across all 3 major ML access patterns we studied. It also scales well with many groups and is often significantly faster than prior art without sacrificing on accuracy.

5.2.1 Algebraic (LBFGS) on Criteo-Country. Figure 8 shows the average CPU usage, max memory usage per worker, and network usage for all compared approaches. Compute (CPU). Kingpin saturates the CPU most of the time while having about 10 falloffs overall. These falloffs are from the synchronization among all workers and occur at our epoch-level scheduling boundaries. Falloffs are more frequent in DP due to its exchanging of gradients after every group, config, and epoch. GTP makes the fullest use of CPU but gradually it sees some workers becoming idle. MTP has the strangest CPU usage with 12 spikes. After poring over the detailed logs, we find that these spikes correspond to 12 sets of hyperparameter configs. MTP evenly distributes configs instead of groups to workers. Training one group and config at a time is not able to saturate all CPU cores in this experiment.

Memory. We see that Kingpin ramps up memory at the beginning and keeps needed data in memory till the end. DP does the same, except that it loads one group at a time, leading to less memory used per worker. GTP loads all data to Ray's shared memory store and trains multiple configs by reading data from the store. Thus, we see its memory usage is also increasing sharply in the beginning and stays there till the end. MTP loads groups' data to memory repeatedly; so, we see 12 spikes on this plot too. Note that GTP tends to consume more memory due to extra serialization in Ray's shared memory store.

Network. These plots concur with our analytical cost models. Kingpin first reads data and then keeps gradients flowing around. DP also loads data and exchanges gradients after each epoch. GTP and MTP only load data over the network once; no more network usage is triggered during training, except MTP copies almost all groups' data to each worker. One unexpected observation here is that Kingpin triggers more network communications than DP but GAP's asynchronous nature amortizes them to yield significantly lower end-to-end runtimes anyway. 5.2.2 Sequential (SGD) on Cityspaces. We compare compute (GPU) and network usages, and Figure 9 shows the details.

Compute (GPU). Kingpin keeps GPU busy most of the time while having lots of falloffs in-between. These falloffs come from the overhead of switching models to train in GAP, leading to Kingpin making worse use of GPU than GTP. Alas, GTP leaves some workers idle towards the end. More skew in the dataset would magnify the idleness. MTP also has lots of falloffs when loading data repeatedly to memory. DP surprisingly saturates the GPU most of the time. We figure the reason being that NVIDIA's GPU utility library (NCCL) is in charge of communicating gradients.

Network. Overall, Kingpin and GTP have the minimum network usages. Both do not explicitly trigger network traffic except loading data and hopping models, though we still see about 100GB of additional network usage started by the underlying Ray infrastructure. MTP essentially copies the whole dataset to all workers, topping about 200GB of network usage. DP is the only approach that requires constant network communication within one epoch, as it uses all-reduce to exchange gradients for each mini-batch.

Worker Scalability. We now vary the number of workers for Algebraic on Criteo-Country to show their speedup behaviors. Figure 10 shows the results. The curves validate our analytical cost modelbased discussion for all approaches, with one difference being DP does not scale as well as we expected (its curve is flat here). As the number of workers increases, Kingpin and MTP see speedups, while MTP benefits the most from more workers. This is because MTP can easily parallelize any group and config on any worker. DP's runtimes do not show any speedup at all mainly due to its high network communication costs. GTP's runtimes also barely change no matter how many workers there are; its total runtime is bounded by the time spent on the largest group.

Group Scalability. Our analytical cost models showed that the runtimes of all approaches except DP are agnostic to the number of groups. DP faces more overhead when training more groups. Table 5 confirmed this on Criteo, with Partner having 15x as many groups as Country. The differences in runtimes on these two attributes are marginal for Kingpin, GTP, and MTP, although we do see minor slowdown on Partner. This minor slowdown comes from the overhead of managing (create, read, and write) more models for all groups. DP suffers from a lot due to its communication costs and idleness as the number of groups went up.

Skew Scalability. We run GBDT (Sampling) and a three-layer MLP (Sequential) on Criteo-Device, a new dataset highly skewed to 2 dominant groups. Figure 11 shows the data distribution and end-to-end runtimes. Kingpin and DP are agnostic to skew, as they both shard data across workers. GTP and MTP step into outof-memory (OOM) errors because the largest group cannot fit in memory. A backup strategy is to swap data to disk, which can be slow and error-prone. Even if the data does fit in memory, GTP cannot match Kingpin in this case. It will only keep 2 workers saturated, exacerbating the idleness seen in Figure 8 and 9.

In-RDBMS ML. There is much prior work on integrating ML with data systems, including RDBMSs (e.g., MADlib [18,27]), cloudnative DBMSs [2, 59], and dataflow systems (e.g., MLlib [45]). All these systems are complementary to our work; none of them optimize learning over groups. Our goal is to study the fundamental systems tradeoffs of this workload at scale, not to integrate ML with data systems. The "keyed models" API in Spark Scikit-learn [5] is the closest to our work in that it trains models in UDFs for each value of a grouping attribute. But it does not fully support model selection for groups and its execution approach is GTP, which suffers from the issues we explained in Section 3.3.

Higher-Level ML Systems. Our work is inspired in part from the recent line of in the DB world on "factorized ML" [12,32,35,51,52,54,61]. They optimize ML over joins of tables. For instance, [12,39] push linear algebra operators down through joins to avoid materializing the join output. In contrast, our work enables and optimizes a GROUP BY abstraction for scalable ML model selection. Ease.ml [29,42] and SystemML/SystemDS [7][8][9][10] are higher-level end-to-end ML platforms. SystemML also uses some hybrid-parallel execution schemes [10] but it is aimed at bulk linear algebra. All of these works are orthogonal to ours. Our goal is not to build a general high-level ML system from scratch. Instead, we enable and optimize a GROUP BY abstraction for scalable ML model selection.

Toward that goal, we devise a novel parallel ML execution scheme and unify the right mix of data partitioning, scheduling, and other system design decisions in a novel and effective manner.

Model Selection Heuristics and Systems. ASHA [37], Hyperband [36], Hyperopt and PBT [28] are recent examples of methods that scale hyperparameter tuning. They are all orthogonal to our work; Kingpin can support them on top in the future. Ray Tune [43], Google Vizier [21], and Dask Hyperband [58] are model selection systems that implement some of the above heuristics; they are all task-parallel systems and thus suffer from the issues we explained in Section 3.3. Horovod [56] and SparkDL [3] offer some model selection support but they are data-parallel systems and thus suffer from the issues we explained in Section 3.3. Overall, none of these systems aim to holistically optimize learning over groups for model selection at scale. That said, they can adopt our techniques in the future to improve their efficiency.

Model Hopper Parallelism (MOP). We view GAP as a sibling of MOP [34,47,48] in that both hybridize TP and DP. But they are fundamentally different and complementary in rationale, generality, and technique. MOP is aimed at Sequential (SGD) on a single large dataset that may not fit on one worker. GAP is aimed at learning over groups, not a single dataset, and supports Algebraic and Sampling access patterns too, not just SGD. MOP-based systems such as Cerebro [49] can be extended to support learning over groups via a lazy group-based subselection on the fly during model selection. However, such an approach will become too slow due to the overhead caused by the relatively low selectivity of most groups in practice, e.g., see the distributions in Figure 6.

Multi-Task Scheduling. Gandiva [60], Tiresias [24], OASiS [6] and SLAQ [62] are cluster scheduling frameworks for deep learning in multi-user scenarios. They focus on lower-level primitives, such as intra-server locality, to reduce average task completion time.

Kingpin is complementary because it exists as a higher-level logical abstraction for learning over groups. In addition, GAP works with ML algorithms besides deep learning. How to allocate computing resources is beyond our scope. There is much prior work in operations research and systems literature about such scheduling algorithms [20,25]. Our goal in Kingpin is to only apply and adapt known techniques to new ML system settings.

Learning over groups is becoming a common practice among practitioners of large-scale ML. Existing approaches to scale this workload using task-or data-parallelism fail to view this process holistically and treat each group as an individual task, which results in poorer resource efficiency, lower model building throughput, and higher total runtimes. We compare existing approaches in depth analytically and empirically and then design a novel approach we call grouped learning to orchestrate learning over groups for large-scale ML model selection holistically and efficiently. We devise a novel parallel ML execution approach we call GAP and support 3 popular classes of ML models: generalized linear models, neural networks, and gradient-boosted decision trees. We adopt a new non-uniform data partitioning scheme suitable for this workload. All of our ideas are easy to integrate with existing ML tools, as we show by building our system, Kingpin, on top Ray and popular ML tools. Both analytically and empirically we find that Kingpin is often substantially faster than the alternative approaches. We hope our work helps ML practitioners interested in training ever more numerous models on group-specific data at scale to achieve more customization for ML applications. As for future work, we aim to extend Kingpin to support AutoML search heuristics, integrate it with Cerebro, and also enable cloud-native scalable execution.