EXPLORING THE LIMITS OF CONCURRENCY IN ML

Exploring the limits of concurrency in ML Training on Google TPUs2018) maps language models along batch and modeldimensions that are then mapped to the physical 2-Dmesh of TPU-v3. We found this approach had significantly high communication overheads as the gradientall-reduce step is executed on a 1-D ring. We presenta novel strided communication optimization schemethat enables high throughput in both the forward andthe gradient reduction steps, that results in the MLPerfTransformer model training in 16 seconds. We scale weight update sharding (distributed optimizer) in a complex hybrid of data and model parallelism scenario via model parallelism and spatialpartitioning. Analysis of the JAX programming model and comparison with TensorFlow. This is the first paper thatstudies JAX at scale, uses JAX on TPU (multi)pods,and uses model parallelism techniques (SPMD partitioning and weight update sharding) in JAX. The JAXresults demonstrate the generality of TPUs and the enhancements added to XLA, and provide a useful comparison for multi- vs. single-controller design points indistributed ML systems. Multipod Performance Results. Four models finishtraining in under 30 seconds. BERT and DLRM, themodels recently added to MLPerf-v0.7, are optimizedat a TPU Multipod scale for the first time.2M ULTIPLE F RAMEWORKSWhile the primary frontend for TPUs has historically beenTensorFlow (Abadi et al., 2016), the hardware and XLAcompiler are general enough to support other programmingenvironments. Therefore in this paper, we chose to benchmark both TensorFlow and JAX (Frostig et al., 2018), anew, research-oriented numerical computing system basedon XLA (TensorFlow.org, 2020). Both systems requiredadditional software engineering to scale effectively to theMultipod, but they ultimately achieved similar benchmarkresults.As shown in Figure 3, two architectural differences betweenTensorFlow and JAX differentiate their performance at scale.First, they have different staging approaches. TensorFlowembeds an expressive and dynamic intermediate language(TensorFlow graphs that can span both accelerators andCPU hosts) in Python, and then JIT-compiles subsets ofthese graphs with XLA. Meanwhile, JAX has one fewerstage: it is a staged programming environment that embedsJIT-compiled XLA programs (for static compiled performance on accelerators and parallelism on the acceleratornetwork) in the Python host language (used for dynamic andon-accelerated computations). As a consequence, TensorFlow has additional compilation steps, which we acceler-Neural Network Model(TPU Estimator)TensorFlow ClientGoogle Compute Engine VMComputational Graph (gRPC)TF StackJAX StackJAX User CodeTensorFlow ServerHostsJAXXLAJust-in-time CompilerTPU Binary (PCIe)Figure 3. Stack view of the TF and JAX frameworks on the TPUv3 machines.ated using multithreading, while JAX requires more carefulmanagement of Python bottlenecks (for instance, movingblocking tasks like data infeed off of the main thread).Second, they enable different distributed programming models. JAX adopts a multi-client approach to distributed programming, running a separate copy of the same JAX code(including the Python interpreter) on each host in the pod.The programs communicate with each other in only twoways: at startup time, to coordinate TPU mesh setup, and inXLA-compiled collectives such as all-reduce that operateover the dedicated TPU network during model training. Onthe other hand, TensorFlow programs TPUs with a singleclient approach, giving one Python process (running eitheron one of the hosts in the pod or elsewhere) global visibilityand control over the entire distributed system. The rest ofthe TPU hosts run a TensorFlow server that executes partitioned subsets of TensorFlow graphs sent via RPCs fromthe client over the datacenter network.These two approaches differ in usability and performancecharacteristics. While TensorFlow’s single-client distributedsystem enables user code that directly reflects the overallworkload, JAX’s multi-client approach enables more directcontrol of the code that runs on each worker. JAX invokesthe XLA compiler independently on each host—relyingon deterministic compilation to avoid incompatibilities between the resulting programs—while TensorFlow compilesonce and distributes the binaries to the workers. The TensorFlow representation of multi-device graphs can also causeAmdahl’s law bottlenecks, as the client process incurs graphconstruction and optimization time proportional to the number of workers, while JAX setup times (other than TPUtopological mesh initialization) do not change significantlywith an increase in the number of workers.

