Flashlight: Enabling Innovation in Tools for Machine Learning
Authors: Jacob Kahn
Vineel Pratap
Tatiana Likhomanenko
Qiantong Xu
Awni Hannun
Jeff Cai
Paden Tomasello
Ann Lee
Edouard Grave
Gilad Avidov
Benoit Steiner
Vitaliy Liptchinsky
Gabriel Synnaeve
Ronan Collobert
Venue: ICLR 2022 Submitted (2022)
Keywords: machine learning
deep learning
systems
frameworks
autograd library
tensor library

Abstract: As the computational requirements for machine learning systems and the size and complexity of machine learning frameworks increases, essential framework innovation has become challenging. While computational needs have driven recent compiler, networking, and hardware advancements, utilization of those advancements by machine learning tools is occurring at a slower pace. This is in part due to the difficulties involved in prototyping new computational paradigms with existing frameworks. Large frameworks prioritize machine learning researchers and practitioners as end users and pay comparatively little attention to systems researchers who can push frameworks forward --- we argue that both are equally-important stakeholders. We introduce Flashlight, an open source library built to spur innovation in machine learning tools and systems by prioritizing open, modular, customizable internals and state-of-the-art, research-ready models and training setups across a variety of domains. Flashlight enables systems researchers to rapidly prototype and experiment with novel ideas in machine learning computation and has low overhead, competing with and often outperforming other popular machine learning frameworks. We see Flashlight as a tool enabling research that can benefit widely-used libraries downstream and bring machine learning and systems researchers closer together.

Under review as a conference paper at ICLR 2022
FLASHLIGHT: ENABLING INNOVATION IN TOOLS FOR MACHINE LEARNING
Anonymous authors Paper under double-blind review
ABSTRACT
As the computational requirements for machine learning systems and the size and complexity of machine learning frameworks increases, essential framework innovation has become challenging. While computational needs have driven recent compiler, networking, and hardware advancements, utilization of those advance- ments by machine learning tools is occurring at a slower pace. This is in part due to the difficulties involved in prototyping new computational paradigms with existing frameworks. Large frameworks prioritize machine learning researchers and practitioners as end users and pay comparatively little attention to systems researchers who can push frameworks forward we argue that both are equally- important stakeholders. We introduce Flashlight, an open source library built spur innovation in machine learning tools and systems by prioritizing open, modular, customizable internals and state-of-the-art, research-ready models and training setups across a variety of domains. Flashlight enables systems researchers to rapidly prototype and experiment with novel ideas in machine learning computa- tion and has low overhead, competing with and often outperforming other popular machine learning frameworks. We see Flashlight as tool enabling research that can benefit widely-used libraries downstream and bring machine learning and systems researchers closer together.
INTRODUCTION
The recent rise of deep dearning-based techniques has been accompanied and sustained by the wide availability of dedicated frameworks such as TensorFlow (Abadi et al.) 2016) and PyTorch (Paszke) et al. 2019). These frameworks have enabled the democratization of machine learning research by providing extensive collections high level primitives to support common use cases. Lowering the barrier to entry for end users has boosted the popularity of both neural networks and the frameworks in which they are implemented. However, order to support what are now vast ecosystems and diverse user base, framework size and complexity have increased dramatically result, deep, groundbreaking framework research has become extremely onerous and time consuming, precluding rapid innovation. Given these barriers, major deep learning frameworks have become stuck in their existing operating modes.
time. As
Innovation in this area remains important as ever. Indeed, framework innovation accelerates machine learning (ML) and artificial intelligence (Al) research. Frameworks that are easier to use reduce the engineering burden researchers, and frameworks that are higher-performance decrease the time required to iterate experimental work and validate hypotheses. Even more critically, tooling plays a fundamental role in deciding which ideas succeed or fail. For example, [eCun (eample,,LeCum et al.) (1989) pioneered the use of convolutional neural networks (CNNs) (Fukushima & Miyake 1982) trained using backpropagation for computer vision tasks the late 1980s, which was subsequently applied handwriting recognition. However, widespread success for CNNs was achieved two decades later when Krizhevsky et al. (2012a) leveraged the CUDA programming model to take advantage of graphies processing units (GPUs) to train much deeper model (AlexNet).
While deep learning frameworks have been optimized to leverage existing hardware paradigms for common neural network architectures, they often fail to deliver similar efficiencies designs that diverge from the mainstream. For example, BBrramm & Isard (2019) explain how the design of these frameworks results in poor hardware utilization for novel type of neural network, known

---

Under review as a conference paper at ICLR 2022
capsule network (Hinton et al. 2018). that leverages new components such as squashing operations and routing by agreement. More generally, what are now unconventional approaches to modern problems in machine learning require highly-specialized additions to popular frameworks. As result of harrowly-optimize systems, research beyond deep learning may be discounted due to purported computational infeasibility given modern frameworks' capabilities.
Furthermore, the waning of Moore's law (Theis & Wong 2017) coupled with the ever-growing computational demands of deep learning are prompting several shifts hardware. Massive-scale distributed computing is now required to train leading models - a process that established frameworks remain unable to handle truly automatically. In parallel, multiple specialized hardware products are now available to better support deep learning applications: Nvidia's TensorCores (Markidis et al. 2018), Google's 's TPUs (Jouppi et al. 2017), Grapheore's IPUs (Jia al. 2019), Apple's Neural Engine and others have been developed to improve total floatpointing operations (FI.OPs), cost per-FLOP,, or energy consumption. Additionally, numerous efforts are underway to move away from conventional von Neumann computing architectures in which memory and processing units are physically separated, either by storing data closer to compute units or by switching to in-memory computing altogether.
While tooling innovation is alive and well given these incentives for progress, working within large, well-established frameworks has become more and more challenging framework size and scope grows. As a result, many recent innovations have required the development of ad-hoc tools. For example, efforts in machine learning-driven compilation of neural networks are largely built on top of Halide (Adams et al., 2019 Steiner et al. (022)) and TVM (Chen et al.) 2018 Zheng et al. 2020); FlexFlow (Jia et al. 2018 2020 underpins recent work aimed at improving the use of distributed computing to accelerate the training of large neural networks; and PET (Wang et al. 2021) provides a framework that enables graph-level neural network optimizations. With ad-hoc approaches, researchers are required to start from scratch for new directions or adapt their ideas to fit into the scaffolding these frameworks provide - resulting significant technical burdens.
To sustain framework innovation, we introduce Flashlight, an open source minimalist ML library designed to support research in machine learning frameworks, facilitate rapid iteration on ideas, reduce the engineering burden on researchers, and remove the need for new tools. Flashlight includes:
A modular, component-ben architecture that makes every aspect of the implementation fully customizable with simple internal APIs.
A compact yet highly-performant reference implementation each component.
A comprehensive set of benchmarks representative of the state-of-the-art in machine learning on which to evaluate alternative implementations.
RELATED WORK
Numerous frameworks have been implemented in recent years to support machine learning, including
Lush (Bottou & LeCun 2002). Theano (Bergstra et al. 2010). Torch (Collobert et al. 201T), Caffe
(Tia et al. 2014), MXNet CChn et al. 2015), deeplearningj (Team 2016), TensorFlow (Abadi et al. 2016). Flux (Inees] 2018). Jax (Bradbury et al. 2018). PyTorch (Paszke et al. 2019), Chainer (Tokui al. 2019), and PaddlePaddle (Ma al. 2019). These frameworks offer programming models designed around multidimensional arrays (TENSORS), modeled as first-class objects and supported by comprehensive set of mathematical primitives (or operators) manipulate them. To provide the computing power required by deep learning- based methods, most natively support hardware accelerators such as general- purppoe GPUs or custom-designed ASICs such as TPUs.
Generally, framework implementations follow one of two computational models:
In the deferred execution model, the neural network to be trained is encoded dataflow graph which can be optimized for a specific set of target hardware devices. The neural network is then executed in a distinct second phase. Since the dataflow graph represents the entire computation, both local and global optimizations can be applied, making the subsequent execution very efficient. However, only programs that can be represented as
pttts:://nr.apl111 com/dx991
2

