PyTorchAlpaka: Batched inference with TensorCollections, eval and frozen model, FP16 convsersion support

[EmanueleCoradin]

This PR results from a collaboration with @valsdav. Thanks also to @Electricks94, who helped implement batched inference support and gave me great suggestions to improve the code.

PR description:

This PR extends the PyTorchAlpaka interface with support for batched inference, half-precision execution, and improved model lifecycle handling.

Main features

• Batched inference via TensorCollection: Introduces native support for splitting a single SoA into multiple logical batches. TensorCollection now supports (batch_size, total_size) construction and computes per-batch offsets internally. Constraint: total_size must be divisible by batch_size (no support for partial batches)

• FPX abstraction (FP32 / FP16 portability): Inspired by the definition used by LSTCore, it adds a portable floating-point type (FPX) which

  • maps to float on CPU
  • maps to __half on CUDA backends
  • falls back to float on ROCm

A PyTorch bridge is included to map __half → torch::kHalf.

• Optional half-precision inference: Introduces a to_half flag in the inference path, such that Inputs can be converted to FP16. Similarly, the model to() function is overloaded, supporting dtype-aware device transfers.
• Model lifecycle improvements:
  • Models are automatically set to eval() mode at construction.
  • Optional automatic freezing (torch::jit::freeze) after loading or device transfer.

Validation

Validation is performed via dedicated mini-batch inference producers:

• SimpleNetMiniBatch
• TinyResNetMiniBatch

The latter reproduces the same inference as the non-batched implementations using the new batching API and the FPX data type.

InspectionSink has been extended to compare outputs between the standard and the batched inference.
A relative difference check is enforced: |ref - batched| / ref < 1e-5

It can be run by:

cd src/PhysicsTools/PyTorchAlpakaTest/test/
cmsRun runPyTorchAlpakaTest.py --compareBatch -e 3

Notes / limitations

• Batch size must be chosen to satisfy alignment and divisibility constraints.
• No support for partial batches.
• Using FPX, FP16 path is currently meaningful only on CUDA backends. ROCm currently falls back to FP32.

[cmsbuild]

cms-bot internal usage

[cmsbuild]

+code-checks

Logs: https://cmssdt.cern.ch/SDT/code-checks/cms-sw-PR-50498/48664

[mmusich]

type ngt

[cmsbuild]

-code-checks

Logs: https://cmssdt.cern.ch/SDT/code-checks/cms-sw-PR-50498/48680

Code check has found code style and quality issues which could be resolved by applying following patch(s)

• code-format:
  https://cmssdt.cern.ch/SDT/code-checks/cms-sw-PR-50498/48680/code-format.patch
  e.g. curl -k https://cmssdt.cern.ch/SDT/code-checks/cms-sw-PR-50498/48680/code-format.patch | patch -p1
  You can also run scram build code-format to apply code format directly

[cmsbuild]

-code-checks

Logs: https://cmssdt.cern.ch/SDT/code-checks/cms-sw-PR-50498/48687

Code check has found code style and quality issues which could be resolved by applying following patch(s)

• code-format:
  https://cmssdt.cern.ch/SDT/code-checks/cms-sw-PR-50498/48687/code-format.patch
  e.g. curl -k https://cmssdt.cern.ch/SDT/code-checks/cms-sw-PR-50498/48687/code-format.patch | patch -p1
  You can also run scram build code-format to apply code format directly

[cmsbuild]

+code-checks

Logs: https://cmssdt.cern.ch/SDT/code-checks/cms-sw-PR-50498/48691

[cmsbuild]

Pull request #50498 was updated.

[cmsbuild]

+code-checks

Logs: https://cmssdt.cern.ch/SDT/code-checks/cms-sw-PR-50498/49074

• There are other open Pull requests which might conflict with changes you have proposed:
  • File DataFormats/PortableTestObjects/src/classes_def.xml modified in PR(s): SoA Schema evolution #50691

[cmsbuild]

Pull request #50498 was updated.

[cmsbuild]

+code-checks

Logs: https://cmssdt.cern.ch/SDT/code-checks/cms-sw-PR-50498/49075

• There are other open Pull requests which might conflict with changes you have proposed:
  • File DataFormats/PortableTestObjects/src/classes_def.xml modified in PR(s): SoA Schema evolution #50691

[cmsbuild]

Pull request #50498 was updated.

[mmusich]

@cmsbuild, please test

[cmsbuild]

A new Pull Request was created by @EmanueleCoradin for master.

It involves the following packages:

