Implementation:Microsoft Onnxruntime TrainingModule

Knowledge Sources	Microsoft_Onnxruntime
Domains	Training, API, Module
Last Updated	2026-02-10 04:00 GMT

Overview

Defines the Parameter and Module structs for running training forward and backward passes within the ORT Training API.

Description

This header provides two principal types. The `Parameter` struct represents a single model parameter (trainable or non-trainable) with its name, data (`OrtValue`), gradient, and a `requires_grad` flag. It supports copying to/from different devices via the `DataTransferManager`, resetting gradients, and setting gradient values. The `ModuleCheckpointState` struct aggregates all named parameters and their data transfer manager reference.

The `Module` struct is the main training execution unit. It holds two `InferenceSession` objects -- one for the training graph (forward + backward) and one for the eval graph (forward only) -- while sharing parameters from a `CheckpointState`. Key capabilities include:

`TrainStep`: Executes forward and backward passes, accumulating gradients in `Parameter` objects.
`EvalStep`: Executes forward-only inference using the eval model.
Parameter buffer operations: `CopyParametersToBuffer` and `CopyBufferToParameters` for efficient parameter serialization.
Model export: `ExportModelForInferencing` strips training-specific nodes for deployment.
Input/output introspection: methods to query training and eval model input/output counts and names.

The Module does not own parameters; it holds a non-owning pointer to `CheckpointState`. During initialization, parameters are moved to the target device if needed.

Usage

Use this header to create a training module that executes forward/backward passes on ONNX training graphs. It is the core compute component wrapped by `TrainingSession`.

Code Reference

Source Location

Repository: Microsoft_Onnxruntime
File: orttraining/orttraining/training_api/module.h
Lines: 1-207

Signature

struct Parameter {
  Parameter(const std::string& name, const OrtValue& data, const bool requires_grad);
  OrtValue& Data();
  Status CopyTo(const DataTransferManager* data_transfer_manager, OrtValue& data) const;
  Status CopyFrom(const DataTransferManager* data_transfer_manager, const OrtValue& data);
  const std::string& Name() const;
  bool RequiresGrad() const;
  OrtValue& Gradient();
  Status ResetGrad();
  Status SetGrad(const std::string& gradient_name, const OrtValue& param_grad);
};

struct Module {
  Module(const ModelIdentifiers& model_identifiers,
         CheckpointState* state,
         const onnxruntime::SessionOptions& session_options,
         const Environment& env,
         const std::vector<std::shared_ptr<IExecutionProvider>>& providers,
         gsl::span<OrtCustomOpDomain* const> op_domains = {});
  ~Module();

  std::vector<std::shared_ptr<Parameter>> Parameters() const;
  std::unordered_map<std::string, std::shared_ptr<Parameter>> NamedParameters() const;
  Status LazyResetGrad();
  Status TrainStep(const std::vector<OrtValue>& inputs, std::vector<OrtValue>& outputs);
  Status EvalStep(const std::vector<OrtValue>& inputs, std::vector<OrtValue>& outputs);
  size_t GetTrainingModelOutputCount() const noexcept;
  size_t GetEvalModelOutputCount() const noexcept;
  std::string GetTrainingModelOutputName(size_t index) const;
  std::string GetEvalModelOutputName(size_t index) const;
  size_t GetParametersSize(const bool trainable_only = true) const;
  Status CopyParametersToBuffer(OrtValue& parameters_buffer, const bool trainable_only = true);
  Status CopyBufferToParameters(OrtValue& parameters_buffer, const bool trainable_only = true);
  Status ExportModelForInferencing(const std::string& inference_model_path,
                                   gsl::span<const std::string> graph_output_names) const;
  size_t GetTrainingModelInputCount() const noexcept;
  size_t GetEvalModelInputCount() const noexcept;
  std::string GetTrainingModelInputName(size_t index) const;
  std::string GetEvalModelInputName(size_t index) const;
};

Import

#include "orttraining/training_api/module.h"

I/O Contract

Method	Inputs	Outputs	Description
TrainStep	vector<OrtValue> inputs	vector<OrtValue> outputs, Status	Runs forward + backward, accumulates gradients
EvalStep	vector<OrtValue> inputs	vector<OrtValue> outputs, Status	Runs forward-only evaluation
CopyParametersToBuffer	OrtValue& buffer, bool trainable_only	Status	Copies all parameters into a contiguous buffer
CopyBufferToParameters	OrtValue& buffer, bool trainable_only	Status	Copies parameter values from a contiguous buffer back to parameters
ExportModelForInferencing	string path, span<string> output_names	Status	Exports eval model for inference deployment

Usage Examples

#include "orttraining/training_api/module.h"

using namespace onnxruntime::training::api;

// Create module
Module module(model_ids, &checkpoint_state, session_options, env, providers);

// Training step
std::vector<OrtValue> inputs = PrepareInputs();
std::vector<OrtValue> outputs;
ORT_THROW_IF_ERROR(module.TrainStep(inputs, outputs));

// Reset gradients
ORT_THROW_IF_ERROR(module.LazyResetGrad());

// Export for inference
module.ExportModelForInferencing("model_inference.onnx", output_names);

Related Pages

Page Connections

Double-click a node to navigate. Hold to expand connections.

Principle

Implementation

Heuristic

Environment