---

Under review as a conference paper at ICLR 2022
dataflow graphs can be processed with this approach, thus limiting flexibility. Frameworks such as Theano, TensorFlow P] Caffe, or MXNet fall into this category.
In the eager model, an interpreter (such as Python) is extended with the high level kernel- based operations needed to train a neural network. These operations are executed im- mediately when called, though this precludes many optimizations. By weaving neural network-related operations into a Turing complete programming language, this approach is extremely flexible. Furthermore, the imperative nature of the underlying programming language allows for fine- aineggrined control over the execution order and memory utilization, which enables more specific user-driven optimization. Frameworks such as Torch, Py Torch. or Chainer exemplify this approach.
3 PRINCIPLES
The aforementioned frameworks are designed and implemented to best-serve their user bases namely, machine learning researchers and practitioners. They rely on large, internally complex codebases to provide comprehensive solutions, as is further discussed in Section 5
In contrast, Flashlight targets an audience of researchers interested in experimenting with new designs and implementations of machine learning tools or broader computational or modeling paradigms. To foster this type of innovation, Flashlight balances simplicity and nimbleness with the need to provide enough functionality to support real use cases. Internal and external simplicity is the key design principle of Flashlight; the ability to dramatically modify software and drive in new directions inversely correlated with codebase size and complexity (Gill & Kemerer 1990). More specifically:
Flashlight is built on a shallow stack of idiomatic, modular, and customizable abstractions. Framework components interact through small, well-defined, stable APIs, which expose most internal aspects its implementation. This ensures that every component of Flashlight can be modified replaced with new custom implementations, even e.g. its memory manager and tensor implementation. To support the exploration of wide array of alternative approaches, Flashlight interfaces are flexible and unopiniated by design. This in contrast to other frameworks, which impose stricter implementation requirements based on tight design constraints for their computation models and support requirements across hardware, downstream frameworks, and other ecosystem members.
Flashlight provides
pntterattllycooppat default implementations of its APIs. This re- duces out-of-the-gate engineering burden and the need for modifications, and enables fast compilation and rapid iteration when experimenting. Furthermore, to mitigate premature op- timization, Flashlight deliberately abstains from adding small efficiency improvements if they conflict with the goals of keeping the codebase simple and A PIs clean.
Flashlight is research-first framework, and is not intended for out of the box production use. To keep codebase size small, it forgoes features such as model servers for deployment and integration with cluster management tools.
Flashlight is viable solution for machine learning research, shipping with a comprehensive set of benchmarks and research setups for toatte-ff-hheeeeee art neural network architectures such as convolutional neural networks (CNNs) (Krizhessky et al.) 2012b) and Transformers (Vaswani et al. (017). as well
task-specific models such VTT (Dosovisskiy et al. 2020), DETR (Carion al. 2020). BERT (Devlin et al. 2018). The speech recognition system wav2letter (Pratap et al. 2019), is also built entirely on Flashlight.
Benchmarks built on these state-of-the-art models make Flashlight a turn key solution for system researchers who to build test benches from the ground-up. More importantly, Flashlight makes possible end-to-end benchmarking on real models rather than microbenchmarks or small-scabe tests.
quickly evaluate their design and implementation choices without needing
FensorFlow 2.0 adds support for eager exccution semanties well.
3

---

Under review as a conference paper at ICLR 2022
4 DESIGN
Flashlight's design centered around internal APIs for framework components which form building blocks for domain-specifee MI puckages and applications this structure is outlined in Figure ! Flashlight implemented as a C++ library and follows a Tensor-based programming methodology, with neural network building blocks that derive from a MODULE interface, communi- cate by exchanging Tensor data, and are composed functionally or imperatively to form complete neural network architectures. Tensor programming in Flashlight is fundamentally dynamic, but given that C++ compiled language, code describing models in Flashlight is compiled. This approach promotes type safety, foregoes the runtime overheads associated with interpreters, and, unlike eager-based approaches, enables global optimizations where possible.
Applications
Speech Racognition
Object Detection
Image Clascification
Language Modkling
Benchmarks
Packages
Core
Speech
Vision
Text
Optimizess
Data Naural Net Operators
Autograd and Losses
Foundation
Tensor
Distributed
Management
Implementation
AmayFr
cLNNN
NCCL
Gloo
Figure 1: Components of the Flashlight library.
4.1 OPEN FOUNDATIONAL INTERFACES
Flashlight is built on top of three open foundational APIs, each addressing design and implementation challenges faced by machine and deep learning tools: Tensor interface, memory management sub- system, and a distributed computing interface. These APIs are backed by reference implementations that enable Flashlight efficiently target CPUs, GPUs, and other accelerators. These include code generation and dedicated kernels for Intel, AMD, OpenCL, and CUDA devices, and leverage libraries such as CuDNN (Chetlur al., 2014), MKL [Intel, 2020a)) oneDNN (Intel 2020b), ArrayFire (Malcolm et al. 2012), and MiOpen (Khan et al. 2019).
4.1.1 TENSOR INTERFACE
Modern deep learning frameworks feature tensor library internals which sit under deep layers of abstractions, requiring numerous framework modifications in order to iterate on tensor stack design. Flashlight's TENSOR abstraction defined in terms of existing tensor libraries via simple, extensible interface and high-level API that mirrors numpy (Harris et al.) (020)) rather than using specific, opinionated intermediate representations (IRs) or large operator sets.
Flashlight TENSOR backend implementations need not follow any particular computation mode as outlined Section 2 and shown Figure 2 Tensor values need only be materialized upon user request typically when extracting the output values of model inspecting intermediary state. This provides a flexibility unique amongst deep learning frameworks to either defer or eagerly- compute intermediate values to experiment with computation paradigms altogether.
4

---

Under review as a conference paper at ICLR 2022
Computation Model
Eager
Op call
Execute
Op call
Execute
Op call
Execute
Ready
Deferred
Op call
Op call
Op call
Data Request
Compile Launch
Execute
Ready
Data Request
Hybrid
Execule
Ready
call
Compile
Execute
Static
op to
to
Launch Compuration
Ready
Time
Figure 2: Flashlight's Tensor API supports backend implementations with any of the above computa- tion modes.
Implementing TENSOR backend in Flashlight involves fulfilling small of implementation requirements. Users have full control their implementations after subclassing two interfaces:
the TENSOR interface (Listing 1 which allows backend attach custom stateful information and metadata to each tensor. This includes shape, type, and memory information which may be implemmentation-deeenddd..
the TENSoRBACKENT interface (Listing 2 which allows backends to store global state as needed (e.g device compute streams, dataflow graphs) and implement small set of primitive tensor operations including unary and binary operations (e.g. arithmetic ops), reductions, matrix multiplication, and convolution.
Rather than require implementations of large, highly-specialized operator sets or interoperability with complex dispatch mechanisms or intermediate representations (IRs) do other frameworks, Flashlight operators outside of the small TENSOR BACKEND API are derived by composition. For example: the ReLUU activation implemented by leveraging the MAX operator. Flashlight's reference TENSOR implementation uses hybrid approach, offloading computation to highly-optimized vendor libraries when advantageous and relying on deferred, on-the-fly code generation via Array Fire for all increase kernel arithmetic intensity.
other operations
class MyTenserInp public TensorAdapter
// State information goes here (e.g. buffers, shape) public:
// Metadata
const Shapes shape () override;
dtype type () override;
};
Listing 1:
The TensorAdapter interface for implementing operations tensor metadata and storing tensor state for individual tensor instances.
class MyTenscrRackend public TenscrRackend
// State information goes here (e.g. compute streams, compiler state)
public:
// Tensor operation primitives
Tensor add(const Tensor& 1hs, const Tensorá rhs) override; Tensor minimum (const Tensors 1hs, const Tensors rhs) override;
//
};
Listing 2: The TensorBackend interface for implementing operations on tensors and storing global backend state.