Exploring the limits of concurrency in ML Training on Google TPUs3S CALABILITY T ECHNIQUESIn this section, we describe the optimization techniques required to scale MLPerf-v0.7 models implemented in bothframeworks to the 4096-chip TPU-v3 Multipod machine.Optimization of MLPerf-v0.6 models to a single TPU-v3pod is presented in (Kumar et al., 2019). To achieve higherscale on the Multipod, we next present novel all-reduce optimizations, aggressive model parallelism and input pipelineoptimizations.3.1Model ParallelismIn models where data parallelism is limited, we use modelparallelism to achieve higher concurrency on TPU-v3 Multipod. We leverage XLA’s Single Program Multiple Data(SPMD) partitioner (Lepikhin et al., 2020) to automaticallypartition model graphs based on light-weight annotations.In the segmentation models, SSD and MaskRCNN, we implement spatial partitioning by annotating input images.The SPMD partitioner can automatically parallelize computation along the spatial dimensions. These models haverelatively large spatial dimensions (8000x1333 for MaskRCNN and 300x300 for SSD). The SPMD partitioner insertshalo exchange communication operations to compute theactivations for the next step from spatially partitioned computations. Both of these models enable spatial partitioningalong 8 cores to achieve the highest level of concurrency.Communication optimization and elimination of Amdahlbottlenecks via the XLA compiler SPMD approach (Lepikhin et al., 2020) enabled higher concurrency in spatialpartitioning. For example, in MaskRCNN the largest batchsize is 256, but we were able to parallelize the training onup to 1024 accelerator cores.In the language models such as the MLPerf transformerbenchmark, where the spatial dimensions are small, weexplore partitioning the feature dimension as describedin (Shazeer et al., 2018), but implemented as annotations forthe SPMD partitioner. In this approach, the model weightsand activations are split on a tile of the TPU mesh. In theforward pass, partial matrix multiplication operations arecomputed on each core of the tiled sub-mesh. The activationcontributions from each core are reduced via an all-reduceoperation on the tiled submesh to execute the next layer ofthe model. The backward pass has a similar partial matrixmultiplication followed by all-reduce producing both activations and gradients. As the weights are also partitioned,the gradients are summed between a partitioned core andits corresponding peer on every other tiled sub-mesh of theTPU machine. Techniques to optimize gradient summationon Multipod are presented in Section 3.3.3.2Weight Update ShardingIn traditional data parallelism, model weights are replicatedand updated by the optimizer at the end of each trainingstep. However, this computation can become significantwhen the mini batch size per core is small. For example,we measured in the MLPerf BERT model, the LAMB optimizer weight-update time is about 18% of the step timeon 512 TPU-v3 chips. The weight-update-sharding technique (Xu et al., 2020) distributes this computation by firstexecuting a global reduce-scatter after which each accelerator has a shard of summed gradients. This is used tocompute a shard of updated weights. In the next step, theshard of updated weights is globally broadcast to updateall replicas. To achieve higher speedups we enable weightupdate sharding in both data and model parallelism. In thesegmentation models, where the weights are replicated, theweight-update-sharing scheme is similar to data parallelism.However, when the weights are distributed, we execute multiple concurrent weight-update-sharding steps in each modelparallel core across all the replicas.3.3Optimized Global SummationThe gradient summation step is critical to achieve strongscaling with MLPerf benchmarks (Mattson et al., 2019). Inorder to optimize gradient summation on the large TPU-v3Multipod, we take advantage of the torus wrap links alongthe Y-dimension. A bidirectional ring is used to executea reduce-scatter operation along the Y-dimension with theoutput being a shard of the summed gradients along theY-ring. Next, a reduce-scatter is executed along the Xdimension. This is followed by a weight update computationwith the gradient shard as the input. The updated weightsare broadcast first along X and then Y in two steps. Note,in data parallelism, the payload transferred along the Xdimension is 32 times less than the data transferred alongthe Y-dimension.In the MLPerf transformer benchmark, we execute distributed matrix multiplication operations by sharding modelweights on up to 4 neighboring TPU cores. These coresare placed along a line on the X-dimension. In the forwardpass of ML training, all-reduce calls are executed alongshort rings of X-neighbors. The gradient summation onthe Y-dimension stays unchanged as with data parallelism.However, the gradient summation along the X-dimensionhops over peers that are model parallelism neighbors. Thedifferent ring reductions in the Transformer benchmark areillustrated in Figure 4. In the BERT and transformer models,we also used the brain-float 16-bit floating point precision(bfloat16) to further reduce gradient summation overheads.

