From 100613cbef491eecadafdc76f5beb3d4250f8ae4 Mon Sep 17 00:00:00 2001 From: bolunz Date: Mon, 6 Jul 2026 16:23:48 +0800 Subject: [PATCH 1/2] feat: save forward autocast context in FunctionCtx --- infini_train/include/autocast.h | 7 ++ infini_train/include/autograd/function.h | 7 ++ infini_train/src/autograd/function.cc | 11 +++ tests/autograd/test_autograd.cc | 107 ++++++++++++++++++++++- 4 files changed, 130 insertions(+), 2 deletions(-) diff --git a/infini_train/include/autocast.h b/infini_train/include/autocast.h index cbb97c3f..0ce83e98 100644 --- a/infini_train/include/autocast.h +++ b/infini_train/include/autocast.h @@ -143,9 +143,16 @@ struct AutocastContext { // Global thread-local storage for autocast context inline thread_local AutocastContext tls_autocast_context; +inline AutocastContext GetCurrentAutocastContext() { return tls_autocast_context; } + // RAII guard to enable/disable autocast in a scope class AutocastGuard { public: + explicit AutocastGuard(const AutocastContext &context) { + saved_context_ = tls_autocast_context; + tls_autocast_context = context; + } + AutocastGuard(Device::DeviceType device_type, DataType autocast_dtype) { saved_context_ = tls_autocast_context; tls_autocast_context.enabled = true; diff --git a/infini_train/include/autograd/function.h b/infini_train/include/autograd/function.h index 4f1274b8..806cbf0d 100644 --- a/infini_train/include/autograd/function.h +++ b/infini_train/include/autograd/function.h @@ -2,10 +2,13 @@ #include #include +#include #include #include #include +#include "infini_train/include/autocast.h" + namespace infini_train { class Tensor; class HookHandle; @@ -46,6 +49,8 @@ class FunctionCtx { const std::vector &needs_input_grad() const; + const AutocastContext &GetForwardAutocastContext() const; + private: struct SavedTensorEntry { std::shared_ptr tensor; @@ -56,6 +61,7 @@ class FunctionCtx { friend class Function; void set_needs_input_grad(std::vector needs_input_grad); + void SaveForwardAutocastContext(const AutocastContext &context); void SaveVariables(const std::vector> &outputs); void ReleaseVariables(); @@ -66,6 +72,7 @@ class FunctionCtx { std::vector saved_tensor_entries_; std::vector needs_input_grad_; std::vector non_differentiable_; + std::optional forward_autocast_context_; }; class Function : public std::enable_shared_from_this { diff --git a/infini_train/src/autograd/function.cc b/infini_train/src/autograd/function.cc index 731bfc0b..4c7b168a 100644 --- a/infini_train/src/autograd/function.cc +++ b/infini_train/src/autograd/function.cc @@ -57,6 +57,11 @@ const std::vector &FunctionCtx::needs_input_grad() const { return needs_in void FunctionCtx::SaveForBackward(const std::vector> &tensors) { to_save_ = tensors; } +const AutocastContext &FunctionCtx::GetForwardAutocastContext() const { + CHECK(forward_autocast_context_.has_value()) << "Forward autocast context has not been saved"; + return *forward_autocast_context_; +} + void FunctionCtx::MarkNonDifferentiable(const std::vector> &outputs) { non_differentiable_.clear(); non_differentiable_.reserve(outputs.size()); @@ -71,6 +76,8 @@ void FunctionCtx::set_needs_input_grad(std::vector needs_input_grad) { needs_input_grad_ = std::move(needs_input_grad); } +void FunctionCtx::SaveForwardAutocastContext(const AutocastContext &context) { forward_autocast_context_ = context; } + void FunctionCtx::SaveVariables(const std::vector> &outputs) { saved_tensor_entries_.clear(); saved_tensor_entries_.reserve(to_save_.size()); @@ -106,6 +113,7 @@ void FunctionCtx::ReleaseVariables() { saved_tensor_entries_.clear(); needs_input_grad_.clear(); non_differentiable_.clear(); + forward_autocast_context_.reset(); } bool FunctionCtx::IsNonDifferentiable(const std::shared_ptr &output) const { @@ -168,6 +176,9 @@ std::vector> Function::Apply(const std::vector AccumulateGrad / non-leaf -> grad_fn). + // Also, save the autocast context in FunctionCtx so that custom backward can + // explicitly restore the forward autocast context from FunctionCtx if needed. + ctx_.SaveForwardAutocastContext(GetCurrentAutocastContext()); auto compute_inputs = input_tensors; for (auto &t : compute_inputs) { tls_autocast_context.Autocast(type_, t); } diff --git a/tests/autograd/test_autograd.cc b/tests/autograd/test_autograd.cc index 8ed27acf..b7322420 100644 --- a/tests/autograd/test_autograd.cc +++ b/tests/autograd/test_autograd.cc @@ -3,13 +3,14 @@ #include "gtest/gtest.h" +#include "infini_train/include/autocast.h" #include "infini_train/include/autograd/activations.h" #include "infini_train/include/autograd/elementwise.h" #include "infini_train/include/autograd/function.h" #include "infini_train/include/autograd/linear.h" #include "infini_train/include/autograd/matmul.h" -#include "infini_train/include/autograd/normalization.h" #include "infini_train/include/autograd/no_op.h" +#include "infini_train/include/autograd/normalization.h" #include "infini_train/include/autograd/outer.h" #include "infini_train/include/autograd/reduction.h" #include "infini_train/include/autograd/softmax.h" @@ -107,6 +108,41 @@ class MarkNonDifferentiableFunction : public autograd::Function { } }; +class ObserveAutocastContextFunction : public autograd::Function { +public: + static constexpr char kType[] = "ObserveAutocastContextFunction"; + + ObserveAutocastContextFunction() : autograd::Function(kType) {} + + std::vector> Forward(const std::vector> &input_tensors) override { + const auto &input = input_tensors[0]; + auto output = std::make_shared(input->Dims(), input->Dtype(), input->GetDevice()); + output->CopyFrom(input); + return {output}; + } + + void SetupContext(const std::vector> &input_tensors, + const std::vector> &) override { + ctx_.SaveForBackward({input_tensors[0]}); + observed_forward_context_ = ctx_.GetForwardAutocastContext(); + } + + std::vector> Backward(const std::vector> &grad_outputs) override { + return {grad_outputs[0]}; + } + + const AutocastContext &ObservedForwardContext() const { return observed_forward_context_; } + const AutocastContext &ForwardAutocastContext() const { return ctx_.GetForwardAutocastContext(); } + + AutocastContext SnapshotWithForwardAutocastContextRestored() const { + AutocastGuard guard(ctx_.GetForwardAutocastContext()); + return GetCurrentAutocastContext(); + } + +private: + AutocastContext observed_forward_context_; +}; + TEST_P(AutogradForwardTest, SavedOutputIsPackedWithoutAutogradMeta) { auto input = std::make_shared(std::vector{2, 2}, DataType::kFLOAT32, GetDevice(), true); input->Fill(1.0f); @@ -123,7 +159,8 @@ TEST_P(AutogradForwardTest, SavedOutputIsPackedWithoutAutogradMeta) { } TEST_P(AutogradForwardTest, FunctionCtxNeedsInputGradAndSaveForBackward) { - auto requires_grad_input = std::make_shared(std::vector{2, 2}, DataType::kFLOAT32, GetDevice(), true); + auto requires_grad_input + = std::make_shared(std::vector{2, 2}, DataType::kFLOAT32, GetDevice(), true); auto no_grad_input = std::make_shared(std::vector{2, 2}, DataType::kFLOAT32, GetDevice(), false); requires_grad_input->Fill(1.0f); no_grad_input->Fill(2.0f); @@ -193,6 +230,72 @@ TEST_P(AutogradForwardTest, FunctionCtxSavedTensorHooksPackAndUnpack) { EXPECT_EQ(saved.get(), packed_tensor.get()); } +TEST_P(AutogradForwardTest, FunctionCtxRecordsForwardAutocastContext) { + auto input = std::make_shared(std::vector{2, 2}, DataType::kFLOAT32, GetDevice(), true); + input->Fill(1.0f); + + auto fn = std::make_shared(); + const auto forward_dtype = kDeviceDefaultDtype[static_cast(GetDevice().type())]; + { + AutocastGuard autocast_guard(GetDevice().type(), forward_dtype); + auto outputs = fn->Apply({input}); + ASSERT_EQ(outputs.size(), 1); + } + + const auto &state = fn->ObservedForwardContext(); + EXPECT_TRUE(state.enabled); + EXPECT_EQ(state.device_type, GetDevice().type()); + EXPECT_EQ(state.autocast_dtype, forward_dtype); +} + +TEST_P(AutogradForwardTest, FunctionCtxRestoresForwardAutocastContextForRecompute) { + auto input = std::make_shared(std::vector{2, 2}, DataType::kFLOAT32, GetDevice(), true); + input->Fill(1.0f); + + auto fn = std::make_shared(); + const auto forward_dtype = kDeviceDefaultDtype[static_cast(GetDevice().type())]; + { + AutocastGuard autocast_guard(GetDevice().type(), forward_dtype); + auto outputs = fn->Apply({input}); + ASSERT_EQ(outputs.size(), 1); + } + + EXPECT_FALSE(GetCurrentAutocastContext().enabled); + const auto restored = fn->SnapshotWithForwardAutocastContextRestored(); + EXPECT_TRUE(restored.enabled); + EXPECT_EQ(restored.device_type, GetDevice().type()); + EXPECT_EQ(restored.autocast_dtype, forward_dtype); + EXPECT_FALSE(GetCurrentAutocastContext().enabled); +} + +TEST_P(AutogradForwardTest, AutocastGuardRestoresCallerContextAfterForwardAutocastContext) { + auto input = std::make_shared(std::vector{2, 2}, DataType::kFLOAT32, GetDevice(), true); + input->Fill(1.0f); + + auto fn = std::make_shared(); + const auto forward_dtype = kDeviceDefaultDtype[static_cast(GetDevice().type())]; + { + AutocastGuard autocast_guard(GetDevice().type(), forward_dtype); + auto outputs = fn->Apply({input}); + ASSERT_EQ(outputs.size(), 1); + } + + const auto backward_dtype = DataType::kFLOAT32; + AutocastGuard backward_autocast_guard(GetDevice().type(), backward_dtype); + { + AutocastGuard restore_guard(fn->ForwardAutocastContext()); + const auto restored = GetCurrentAutocastContext(); + EXPECT_TRUE(restored.enabled); + EXPECT_EQ(restored.device_type, GetDevice().type()); + EXPECT_EQ(restored.autocast_dtype, forward_dtype); + } + + const auto current = GetCurrentAutocastContext(); + EXPECT_TRUE(current.enabled); + EXPECT_EQ(current.device_type, GetDevice().type()); + EXPECT_EQ(current.autocast_dtype, backward_dtype); +} + TEST_P(AutogradForwardTest, AddForward) { auto a = std::make_shared(std::vector{2, 3}, DataType::kFLOAT32, GetDevice(), true); a->Fill(1.0f); From 92d72a40ea5128b101ac313b03030438d9e1f156 Mon Sep 17 00:00:00 2001 From: bolunz Date: Wed, 8 Jul 2026 15:27:33 +0800 Subject: [PATCH 2/2] fix: rename getter --- infini_train/include/autograd/function.h | 4 ++-- infini_train/src/autograd/function.cc | 6 +++--- tests/autograd/test_autograd.cc | 6 +++--- 3 files changed, 8 insertions(+), 8 deletions(-) diff --git a/infini_train/include/autograd/function.h b/infini_train/include/autograd/function.h index 806cbf0d..9e05a7b8 100644 --- a/infini_train/include/autograd/function.h +++ b/infini_train/include/autograd/function.h @@ -49,7 +49,7 @@ class FunctionCtx { const std::vector &needs_input_grad() const; - const AutocastContext &GetForwardAutocastContext() const; + const AutocastContext &forward_autocast_context() const; private: struct SavedTensorEntry { @@ -61,7 +61,7 @@ class FunctionCtx { friend class Function; void set_needs_input_grad(std::vector needs_input_grad); - void SaveForwardAutocastContext(const AutocastContext &context); + void set_forward_autocast_context(const AutocastContext &context); void SaveVariables(const std::vector> &outputs); void ReleaseVariables(); diff --git a/infini_train/src/autograd/function.cc b/infini_train/src/autograd/function.cc index 4c7b168a..af23c163 100644 --- a/infini_train/src/autograd/function.cc +++ b/infini_train/src/autograd/function.cc @@ -57,7 +57,7 @@ const std::vector &FunctionCtx::needs_input_grad() const { return needs_in void FunctionCtx::SaveForBackward(const std::vector> &tensors) { to_save_ = tensors; } -const AutocastContext &FunctionCtx::GetForwardAutocastContext() const { +const AutocastContext &FunctionCtx::forward_autocast_context() const { CHECK(forward_autocast_context_.has_value()) << "Forward autocast context has not been saved"; return *forward_autocast_context_; } @@ -76,7 +76,7 @@ void FunctionCtx::set_needs_input_grad(std::vector needs_input_grad) { needs_input_grad_ = std::move(needs_input_grad); } -void FunctionCtx::SaveForwardAutocastContext(const AutocastContext &context) { forward_autocast_context_ = context; } +void FunctionCtx::set_forward_autocast_context(const AutocastContext &context) { forward_autocast_context_ = context; } void FunctionCtx::SaveVariables(const std::vector> &outputs) { saved_tensor_entries_.clear(); @@ -178,7 +178,7 @@ std::vector> Function::Apply(const std::vector AccumulateGrad / non-leaf -> grad_fn). // Also, save the autocast context in FunctionCtx so that custom backward can // explicitly restore the forward autocast context from FunctionCtx if needed. - ctx_.SaveForwardAutocastContext(GetCurrentAutocastContext()); + ctx_.set_forward_autocast_context(GetCurrentAutocastContext()); auto compute_inputs = input_tensors; for (auto &t : compute_inputs) { tls_autocast_context.Autocast(type_, t); } diff --git a/tests/autograd/test_autograd.cc b/tests/autograd/test_autograd.cc index b7322420..e7ca5d81 100644 --- a/tests/autograd/test_autograd.cc +++ b/tests/autograd/test_autograd.cc @@ -124,7 +124,7 @@ class ObserveAutocastContextFunction : public autograd::Function { void SetupContext(const std::vector> &input_tensors, const std::vector> &) override { ctx_.SaveForBackward({input_tensors[0]}); - observed_forward_context_ = ctx_.GetForwardAutocastContext(); + observed_forward_context_ = ctx_.forward_autocast_context(); } std::vector> Backward(const std::vector> &grad_outputs) override { @@ -132,10 +132,10 @@ class ObserveAutocastContextFunction : public autograd::Function { } const AutocastContext &ObservedForwardContext() const { return observed_forward_context_; } - const AutocastContext &ForwardAutocastContext() const { return ctx_.GetForwardAutocastContext(); } + const AutocastContext &ForwardAutocastContext() const { return ctx_.forward_autocast_context(); } AutocastContext SnapshotWithForwardAutocastContextRestored() const { - AutocastGuard guard(ctx_.GetForwardAutocastContext()); + AutocastGuard guard(ctx_.forward_autocast_context()); return GetCurrentAutocastContext(); }