---

Under review as a conference paper at ICLR 2022
4.1.2 MEMORY MANAGEMENT
Robust memory management an important research area model size increases. While indi- vidual TENSOR backends in Flashlight can perform their own memory management as defined by implementers, Flashlight's is default TENSOR backend also provides a generic API for defining custom memory management schemes. By default, memory is only allocated when needed for just-in-ting iistinntimee compilation. A sample of this API is shown in Listing S To support the lazy computation model as well as just in time code generation, memory allocations only occur when tensors need to be material- ized per the compute graph. Buffers are used asynchronously after they are requested depending on the timing of kernel launches but are not freed until computation is complete.
class GachingMemoryManagee // Store state needed publ ic:
public KemoryWanagerAddptte:
void* alloc (bool userllock, const unsigned ndims, dim_tr tx dims, const unsigned elSize) override;
void unlock (void* ptr, bool userLeck) override; // free nemory
};
Listing 3: An implementation of memory manager using the memory management API.
4.1.3 DISTRIBUTED TRAINING
Flashlight provides a low-level API for distributed training primitives with performant default implementations in GPU and CPU settings using NCCL (2019) and Gloo (2019), respectively. Users can add new backends or custom methods of performing distributed computation, and can use primitives in other framework components as needed. The API is unopinionated and supports both synchronous and asynchronous communication and arbitrary distributed computation schemes.
4.2 THE FLASHLIGHT CORE
The TENSOR API, together with memory management and distributed computation abstractions provide a foundation on which to build other core machine learning algorithms and applications. These other components outlined below. Section A.4 provides code samples and linking documentation for the below components.
Neural Network Primitives To facilitate the implementation of neural networks, Flashlight ships with common neural building-blocks encompassing activation functions, normalization layers, regu- larizers, losses, and more. These derive from the MODULE abstraction, discussed above, which provides method of chaining and nesting Secrioio A.4. contains more details and sample implementations.
Automatic Differentiation Automatic differentiation (autograd) is implemented via simple VARIABLE abstraction. A Variable takes TENSOR argument when created, and its underlying Tensor (or gradient Tensor) be accessed any time. Variables feature operators which call underlying Tensor operations and record those operations to a dynamic tape in a design similar to Paszke et al. (2017) while being lightweight enough to allow implementations of other autograd paradigms. In keeping with Flashlight's modularity, TENSOR and VARIABLE are separated order avoid performance and implementation overhead in non-gradient-base MI algorithms.
Variable cos (const Variables input)
auto result fl::c0s // Called with backward () to compute gradients for this op's inputs auto gradFunc = [1 (std: vectorrarriabblee inputs, const Variables gradOutput)
()); // get a Tensor from a Variable
inputs[0] adddraa // Add gradient the input Variable (gradOutput + negate (sin (inputs [0] tensor ())), falso) )
};
// Construct Variable from Tensor and gradient- comppttng function
recacteddooranonyiity

---

