Google Summer of Code 2024 Work Submission - Machine Learning in Julia for Calorimeter Showers @ CERN-HSF

GSoC'24 Final Report.md

Machine Learning in Julia for Calorimeter Showers

Organization: CERN-HSF

Contributor: Daniel Regado (guiyrt)

Code Repository: MLinJulia

Presentation at JuliaHEP 2024: Machine Learning in Julia for Calorimeter Showers

Project Proposal: ML in Julia for Calorimeter Showers

Mentors:

• Graeme Stewart (graeme-a-stewart)
• Pere Mato (peremato)
• Piyush Raikwar (Piyush-555)

Description: The calorimeter, a key detector in Large Hadron Collider (LHC) experiments, measures the energy of particles interacting with detector materials. Particles emerging from collisions create cascades of secondary particles, or showers. Describing these processes requires precise simulation methods like the Geant4 toolkit. Machine learning (ML) techniques, such as generative modeling, offer a promising alternative. Recently, the Fast Calorimeter Simulation Challenge (CaloChallenge) spurred the development and evaluation of different models. Also, in High-Energy Physics (HEP), there has been an increasing interest in using Julia as a language for software development, for combining the ease of programming of interactive languages, e.g. Python, with the speed of compiled languages, e.g. C++. This project targets to assess Julia's machine learning ecosystem's maturity, in terms of availability and robustness of libraries and tools for deep learning. To this end, a selection of models from the CaloChallenge will be chosen for implementation in Julia, as well as the required functionality to enable training and evaluation of the networks. Ultimately, a comprehensive analysis will be conducted between Python/C++ and Julia implementations, considering training time, inference time and evaluation metrics from CaloChallenge.

Work Completed

The initial milestone for this project was comparing the different approaches from the CaloChallenge, and deciding which to implement in Julia. For this purpose, the CaloChallenge Summary document was created, on which key details from each implementation are presented, such as architecture, ML framework and sample generation fidelity. From this investigation, the denoising diffusion model CaloDiffusion stood out as the best performing model for every dataset, albeit with the drawback of high sampling latency comparing with other architectures. For this reason, CaloDiffusion was chosen as the first model to implement, with CaloINN as the most probable implementation to follow. However, as it was identified a significant difference in terms of performance between PyTorch and Flux.jl implementations, we decided to favor understanding this unexpected slowdown by profiling the CaloDiffusion implementations, rather than implementing more models that would be equally slower against their PyTorch counterpart. Considering this change to the original planning, the following objectives were achieved:

Data

• Pre-processing: Includes data normalization and other transformations in the same manner as in the original Python implementation. The data is loaded and kept in Float32 precision, a process that required careful attention on operations performed, to avoid casting to Float64 in intermediary steps.

Model and Training

• Implemented base layers: CaloDiffusion uses convolutional layers with circularly padded input, which more accurately represents the calorimeter geometry. This required implementing CylindricalConv and CylindricalConvTranspose layers, as well as LinearAttention using CylindricalConv layers. Additional support layers for LinearAttention were also implemented, namely PreNorm and Residual layers.
• Add output padding for ConvTranspose #2462: Contribution on Flux.jl to match PyTorch's ConvTranspose implementation, on which the output shape for stride > 1 can be increased according to an output padding argument. This change was merged and released in Flux.jl v0.14.17.
• Implemented block layers: Using layers from Flux.jl and the base layers previously mentioned, the blocks required to create the model were implemented. These include ResNetBlocks and ConvBlocks, which are part of each Up and Down sampling blocks.
• CaloDiffusion Model: With all the layers and blocks available, the model creation and forward pass were defined. The model architecture can be adjusted based on a configuration file. This facilitates training with different datasets, which have different calorimeter geometries. Additionally, it also allows to have different model sizes, according to the dataset complexity.
• Training script: Based on a configuration file, the model is instantiated and starts training, with a validation loop after each epoch. A checkpoint is saved for the lowest loss value, and a final evaluation loop is performed when the training is finished.
• GPU support with CUDA: Training with a GPU was a highly desired functionality, as it considerably speeds up training. This required changes on the forward pass code of some layers, as the automatic differentiation has GPU-specific incompatibilities.

Validation Tests

To assert that this implementation is consistent with the original, numerous tests were created comparing Python and Julia output. Using the PyCall Julia package, one can import Python packages, call Python functions, and interact with their output in Julia native types, when available. For validation purposes, the tests followed this general flow:

1. Random input data: Create random Julia Float32 multi-dimensional Array, which is used to create an equivalent Torch Tensor object.
2. Instantiate Python object and Julia struct: First, the Python object of interest is created. From this object, a previously implemented constructor is used to instantiate the equivalent Julia struct. This constructor maps the properties of the Python object (generally layer parameters, weights, and biases) to the Julia struct.
3. Generate output and test equivalence with ≈ operator: Both outputs are compared with the binary operator ≈, which only returns True if the difference is less than square root of the machine epsilon for the given type.

Since usually the Python output is a Torch Tensor, the function fromtorchtensor was created to cast a Torch Tensor to a NumPy Array, which is then automatically usable as a Julia Array. It also reverses the array dimensions, as where Python is row-major, Julia is column-major and therefore this adjustment is needed.

As an example, below is the test for validating that the Julia and Python ResNetBlock implementations are equivalent:

@testset "ResNetBlock" begin
    data, torchdata = rand32tensors(9, 16, 45, 32, 128)
    time, torchtime = rand32tensors(128, 1)

    torchrnb = pymodels.ResNetBlock(dim_in=32, dim_out=16, cond_dim=128, groups=8, cylindrical=true)
    rnb = ResNetBlock(torchrnb)

    @test rnb(data, time) ≈ torchrnb(torchdata, torchtime) |> fromtorchtensor
