Converting PyTorch Models to StreamIt Pipelines
Name
rajvee-muhender-meng-eecs-2025-thesis.pdf
Description
Thesis PDF
Size
928.42 KB
Format
Adobe PDF
Checksum (MD5)
87bd6019e6cf1f313208546d08afe866
Author(s)
Rajvee, Muhender Raj
Advisor(s)
Amarasinghe, Saman
Date Issued
May 2025
Publisher
Massachusetts Institute of Technology
Abstract
With the rise of large language models, there have been efforts to optimize machine learning inference to support a large volume of queries. Currently, the two main ways to do this are running optimized kernels for computing the forward inference pass and distributing computation across multiple GPUs or different cores in a GPU. Machine learning libraries such as PyTorch produce dynamic computation graphs in order to represent the forward pass of the model. PyTorch allows conversion of these dynamic graphs into static ones through just-in-time (JIT) compilation. These graphs can then be optimized further by the compiler. We propose an alternate way of optimizing these dynamic graphs. We convert the dynamic computation graph of PyTorch to pipelines in StreamIt, a domain specific language (DSL) for streaming applications, and use the multi-stage compilation property of BuildIt to compile this pipeline in stages to inference code. We found that, while the inference latencies of models compiled in this way are slightly higher, they are still comparable to those of PyTorch models and are open to future optimizations.
MIT Department
Massachusetts Institute of Technology. Department of Electrical Engineering and Computer Science
Terms of Use
In Copyright - Educational Use Permitted
Copyright retained by author(s)
Persistent DSpace Link