Under review as a conference paper at ICLR 2022
return Variable (result, (input): gradFunc)
Listing 4: Defining cosine autograd operator in Flashlight using TENSOR operations and VARIABLE.
Optimizers Flashlight provides implementations of most common first-order stochastic optimiza- tion algorithms, as included in other frameworks. These are defined in terms of Variable and Tensor operations, allowing for open-ended experimentation (e.g. with distributed computation, in in-place operations, etc).
Data Loaders Flashlight provides a simple DATASET class which abstracts away the notion of sample in ML algorithms. A sample is viewed here as a TENSOR or vector TENSORS. Datasets are trivially composable to create pipelines to transform, resample, or parallelize (via native C++ threads) the construction of such samples. While Flashlight core dataset abstractions are agnostic from the end-user task, data-specific datasets are provided in higher-level packages, efficiently load from disk structured data (e.g. images, audio or text).
4.3 PACKAGES AND APPLICATIONS
Flashlight contains additional domain-speciff abstractions leveraging both core components as well stand-alone implementations. These abstractions allow end-users ML researchers quickly get started on various MI applications. The package module provides building blocks for common ML tasks, Nomain-speeiffe algorithms, and helpers. The application module leverages these building blocks to provide complete, ready to use solutions (c.g. models, training loops, evaluation pipelines). When original to Flashlight, implementations reproduce the task performance those they reference. We leverage several of these applications to evaluate Flashlight's performance in Section
Speech. Flashlight provides an implementation of classical featurization (spectogram, log- mel filterbanks, etc.) that can run on-the-fly with minimal overhead. It also provides a collection of data- augmentation techniques (including additive noise and reverberation). as well as implementations of speech-specific sequential criteria and model architectures. Flashlight contains fast beam- search decoder (which can interface any language model) and beam rescorers (Collobert et al. 2016) Pratap al. 2019). Research performed with the speech application have reached and are competitive with state-of-the-art results (Synnaeve et al., 2019 Likhomanenko et al. 2020).
Vision. Flashlight offers built-in data loaders for standard computer vision benchmarks (such as ImageNet (Deng et al.) 2009) and COCO (Lin et al.) 2014)) along with large set of efficient data-augmentations and transtormations. includes mainstream image classification models: convo- lutional (e.g. ResNet He et al. (2016)) and transformer-ban based architectures (e.g. ViT Dosovitskiy al. (0200), well modern, transformer-bigg object detection model (DETR Carion al. (2020) and helpers (e.g. Hungarian matching and object detection evaluation).
Text. Flashlight ships with support for text dataset manipulation and tokenization along with language modeling training pipelines for a variety of neural language models, including transformer (Vaswani et al. 2017) and CNN-based (Dauphin et al., (2017). Both autoregressive and masked, e.g. BERT, language modeling tasks supported. These language models can be combined with other domain-specific packages such as speech.
EVALUATION
In the sections that follow, evaluate Flashlight in context with two widely-used deep learning frameworks - Torch and Tensorflow - with metrics relevant to framework research velocity. We also evaluate framework performance proxy for overhead and the quality of default imple- mentations. We outline the steps needed to reproduce all our results included in this section in the Appendix.
7

---

Under review as a conference paper at ICLR 2022
Table 1: Complexity of various frameworks based high-level metrics.
Metric
Py Py/orchh
TensorFlow
(Ours) Flashlight
Binary Size (MB)
527
768
9 27,173
Lines of Code
1,798,292
1,306,159
Number of Operators
2,166
1,423
110
Approx num. ops. that perform:
ADD
55
20
1
CONV SUM
85
30
2
25
10
Table 2: Compile times in CPU minutes across frameworks.
Platform
From Scratch
Incremental
PyTorch
754
132
Tensorflow
2061
371
(Ours) Flashlight
27
0.6
5.1 CODE COMPLEXITY
Flashlight is built to minimize complexity and operating surface. As frameworks grow and are combined with other frameworks or take new platform-specific requirements, internal modifiability decreases. Table T comparess frameworks across binary size, lines of code and number operators and operator implementations. Flashlight's small surface facilitates easily exploring new designs and prototyping on new hardware - having few sources of truth simplifies the process of replacing core components and ensures end-to-end tests don't opaquely fall back existing implementations.
5.1.1 COMPILATION TIME
When modifying adding significant new research code to framework internals, recompilation can be costly. Large frameworks depend on code generation for broad platform suppor increasing compilation time. Further, expensive incremental rebuilds can slow iteration speed.
Flashlight suffiiiettl-lightweiiee and modular so as to enable from-source build times that are orders of magnitude faster than other frameworks. shown in Table 2 Times were measured for both from-scratch and incremental builds with Intel Xeon Gold 6138 CPUs with 80 cores and 750 GB of memory. To estimate incremental build performance, randomly sample 100 source files without replacement, make trivial modifications that force recompilation, and time the resulting rebuild.
5.1.2 PERFORMANCE
When improving framework components or modifying internals, framework overhead makes difficult disambiguate performance changes due to in in-ftight modifications from existing bottlenecks or overhead due to other framework components as discussed in Section 3 performance of Flashlight 0. PyTorch 1.8. TensorFlow 2.4 on 6 common large-scale deep neural networks. For each configuration, we benchmark 100 iterations of data loading preprocessing, and forward/backwaar passes, with data aata-parallll gradient synchronization in distributed settings. Benchmarks are performed on Intel ES-2698 CPUs with 512GB of RAM, and NVIDIA V 1100-32GB GPUs in DGX-1 Station. Inter-GPP interconnects in the 8 GPUs node) setting and the multi-node (64 GPUs) setting are Nvidia NVLink and InfiniBand- based, respectively.
compares the
*In attempt not include auxiliary framework components, count only C, C++, Python, YAML, CUDA, and CMake files from relevant subset of components for each project.
'Examples PyTooch: https: ///tt. it. io/ 9. TensorFlow: https: / //git ic/JzeRw 'For consistency, BERT-like models random data in-memory; ViT models exclude data augmentation.
8

---

Under review as a conference paper at ICLR 2022
Table 3: Performance common state-of-the-art models across frameworks. Values the number of seconds needed to perform 100 iterations of the forward and backwards passes, with data loading (unless indicated). Number parameters in millions (#param) and batch size (bsz) are specified. Framework labels: PT= = PyTorch, TFF TensorFlow, and FL Flashlight.
GPU
8 GPUs
Model
Num. Params (M)
Batch Size
PT
TF
FL
PT
TF
FL
AlexNet
32
2.0
4.0
1.4
6.0
6.5
2.1
VGG16
138
32
14.8
12.6
13.2
16.3
17.9
14.9
ResNet-50
25
32
11.1
12.4
10.3
12.3
15.9
11.9
BERT-like
406
128
19.6
19.8
17.5
22.7
23.6
19.2
ASR Tr.
263
10
58.5
63.7
53.6
63.7
69.7
57.5
VIT
87
128
137.8
140.3
129.3
143.1
169.6
141.0
Flashlight is competitive and can exceed the performance of other frameworks, especially on archi- tectures which are of lower arithmetic intensity and spend less compute time in vendor-optimized librarics, such as AlexNet. Given strong performance with simple reference implementations that have undergone far less optimization than have large frameworks, we see exciting potential for improvement with future research done in Flashlight.
CONCLUSION
We presented Flashlight, a modular machine learning library supporting modern, state-of-the-art baselines that features orders of magnitude less code and binary size as compared frameworks such as Py Torch and TensorFlow. These large frameworks come fully-featured; Flashlight aims complement them in providing tool to do machine learning framework research. To this end, Flashlight features lightweight, modular design, well as full implementations of mainstream models across a variety of domains, making it easy for researchers to implement new internal tensor, memory management, or distributed computation backends. Flashlight includes variety of reference APIs that compete with and often outperforms that of popular machine learning frameworks, thus demonstrating the viability of our approach.
REFERENCES
Martín Abadi, Paul Barham, Jianmin Chen, Zhifeng Chen, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Geoffrey Irving, Michael Isard, et al. Tensorflow: A system for large-scale machine learning. In 12th USENIX symposium on operating systems design and implementation (OSDI 16), pp. 265-283, 2016.
Andrew Adams, Karima Ma, Luke Anderson, Riyadh Baghdadi, Tzu-Mao Li, Michaël Gharbi, Benoit Steiner, Steven Johnson, Kayvon Fatahalian, Frédo Durand, and Jonathan Ragan- Learning optimize halide with tree search and random programs. ACM Trans. Graph., 38(4), July 2019. ISSN 0730-0301. 1145/330634.11 3322999663336
Ragan-Kelley Kelley.
doi: 10. 00.114//33663443322966 URL httpt::t/ h/op:://do... org/10/
Paul Barham and Michael Isard. Machine learning systems are stuck in a rut. In Proceedings of the Workshop on Hot Topics in Operating Systems, HotOS '19, 177-183, New York, NY, USA, 2019. Association for Computing Machinery. ISBN 978145036727L. doi: 10.114523117550.3214441 3321441
URL wttpt:tt///// ULL.:
James Bergstra, Olivier Breuleux, Frédéric Bastien, Pascal Lamblin, Razvan Pascanu, Guillaume Desjardins, Joseph Turian, David Warde-Farley. and Yoshua Bengio. Theano: cpu and gpu math expression compiler. Proceedings of the Python for scientific computing conference (SciPy), volume 4, pp. 1-7. Austin, TX, 2010.
Léon Bottou and Yann LeCun. Lush. 2002. URL

---

Under review as a conference paper at ICLR 2022
James Bradbury. Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Miine and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL
http:: http//ggttuuie com
Nicolas Carion, Francisco Massa, Gabriel Synnaeve, Nicolas Usumier, Alexander Kirillov, and Sergey Zagoruyko. End-to-end object detection with transformers. In European Conference on Computer Vision, pp. 213-229. Springer, 2020.
Tianqi Chen, Mu Li, Yutian Li, Min Lin, Naiyan Wang, Minjie Wang, Tianjun Xiao, Bing Xu, Chiyuan Zhang, and Zheng Zhang. Mxnet: A flexible and efficient machine learning library for heterogencous distributed systems. arXiv preprint arXiv:1512.01274, 2015.
Tianqi Chen, Lianmin Zheng, Eddie Yan, Ziheng Jiang. Thierry Moreau, Luis Ceze, Carlos Guestrin, and Arvind Krishnamurtthy. Learning to optimize tensor programs. In Proceedings of the 32nd International Conference on Neural Information Processing Systems, NIPS' 18, pp. 3393-3404, Red Hook, NY, USA, 2018. Curran Associates Inc.
Sharan Chetlur, Cliff Woolley, Philippe Vandermersch, Jonathan Cohen, John Tran, Bryan Catanzaro, and Evan Shelhamer. cudnn: Efficient primitives for deep learning. arXiv preprint arXiv: 1410.0759 2014.
Ronan Collobert, Koray Kavukcuoglu, and Clement Farabet. Torch7: matlab-like environment for machine learning. 01 2011.
Ronan Collobert, Christian Puhrsch, and Gabriel Synnaeve. Wav2letter: an end-to-end convnet-based speech recognition system, 2016.
Yann N. Dauphin, Angela Fan, Michael Auli, and David Grangier. Language modeling with gated convolutional networks. In Proceedings of the 34th International Conference on Machine Learning Volume 70, ICML' 17, pp. 933-941, JMLR.org, .org 2017.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 JEEE conference on computer vision and pattern recognition, pp. 248-255. Icee, 2009.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805. 2018.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly. et al. An image is worth 16x16 words: Transformers for image recognition scale. arXiv preprint arXiv:2010. 11929, 2020.
Kunihiko Fukushima and Sei Miyake. Neocognitron: A new algorithm for pattern recognition tolerant of deformations and shifts in position. Pattern Recognition, 15(6):455 469, 1982. ISSN 0031-3203. doi: http:://...oog////...00001111111 com/ http://cccoceeaaraaaee//////1122
htt.h:tt::/// //wp::://wwww../. sci www.wwiiiincncerett..ccc..c
G.K Gill and C.F Kemerer. Productivity impacts of software complexity and developer experience.
MIT Sloan, 1990.
Gloo. Gloo: collective communications library.
https://github com/
faacbookincubator/gl/ 2019.
Charles R. Harris, K. Jarrod Millman, Stéfan J. van der Walt, Ralf Gommers, Pauli Virtanen, David Cournapeau, Eric Wieser, Julian Taylor, Sebastian Berg, Nathaniel J. Smith, Robert Kern, Matti Picus, Stephan Hoyer, Marten H. van Kerkwijk, Matthew Brett, Allan Haldane, Jaime Fernández del Río, Mark Wiebe, Pearu Peterson, Pierre Warren Weckesser, Hameer Abbasi, Christoph Gohlke, and Travis E. Oliphant. Array programming with NumPy. Nature, 585 URL https::::: httttp:://doi..ogg111.o /
Kevin Sheppard, Tyler Reddy,
September 2020. doi: 10.1033//41588-0222 -20022966-2-922 2649-2.
10

---

Under review as a conference paper at ICLR 2022
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Geoffrey Hinton, Sara Sabour, and Nicholas Frosst. Matrix capsules with em routing. 2018. URL https:// http:s///ppeeeweeww wwtpeanrveeww.eeeeeeeeeeeeeeeeee net http.e/e/pdd.i---BRRooo.o..
Mike Innes. Flux: Elegant machine learning with julia. Journal of Open Source Software, 3(25):602,
2018.
Intel. Mkl httpr//www.rrrr/tettttaaaaaaaooc
httpp///oociiiataaioiaiereeeeeea 2020b.
2020a.
Intel. Onednn. http:s//ggttttt.com///eeeeessssp
Yangqing Jia, Evan Shelhamer, Jeff Donahue, Sergey Karayev, Jonathan Long, Ross Girshick, Sergio Guadarrama, and Trevor Darrell, Caffe: Convolutional architecture for fast feature embedding. In Proceedings of the 22nd ACM international conference Multimedia, pp. 675-678, 2014.
Zhe Jia, Blake Tillman, Marco Maggioni, and D. Scarpazza. Dissecting the graphcore ipu architecture via microbenchmarking. ArXiv, abs/ 1ttps//912.3311,,, 2019.
Zhihao Jía, Matei Zaharia, and Alex Aiken. Beyond data and model parallelism for deep neural networks. CoRR, abs/1807.05358, 2018. (ttLittt://aaat org. abs//as1111 05358
Zhihao Jia, Sina Lin, Mingyu Gao, Matei Zaharia, and Alex Aiken. Improving the accu- racy, scalability, and performance of graph neural networks with roc. In I. Dhillon, D. Pa- pailiopoulos, and V. Sze (eds.), Proceedings of Machine Learning and Systems, volume pp. 187-198, 2020. URL 1EBF12990EE0011222beeee332229223
inw.innsmmmm..... ml mysss..... Paperr. pdf
Norman P. Jouppi, Cliff Young, Nishant Patil, David Patterson, Gaurav Agrawal, Raminder Bajwa, Sarah Bates, Suresh Bhatia, Nan Boden, Al Borchers, Rick Boyle, Pierre-lue Cantin, Clifford Chao, Chris Clark, Jeremy Coriell, Mike Daley, Matt Dau, Jeffrey Dean, Ben Gelb, Tara Vazir Ghaemmaghami, Rajendra Gottipati, William Gulland, Robert Hagmann, C. Richard Ho, Doug Hogberg, John Hu, Robert Hundt, Dan Hurt, Julian Ibarz, Aaron Jaffey, Alek Jaworski, Alexander Kaplan, Harshit Khaitan, Daniel Killebrew, Andy Koch, Naveen Kumar, Steve Lacy, James Laudon, James Law, Diemthu Le, Chris Leary, Zhuyuan Liu, Kyle Lucke, Alan Lundin, Gordon MacKean, Adriana Maggiore, Maire Mahony, Kieran Miller, Rahul Nagarajan, Ravi Narayanaswami, Ray Ni. Kathy Nix, Thomas Norrie, Mark Omernick, Narayana Penukonda, Andy Phelps, Jonathan Ross, Matt Ross, Amir Salek, Emad Samadiani, Chris Severn, Gregory Sizikov, Matthew Snelham, Jed Souter, Dan Steinberg, Andy Swing, Mercedes Tan, Gregory Thorson, Bo Tian, Horia Toma, Erick Tuttle, Vijay Vasudevan, Richard Walter, Walter Wang, Eric Wilcox, and Doe Hyun Yoon. In-datacenter performance analysis of tensor processing unit. SIGARCH Comput. Archit. News, doi: 14.114573140559300022222 URL https
412 June 2017. ISSN 0163-5964. 3080246
//doi.org// 10.
Jehandad Khan, Paul Fultz, Artem Tamazov, Daniel Lowell, Chao Liu, Michael Melesse, Murali Nandhimandalam, Kamil Nasyrov, Ilya Perminov, Tejash Shah, et al. Miopen: An open source library for deep learning primitives. arXiv preprint arXiv:1910.0000 2019.
Alex Krizhevsky, Ilya Sutskever, and Geoffrey Hinton. Imagenet classification with deep convolu- tional neural networks. In Proceedings of the 25th International Conference on Neural Information Processing Systems Volume 1, NIPS' 12, pp. 1097-1105, Red Hook, NY, USA, 2012a. Curran Associates Inc.
Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. Imagenet classification with deep convolu- tional neural networks. In Advances in neural information processing systems, pp. 1097-1105, 2012b.
Y. LeCun, B. Boser, J. S. Denker, D. Henderson, R. E. Howard, W. Hubbard, and L. D. Jackel. Backpropagation applied to handwritten zip code recognition. Neural Computation, 1(4):541-551, 1989. doi: 10. 162/neco. 1989.1.4.5541.
11

---

Under review as a conference paper at ICLR 2022
Tatiana Likhomanenko. Qiantong Xu, Vineel Pratap, Paden Tomasello, Jacob Kahn, Gilad Avidov, Ronan Collobert, and Gabriel Synnaeve. Rethinking evaluation in asr: Are our models robust enough? arXiv preprint arXiv:2010.11745, 2020.
Tsung-Yi Yi Lin, Michael Maire, Serge Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and C Lawrence Zitnick. Microsoft coco: Common objects in context. In European conference on computer vision, pp. 740-755. Springer, 2014.
Yanjun Ma, Dianhai Yu, Tian Wu, and Haifeng Wang. Paddlepaddle: An open-source deep learning platform from industrial practice. Frontiers of Data and Computing, 1(I):105-115, 2019
James Malcolm, Pavan Yalamanchili, Chris McClanahan, Vishwanath Venugopalakrishman. Krunal Patel, and John Melonakos. Arrayfire: gpu acceleration platform. In Modeling and simulation for defense systems and applications VII, volume 8403, pp. 84030A. International Society for Optics and Photonics, 2012.
S. Markidis, S. W. Chien, E. Laure, I. Peng, and J. Vetter. Nvidia tensor core programmability. performance precision. 2018 IEEE International Purallel and Distributed PaclteddPPPooee Processing Symposium Workshops (IPDPSW), pp. 522-531, 2018.
NCCL. Nvidia collective communications library (ncel). bttp:://github..cmmm
nec1 2019.
Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zach DeVito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in pytorch. 2017.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32: 8026-8037, 2019.
Vineel Pratap, Awni Hannun, Qiantong Xu, Jeff Cai, Jacob Kahn, Gabriel Synnaeve, Vitaliy Liptchin- sky, and Ronan Collobert. Wav2letter++: A fast open-source speech recognition system. In ICASSP 2019- 2019 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 6460 6464, 2019.
doi: 10. 10.1109/CCASSP22199.60 1.683533.
Benoit Steiner, Chris Cummins, Horace He, and Hugh Leather. Value learning for throughput optimization of deep learning workloads. In A. Smola, A. Dimakis, and Stoica (eds.), Proceedings of Machine Learning and Systems. volume 3, pp. 323 334, 2021. URL htttps:/pproceedingg wtwp:eewwiiin..........o.m..m
Gabriel Synnaeve, Qiantong Xu, Jacob Kahn, Tatiana Likhomanenko, Edouard Grave, Vineel Pratap, Anuroop Sriram, Vitaliy Liptchinsky, and Ronan Collobert. End-to-end asr: from supervised to semi-supervised learning with modern architectures. arXiv preprint arXiv:1911.08460, 2019.
Eclipse Deeplearning4j Development Team. Deepleaming44j: Open- distributed deep learning for the jvm. 2016. URL /ttpttt:///ttttttt
com www.////ww////eeooccccococtpeeee hetp:ppappeeeppleeainingggg
Thomas N. Theis and H.-S. Philip Wong. The end of moore's law: A new beginning for information technology. Computing in Science Engineering, 19(2):41-50) 2017. doi: 10.1109/MCSE.201..22
Seiya Tokui, Ryosuke Okuta, Takuya Akiba, Yusuke Niitani, Toru Ogawa, Shunta Saito, Shuji Suzuki, Kota Uenishi, Brian Vogel, and Hiroyuki Yamazaki Vincent. Chainer: A deep learning framework for accelerating the research cycle. In Proceedings of the 25th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, pp. 2002-2011, 2019.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Proceedings of the 31 International Conference on Neural Information Processing Systems, NIPS' 17, pp. 6000-6010, 6010, Red Hook, NY, USA, 2017. Curran Associates Inc. ISBN 518169066
12

---

Under review as a conference paper at ICLR 2022
Haojie Wang, Jidong Zhai, Mingyu Gao, Zixuan Ma, Shizhi Tang, Liyan Zheng, Yuanzhi Li, Kaiyuan Rong, Yuanyong Chen, and Zhihao Jia. PET: Optimizing tensor programs with par- tially equivalent transformations and automated corrections. In 15th USENIX Symposium Operating Systems Design and Implementation (OSDI 21), pp. 37-54. USENIX Association, URL httph:htt//wtt useni i wwt.o//////ooooreeoeeoo
July 2021. ISBN
wang
Lianmin Zheng, Chengfan Jía, Minmin Sun, Zhao Wu, Cody Hao Yu, Ameer Haj-Ali, Yida Wang, Jun Yang, Danyang Zhuo, Koushik Sen, Joseph E. Gonzalez, and Ion Stoica. Ansor: Generating high-performance tensor programs for deep learning. In 14th USENIX Symposium on Operating Systems Design and Implementation (OSDI 20), pp. 863-879. USENIX Association, Novem- ber 2020. ISBN 978 19999993933339
URL https: www useni x org/conferen.ce)
osaii20/preeeetat..onm// zheng
A APPENDIX
A.1 REPRODUCIBILITY
All tools and code used in our evaluation are open source and available on Github this includes scripts to reproduce benchmarks and other quantitative codebase metrics. Code can be found at URL redacted to preserve anonymity.
Flashlight is open source on Githut as stated in Section 5 Flashlight v0.3.1 is used to reproduce results using ArrayFire 3.8 as the underlying tensor backend. No other specialized configuration used for either Flashlight or PyTorch or TensorFlow.
A.1.1 FROM-SOURCE COMPILATION
Flashlight is built in CMake release mode, which is also the default for both PyTorch and TensorFlow builds. Build settings are kept default for all frameworks; Flashlight uses Ninja with CMake in accordance with PoTochh's build, while Tensorflow employs Bazel.
Exact build step reproduction be found in the aforementioned repository.
A.1.2 INCREMENTAL COMPILATION
Incremental compilation benchmarks are performed using similar build setups per Section A.1.I To test incremental rebuilds, source files are randomly selected without replacement from a distribution constructing by weighting each file by the number of lines in the file, and using that number to determine the probability of selecting for modification.
Scripts perform and time this incremental compilation can be found in the aforementioned public repository.
A.2
CODE COMPLEXITY
A.2.1 OPERATOR COUNTING
To count the number of operators for each framework, we utilize operator schemas for PyTorch and Tensorflow (which generate code from those schemas, accordingly) written in YAML and Protobuf, respectively. For Flashlight, we count the number functions in the Flashlight TENSOR interface and autograd interfaces, as these form the full implementation requirements for full tensor backend that functions on all platforms. The scripts released on Github detail the files and filtering techniques used reproduce the number of results.
To count the number of operators for each implementation, we use the above operator lists, then count the number of operators that perform the specified function, even if those operators perform other functions. For example: an operator called ADDMM, which performs an addition operation followed
Precactedfooannnnyyne httpss//minja-hulld.om
13

---

Under review as a conference paper at ICLR 2022
by matrix- matrix multiplication, performs addition operation, and would thus be counted when tallying the number of ADD operators.
A.3 PERFORMANCE
The aforementioned public repository provides scripts required fully-reproduce all performance measures.
A.4
DESIGN DETAILS AND CODE SAMPLES
In the following sections, we show brief code samples expounding on those Section 4
A.4.1 MODULES
Flashlight's MODULE abstraction is similar to that of frameworks such Torch and PyTorch, can recursively store other modules and interoperate with more sophisticated abstractions including CONTAINER, which wraps multiple modules, SEQUENTIAL, which stores sequences of modules and forwards data through them sequentially, and user-defined abstractions. Listing 7 in Section A.4.2 shows an example of Sequential usage.
Listing 5
shows small Dropout module implementation that calls into the dropout autograd primitive, stores and serializes a small amount of state, and defines simple forward function on a Variable.
class Dropout public Module
private:
double ratic_;
(Module, ratio_) serialization
public:
Dropout (dcuble drop_ratio = 0.5); Variable forward (const Variables input) override { it (train_)
10
return dropout (input, ratic_); // autograd primitive else
return input;
13
14
15
//
16
};
Listing 5: Dropout layer implemented a Flashlight module.
The SLL_EAVVELLON WITH_BASH macro defines serialization of the Dropout class as a module, including any fields to be serialized (in this case, only the dropout ratio).
A.4.2 AN END-TO-EN END EXAMPLE: MNIST
Below, we detail simple end-to-end training setup.
First, data loaded using the BATCHDATASET abstraction in Listing 6
const int KTrainSize = 60000; const int kValSize 5000;
autoś [train_x, train_y] load_datased (data_dir);
// Hold out a dev set auto val_x
train_x(span) span, range kValSize)); train_x train_x span, range (kValSize, kTrainSize));
(range (0, kValsize)); 10 train_y = train_y (kValsize, kTrainSize));
auto val_y train (range
11
12 // Make the training batch dataset. Hatchlataset trainset
14

---

Under review as a conference paper at ICLR 2022
x, train_y
}),
15
batch_size);
16
Make the validation batch dataset
18
BatchDataset valset
19
std: mar.aeeeeesrrreceeesssttttassss (std: (ector<<ensor>>vaall val_yi),
20
match_sized
Listing 6: Loading MNIST data into train and evaluation set.
A full description of the LOAD_ Githut
LOAD_DATASEE DATASET function be found in the MNIST training example
We construct the model using simple SEQUENTIAL. in Listing 7
const int k klmagellim 28;
auto pac PaddingMode:: SAME;
Sequential model;
(view (K Imaggeiim, kImageDim, 1, -1)1); // WHCN (col major)
model add (Conv2D
1 /* input channels +/,
32 /* output channels +/,
5 /* kernel width +/,
/* kernel height */, stride x */,
11
12
1 /* stride y +/,
13
pad /+ padding mode +/,
14
pad /+ padding mode +/1);
15 model. .add ())IU());
16
model add (Poc120(
17
2 /x kernel width */,
18
2 /* kernel height x/,
/* stride */,
20
stride y +/));
21 model. (Conv2D (32, 64, 5, 5, 1, 1, pad, pad) );
22 model. ());
model.ace (Poc12D(2 (2, 2, 2, 2)); model. (V7
64, -1});;
25 model add (Linear (7
64, 1024));
26 model acc (ReLU());
27 model acc. (Dropout (0.5)) 51
model (Linear (1024, 10));
model acc (LegSoitmax )
Listing 7: Constructing CNN for MNIST training.
we create a simple custom training loop. This uses optimizer, loss function, and meter abstractions as provided by default by Flashlight. We perform the forward and backward pass, step the optimizer to update parameters, and zero out gradients before moving to the next batch. We evaluate the model using the function defined in Listing pulls out the max prediction and comparing it against the ground truth, updating the loss meter we go, then returning the final loss values.
In Listing 8
// Make the optimizer sGnoptimize opt (model.3 modeel.prram. params (), learning_rate)
// The main training loop for (int e 0; e epochs; e++) AverageValueMeter train loss meter;
// Gct an iterator over the data for (auto&& example dataset) { auto inputs = noGrad (example (INNUUTTDX]]]
recactedfooranonmiity
15