end

There are tests in place for the following structs and functions:

• Pre-processing
• Layers
• Blocks
• Model
• Loss

As a fortunate consequence of implementing these tests, the strecth goal initially listed in the proposal was achieved, which refers to the ability of loading the PyTorch model in Julia. To validate that both model implementations produce the same output, they must have not only the same input data, but also the exact same parameters.

Benchmarks and Profiling

As initial results pointed to significant differences in performance, we decided to invest development time into collecting and analyzing information about execution times and resources utilized. This analysis targeted the training step, which starts with loading a batch of data, and ends with updating the weights.

• Benchmarking: Analyses the training step without introspection, providing values regarding the entire execution step. In Julia, BenchmarkingTools.jl was used to provide statistics about execution time, memory allocated and garbage collection (GC) time. In Python, PyTorch's internal util.benchmark was used, which only provides mean execution time.
• CUDA profiler: Unlike with benchmarks, the profiler used enabled a fine-grained analysis into function calls, system calls and GPU memory usage in a timeline. Using an external profiler, NVIDIA® Nsight Systems, it was possible to visualize and analyze both implementations with minimal diferences. Not only does this make it easier for comparisons, but it is also safe to assume that the overhead should be similar, independently of application profiled.

Current State

By the end of the project, CaloDiffusion is trainable both on CPU and NVIDIA GPUs. Additionaly, all the validation test are passing, which assures us that if there are numerical differences in the output, these are not significant. By employing the benchmarker and profilier, we can measure how differently both approaches perform. All of the results below are using CaloChallenge's dataset 2, and values represent the mean over 20 samples. In terms of hardware, the machine on which the data was extracted has a Intel® Core™ i5-13600KF, NVIDIA® GeForce RTX™ 4070 Ti SUPER 16GB GPU, and 32GB of RAM. The data from benchmarking and profiling is available in the Links section below.

CPU Training

Using data from the benchmarks, we can see that Julia's implementation with Flux.jl takes roughly 10x to perform the same training step. Additionaly, as BenchmarkTools.jl provides information regarding garbage collection, we noticed that a big part of the execution time is courtesy of GC.

Batch size	Implementation	Step time	Memory usage	GC time (%)
4	PyTorch	0.31 s	-	-
	Flux.jl	3.09 s	3.46 GiB	44.85%
16	PyTorch	1.30 s	-	-
	Flux.jl	11.54 s	12.54 GiB	40.30%
32	PyTorch	2.75 s	-	-
	Flux.jl	22.91 s	24.65 GiB	39.78%

GPU Training

Using the GPU profiler, similar results are obtained but for batch size 4. With this smaller batch size, GC is not so frequent and it significantly reduces the performance gap. However, even in this scenario, the Flux.jl implementation remains over 3x slower. In terms of memory usage, there is also a big difference, with Julia being more memory hungry. Additionaly, it is worth highlighting that memory usage with PyTorch is constant throughout the training steps, and it appears that all the memory required is allocated before-hand. This can be verified in the screenshots below.

Batch size	Implementation	Step time	Memory usage	GC time (%)
4	PyTorch	24.02 ms	282.60 MiB
	Flux.jl	87.95 ms	3.07 GiB	13.95%
16	PyTorch	37.72 ms	724.60 MiB	-
	Flux.jl	333.53 ms	12.39 GiB	62.56%
32	PyTorch	53.67 ms	1.22 GiB	-
	Flux.jl	651.30 ms	24.69 GiB	66.23%

GPU Profiling

Using batch size 32 as example, we first a baseline PyTorch training step. It lasted 56.22 ms with constant memory usage of 1.22 GiB. We can also see above the Memory usage green chart that kernels are constantly running, keeping the GPU busy.

Now in Julia, we have a 498.96 ms step reaching maximum GPU memory utilization at 14.33 GiB. This triggers a full GC call that lasts 245.31 ms and interrupts kernel executions. This is the main cause for the performance degradation on the Julia side, where the GPU sits idle with full memory.

However, the previous example was the best case scenario in terms of GC. In this image, we have a step that lasted 834.74 ms with 3 GC calls, which sums to 634.68 ms of waiting time!

Based on this, we can confidently say that the main issue on the Julia side is related to memory management, whether by sub-optimal implementation or by the automatic differentiation for calculating gradients. For reference, there is a slim vertical blue line on the previous Julia profiler images, which represents the point at which the loss is obtained as the last step of the forward pass. After that line, a majority of the time is spent in the backwards pass and about 10 ms updating the weights.

Future Work

With the pivot to performance analysis, not all the original objectives were achieved. The following list identifies tasks or features of interest, ordered by importance:

• ☐ Further identify performance bottlenecks with layer-level analysis.
• ☐ Implement additional models from CaloChallenge.
• ☐ Implement evaluation metrics in Julia.
• ☐ Implement high level feature visualizations in Julia.

Links

https://arxiv.org/abs/2308.03876 - CaloDiffusion - "Denoising diffusion models with geometry adaptation for high fidelity calorimeter simulation" arXiv Paper.

https://github.com/guiyrt/CaloDiffusion - Fork from CaloDiffusion original implementation, refactored and with additional script.

https://github.com/guiyrt/MLinJulia - CaloDiffusion implementation in Julia, main project repository.

Benchmarking / Profiling - Data from CaloDiffusion performance analysis on dataset 2 size model.