diff --git a/orttraining/orttraining/core/graph/gradient_builder.cc b/orttraining/orttraining/core/graph/gradient_builder.cc index 5576e8ccb1..ee9707f77c 100755 --- a/orttraining/orttraining/core/graph/gradient_builder.cc +++ b/orttraining/orttraining/core/graph/gradient_builder.cc @@ -1994,5 +1994,14 @@ IMPLEMENT_GRADIENT_BUILDER(GetLSTMGradient) { return {NodeDef(OpDef{"LSTMGrad", kMSDomain, 1}, input_args, output_args, SrcNodeAttributes())}; } +IMPLEMENT_GRADIENT_BUILDER(GetReciprocalGradient) { + // y = 1 / x + // dy/dx = -1 / x^2 + // dL/dx = dL/dy * dy/dx = dL/dy * (-1 / x^2) + return {NodeDef("Mul", {O(0), O(0)}, {IA("Square_O0")}), + NodeDef("Neg", {IA("Square_O0")}, {IA("Neg_Square_O0")}), + NodeDef("Mul", {GO(0), IA("Neg_Square_O0")}, {GI(0)})}; +} + } // namespace training } // namespace onnxruntime diff --git a/orttraining/orttraining/core/graph/gradient_builder.h b/orttraining/orttraining/core/graph/gradient_builder.h index 30da6e53ba..f3b5be06f5 100755 --- a/orttraining/orttraining/core/graph/gradient_builder.h +++ b/orttraining/orttraining/core/graph/gradient_builder.h @@ -85,6 +85,7 @@ DECLARE_GRADIENT_BUILDER(GetScatterElementsGradient) DECLARE_GRADIENT_BUILDER(GetTriluGradient) DECLARE_GRADIENT_BUILDER(GetFakeQuantGradient) DECLARE_GRADIENT_BUILDER(GetLSTMGradient) +DECLARE_GRADIENT_BUILDER(GetReciprocalGradient) DECLARE_GRADIENT_BUILDER(GetExternalGradient) diff --git a/orttraining/orttraining/core/graph/gradient_builder_registry.cc b/orttraining/orttraining/core/graph/gradient_builder_registry.cc index ca0a397ced..bbb7f174e5 100755 --- a/orttraining/orttraining/core/graph/gradient_builder_registry.cc +++ b/orttraining/orttraining/core/graph/gradient_builder_registry.cc @@ -117,6 +117,7 @@ void GradientBuilderRegistry::RegisterGradientBuilders() { REGISTER_GRADIENT_BUILDER("Trilu", GetTriluGradient); REGISTER_GRADIENT_BUILDER("FakeQuant", GetFakeQuantGradient); REGISTER_GRADIENT_BUILDER("LSTMTraining", GetLSTMGradient); + REGISTER_GRADIENT_BUILDER("Reciprocal", GetReciprocalGradient); REGISTER_GRADIENT_BUILDER("ExternalGradient", GetExternalGradient); }; diff --git a/orttraining/orttraining/test/gradient/gradient_ops_test.cc b/orttraining/orttraining/test/gradient/gradient_ops_test.cc index b6f2639512..4834451616 100644 --- a/orttraining/orttraining/test/gradient/gradient_ops_test.cc +++ b/orttraining/orttraining/test/gradient/gradient_ops_test.cc @@ -3027,6 +3027,12 @@ TEST(GradientCheckerTest, PadAndUnflattenGrad) { } #endif +TEST(GradientCheckerTest, ReciprocalGrad) { + // Avoid division by 0 by using the transformer. + std::function transformer = [](float x) { return x > 0 ? x + 0.2f : x - 0.2f; }; + UnaryOpGradientTest("Reciprocal", kOnnxDomain, 12, nullptr, &transformer); +} + } // namespace test } // namespace onnxruntime diff --git a/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py b/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py index 335b3f051a..e0dbc44741 100644 --- a/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py +++ b/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py @@ -6140,3 +6140,34 @@ def test_cache_exported_model(): torch.onnx.export.reset_mock() del os.environ["ORTMODULE_CACHE_DIR"] + + +def test_reciprocal_gradient(): + class ReciprocalModel(torch.nn.Module): + def __init__(self): + super().__init__() + + def forward(self, x): + return 1 / x + + def run_step(model, x): + prediction = model(x) + loss = prediction.sum() + loss.backward() + return prediction, loss + + device = "cuda" + pt_model = ReciprocalModel().to(device) + ort_model = ORTModule(copy.deepcopy(pt_model)) + + pt_x = torch.zeros(3, 224, 224, requires_grad=True, device=device) + with torch.no_grad(): + pt_x[pt_x <= 0] -= 0.2 + pt_x[pt_x > 0] += 0.2 + ort_x = copy.deepcopy(pt_x) + + pt_prediction, pt_loss = run_step(pt_model, pt_x) + ort_prediction, ort_loss = run_step(ort_model, ort_x) + _test_helpers.assert_values_are_close(pt_prediction, ort_prediction) + _test_helpers.assert_values_are_close(pt_loss, ort_loss) + _test_helpers.assert_values_are_close(pt_x.grad, ort_x.grad)