• DataFormats/Common (core)
• DataFormats/PortableTestObjects (heterogeneous)
• PhysicsTools/PyTorch (ml)
• PhysicsTools/PyTorchAlpaka (heterogeneous, ml)
• PhysicsTools/PyTorchAlpakaTest (heterogeneous, ml)

@Dr15Jones, @fwyzard, @hjkwon260, @makortel, @smuzaffar, @valsdav, @y19y19 can you please review it and eventually sign? Thanks.
@makortel, @missirol, @mmusich, @rovere, @wddgit this is something you requested to watch as well.
@ftenchini, @mandrenguyen, @sextonkennedy you are the release manager for this.

cms-bot commands are listed here

[makortel]

I think DataFormats/Math (or HeterogeneousCore/AlpakaMath) would be a better place for this code.

[makortel]

The use of FPX here introduces an ODR violation, no?

The code compiled with ALPAKA_ACC_GPU_CUDA_ENABLED sees the version with __half, while other code sees float.

[cmsbuild]

-1

Failed Tests: UnitTests
Size: This PR adds an extra 16KB to repository
Summary: https://cmssdt.cern.ch/SDT/jenkins-artifacts/pull-request-integration/PR-f804d2/52802/summary.html
COMMIT: 9f505be
CMSSW: CMSSW_17_0_X_2026-04-21-1100/el8_amd64_gcc13
User test area: For local testing, you can use /cvmfs/cms-ci.cern.ch/week0/cms-sw/cmssw/50498/52802/install.sh to create a dev area with all the needed externals and cmssw changes.

Failed Unit Tests

I found 1 errors in the following unit tests:

---> test testSoADataTypesSerialSync had ERRORS

Comparison Summary

Summary:

• No significant changes to the logs found
• Reco comparison results: 6 differences found in the comparisons
• DQMHistoTests: Total files compared: 53
• DQMHistoTests: Total histograms compared: 4186853
• DQMHistoTests: Total failures: 3
• DQMHistoTests: Total nulls: 0
• DQMHistoTests: Total successes: 4186830
• DQMHistoTests: Total skipped: 20
• DQMHistoTests: Total Missing objects: 0
• DQMHistoSizes: Histogram memory added: 0.0 KiB( 52 files compared)
• Checked 227 log files, 197 edm output root files, 53 DQM output files
• TriggerResults: no differences found

[fwyzard]

Let me expand on @makortel's comment.

The approach to using half precision explored in this PR is not really workable: much of the work done implementing the SoA interface and the PortableCollection approach is based on the fact that we want the exact same memory layout in memory for all device types: CPUs, NVIDIA GPUs, AMD GPUs.

Using a type with a different size on different back-ends breaks this assumption. This makes it impossible to transparently copy the data between host and device, and to automatically convert between a host and device collection.

More in general, why do you propose to limit the use of fp16 to the CUDA backend ?

[fwyzard]

-heterogeneous

[fwyzard]

I would propose to split this PR in two parts, one that adds support for batches, one that adds support for fp16.

For the batching part, I would suggest to investigate if it's possible to support a tensor size that is not a multiple of the batch size.

For the fp16 part, if usage of fp16 should only ever be internal to Torch, it should not be exposed as a data type to CMSSW.
One possibility would be to define only the "hidden" buffer used for the extra copy to use fp16 and convert the data back and forth during the copy, although this would not be very efficient if one model's output is used as the input for another model.

If usage of fp16 should be exposed to CMSSW, please define or use an fp16 type that is available and behaves in the same way on all devices.
C++23 provides std::float16_t, but unfortunately we are still stuck with c++ 20 (in part because CUDA does not support c++ 23 yet...)
In the meantime you could do something like

namespace cms {

#if defined(ALPAKA_ACC_GPU_CUDA_ENABLED) or defined(ALPAKA_ACC_GPU_HIP_ENABLED)
  // on NVIDIA or AMD GPUs
  using float16_t = ::__half;
#else
  // on CPU
  using float16_t = ::_Float16;
#endif

}  // namespace cms

[EmanueleCoradin]

Thanks @makortel and @fwyzard for the feedback.

My intention with FPX was to provide a way for users to store data in SoA format using FP16 and pass it to PyTorch.
Your suggestion of defining a float16_t abstraction fits well with the original goal; I will test it.
The backend-dependent implementation I used was mainly due to my lack of awareness of _Float16 on CPU, so thanks for pointing that out and for explaining the issues with the previous design.

Following @fwyzard, I will close this PR and split it into two updated parts.