feat: save forward autocast context in FunctionCtx

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;

infini_train/include/autograd/function.h

@@ -2,10 +2,13 @@
 
 #include <functional>
 #include <memory>
+#include <optional>
 #include <string>
 #include <utility>
 #include <vector>
 
+#include "infini_train/include/autocast.h"
+
 namespace infini_train {
 class Tensor;
 class HookHandle;
@@ -46,6 +49,8 @@ class FunctionCtx {
 
     const std::vector<bool> &needs_input_grad() const;
 
+    const AutocastContext &forward_autocast_context() const;
+
 private:
     struct SavedTensorEntry {
         std::shared_ptr<Tensor> tensor;
@@ -56,6 +61,7 @@ class FunctionCtx {
     friend class Function;
 
     void set_needs_input_grad(std::vector<bool> needs_input_grad);
+    void set_forward_autocast_context(const AutocastContext &context);
 
     void SaveVariables(const std::vector<std::shared_ptr<Tensor>> &outputs);
     void ReleaseVariables();
@@ -66,6 +72,7 @@ class FunctionCtx {
     std::vector<SavedTensorEntry> saved_tensor_entries_;
     std::vector<bool> needs_input_grad_;
     std::vector<Tensor *> non_differentiable_;
+    std::optional<AutocastContext> forward_autocast_context_;
 };
 
 class Function : public std::enable_shared_from_this<Function> {

infini_train/src/autograd/function.cc

@@ -57,6 +57,11 @@ const std::vector<bool> &FunctionCtx::needs_input_grad() const { return needs_in
 
 void FunctionCtx::SaveForBackward(const std::vector<std::shared_ptr<Tensor>> &tensors) { to_save_ = tensors; }
 
+const AutocastContext &FunctionCtx::forward_autocast_context() const {
+    CHECK(forward_autocast_context_.has_value()) << "Forward autocast context has not been saved";
+    return *forward_autocast_context_;
+}
+
 void FunctionCtx::MarkNonDifferentiable(const std::vector<std::shared_ptr<Tensor>> &outputs) {
     non_differentiable_.clear();
     non_differentiable_.reserve(outputs.size());
@@ -71,6 +76,8 @@ void FunctionCtx::set_needs_input_grad(std::vector<bool> needs_input_grad) {
     needs_input_grad_ = std::move(needs_input_grad);
 }
 
+void FunctionCtx::set_forward_autocast_context(const AutocastContext &context) { forward_autocast_context_ = context; }
+
 void FunctionCtx::SaveVariables(const std::vector<std::shared_ptr<Tensor>> &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<Tensor> &output) const {
@@ -168,6 +176,9 @@ std::vector<std::shared_ptr<Tensor>> Function::Apply(const std::vector<std::shar
     // tensors already in the compute dtype. The shared_ptr copies are local; we keep
     // the caller's `input_tensors` untouched so next_functions_ wires up to the
     // original autograd graph (leaf -> 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_.set_forward_autocast_context(GetCurrentAutocastContext());
     auto compute_inputs = input_tensors;
     for (auto &t : compute_inputs) { tls_autocast_context.Autocast(type_, t); }
 

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<std::shared_ptr<Tensor>> Forward(const std::vector<std::shared_ptr<Tensor>> &input_tensors) override {
+        const auto &input = input_tensors[0];
+        auto output = std::make_shared<Tensor>(input->Dims(), input->Dtype(), input->GetDevice());
+        output->CopyFrom(input);
+        return {output};
+    }
+
+    void SetupContext(const std::vector<std::shared_ptr<Tensor>> &input_tensors,
+                      const std::vector<std::shared_ptr<Tensor>> &) override {
+        ctx_.SaveForBackward({input_tensors[0]});
+        observed_forward_context_ = ctx_.forward_autocast_context();
+    }
+
+    std::vector<std::shared_ptr<Tensor>> Backward(const std::vector<std::shared_ptr<Tensor>> &grad_outputs) override {
+        return {grad_outputs[0]};
+    }
+
+    const AutocastContext &ObservedForwardContext() const { return observed_forward_context_; }
+    const AutocastContext &ForwardAutocastContext() const { return ctx_.forward_autocast_context(); }
+
+    AutocastContext SnapshotWithForwardAutocastContextRestored() const {
+        AutocastGuard guard(ctx_.forward_autocast_context());
+        return GetCurrentAutocastContext();
+    }
+
+private:
+    AutocastContext observed_forward_context_;
+};
+
 TEST_P(AutogradForwardTest, SavedOutputIsPackedWithoutAutogradMeta) {
     auto input = std::make_shared<Tensor>(std::vector<int64_t>{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<Tensor>(std::vector<int64_t>{2, 2}, DataType::kFLOAT32, GetDevice(), true);
+    auto requires_grad_input
+        = std::make_shared<Tensor>(std::vector<int64_t>{2, 2}, DataType::kFLOAT32, GetDevice(), true);
     auto no_grad_input = std::make_shared<Tensor>(std::vector<int64_t>{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<Tensor>(std::vector<int64_t>{2, 2}, DataType::kFLOAT32, GetDevice(), true);
+    input->Fill(1.0f);
+
+    auto fn = std::make_shared<ObserveAutocastContextFunction>();
+    const auto forward_dtype = kDeviceDefaultDtype[static_cast<size_t>(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<Tensor>(std::vector<int64_t>{2, 2}, DataType::kFLOAT32, GetDevice(), true);
+    input->Fill(1.0f);
+
+    auto fn = std::make_shared<ObserveAutocastContextFunction>();
+    const auto forward_dtype = kDeviceDefaultDtype[static_cast<size_t>(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<Tensor>(std::vector<int64_t>{2, 2}, DataType::kFLOAT32, GetDevice(), true);
+    input->Fill(1.0f);
+
+    auto fn = std::make_shared<ObserveAutocastContextFunction>();
+    const auto forward_dtype = kDeviceDefaultDtype[static_cast<size_t>(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<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
     a->Fill(1.0f);