Exploring the limits of concurrency in ML Training on Google TPUsthe available memory capacity in the system is sufficient.The Multipod has about a thousand CPU host servers andthe input is sharded across all these. Using uncompressedimages does not incur extra memory throughput overhead,since decompressing images in host memory results in morememory transfers.Figure 4. Figure shows a 16(mesh) 8 (torus) with model parallelism along 4 chips. Three different ring reductions are shownhere. i) a black ring reduction for the model parallel forward pass.ii) red rings do the bulk of the gradient reduce scatter along the Ydimension iii) the dotted blue line shows gradient reduction amongmodel peers (only peer id 0 is shown).3.4Distributed Computation of Evaluation MetricsThe train and evaluation computations are executed in a tightloop on the TPU accelerators. The result of the train loopupdates model weights in the HBM storage on each TPUaccelerator. The updated weights are then used to evaluatethe output metrics for the number of epochs specified inthe MLPerf rules. In benchmarks where the evaluationbatch size is larger than the number of examples in theevaluation dataset, the evaluation dataset is padded withdummy examples. In the TensorFlow implementation, theeval output tensors are used to compute the evaluation metric(for example top-1 accuracy in the Resnet-50 benchmark) onthe TPU master host. However, in the JAX implementationthe computation of the evaluation quality metric is fullydistributed via global summation calls.3.5Input pipeline optimizationsOne of the challenges of scaling the ResNet-50 model isload-imbalance in the host input pipeline. With the massivescale of a Multipod, some host input pipelines have highoverheads of decompressing large JPEG images. Our solution is to store uncompressed images in the host memoryso that the host input pipelines execute only i) random crop,ii) random flip and iii) image normalization with a constantmean and variance as specified in the MLPerf reference.This significantly increases the throughput of the host inputpipeline allowing it to create a large prefetch buffer. So,when the host pipeline preprocesses a large input image itcan feed TPUs with images in the prefetch buffer, thus, eliminating the input pipeline load-imbalance on the Multipodsystem. This optimization increases the training throughputof ResNet-50 by 35% on a Multipod. With uncompressedimages, although the need for memory capacity increases,For BERT, one of the key techniques to improve convergence is to guarantee randomness and coverage in data shuffling. We find two things very helpful for BERT, using thetf.data.shuffle function before the tf.data.repeat function atfile level, and increasing the shuffle buffer size at sequencelevel. At file level, proper data shuffling is especially important as the system scale increases, where every host hasfewer data files to work with. For example, the 500 filesin the BERT reference model will result in a medium-scalesystem with 128 hosts having only about 4 files per host.Executing a tf.data.repeat before tf.data.shuffle gives betterrandomness and coverage of the whole dataset, where theformer guarantees the stochasticity, and the latter guaranteesthe model catches all information available in the dataset.At sequence level, shuffling with small buffer size incurslarge run-to-run convergence difference, which originatesfrom the difference of biased training batch at each trainingiteration, leading to very different convergence trajectoriesof different runs. With larger buffer sizes, every trainingbatch of different runs can be more uniformly sampled fromthe whole dataset, which therefore reduces run-to-run difference.DLRM, like many other recommendation engines, canquickly become input bound as the model accommodates alarge per-core batch size while having a small step latency.One key input pipeline optimization for such models is touse host parallel processing to parse data at batch granularity, instead of per-sample. In the case of the dataset usedfor this model, each training sample is composed of about40 input features. An additional optimization is to transmitinput features over the PCIe bus in a stacked form, reducingthe overhead of transmitting many features separately. Finally, batching overhead can be mitigated by shuffling andpre-serializing data in batch form.4M ODEL O PTIMIZATIONSIn addition to the optimizations mentioned previously, inthis section we present the optimizations applied to eachMLPerf model. With the exception of MaskRCNN andDLRM, all other models are implemented in both TF andJAX. Note, the JAX implementations use the same scalability and convergence techniques as TF models, resulting invery similar step times as well as number of convergencesteps. There are subtle differences w.r.t. TF implementations due to JAX’s multi-client design. For example, JAXallows initializing datasets and input pipelines concurrently

