cms-sw · cmsbuild · May 22, 2026 · Apr 22, 2026
diff --git a/PhysicsTools/PyTorch/interface/Model.h b/PhysicsTools/PyTorch/interface/Model.h
@@ -24,12 +24,17 @@ namespace cms::torch {
     }
 
     // Move model to specified device memory space. Async load by specifying `non_blocking` (in default stream if not overridden by the caller)
-    void to(::torch::Device dev, const bool non_blocking = false) {
+    void to(::torch::Device dev, const bool non_blocking = false, std::optional<::torch::Dtype> dtype = std::nullopt) {
       if (dev == device_)
         return;
 
       TORCH_CHECK(!is_frozen_ && "Model is frozen, cannot be moved to another device!");
-      model_.to(dev, non_blocking);
+
+      if (dtype)
+        model_.to(dev, *dtype, non_blocking);
+      else
+        model_.to(dev, non_blocking);
+
       device_ = dev;
       if (auto_freeze_) {
         freeze();
@@ -43,6 +48,8 @@ namespace cms::torch {
       }
     }
 
+    void to(::torch::Dtype dtype) { model_.to(dtype); }
+
     // Forward pass (inference) of model, returns torch::IValue (multi output support). Match native torchlib interface.
     ::torch::IValue forward(std::vector<::torch::IValue> &inputs) {
       // Disabling autograd

diff --git a/PhysicsTools/PyTorchAlpaka/README.md b/PhysicsTools/PyTorchAlpaka/README.md
@@ -8,6 +8,7 @@ All Pytorch based modules should add `PyTorchService` to disable internal torchl
 
 Examples demonstrating the interoperability of PyTorch with Alpaka in the CMSSW environment can be found in the [PyTorchAlpakaTest](../PyTorchAlpakaTest) directory. The basic test pipeline includes:
 - *SimpleNet* composed with few Dense layers, that operate on SoA style portable data structures
+- *SimpleNet* composed with few Dense layers, that operate on SoA style portable data structures. It provides also an example for Runtime FP16 conversion.
 - *SimpleNetMiniBatch*, providing and example of inference perfomed in mini-batches
 - *MaskedNet* shows how to use multiple input data with `Eigen::Vector` and `SOA_SCALAR`
 - *TinyResNet* emulate more complex scenario with `Eigen::Matrix` and how one can implement image-like Tensor implementation
@@ -22,7 +23,6 @@ By default, the model is automatically frozen using the `torch::jit::freeze()` f
 You can skip this optimization step by setting `auto_freeze=false` when calling the model constructor.
 **Important:** Once a model is frozen, it cannot be moved to another device. Attempting to do so will trigger a runtime assertion.
 
-
 ## Direct Inference on SoA 
 The interface provides a converter to dynamically wrap SoA data into one or more `torch::tensors` without the need to copy data (or minimal copy overhead).
 
@@ -100,6 +100,36 @@ These checks rely on `assert`.
 
 Look at [SimpleNetMiniBatch](PhysicsTools/PyTorchAlpakaTest/plugins/alpaka/SimpleNetMiniBatch.cc) to have an example.
 
+## FP16 Inference Support
+
+FP16 (half precision) inference is supported alongside the default FP32 execution path. The goal is to enable reduced memory usage while preserving numerical compatibility with FP32 results.
+
+### 1. Runtime FP16 conversion
+
+FP32 data is stored in SoA format and explicitly converted to FP16 at inference time.
+In this case, you just need to pass `torch::kHalf` to the forward call; the model and input tensors are converted to FP16 under the hood using the PyTorch API.
+
+```cpp
+// SoA with input features and output
+GENERATE_SOA_LAYOUT(SimpleNetLayout, SOA_COLUMN(float, reco_pt))
+GENERATE_SOA_LAYOUT(ParticleLayout, SOA_COLUMN(float, pt), SOA_COLUMN(float, eta), SOA_COLUMN(float, phi))
+
+TensorCollection<Queue> inputs(batch_size);
+inputs.add<ParticleSoA>(
+    "particles",
+    input_records.pt(),
+    input_records.eta(),
+    input_records.phi()
+);
+TensorCollection<Queue> outputs(batch_size);
+outputs.add<SimpleNetSoA>("regression_head", output_records.reco_pt());
+
+// Runtime FP16 inference
+model.forward(queue, inputs, outputs, torch::kHalf);
+```
+
+FP16 and FP32 outputs may differ slightly due to reduced precision and floating-point accumulation effects. Users are encouraged to check the output compatibility.
+
 ## Limitations
 - Current implementation supports `SerialSync` and `CudaAsync` backends only. `ROCmAsync` backend is supported via SerialSync fallback mechanism due to missing `pytorch-hip` library in CMSSW (see: https://github.com/pytorch/pytorch/blob/main/aten/CMakeLists.txt#L75), with explicit `alpaka::wait()` call to copy data to host and back to device.
 - Const correctness and thread-safety relies on `torch::from_blob()` mechanism which currently does not ensure that data will not be modified internally. There is ongoing work to support COW tensors but until this support will be integrated in mainstream PyTorch the provided solution materialises (copies) the tensors if passed registry points to `const` memory. For more information please check [Const correctness and thread-safety of torch::from_blob with external memory](https://discuss.pytorch.org/t/const-correctness-and-thread-safety-of-torch-from-blob-with-external-memory/223521) and [pytorch:#97856](https://github.com/pytorch/pytorch/issues/97856)

diff --git a/PhysicsTools/PyTorchAlpaka/interface/SoAConversion.h b/PhysicsTools/PyTorchAlpaka/interface/SoAConversion.h
@@ -24,10 +24,15 @@ namespace cms::torch::alpakatools::detail {
   }
 
   template <typename TQueue>
-  inline std::vector<::torch::IValue> convertInput(TensorCollection<TQueue>& inputs, ::torch::Device device) {
+  inline std::vector<::torch::IValue> convertInput(TensorCollection<TQueue>& inputs,
+                                                   ::torch::Device device,
+                                                   std::optional<::torch::Dtype> dtype = std::nullopt) {
     std::vector<::torch::IValue> tensors(inputs.size());
     for (size_t i = 0; i < inputs.size(); i++) {
-      tensors[i] = cms::torch::alpakatools::detail::arrayToTensor(device, inputs[i]);
+      if (dtype)
+        tensors[i] = cms::torch::alpakatools::detail::arrayToTensor(device, inputs[i]).to(*dtype);
+      else
+        tensors[i] = cms::torch::alpakatools::detail::arrayToTensor(device, inputs[i]);
     }
     return tensors;
   }

diff --git a/PhysicsTools/PyTorchAlpaka/interface/alpaka/AlpakaModel.h b/PhysicsTools/PyTorchAlpaka/interface/alpaka/AlpakaModel.h
@@ -36,7 +36,8 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE::torch {
     // Refer: PhysicsTools/PyTorch/interface/SoAConversion.h for details about wrapping memory layouts.
     void forward(Queue &queue,
                  cms::torch::alpakatools::TensorCollection<Queue> &inputs,
-                 cms::torch::alpakatools::TensorCollection<Queue> &outputs) {
+                 cms::torch::alpakatools::TensorCollection<Queue> &outputs,
+                 std::optional<::torch::Dtype> dtype = std::nullopt) {
 #ifdef ALPAKA_ACC_GPU_HIP_ENABLED
       inputs.copy(queue, cms::torch::alpakatools::detail::MemcpyKind::DeviceToHost);
       outputs.copy(queue, cms::torch::alpakatools::detail::MemcpyKind::DeviceToHost);
@@ -45,10 +46,10 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE::torch {
 #endif  // ALPAKA_ACC_GPU_HIP_ENABLED
       cms::torch::alpakatools::QueueGuard<Queue> guard(queue);
       if (cms::torch::alpakatools::getDevice(queue) != this->Model::device()) {
-        to(queue);
+        to(queue, dtype);
       }
 
-      auto input_tensor = cms::torch::alpakatools::detail::convertInput(inputs, device_);
+      auto input_tensor = cms::torch::alpakatools::detail::convertInput(inputs, device_, dtype);
       if (outputs.size() > 1) {
         auto output_tensors = model_.forward(input_tensor);
         cms::torch::alpakatools::detail::convertOutput(output_tensors, outputs, device_);
@@ -61,24 +62,26 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE::torch {
 #endif  // ALPAKA_ACC_GPU_HIP_ENABLED
     }
 
-    // Move model to specified device memory space. Async load (in default stream if not overridden by the caller)
+    // Move model to specified device memory space and with the specified dtype. Async load (in default stream if not overridden by the caller)
     // The caller should ensure the QueueGuard is instantiated and PyTorch stream context is properly set.
-    void to(const Device &dev) {
+    void to(const Device &dev, std::optional<::torch::Dtype> dtype = std::nullopt) {
       if constexpr (std::is_same_v<::alpaka::Dev<Device>, ::alpaka::DevCpu>) {
-        this->Model::to(cms::torch::alpakatools::getDevice(dev));
+        this->Model::to(cms::torch::alpakatools::getDevice(dev), false, dtype);
         return;
       }
 #ifdef ALPAKA_ACC_GPU_HIP_ENABLED
       // ROCm/HIP not yet directly supported → fallback to CPU inference
-      this->Model::to(cms::torch::alpakatools::getDevice(dev));
+      this->Model::to(cms::torch::alpakatools::getDevice(dev), false, dtype);
       return;
 #endif  // ALPAKA_ACC_GPU_HIP_ENABLED
       // CUDA → keep async execution
-      this->Model::to(cms::torch::alpakatools::getDevice(dev), true);
+      this->Model::to(cms::torch::alpakatools::getDevice(dev), true, dtype);
     }
 
     // Overload for Queue to simplify the interface for the common case of async execution.
-    void to(const Queue &queue) { this->AlpakaModel::to(::alpaka::getDev(queue)); }
+    void to(const Queue &queue, std::optional<::torch::Dtype> dtype = std::nullopt) {
+      this->AlpakaModel::to(::alpaka::getDev(queue), dtype);
+    }
   };
 
 }  // namespace ALPAKA_ACCELERATOR_NAMESPACE::torch

diff --git a/PhysicsTools/PyTorchAlpakaTest/plugins/InspectionSink.cc b/PhysicsTools/PyTorchAlpakaTest/plugins/InspectionSink.cc
@@ -33,6 +33,7 @@ namespace torchtest {
           particles_token_{consumes(params.getParameter<edm::InputTag>("particles"))},
           simple_net_token_{consumes(params.getParameter<edm::InputTag>("simple_net"))},
           simple_net_minibatch_token_{consumes(params.getParameter<edm::InputTag>("simple_net_minibatch"))},
+          simple_net_runtimeFP16_token_{consumes(params.getParameter<edm::InputTag>("simple_net_runtimeFP16"))},
           masked_net_token_{consumes(params.getParameter<edm::InputTag>("masked_net"))},
           multi_head_net_token_{consumes(params.getParameter<edm::InputTag>("multi_head_net"))},
           images_token_{consumes(params.getParameter<edm::InputTag>("images"))},
@@ -55,6 +56,7 @@ namespace torchtest {
       desc.add<edm::InputTag>("particles");
       desc.add<edm::InputTag>("simple_net");
       desc.add<edm::InputTag>("simple_net_minibatch");
+      desc.add<edm::InputTag>("simple_net_runtimeFP16");
       desc.add<edm::InputTag>("masked_net");
       desc.add<edm::InputTag>("multi_head_net");
       desc.add<edm::InputTag>("images");
@@ -77,6 +79,7 @@ namespace torchtest {
       auto particles_handle = event.getHandle(particles_token_);
       auto simple_net_handle = event.getHandle(simple_net_token_);
       auto simple_net_minibatch_handle = event.getHandle(simple_net_minibatch_token_);
+      auto simple_net_runtimeFP16_handle = event.getHandle(simple_net_runtimeFP16_token_);
       auto masked_net_handle = event.getHandle(masked_net_token_);
       auto multi_head_net_handle = event.getHandle(multi_head_net_token_);
       auto images_handle = event.getHandle(images_token_);
@@ -110,6 +113,22 @@ namespace torchtest {
                 static_cast<cms::alpakatools::Backend>(event.get(simple_net_minibatch_backend_));
             print(simple_net_minibatch.const_view(), cms::alpakatools::toString(simple_net_minibatch_backend));
           }
+          if (simple_net_runtimeFP16_handle.isValid()) {
+            auto const& simple_net_runtimeFP16 = *simple_net_runtimeFP16_handle;
+            print(simple_net_runtimeFP16.const_view(), "runtimeFP16", "SimpleNetCollection");
+          }
+          // assert the FP16 precision si producing compatible results
+          if (simple_net_handle.isValid() && simple_net_runtimeFP16_handle.isValid()) {
+            auto const& ref = *simple_net_handle;
+            auto const& FP16 = *simple_net_runtimeFP16_handle;
+
+            assert(ref.const_view().metadata().size() == FP16.const_view().metadata().size());
+            for (auto i = 0; i < ref.const_view().metadata().size(); i++) {
+              auto diff = std::abs(ref.const_view()[i].reco_pt() - FP16.const_view()[i].reco_pt()) /
+                          ref.const_view()[i].reco_pt();
+              assert(diff < 1e-2 && "Results from simple_net and simple_net_runtimeFP16 do not match!");
+            }
+          }
           // masked_net
           if (masked_net_handle.isValid()) {
             auto const& masked_net = *masked_net_handle;
@@ -192,6 +211,7 @@ namespace torchtest {
     const edm::EDGetTokenT<portabletest::ParticleHostCollection> particles_token_;
     const edm::EDGetTokenT<portabletest::SimpleNetHostCollection> simple_net_token_;
     const edm::EDGetTokenT<portabletest::SimpleNetHostCollection> simple_net_minibatch_token_;
+    const edm::EDGetTokenT<portabletest::SimpleNetHostCollection> simple_net_runtimeFP16_token_;
     const edm::EDGetTokenT<portabletest::SimpleNetHostCollection> masked_net_token_;
     const edm::EDGetTokenT<portabletest::MultiHeadNetHostCollection> multi_head_net_token_;
     const edm::EDGetTokenT<portabletest::ImageHostCollection> images_token_;
@@ -318,9 +338,8 @@ namespace torchtest {
       fmt::print("{}\n", fmt::to_string(buffer));
     }
 
-    void print(const portabletest::SimpleNetHostCollection::ConstView& simple_net,
-               const std::string_view simple_net_backend,
-               const std::string& label = "SimpleNetCollection") {
+    template <typename ViewT>
+    void print_view(const ViewT& simple_net) {
       constexpr auto line = "+-------+---------+\n";
       const auto size = simple_net.metadata().size();
       if (size == 0) {
@@ -329,20 +348,15 @@ namespace torchtest {
       }
       fmt::memory_buffer buffer;
 
-      // Header message
-      fmt::format_to(std::back_inserter(buffer), "[DEBUG] {}[{}] ({}):\n", label, size, simple_net_backend);
-      fmt::format_to(std::back_inserter(buffer), "{}", line);
-      fmt::format_to(std::back_inserter(buffer), "| {:>5} | {:>7} |\n", "index", "reco_pt");
-      fmt::format_to(std::back_inserter(buffer), "{}", line);
-
-      // Table rows (preview)
       int32_t range = (environment_ >= Environment::kTest) ? size : std::min<int32_t>(kMaxView, size);
+
       for (int32_t i = 0; i < range; ++i) {
-        fmt::format_to(
-            std::back_inserter(buffer), "| {:5d} | {:7.2f} |\n", static_cast<int>(i), simple_net[i].reco_pt());
+        fmt::format_to(std::back_inserter(buffer),
+                       "| {:5d} | {:7.2f} |\n",
+                       static_cast<int>(i),
+                       static_cast<float>(simple_net[i].reco_pt()));
       }
 
-      // Ellipsis row if truncated
       if (range < kMaxView) {
         fmt::format_to(std::back_inserter(buffer), "| {:>5} | {:>7} |\n", "...", "...");
       }
@@ -351,6 +365,23 @@ namespace torchtest {
       fmt::print("{}\n", fmt::to_string(buffer));
     }
 
+    template <typename ViewT>
+    void print(const ViewT& simple_net, std::string_view backend, const std::string& label = "SimpleNetCollection") {
+      constexpr auto line = "+-------+---------+\n";
+      const auto size = simple_net.metadata().size();
+      fmt::memory_buffer buffer;
+
+      fmt::format_to(std::back_inserter(buffer), "[DEBUG] {}[{}] ({}):\n", label, size, backend);
+
+      fmt::format_to(std::back_inserter(buffer), "{}", line);
+      fmt::format_to(std::back_inserter(buffer), "| {:>5} | {:>7} |\n", "index", "reco_pt");
+      fmt::format_to(std::back_inserter(buffer), "{}", line);
+
+      fmt::print("{}\n", fmt::to_string(buffer));
+
+      print_view(simple_net);
+    }
+
     void print(const portabletest::ParticleHostCollection::ConstView& particles,
                const std::string_view particles_backend) {
       constexpr auto line = "+-------+---------+---------+---------+\n";

diff --git a/PhysicsTools/PyTorchAlpakaTest/plugins/alpaka/SimpleNet.cc b/PhysicsTools/PyTorchAlpakaTest/plugins/alpaka/SimpleNet.cc
@@ -23,11 +23,18 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE::torchtest {
           particles_token_(consumes(params.getParameter<edm::InputTag>("particles"))),
           simple_net_token_{produces()},
           model_(params.getParameter<edm::FileInPath>("model").fullPath()),
-          environment_{static_cast<::torchtest::Environment>(params.getUntrackedParameter<int>("environment"))} {}
+          convertToFP16_(params.getParameter<bool>("convertToFP16")),
+          environment_{static_cast<::torchtest::Environment>(params.getUntrackedParameter<int>("environment"))} {
+      // Cast the model in half precision if required.
+      // Note: this passage can be skipped if you exported the model in FP16 precision in the .pt file
+      if (convertToFP16_)
+        model_.to(::torch::kHalf);
+    }
 
     static void fillDescriptions(edm::ConfigurationDescriptions &descriptions) {
       edm::ParameterSetDescription desc;
       desc.add<edm::FileInPath>("model");
+      desc.add<bool>("convertToFP16");
       desc.add<edm::InputTag>("particles");
       desc.addUntracked<int>("environment", static_cast<int>(::torchtest::Environment::kProduction));
       descriptions.addWithDefaultLabel(desc);
@@ -49,7 +56,10 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE::torchtest {
       cms::torch::alpakatools::TensorCollection<Queue> outputs(total_size);
       outputs.add<portabletest::SimpleNetSoA>("regression_head", output_records.reco_pt());
 
-      model_.forward(event.queue(), inputs, outputs);
+      if (convertToFP16_)
+        model_.forward(event.queue(), inputs, outputs, ::torch::kHalf);
+      else
+        model_.forward(event.queue(), inputs, outputs);
       // put device-side product into event
       event.emplace(simple_net_token_, std::move(regression_collection));
     }
@@ -60,6 +70,7 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE::torchtest {
     const device::EDPutToken<portabletest::SimpleNetDeviceCollection> simple_net_token_;
     // model
     torch::AlpakaModel model_;
+    const bool convertToFP16_;
     // debug mode flag
     const ::torchtest::Environment environment_;
   };

diff --git a/PhysicsTools/PyTorchAlpakaTest/test/runPyTorchAlpakaTest.py b/PhysicsTools/PyTorchAlpakaTest/test/runPyTorchAlpakaTest.py
@@ -44,24 +44,38 @@
     from PhysicsTools.PyTorchAlpakaTest.modules import torchtest_SimpleNet_alpaka, torchtest_SimpleNetMiniBatch_alpaka
     process.SimpleNet = torchtest_SimpleNet_alpaka(
         model = cms.FileInPath(args.simpleNet),
+        convertToFP16 = cms.bool(False),
         particles = 'DataSource',
         alpaka = cms.untracked.PSet(
             backend = cms.untracked.string("serial_sync")  # force serial backend to emulate heterogeneous pipeline
         ),
         environment = cms.untracked.int32(args.environment)
     )
     process.path += process.SimpleNet
-if "SimpleNetMiniBatch" in args.only:
-    process.SimpleNetMiniBatch = torchtest_SimpleNetMiniBatch_alpaka(
+
+    process.SimpleNetRuntineFP16 = torchtest_SimpleNet_alpaka(
         model = cms.FileInPath(args.simpleNet),
-        batchSize = cms.int32(args.batchSize),
+        convertToFP16 = cms.bool(True),
         particles = 'DataSource',
         alpaka = cms.untracked.PSet(
-            backend = cms.untracked.string("serial_sync")
+            backend = cms.untracked.string("serial_sync")  # force serial backend to emulate heterogeneous pipeline
         ),
         environment = cms.untracked.int32(args.environment)
     )
-    process.path += process.SimpleNetMiniBatch
+    process.path += process.SimpleNetRuntineFP16
+
+    if "SimpleNetMiniBatch" in args.only:
+        process.SimpleNetMiniBatch = torchtest_SimpleNetMiniBatch_alpaka(
+            model = cms.FileInPath(args.simpleNet),
+            batchSize = cms.int32(args.batchSize),
+            particles = 'DataSource',
+            alpaka = cms.untracked.PSet(
+                backend = cms.untracked.string("serial_sync")
+            ),
+            environment = cms.untracked.int32(args.environment)
+        )
+        process.path += process.SimpleNetMiniBatch
+
 # --only MultiHeadNet
 if "MultiHeadNet" in args.only:
     from PhysicsTools.PyTorchAlpakaTest.modules import torchtest_MultiHeadNet_alpaka
@@ -115,6 +129,7 @@
     particles = 'DataSource',
     simple_net = 'SimpleNet',
     simple_net_minibatch = 'SimpleNetMiniBatch',
+    simple_net_runtimeFP16 = 'SimpleNetRuntineFP16',
     masked_net = 'MaskedNet',
     multi_head_net = 'MultiHeadNet',
     images = 'DataSource',