diff --git a/PhysicsTools/PyTorch/interface/Model.h b/PhysicsTools/PyTorch/interface/Model.h index c0efd2df4fab0..15156a7706190 100644 --- 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 index 91ecd371d27e5..89c56f88189e8 100644 --- 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 inputs(batch_size); +inputs.add( + "particles", + input_records.pt(), + input_records.eta(), + input_records.phi() +); +TensorCollection outputs(batch_size); +outputs.add("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 index dd46a9a67dac4..c1230a1fa05c9 100644 --- a/PhysicsTools/PyTorchAlpaka/interface/SoAConversion.h +++ b/PhysicsTools/PyTorchAlpaka/interface/SoAConversion.h @@ -24,10 +24,15 @@ namespace cms::torch::alpakatools::detail { } template - inline std::vector<::torch::IValue> convertInput(TensorCollection& inputs, ::torch::Device device) { + inline std::vector<::torch::IValue> convertInput(TensorCollection& 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 index f69fedd690c80..3c643cec49a36 100644 --- 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 &inputs, - cms::torch::alpakatools::TensorCollection &outputs) { + cms::torch::alpakatools::TensorCollection &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 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, ::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 index 783f9768798a4..5d163c2c10674 100644 --- a/PhysicsTools/PyTorchAlpakaTest/plugins/InspectionSink.cc +++ b/PhysicsTools/PyTorchAlpakaTest/plugins/InspectionSink.cc @@ -33,6 +33,7 @@ namespace torchtest { particles_token_{consumes(params.getParameter("particles"))}, simple_net_token_{consumes(params.getParameter("simple_net"))}, simple_net_minibatch_token_{consumes(params.getParameter("simple_net_minibatch"))}, + simple_net_runtimeFP16_token_{consumes(params.getParameter("simple_net_runtimeFP16"))}, masked_net_token_{consumes(params.getParameter("masked_net"))}, multi_head_net_token_{consumes(params.getParameter("multi_head_net"))}, images_token_{consumes(params.getParameter("images"))}, @@ -55,6 +56,7 @@ namespace torchtest { desc.add("particles"); desc.add("simple_net"); desc.add("simple_net_minibatch"); + desc.add("simple_net_runtimeFP16"); desc.add("masked_net"); desc.add("multi_head_net"); desc.add("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(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 particles_token_; const edm::EDGetTokenT simple_net_token_; const edm::EDGetTokenT simple_net_minibatch_token_; + const edm::EDGetTokenT simple_net_runtimeFP16_token_; const edm::EDGetTokenT masked_net_token_; const edm::EDGetTokenT multi_head_net_token_; const edm::EDGetTokenT 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 + 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(kMaxView, size); + for (int32_t i = 0; i < range; ++i) { - fmt::format_to( - std::back_inserter(buffer), "| {:5d} | {:7.2f} |\n", static_cast(i), simple_net[i].reco_pt()); + fmt::format_to(std::back_inserter(buffer), + "| {:5d} | {:7.2f} |\n", + static_cast(i), + static_cast(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 + 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 index bdb23913edf69..6d8597545e20a 100644 --- 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("particles"))), simple_net_token_{produces()}, model_(params.getParameter("model").fullPath()), - environment_{static_cast<::torchtest::Environment>(params.getUntrackedParameter("environment"))} {} + convertToFP16_(params.getParameter("convertToFP16")), + environment_{static_cast<::torchtest::Environment>(params.getUntrackedParameter("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("model"); + desc.add("convertToFP16"); desc.add("particles"); desc.addUntracked("environment", static_cast(::torchtest::Environment::kProduction)); descriptions.addWithDefaultLabel(desc); @@ -49,7 +56,10 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE::torchtest { cms::torch::alpakatools::TensorCollection outputs(total_size); outputs.add("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 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 index d2d2e861548cf..f2c4f929224cc 100755 --- a/PhysicsTools/PyTorchAlpakaTest/test/runPyTorchAlpakaTest.py +++ b/PhysicsTools/PyTorchAlpakaTest/test/runPyTorchAlpakaTest.py @@ -44,6 +44,7 @@ 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 @@ -51,17 +52,30 @@ 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',