Exploring the limits of concurrency in ML Training on Google TPUsin each TPU worker. Also, the global evaluation accuracy iscomputed via a global all reduce operation on TPUs, compared to TF, where the coordinator CPU process computesthe sum after gathering all local accuracy metrics via hostRPC calls.4.1BERTBER (Devlin et al., 2018) with the wikipedia dataset is newlyadded in MLPerf-v0.7. It is a pre-training task for languageunderstanding with bi-directional transformer architecture.Thanks to the LAMB optimizer (You et al., 2019), BERTcan scale very well to large batch sizes, and we are able touse data parallelism at a 4096-chip scale. Scaling BERT inlarge systems involves optimizations of two aspects, steptime and steps to converge.Other than the optimizations in Section 3, at model level,we optimize the step time of BERT by reducing the stresson architectural bottlenecks, including memory bandwidth,vector units, and registers allowing computation to be executed on the TPU-v3 matrix units with minimal pipelinebottlenecks. To reduce memory bandwidth, we utilize thebfloat16 data type (Wang & Kanwar, 2019) for model activations and gradients aggregation, which largely improvesthe step time and does not have negative effects on modelconvergence. To reduce the stress on vector units, we movethe scalar multiplications and divisions to the smaller sideof matrix multiplication by leveraging the commutativity ofscalar multiplication and matrix multiplication. To reduceregister spilling, we combine small variables, such as layernorm variables, into one large TensorFlow tensor. Thislargely reduces the number of variable addresses to store inthe registers and therefore speeds up the training step time.To reduce the steps to converge, we optimize hyperparameters and data shuffling in the input pipeline. First, weuse Google Vizier (Golovin et al., 2017) to fine tune thehyperparameters for large batch training, enabled by thescalability of LAMB optimizer. This allows us to leveragemaximum data parallelism, which gives better time to accuracy compared to model parallelism. Second, as detailed inSection 3.5, we optimize the way data is shuffled in the datainput pipeline in order to ensure the convergence in largesystems. This is critical for guaranteeing the stochasticityof the optimizer because large systems with thousands ofhosts typically assign less data to each host.4.2ResNet-50ResNet-50 (He et al., 2016) is one of the most widely-usedmodels for ML benchmarking. MLPerf uses the ResNet-50model with the ImageNet-1K (Russakovsky et al., 2015)dataset as the image classification benchmark. Specifically,MLPerf uses the variant termed “version 1.5” (Goyal et al.,2017) to indicate a slight modification to the original modelarchitecture which is commonly found in practice. In orderto scale the ResNet-50 benchmark to the TPU-v3 Multipod system, we apply optimizations including distributedevaluation, distributed batch normalization, weight updatesharding, and gradient summation. The MLPerf-v0.7 reference model uses the LARS optimizer (You et al., 2017)that adaptively scales learning rates, which enables training ResNet-50 with data parallelism on large batch sizes.After the momentum hyperparameters are tuned, we areable to finish training in 88 epochs with batch 65536 on theMultipod.4.3TransformerTransformer represents the state-of-the-art language translation in the MLPerf suite and is one of the two translationmodels. Trained on the WMT English to German dataset,Transformer uses an attention-based model which differentiates it from the other language model in MLPerf, GNMT. Ithas been observed that it is hard to scale Transformer witha fixed epoch budget beyond a global batch size thresholdgiven the current dataset (Shallue et al., 2018). Thereforeboth data and model parallelism are applied to scale theTransformer model to a TPU-v3 Multipod system. Withmodel parallelism, the model is able to run with fewer thanbatch one per core, using a fixed global batch size of 2048where the hyperparameters have been well tuned.SPMD sharding is employed to enable model parallelism.Unlike spatial partitioning (sharding the images) norGshard (Lepikhin et al., 2020) (which has sparse components and all-to-all communications), dense sharding is applied to the Transformer model. Shared embedding layers,multi-heads attention projection layers and feed-forward layers are sharded, along with vocab, num heads, and hiddendimensions, respectively. To speed up gradient all-reduce,2D cross-replica all-reduce is enabled for SPMD shardingwith X-dimension hops over model parallelism neighborreplica. The all-reduce communication is performed inbfloat16 floating point precision to further improve the performance.4.4SSDSingle Shot Detection (SSD) is one of two image segmentation models in MLPerf; SSD is intended to reflect a simplerand lower latency model for interactive use cases such asin end-point and non-server situations. Notably, SSD usesa pre-trained ResNet-34 backbone as part of the architecture. The MLPerf SSD benchmark is trained on the COCOdataset (Lin et al., 2014). In the MLPerf-v0.6 submission wehad used a global batch size of 2048 and 4-way model parallelism. In this round of MLPerf submissions, we are able totrain with a batch size of 4096 using new hyperparameters.Note, this is still much smaller than the batch size of 65536

Exploring the limits of concurrency in ML Training on Google TPUsavailable in the ResNet-50 model. We used XLA’s SPMDpartitioner to enable scaling up to eight TPU cores via modelparallelism, replacing XLA’s MPMD spatial partitioner usedin MLPerf-v0.6. SPMD has better scalability in compilation time and enabled us to increase the largest scale forSSD training from 2048 TPU-v3 cores in MLPerf-v0.6 to8192 cores in MLPerf-v0.7. A unique benefit of the SPMDpartitioner is that it enables the weight-update-sharding optimization even with model parallelism, that results in a 10%speedup. It is challenging to get high speedups in the SSDmodel from model parallelism as there are communicationoverheads from halo exchange and load imbalance as different workers may get uneven tiles of work. In addition, theinput image to the SSD model is relatively

mizer) in a complex hybrid of data and model par-allelism scenario via model parallelism and spatial partitioning. Analysis of the JAX programming model and com-parison with TensorFlow. This is the first paper that studies JAX at scale, uses JAX on TPU (multi)pods, and uses