Implementation:Microsoft Onnxruntime CPU SGDv2
| Knowledge Sources | |
|---|---|
| Domains | Training, CPU_Kernels |
| Last Updated | 2026-02-10 04:00 GMT |
Overview
Concrete tool for the SGDOptimizerV2 (TensorSeq-based SGD) kernel on CPU in the ONNX Runtime training framework.
Description
This file implements the SGDOptimizerV2 kernel, which performs stochastic gradient descent on a sequence of weight-gradient tensor pairs. Unlike the legacy SGDOptimizer that operates on single tensors, this version processes TensorSeq inputs containing all model weights and their corresponding gradients in a single kernel call. The base class SGDOptimizerV2Base::PrepareForCompute validates that weights and gradients match in count and shape, and prepares grouped tensor pointers in parallel using the thread pool. The compute method applies the standard SGD update: weight = weight - lr * gradient for each weight-gradient pair. An optional update signal can skip the update step. Weights are updated in-place through aliasing.
Usage
This kernel is invoked at each training step when the SGD optimizer is used with the grouped weight/gradient pattern. It updates all model weights simultaneously.
Code Reference
Source Location
- Repository: Microsoft_Onnxruntime
- File: orttraining/orttraining/training_ops/cpu/optimizer/sgd/sgd.cc
- Lines: 1-102
Signature
Status SGDOptimizerV2Base::PrepareForCompute(OpKernelContext* ctx,
SGDOptimizerV2Base::Prepare& prepare) const;
template <typename T>
Status SGDOptimizerV2<T>::Compute(OpKernelContext* ctx) const;
Import
#include "orttraining/orttraining/training_ops/cpu/optimizer/sgd/sgd.h"
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| learning_rate | Tensor(float) | Yes | Scalar learning rate |
| weights | TensorSeq(float) | Yes | Sequence of weight tensors |
| gradients | TensorSeq(float) | Yes | Sequence of gradient tensors |
| update_signal | Tensor(bool) | No | Optional signal to skip update |
Outputs
| Name | Type | Description |
|---|---|---|
| update_completed | Tensor(bool) | Whether the update was performed |
| updated_weights | TensorSeq(float) | Updated weight tensors (in-place alias) |
Usage Examples
ONNX_OPERATOR_KERNEL_EX(
SGDOptimizerV2, kMSDomain, 1, kCpuExecutionProvider,
KernelDefBuilder()
.Alias(1, 1) // Update weights in-place
.TypeConstraint("T1", DataTypeImpl::GetTensorType<float>())
.TypeConstraint("T_BOOL", DataTypeImpl::GetTensorType<bool>())
.TypeConstraint("S_WEIGHT", DataTypeImpl::AllFixedSizeSequenceTensorTypes())
.TypeConstraint("S_GRAD", DataTypeImpl::AllFixedSizeSequenceTensorTypes()),
SGDOptimizerV2<float>);