---

Under review as a conference paper at ICLR 2022
auto output model (inputs)
12
13
14
auto target noGrad (cxample (TARGET_IDX]]
// Compute and record the less.
17
auto lcss categorical cyt.ossertrrorrrryyoooooooo (output, target); train_loss_meter. add (Loss.tensor t
18
scalar<float>
19
// Backprop, update the weights and then zero the gradients.
loss (ackward
22
ste.
23
opt zer.zerrd
26
double train loss train loss meter. meter. () (0)
scalar<double> ()
27
28
// Evaluate the dev sct. double val_loss, val_error; std::tie (val_loss, val_error) eval eval_loop (model, valset);
31
32
std: :cout
<< "Epoch << e << std: :ettpeccppiiinin (3)
33
<< ": Avg Train Loss: << train_
<< Validation Loss: << val loss
35
<< Validation Error (8): " val << std::endi
36
Listing 8: A simple training loop.
std: std::bble, double> eval_lcop (Sequentials model, BatchDatasets dataset) AverageValueMeter FrameErrorMeter error_meter:
Place the model in eval mode. model.eval ();
for (autic& example : dataset) { auto inputs = noGrad (example INPUT IDX]) auto output model (inputs)
10
11
// Get the predictions in max_ids
12
Tensor max_vals, max_ids; max (max_vals, max_ids, output.tensor. 0);
15
auto target noGrad (example (TARGET_IDX)) (TAA
16
// Compute and record the prediction error. error meter. wdweeeeeeeeeeeeeeeeeeeaeeeeeeeee (deorde (max ids, 1, 0), target sarget.tensor )
19
20
// Compute and record the loss. auto loss categoricallrossent loss_meter. add tensr.
21
(output, target); tensor scalar<float> ());
23
24
Place the model back into train mode.
25
model.train())
double errer error_meter value scalarcdouble> () double loss Loss_meter. value () scaler<double> (); return std: :adkm_ma__pppp pair (loss, error);
28
29
30
Listing 9: Evaluating a training model MNIST.
we evaluate the trained model on the test set by creating a test dataset and using the previously-defined evaluation function.
Finally, in Listing 10
16

---

Under review as a conference paper at ICLR 2022
std: :pair<duubb, double> eval 1cep (Sequential& model, BatchDatasets dataset)
AveragevalucMeter loss_meter; FrameErrorMeter error_meter;
// Place the model in eval mode. model.eval ()
for (aute& example : dataset) [ auto inputs = noGrad (example auto output model (inputs);
// Get the predictions in max_ids Tensor max_vals, max_ids; max vals, max max_ids, output.tensor 0);
12
15
auto target noGrad (example (TARGET__DG
16
17
// Compute and record the prediction error.
error_meter, add (rcorder (max_ids, 1, 0), carget.tensort
20
// Compute and record the lcss. auto lcss - categorica Caroosssoocrrss (output, target); less_meter.at (1oss.
21
22
() scalar<float>
24
// Place the model back into train mode.
25
model.traint)
26
double error error_meter. value (). scalar<double>> (); double less loss_ boss_meter.ve value scalarcdouble> () return std: :make_ :make_pair (loss, error);
29
30
Listing 10: Evaluating training model MNIST.
17


Reviews:
[{"review_id": "vd1fC5pLEEG", "reviewer": "Reviewer_KzrA", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "I have to admit that this paper gives me mixed feelings. Machine learning compiler/framework is my major of study, and I understand how we often have the desire to redesign the framework. It is not only (in most of the cases) harder than anticipated, but also (in most cases) just create \"another\" framework that is not necessary better. It is also very hard to argue that a framework is better, since there isn't a good way to measure the \"quality\" of a framework.\n\nStrengths:\n1. Flashlight captured the \"core\" of the Machine Learning Frameworks: Tensor and Operation, Memory Management, and Distributed Computation. The core APIs are interesting and somewhat innovative. \n\n\nWeakness:\n1. I have to admit that it is tremendously hard to argue (with evaluation) that one machine learning framework is better than another. In this paper, we see too many vague arguments such as \"agile\", \"modular\", and \"nimble\". It is unfortunately that such vague arguments are bad for a scientific paper. On the same cord, the evaluations in the paper doesn't really mean too much. For instance, Table 1 argues that Flashlight is small. Table 2 argues that Flashlight is quick to compile. Table 3 show model performance, but most models are similar (AlexNet is small and small models are easier to be faster in a simple framework). None of the evaluations support the argument of \"agile\", \"modular\", and \"nimble\".\n\n2. The interesting parts (or, we could say the \"pain points\" of machine learning frameworks in general) are memory management and distributed computation. The paper highlighted the importance of them, but the contribution of Flashlight in these aspects seems minimal (just described the minimal APIs and connected to downstream implementations). I agree that other researchers might be able to take advantage of these APIs, but the contribution seems small.\n\nHow could the paper be improved?\n\nHonest, I am not sure. To be fair, Pytorch paper was published years after it was popular. TensorFlow was published when it was already well-known. Recent, MLIR made a big noise in machine learning framework community, but still, the popularity of the tool is way more important that the publication of it.\n\nTaking MLIR as a reference, I think the reason it is loved so much is that it facilities (greatly) the development of IRs and IR transformations. I think it will be helpful to show case how Flashlight facilities other researchers by the \"agile and minimal\" design. That would be a more convincing argument of the value of Flashlight.\n\n", "summary_of_the_paper": "The paper proposed a minimal design API (or mostly API?) of machine learning framework called Flashlight. The key argument of the paper is that Flashlight is modular and agile. The Flashlight captured the key aspects of machine learning frameworks: Tensor and Operation, Memory Management, and Distributed. There are some evaluations to argue the benefit of Flashlight.", "summary_of_the_review": "I empathy the difficult of arguing the \"betterness\" of a machine learning framework. However, I think the paper in general has not made a convincing argument or provide enough scientific contribution. So I recommend rejecting this paper. However, I am very interested to check out the open source implementation of Flashlight, and see how amazing it is.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "gHlXAZDsYv_", "reviewer": "Reviewer_gStW", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The authors have built what appears to be a well thought out, and cleanly implemented neural network framework with researcher modification and extensibility in mind. The paper is well written, with clear background sections on the current state of framework design, and the need for a more research-oriented extensible framework. Indeed a simplistic, easily modifiable/extensible framework that can be used to easily implement new research methods that might require custom forms of memory management, runtime, or tensor functions is immediately appealing.\n\nWhile I commend the authors for what is no small undertaking, I have a few concerns:\n\n1. The first and primary one is whether this work sufficiently distinguishes itself in a crowded field. Between the PyTorch C++ library, Tensorflow Lite, Tensorflow Lite Micro, etc., there are many options for implementing custom built solutions, some of which offer a similar level of control over memory management, and operator behavior as described in Flashlight (albeit somewhat less cleanly perhaps). Via the PyTorch C++ library, custom operators can be implemented, memory allocation behavior can be customized as desired, and hooking into the PyTorch ecosystem carries significant advantages. So the question I'm left asking myself is whether Flashlight provides any significant advantages over the PyTorch C++ library for the majority of use cases. Unless a researcher is implementing functionality that really does involve altering very large portions of framework internals, I'm not convinced that optimizations are worth giving up for the sake of readability.\n\nThat being said, the framework seems nicely designed, with clean interfaces and its use of ArrayFire for (nice) cross-platform tensor calculation, and a small binary size and compile time to boot which is great for development purposes (although again, extending PyTorch with a custom op does not require compiling the library from scratch). I could see using Flashlight in order to roll a very custom solution.\n\nI am willing to believe that the code cleanliness, modularity and practicality of Flashlight could be beneficial to some researchers, but beyond that, I'm unsure as to what Flashlight brings to the table other than that it's more barebones and a good jumping off point for research than other frameworks. The methods it employs are (unless I'm mistaken) not exactly unique to the work, and I'm unsure if a conference paper is necessarily the right venue for this work. \n\n2. I assume the 9MB binary size stated in the paper does not include ArrayFire... Does the PyTorch 528mb include Aten? and similar for Tensorflow?\n\n3. I would like to hear more about why Flashlight performs better than PyTorch or Tensorflow. Is this an Aten v.s. ArrayFire speedup?", "summary_of_the_paper": "This paper proposes a neural network training and inference framework aimed at framework researchers with a focus on modularity, simplicity of design, and extensability by researchers.", "summary_of_the_review": "Overall the paper was interesting, and I commend the authors for their work. It seems the Flashlight framework could be beneficial for researchers looking for a boilerplate template framework on which to build their very highly customized solutions. However, for anything other than an solution requiring rewriting large portions of a neural network framework, the researchers using this library will give up optimizations and other benefits of using a more mature library such as PyTorch C++ (for which custom ops and some other features mentioned in this work are possible without rewriting internals).", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "ZHHPEkgDLog", "reviewer": "Reviewer_1j8s", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Flashlight’s narrow use-case, as a testbed for systems research in machine learning, is important and poorly supported by its major competitors. This comes at the cost of other use-cases like model design research (limited feature set, need to compile models in C++) or usage in production environments (avoidance of low-level optimizations, absence of model servers for deployment).\n\nWhile Flashlight with its native support for efficient memory management, distributed training and end to end benchmarking, shows promise as a research tool. However, it has not demonstrated its capabilities in practice. Whereas a long list of well-established (reference) models in speech, vision, and text applications are provided, the only discussion of novel work enabled by Flashlight is a brief mention of wav2letter in Section 3. The performance study on well-established models gives confidence in the basic implementation of Flashlight, which is of limited use by itself. This would be a strong paper if it contained a case study showing how its design contributed to systems research.\n", "summary_of_the_paper": "The paper describes the design philosophy and structure of the Flashlight deep learning framework. Flashlight is modular, small, and narrowly oriented toward systems researchers. Rather than (or in addition to) high level productivity, Flashlight focuses on internal and external simplicity of ML tools. The authors evaluate Flashlight by training several standard reference models and it achieves slightly better performance than PyTorch or TensorFlow\n", "summary_of_the_review": "Flashlight is well-designed and shows promise as a research tool for systems research in machine learning. However, it has not demonstrated its usefulness in practice.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "P5ZwM9emtep", "reviewer": "Reviewer_DYUa", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "\n**Strength of the work.**\n- The paper provides performance comparison for end-to-end training performance on 6 representative\nworkloads, and the proposed system outperforms PyTorch and TensorFlow in most of the cases;\n- The proposed system allows customized memory manager, which is a missing piece in many of the existing\ndeep learning frameworks;\n- The system is ready to use in production, which makes it outstanding given the robustness requirement\nin the production environment.\n\n**Weakness.**\n- Ablation study is missing from the paper. For example, training ResNet-50 on a single GPU\nhas become a well-studied and highly-optimized task. However, the system is still able to outperform\nexisting large systems like TensorFlow. Therefore, it would be desirable to showcase why the system\nis able to do better. Is it because Flashlight does better latency hiding in data loading, or the\nlow-overhead nature of C++ APIs, or any other reasons?\n- More direct and fair comparison might be required for the reviewer to better understand\nin Table 1 and 2. According to Table 1, FlashLight has far fewer operators than PyTorch (110 vs 2166),\nso it seems nature that it compiles faster than PyTorch (27min vs 754min), and generates smaller binaries (9MB vs 527MB).\n- Claims need more justification with data and comparison with other frameworks. For example, in Section 3,\nit's claimed that the proposed system is research-first. It would be desirable to quantifiably understand the reason\nbehind the claim. Compared with TensorFlow where a lot of research have been done on for distributed training,\nwhy the proposed system is more suitable?\n\n**Correctness.** To the best of the reviewer's knowledge, there is no correctness issue.\n\n**Clarity.** The paper is not hard to read, but might need more data to justify many of it claims.\n\n", "summary_of_the_paper": "The paper presents Flashlight, a C++-based deep learning framework, based on the ArrayFire\ntensor library, and implemented all the necessary functionalities in the core infrastructure,\nincluding memory management, distributed training, data loader, etc. The major contribution includes:\n- An end-to-end framework with solid engineering efforts, which comes with highly optimized models\nincluding speech, vision and text;\n- The system is modularized and customizable with simple C++ APIs, and each component is equipped\nwith highly optimized default implementation;\n- The system is benchmarked with a variety of the state-of-the-art models, and supports end-to-end\npackaging for these applications.\n", "summary_of_the_review": "In summary, the paper presents a solid deep learning framework written in C++, and has been put in production for a wide range of applications. However, in terms of paper writing, it would be more desirable for the readers and reviewers to understand the novelty and quantifiably justify the claims in the paper.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper describes Flashlight, a tool for ML researchers with specific design considerations for conducting systems research. The needs for such tool are significant, and recent advances in this topic have been relatively slow, so this research is timely and important. 

Reviewers are positive about the importance of the problem and the nice design of Flashlight. It seems the tool has been used by researchers with positive feedback. At the time of the original submission, reviewers expressed some concerns about the novelty and the weak arguments for convincingly showing the advantages over other similar tools. 

Authors provided nice replies including specific case studies, but with the short time period to reassess the proposed changes and additions, some reviewers remain hesitant, and thus this paper cannot be accepted at this time. I strongly encourage the authors to incorporate all of the proposed revisions and resubmit to a future venue.
Decision: Reject