diff --git a/example/gpt2/checkpoint_loader.cc b/example/gpt2/checkpoint_loader.cc
index 95e54730b..b2ef92b28 100644
--- a/example/gpt2/checkpoint_loader.cc
+++ b/example/gpt2/checkpoint_loader.cc
@@ -57,7 +57,8 @@ std::tuple<int32_t, infini_train::DataType> DetermineAndCheckVersion(const std::
 
 namespace gpt2 {
 
-std::shared_ptr<nn::TransformerModel> LoadFromLLMC(const std::string &filepath) {
+std::shared_ptr<nn::TransformerModel> LoadFromLLMC(const std::string &filepath,
+                                                   const nn::ActivationRecomputeOptions &recompute_options) {
     if (!std::filesystem::exists(filepath)) {
         LOG(FATAL) << "File not found: " << filepath;
     }
@@ -89,6 +90,7 @@ std::shared_ptr<nn::TransformerModel> LoadFromLLMC(const std::string &filepath)
     gpt2_config.n_head = n_head;
     gpt2_config.n_embd = n_embd;
     gpt2::SanitizeGPT2Config(gpt2_config);
+    nn::ApplyActivationRecomputeOptions(&gpt2_config, recompute_options);
     auto local_gpt2 = std::make_shared<nn::TransformerModel>(gpt2_config);
 
     LOG(INFO) << "magic: " << magic << " version: " << version << " block_size: " << block_size
diff --git a/example/gpt2/checkpoint_loader.h b/example/gpt2/checkpoint_loader.h
index e80c356e3..1ba822f5b 100644
--- a/example/gpt2/checkpoint_loader.h
+++ b/example/gpt2/checkpoint_loader.h
@@ -3,10 +3,13 @@
 #include <memory>
 #include <string>
 
+#include "infini_train/include/nn/modules/transformer/transformer_config.h"
+
 namespace infini_train::nn {
 class TransformerModel;
 } // namespace infini_train::nn
 
 namespace gpt2 {
-std::shared_ptr<infini_train::nn::TransformerModel> LoadFromLLMC(const std::string &filepath);
+std::shared_ptr<infini_train::nn::TransformerModel>
+LoadFromLLMC(const std::string &filepath, const infini_train::nn::ActivationRecomputeOptions &recompute_options = {});
 } // namespace gpt2
diff --git a/example/gpt2/main.cc b/example/gpt2/main.cc
index 67b77f75f..a23984689 100644
--- a/example/gpt2/main.cc
+++ b/example/gpt2/main.cc
@@ -95,6 +95,12 @@ DEFINE_bool(sequence_parallel, false, "Whether to enable Sequence Parallel");
 DEFINE_uint32(pipeline_parallel, 1, "Pipeline Parallel world size, specified the number of PP stages.");
 DEFINE_uint32(virtual_pipeline_parallel, 1, "Number of chunks in PP stage.");
 
+// activation recompute
+DEFINE_bool(activation_recompute, false, "Enable activation recompute to trade compute for memory.");
+DEFINE_string(recompute_granularity, "full", "Activation recompute granularity: none|full|selective");
+DEFINE_string(recompute_method, "none", "Activation recompute method: none|uniform|block");
+DEFINE_uint32(recompute_num_layers, 0, "Number of transformer layers per recompute region for uniform/block methods.");
+
 // precision
 DEFINE_string(dtype, "float32", "precision used in training (float32/bfloat16)");
 DEFINE_uint32(save_interval, 0, "save checkpoint every N steps; 0 disables saving");
@@ -230,22 +236,25 @@ void Train(const nn::parallel::Rank &rank) {
     nn::TransformerConfig model_config = gpt2::GPT2Config();
     std::shared_ptr<nn::Module> model = nullptr;
 
+    const auto recompute_options
+        = nn::ParseActivationRecomputeOptions(FLAGS_activation_recompute, FLAGS_recompute_granularity,
+                                              FLAGS_recompute_method, static_cast<int64_t>(FLAGS_recompute_num_layers));
+
     if (!FLAGS_llmc_filepath.empty()) {
-        model = gpt2::LoadFromLLMC(FLAGS_llmc_filepath);
+        model = gpt2::LoadFromLLMC(FLAGS_llmc_filepath, recompute_options);
     } else if (kModelToConfigs.count(FLAGS_model)) {
         model_config = kModelToConfigs.at(FLAGS_model);
         gpt2::SanitizeGPT2Config(model_config);
+        nn::ApplyActivationRecomputeOptions(&model_config, recompute_options);
         model = std::make_shared<nn::TransformerModel>(model_config);
     }
 
+    CHECK(model) << "GPT2 example expects GPT2 model.";
+
     model->To(device);
 
     utils::PrecisionChecker::BuildNameMap(model.get());
 
-    // Get chunk size before wrapping with LoRA (needed for PipelineParallel)
-    auto gpt2_model = std::dynamic_pointer_cast<nn::TransformerModel>(model);
-    CHECK(gpt2_model) << "GPT2 example expects GPT2 model.";
-
     // Apply LoRA using GetLoRAModel (in-place injection)
     bool lora_enabled = FLAGS_lora_rank > 0;
     if (lora_enabled) {
diff --git a/example/llama3/checkpoint_loader.cc b/example/llama3/checkpoint_loader.cc
index 9c4058479..52a52d357 100644
--- a/example/llama3/checkpoint_loader.cc
+++ b/example/llama3/checkpoint_loader.cc
@@ -40,7 +40,8 @@ constexpr int32_t kLLaMA3FP32Version = 3;
 
 namespace llama3 {
 
-std::shared_ptr<nn::TransformerModel> LoadFromLLMC(const std::string &filepath) {
+std::shared_ptr<nn::TransformerModel> LoadFromLLMC(const std::string &filepath,
+                                                   const nn::ActivationRecomputeOptions &recompute_options) {
     if (!std::filesystem::exists(filepath)) {
         LOG(FATAL) << "File not found: " << filepath;
     }
@@ -95,6 +96,7 @@ std::shared_ptr<nn::TransformerModel> LoadFromLLMC(const std::string &filepath)
     llama3_config.norm_eps = norm_eps;
     llama3_config.max_gen_batch_size = max_gen_bs;
     llama3::SanitizeLLaMA3Config(llama3_config);
+    nn::ApplyActivationRecomputeOptions(&llama3_config, recompute_options);
     auto llama3 = std::make_shared<nn::TransformerModel>(llama3_config);
 
     // ========== pp_size：num_stages; vpp_size: num_chunks_per_stage ==========
diff --git a/example/llama3/checkpoint_loader.h b/example/llama3/checkpoint_loader.h
index d4aea3d04..7abd84d10 100644
--- a/example/llama3/checkpoint_loader.h
+++ b/example/llama3/checkpoint_loader.h
@@ -3,10 +3,13 @@
 #include <memory>
 #include <string>
 
+#include "infini_train/include/nn/modules/transformer/transformer_config.h"
+
 namespace infini_train::nn {
 class TransformerModel;
 } // namespace infini_train::nn
 
 namespace llama3 {
-std::shared_ptr<infini_train::nn::TransformerModel> LoadFromLLMC(const std::string &filepath);
+std::shared_ptr<infini_train::nn::TransformerModel>
+LoadFromLLMC(const std::string &filepath, const infini_train::nn::ActivationRecomputeOptions &recompute_options = {});
 } // namespace llama3
diff --git a/example/llama3/main.cc b/example/llama3/main.cc
index a8760e28b..5d95a0991 100644
--- a/example/llama3/main.cc
+++ b/example/llama3/main.cc
@@ -92,6 +92,12 @@ DEFINE_uint32(tensor_parallel, 1, "Tensor Parallel world size");
 DEFINE_bool(sequence_parallel, false, "Whether to enable Sequence Parallel");
 DEFINE_uint32(pipeline_parallel, 1, "Pipeline Parallel world size, specified the number of PP stages.");
 DEFINE_uint32(virtual_pipeline_parallel, 1, "Number of chunks in PP stage.");
+
+// activation recompute
+DEFINE_bool(activation_recompute, false, "Enable activation recompute to trade compute for memory.");
+DEFINE_string(recompute_granularity, "full", "Activation recompute granularity: none|full|selective");
+DEFINE_string(recompute_method, "none", "Activation recompute method: none|uniform|block");
+DEFINE_uint32(recompute_num_layers, 0, "Number of transformer layers per recompute region for uniform/block methods.");
 // precision
 DEFINE_string(dtype, "float32", "precision used in training (float32/bfloat16)");
 DEFINE_uint32(save_interval, 0, "save checkpoint every N steps; 0 disables saving");
@@ -218,13 +224,19 @@ void Train(const nn::parallel::Rank &rank) {
 
     nn::TransformerConfig model_config = llama3::LLaMA3Config();
     std::shared_ptr<nn::Module> model = nullptr;
+    const auto recompute_options
+        = nn::ParseActivationRecomputeOptions(FLAGS_activation_recompute, FLAGS_recompute_granularity,
+                                              FLAGS_recompute_method, static_cast<int64_t>(FLAGS_recompute_num_layers));
     if (!FLAGS_llmc_filepath.empty()) {
-        model = llama3::LoadFromLLMC(FLAGS_llmc_filepath);
+        model = llama3::LoadFromLLMC(FLAGS_llmc_filepath, recompute_options);
     } else {
+        nn::ApplyActivationRecomputeOptions(&model_config, recompute_options);
         llama3::SanitizeLLaMA3Config(model_config);
         model = std::make_shared<nn::TransformerModel>(model_config);
     }
 
+    CHECK(model) << "LLaMA3 example expects LLaMA3 model.";
+
     model->To(device);
 
 #ifdef INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_WORKAROUNDS
diff --git a/example/mnist/main.cc b/example/mnist/main.cc
index 7744e0947..54e912c20 100644
--- a/example/mnist/main.cc
+++ b/example/mnist/main.cc
@@ -66,7 +66,7 @@ int main(int argc, char *argv[]) {
             auto new_image = std::make_shared<Tensor>(image->To(device));
             auto new_label = std::make_shared<Tensor>(label->To(device));
 
-            auto outputs = network.Forward({new_image});
+            auto outputs = network({new_image});
             optimizer.ZeroGrad();
 
             auto loss = loss_fn.Forward({outputs[0], new_label});
@@ -104,7 +104,7 @@ int main(int argc, char *argv[]) {
         auto new_label = std::make_shared<Tensor>(label->To(device));
 
         auto label_cpu = label->To(cpu_device);
-        auto outputs = network.Forward({new_image});
+        auto outputs = network({new_image});
         auto output_cpu = outputs[0]->To(cpu_device);
         auto loss = loss_fn.Forward({outputs[0], new_label});
         auto loss_cpu = loss[0]->To(cpu_device);
diff --git a/example/qwen3/checkpoint_loader.cc b/example/qwen3/checkpoint_loader.cc
index 94ff782e9..880a630d6 100644
--- a/example/qwen3/checkpoint_loader.cc
+++ b/example/qwen3/checkpoint_loader.cc
@@ -34,7 +34,8 @@ constexpr size_t kQwen3HeaderBytes = 256 * sizeof(int32_t);
 
 namespace qwen3 {
 
-std::shared_ptr<nn::TransformerModel> LoadFromLLMC(const std::string &filepath) {
+std::shared_ptr<nn::TransformerModel> LoadFromLLMC(const std::string &filepath,
+                                                   const nn::ActivationRecomputeOptions &recompute_options) {
     if (!std::filesystem::exists(filepath)) {
         LOG(FATAL) << "File not found: " << filepath;
     }
@@ -108,6 +109,7 @@ std::shared_ptr<nn::TransformerModel> LoadFromLLMC(const std::string &filepath)
     qwen3_config.max_gen_batch_size = max_gen_bs;
     qwen3_config.qk_layernorm = true;
     qwen3::SanitizeQwen3Config(qwen3_config);
+    nn::ApplyActivationRecomputeOptions(&qwen3_config, recompute_options);
     auto qwen3 = std::make_shared<nn::TransformerModel>(qwen3_config);
 
     // ========== pp_size: num_stages; vpp_size: num_chunks_per_stage ==========
diff --git a/example/qwen3/checkpoint_loader.h b/example/qwen3/checkpoint_loader.h
index 5f4880f6e..a56658404 100644
--- a/example/qwen3/checkpoint_loader.h
+++ b/example/qwen3/checkpoint_loader.h
@@ -3,10 +3,13 @@
 #include <memory>
 #include <string>
 
+#include "infini_train/include/nn/modules/transformer/transformer_config.h"
+
 namespace infini_train::nn {
 class TransformerModel;
 } // namespace infini_train::nn
 
 namespace qwen3 {
-std::shared_ptr<infini_train::nn::TransformerModel> LoadFromLLMC(const std::string &filepath);
+std::shared_ptr<infini_train::nn::TransformerModel>
+LoadFromLLMC(const std::string &filepath, const infini_train::nn::ActivationRecomputeOptions &recompute_options = {});
 } // namespace qwen3
diff --git a/example/qwen3/main.cc b/example/qwen3/main.cc
index 435bae301..1eb23c295 100644
--- a/example/qwen3/main.cc
+++ b/example/qwen3/main.cc
@@ -92,6 +92,12 @@ DEFINE_uint32(tensor_parallel, 1, "Tensor Parallel world size");
 DEFINE_bool(sequence_parallel, false, "Whether to enable Sequence Parallel");
 DEFINE_uint32(pipeline_parallel, 1, "Pipeline Parallel world size, specified the number of PP stages.");
 DEFINE_uint32(virtual_pipeline_parallel, 1, "Number of chunks in PP stage.");
+
+// activation recompute
+DEFINE_bool(activation_recompute, false, "Enable activation recompute to trade compute for memory.");
+DEFINE_string(recompute_granularity, "full", "Activation recompute granularity: none|full|selective");
+DEFINE_string(recompute_method, "none", "Activation recompute method: none|uniform|block");
+DEFINE_uint32(recompute_num_layers, 0, "Number of transformer layers per recompute region for uniform/block methods.");
 // precision
 DEFINE_string(dtype, "float32", "precision used in training (float32/bfloat16)");
 DEFINE_uint32(save_interval, 0, "save checkpoint every N steps; 0 disables saving");
@@ -217,9 +223,13 @@ void Train(const nn::parallel::Rank &rank) {
 
     nn::TransformerConfig model_config = qwen3::Qwen3Config();
     std::shared_ptr<nn::Module> model = nullptr;
+    const auto recompute_options
+        = nn::ParseActivationRecomputeOptions(FLAGS_activation_recompute, FLAGS_recompute_granularity,
+                                              FLAGS_recompute_method, static_cast<int64_t>(FLAGS_recompute_num_layers));
     if (!FLAGS_llmc_filepath.empty()) {
-        model = qwen3::LoadFromLLMC(FLAGS_llmc_filepath);
+        model = qwen3::LoadFromLLMC(FLAGS_llmc_filepath, recompute_options);
     } else {
+        nn::ApplyActivationRecomputeOptions(&model_config, recompute_options);
         qwen3::SanitizeQwen3Config(model_config);
         model = std::make_shared<nn::TransformerModel>(model_config);
     }
diff --git a/infini_train/include/autograd/grad_mode.h b/infini_train/include/autograd/grad_mode.h
index 651573875..1a21bff48 100644
--- a/infini_train/include/autograd/grad_mode.h
+++ b/infini_train/include/autograd/grad_mode.h
@@ -8,10 +8,13 @@ class GradMode {
     // Whether to enable Autograd (enabled by default)
     static bool IsEnabled() { return grad_enabled_; }
     static void SetEnabled(bool enabled) { grad_enabled_ = enabled; }
+    static bool PropagateRequiresGrad() { return propagate_requires_grad_; }
+    static void SetPropagateRequiresGrad(bool enabled) { propagate_requires_grad_ = enabled; }
 
 private:
     // grad mode should be thread_local
     static thread_local bool grad_enabled_;
+    static thread_local bool propagate_requires_grad_;
 };
 
 // RAII: Disable grad (align with torch.no_grad)
@@ -34,4 +37,20 @@ class EnableGradGuard {
     bool prev_;
 };
 
+// RAII: Propagate requires_grad metadata while graph construction is disabled.
+// Used by non-reentrant checkpoint recomputation so downstream SetupContext
+// calls see the same needs_input_grad_ pattern as the original forward,
+// without wiring the recompute graph into the engine. Together with grad mode,
+// this gives three behaviors: build a graph, propagate metadata only, or neither.
+class PropagateRequiresGradGuard {
+public:
+    explicit PropagateRequiresGradGuard(bool enabled = true) : prev_(GradMode::PropagateRequiresGrad()) {
+        GradMode::SetPropagateRequiresGrad(enabled);
+    }
+    ~PropagateRequiresGradGuard() { GradMode::SetPropagateRequiresGrad(prev_); }
+
+private:
+    bool prev_ = false;
+};
+
 } // namespace infini_train::autograd
diff --git a/infini_train/include/nn/modules/module.h b/infini_train/include/nn/modules/module.h
index 8570c4768..5b875517f 100644
--- a/infini_train/include/nn/modules/module.h
+++ b/infini_train/include/nn/modules/module.h
@@ -81,7 +81,7 @@ class Module : public std::enable_shared_from_this<Module> {
     // WARNING-only and silently ignored.
     void LoadStateDict(const std::unordered_map<std::string, std::shared_ptr<Tensor>> &state_dict);
 
-    // operator() calls hooks and Forward
+    // Public module invocation entry point.
     std::vector<std::shared_ptr<Tensor>> operator()(const std::vector<std::shared_ptr<Tensor>> &input_tensors);
 
     // Forward to be overridden by subclasses
diff --git a/infini_train/include/nn/modules/transformer/transformer_config.h b/infini_train/include/nn/modules/transformer/transformer_config.h
index 11df977f9..80558d739 100644
--- a/infini_train/include/nn/modules/transformer/transformer_config.h
+++ b/infini_train/include/nn/modules/transformer/transformer_config.h
@@ -2,6 +2,7 @@
 
 #include <cstdint>
 #include <optional>
+#include <string_view>
 
 namespace infini_train::nn {
 
@@ -54,6 +55,24 @@ struct MoEConfig {
     ExpertImpl expert_impl = ExpertImpl::kSequential;
 };
 
+enum class ActivationRecomputeGranularity {
+    kNone,      // Disable activation recompute.
+    kFull,      // Recompute full transformer layers.
+    kSelective, // Recompute selected transformer submodules.
+};
+
+enum class ActivationRecomputeMethod {
+    kNone,    // Recompute every layer when granularity is full.
+    kUniform, // Uniformly divide layers into recompute chunks.
+    kBlock,   // Recompute only the first N layers in the local chunk/stage.
+};
+
+struct ActivationRecomputeOptions {
+    ActivationRecomputeGranularity granularity = ActivationRecomputeGranularity::kNone;
+    ActivationRecomputeMethod method = ActivationRecomputeMethod::kNone;
+    int64_t num_layers = 0;
+};
+
 struct TransformerConfig {
     int64_t block_size = 1024;           // Max seq_len
     int64_t vocab_size = 50304;          // Vocab size
@@ -87,6 +106,11 @@ struct TransformerConfig {
     // Normalization
     float norm_eps = 1e-5f; // epsilon in RMSNorm
 
+    // Activation recomputation, aligned with Megatron-Core TransformerConfig.
+    ActivationRecomputeGranularity recompute_granularity = ActivationRecomputeGranularity::kNone;
+    ActivationRecomputeMethod recompute_method = ActivationRecomputeMethod::kNone;
+    int64_t recompute_num_layers = 0;
+
     // Inference
     bool use_kv = false;            // kv cache
     bool flash = false;             // flash attention
@@ -96,5 +120,10 @@ struct TransformerConfig {
 
     bool UseGQA() const;
     int GetChunkSize() const;
+    bool RecomputeEnabled() const;
 };
+
+ActivationRecomputeOptions ParseActivationRecomputeOptions(bool enabled, std::string_view granularity,
+                                                           std::string_view method, int64_t num_layers);
+void ApplyActivationRecomputeOptions(TransformerConfig *config, const ActivationRecomputeOptions &options);
 } // namespace infini_train::nn
diff --git a/infini_train/include/utils/checkpoint.h b/infini_train/include/utils/checkpoint.h
new file mode 100644
index 000000000..e5870bbe2
--- /dev/null
+++ b/infini_train/include/utils/checkpoint.h
@@ -0,0 +1,23 @@
+#pragma once
+
+#include <functional>
+#include <memory>
+#include <vector>
+
+namespace infini_train {
+class Tensor;
+} // namespace infini_train
+
+namespace infini_train::utils::checkpoint {
+
+using CheckpointForwardFn
+    = std::function<std::vector<std::shared_ptr<Tensor>>(const std::vector<std::shared_ptr<Tensor>> &)>;
+
+// RNG state preservation is not implemented. Random ops such as Dropout cannot
+// be replayed consistently; preserve_rng_state=true is rejected.
+std::vector<std::shared_ptr<Tensor>> Checkpoint(const CheckpointForwardFn &forward_fn,
+                                                const std::vector<std::shared_ptr<Tensor>> &inputs,
+                                                bool use_reentrant = false, bool preserve_rng_state = false,
+                                                bool determinism_check = true, bool early_stop = true);
+
+} // namespace infini_train::utils::checkpoint
diff --git a/infini_train/src/autograd/function.cc b/infini_train/src/autograd/function.cc
index af23c1631..70130822b 100644
--- a/infini_train/src/autograd/function.cc
+++ b/infini_train/src/autograd/function.cc
@@ -162,12 +162,15 @@ std::vector<std::shared_ptr<Tensor>> Function::Apply(const std::vector<std::shar
 
     // Populate needs_input_grad before Forward/SetupContext so that
     // SetupContext can use it for saved-tensor pruning.
-    // Must be done before NoGradGuard since it checks GradMode.
+    // Non-reentrant checkpoint recomputes under no-grad; PropagateRequiresGrad
+    // keeps the same context metadata without wiring the recompute graph.
     ctx_.ReleaseVariables();
-    if (autograd::GradMode::IsEnabled()) {
+    const bool grad_enabled = autograd::GradMode::IsEnabled();
+    const bool record_context = grad_enabled || autograd::GradMode::PropagateRequiresGrad();
+    if (record_context) {
         std::vector<bool> needs_input_grad(input_tensors.size());
         for (size_t idx = 0; idx < input_tensors.size(); ++idx) {
-            needs_input_grad[idx] = input_tensors[idx]->requires_grad();
+            needs_input_grad[idx] = input_tensors[idx] && input_tensors[idx]->requires_grad();
         }
         ctx_.set_needs_input_grad(std::move(needs_input_grad));
     }
@@ -185,12 +188,16 @@ std::vector<std::shared_ptr<Tensor>> Function::Apply(const std::vector<std::shar
     std::vector<std::shared_ptr<Tensor>> output_tensors;
     {
         autograd::NoGradGuard no_grad;
+        // Nested Functions inside Forward/SetupContext must not record context,
+        // including during checkpoint recomputation. Restore the caller mode
+        // before saving this Function's tensors and propagating output metadata.
+        autograd::PropagateRequiresGradGuard no_propagate(false);
         // no_grad in autograd.Function.Forward()
         output_tensors = Forward(compute_inputs);
         SetupContext(compute_inputs, output_tensors);
     }
 
-    if (autograd::GradMode::IsEnabled()) {
+    if (record_context) {
         ctx_.SaveVariables(output_tensors);
     } else {
         ctx_.ReleaseVariables();
@@ -203,51 +210,67 @@ std::vector<std::shared_ptr<Tensor>> Function::Apply(const std::vector<std::shar
         }
     }
 
-    if (!autograd::GradMode::IsEnabled()) {
-        // with no_grad: block graph building operations
+    if (!record_context) {
         return output_tensors;
     }
 
     bool output_requires_grad = false;
-    for (const auto &input_tensor : input_tensors) { output_requires_grad |= input_tensor->requires_grad(); }
+    for (const auto &input_tensor : input_tensors) {
+        output_requires_grad |= input_tensor && input_tensor->requires_grad();
+    }
 
-    grad_outputs_reached_ = 0;
-    grad_outputs_.resize(output_tensors.size(), nullptr);
     std::vector<bool> differentiable_outputs(output_tensors.size(), false);
     bool has_differentiable_output = false;
     for (int output_idx = 0; output_idx < output_tensors.size(); ++output_idx) {
-        differentiable_outputs[output_idx]
-            = output_requires_grad && !ctx_.IsNonDifferentiable(output_tensors[output_idx]);
+        differentiable_outputs[output_idx] = output_tensors[output_idx] && output_requires_grad
+                                          && !ctx_.IsNonDifferentiable(output_tensors[output_idx]);
         has_differentiable_output |= differentiable_outputs[output_idx];
-        if (!differentiable_outputs[output_idx]) {
-            ++grad_outputs_reached_;
-        }
     }
 
-    if (!has_differentiable_output) {
-        next_functions_.clear();
-        return output_tensors;
-    }
+    // Metadata-only recomputation must not change graph dependency counters.
+    if (grad_enabled) {
+        grad_outputs_reached_ = 0;
+        grad_outputs_.resize(output_tensors.size(), nullptr);
+        for (bool differentiable : differentiable_outputs) {
+            if (!differentiable) {
+                ++grad_outputs_reached_;
+            }
+        }
 
-    for (int idx = 0; idx < input_tensors.size(); ++idx) {
-        const auto &input_tensor = input_tensors[idx];
-        if (input_tensor->requires_grad() && input_tensor->is_leaf()) {
-            next_functions_.emplace_back(input_tensor->grad_accumulator(), input_tensor->output_idx());
-            input_tensor->grad_accumulator()->IncreaseDependenciesNumber();
+        if (!has_differentiable_output) {
+            next_functions_.clear();
         } else {
-            next_functions_.emplace_back(input_tensor->grad_fn(), input_tensor->output_idx());
-            if (input_tensor->grad_fn()) {
-                input_tensor->grad_fn()->IncreaseDependenciesNumber();
+            for (int idx = 0; idx < input_tensors.size(); ++idx) {
+                const auto &input_tensor = input_tensors[idx];
+                if (!input_tensor) {
+                    next_functions_.emplace_back(nullptr, 0);
+                    continue;
+                }
+                if (input_tensor->requires_grad() && input_tensor->is_leaf()) {
+                    next_functions_.emplace_back(input_tensor->grad_accumulator(), input_tensor->output_idx());
+                    input_tensor->grad_accumulator()->IncreaseDependenciesNumber();
+                } else {
+                    next_functions_.emplace_back(input_tensor->grad_fn(), input_tensor->output_idx());
+                    if (input_tensor->grad_fn()) {
+                        input_tensor->grad_fn()->IncreaseDependenciesNumber();
+                    }
+                }
             }
         }
     }
 
+    // Both modes take MarkNonDifferentiable into consideration.
+    // Set output metadata only after wiring input edges, since a custom Function may return an input tensor.
     for (int output_idx = 0; output_idx < output_tensors.size(); ++output_idx) {
         auto &output_tensor = output_tensors[output_idx];
+        if (!output_tensor) {
+            continue;
+        }
         const bool differentiable_output = differentiable_outputs[output_idx];
+        const bool attach_grad_fn = grad_enabled && differentiable_output;
         output_tensor->set_requires_grad(differentiable_output);
-        output_tensor->set_grad_fn(differentiable_output ? shared_from_this() : nullptr);
-        output_tensor->set_is_leaf(!differentiable_output);
+        output_tensor->set_grad_fn(attach_grad_fn ? shared_from_this() : nullptr);
+        output_tensor->set_is_leaf(!attach_grad_fn);
         output_tensor->set_output_idx(output_idx);
     }
 
diff --git a/infini_train/src/autograd/grad_mode.cc b/infini_train/src/autograd/grad_mode.cc
index 28a6e693c..f638f2bc7 100644
--- a/infini_train/src/autograd/grad_mode.cc
+++ b/infini_train/src/autograd/grad_mode.cc
@@ -2,4 +2,5 @@
 
 namespace infini_train::autograd {
 thread_local bool GradMode::grad_enabled_ = true;
+thread_local bool GradMode::propagate_requires_grad_ = false;
 } // namespace infini_train::autograd
diff --git a/infini_train/src/nn/modules/module.cc b/infini_train/src/nn/modules/module.cc
index 498bcc075..f81681200 100644
--- a/infini_train/src/nn/modules/module.cc
+++ b/infini_train/src/nn/modules/module.cc
@@ -237,6 +237,7 @@ std::vector<std::shared_ptr<Tensor>> Module::Forward(const std::vector<std::shar
 }
 
 std::vector<std::shared_ptr<Tensor>> Module::operator()(const std::vector<std::shared_ptr<Tensor>> &input_tensors) {
+    // Compared to `Forward()` which will bypass hooks, this is the full module call protocol including hooks.
     // 1. Call global module forward pre-hooks
     utils::GlobalModuleHookRegistry::Instance().CallModuleForwardPreHooks(this, input_tensors);
 
diff --git a/infini_train/src/nn/modules/transformer/transformer.cc b/infini_train/src/nn/modules/transformer/transformer.cc
index 13147bf50..133577d21 100644
--- a/infini_train/src/nn/modules/transformer/transformer.cc
+++ b/infini_train/src/nn/modules/transformer/transformer.cc
@@ -1,5 +1,6 @@
 #include "infini_train/include/nn/modules/transformer/transformer.h"
 
+#include <algorithm>
 #include <cmath>
 #include <cstdio>
 #include <map>
@@ -7,6 +8,7 @@
 
 #include "glog/logging.h"
 
+#include "infini_train/include/autograd/grad_mode.h"
 #include "infini_train/include/nn/functional.h"
 #include "infini_train/include/nn/init.h"
 #include "infini_train/include/nn/modules/container.h"
@@ -22,8 +24,54 @@
 #include "infini_train/include/nn/parallel/tensor_parallel.h"
 #include "infini_train/include/nn/parallel/utils.h"
 #include "infini_train/include/tensor.h"
+#include "infini_train/include/utils/checkpoint.h"
 
 namespace infini_train::nn {
+namespace {
+std::vector<std::shared_ptr<Tensor>> RunTransformerLayers(const std::shared_ptr<nn::ModuleList> &layers, size_t start,
+                                                          size_t end,
+                                                          const std::vector<std::shared_ptr<Tensor>> &inputs) {
+    auto hidden = inputs[0];
+    auto layer_inputs = inputs;
+    for (size_t layer_idx = start; layer_idx < end; ++layer_idx) {
+        layer_inputs[0] = hidden;
+        hidden = (*(*layers)[layer_idx])(layer_inputs)[0];
+    }
+    return {hidden};
+}
+
+std::vector<std::shared_ptr<Tensor>> CheckpointTransformerLayers(const std::shared_ptr<nn::ModuleList> &layers,
+                                                                 size_t start, size_t end,
+                                                                 const std::vector<std::shared_ptr<Tensor>> &inputs) {
+    auto forward_fn = [layers, start, end](const std::vector<std::shared_ptr<Tensor>> &checkpoint_inputs) {
+        return RunTransformerLayers(layers, start, end, checkpoint_inputs);
+    };
+    constexpr bool kUseReentrant = false;
+    // FIXME(zbl): RNG state preservation is not implemented. Checkpointed layers
+    //             must not use Dropout or other random ops that need identical replay.
+    constexpr bool kPreserveRngState = false;
+    return utils::checkpoint::Checkpoint(forward_fn, inputs, kUseReentrant, kPreserveRngState);
+}
+
+bool ShouldUseLayerRecompute(const TransformerConfig &config) {
+    if (!config.RecomputeEnabled() || !autograd::GradMode::IsEnabled()) {
+        return false;
+    }
+    if (config.recompute_granularity == ActivationRecomputeGranularity::kSelective) {
+        LOG(FATAL) << "Selective activation recompute is not implemented yet. Use full layer recompute.";
+    }
+    return config.recompute_granularity == ActivationRecomputeGranularity::kFull;
+}
+
+size_t GetRecomputeNumLayers(const TransformerConfig &config) {
+    if (config.recompute_method == ActivationRecomputeMethod::kNone) {
+        return 1;
+    }
+    CHECK_GE(config.recompute_num_layers, 1)
+        << "recompute_num_layers must be >= 1 when recompute_method is uniform or block.";
+    return static_cast<size_t>(config.recompute_num_layers);
+}
+} // namespace
 
 TransformerFirstStage::TransformerFirstStage(const TransformerConfig &config)
     : CloneableModule(kType), config_(config) {
@@ -138,8 +186,11 @@ TransformerChunk::TransformerChunk(const TransformerConfig &config, int start_la
 
 std::vector<std::shared_ptr<Tensor>> TransformerChunk::Forward(const std::vector<std::shared_ptr<Tensor>> &x) {
     auto x1 = x[0];
+    auto layers = std::dynamic_pointer_cast<nn::ModuleList>(modules_[kHLayerName]);
+    CHECK(layers);
 
     // Check if we need to pass RoPE parameters (for LLaMA3 style models).
+    std::vector<std::shared_ptr<Tensor>> layer_inputs = {x1};
     if (config_.position_embedding_type == PositionEmbeddingType::kRoPE) {
         // For RoPE models, we need to prepare freqs_cis and potentially other parameters
         const auto device = x1->GetDevice();
@@ -147,6 +198,7 @@ std::vector<std::shared_ptr<Tensor>> TransformerChunk::Forward(const std::vector
         // Init freqs_cis on device only once
         if (buffers_[kFreqsCisName] == nullptr) {
             int64_t head_dim = config_.n_embd / config_.n_head;
+            autograd::NoGradGuard no_grad;
             buffers_[kFreqsCisName] = PrecomputeFreqsCis(head_dim, config_.block_size * 2, config_.rope_theta,
                                                          config_.use_scaled_rope, device);
         }
@@ -163,17 +215,51 @@ std::vector<std::shared_ptr<Tensor>> TransformerChunk::Forward(const std::vector
 
         std::shared_ptr<Tensor> start_pos_ptr = nullptr;
 
-        // Pass RoPE parameters to each transformer block
-        for (auto &h : *std::dynamic_pointer_cast<nn::ModuleList>(modules_[kHLayerName])) {
-            x1 = (*h)({x1, freqs_view, start_pos_ptr, mask})[0];
-        }
-    } else if (config_.position_embedding_type == PositionEmbeddingType::kLearnedAbsolute) {
-        // Learned absolute position embedding models (GPT-2 style).
-        for (auto &h : *std::dynamic_pointer_cast<nn::ModuleList>(modules_[kHLayerName])) { x1 = (*h)({x1})[0]; }
-    } else {
+        layer_inputs = {x1, freqs_view, start_pos_ptr, mask};
+    } else if (config_.position_embedding_type != PositionEmbeddingType::kLearnedAbsolute) {
         LOG(FATAL) << "Unsupported position embedding type";
     }
 
+    if (!ShouldUseLayerRecompute(config_)) {
+        return RunTransformerLayers(layers, 0, static_cast<size_t>(layers->end() - layers->begin()), layer_inputs);
+    }
+
+    const size_t num_layers = static_cast<size_t>(layers->end() - layers->begin());
+    const size_t recompute_num_layers = GetRecomputeNumLayers(config_);
+    switch (config_.recompute_method) {
+    case ActivationRecomputeMethod::kNone: {
+        for (size_t layer_idx = 0; layer_idx < num_layers; ++layer_idx) {
+            layer_inputs[0] = x1;
+            x1 = CheckpointTransformerLayers(layers, layer_idx, layer_idx + 1, layer_inputs)[0];
+        }
+        break;
+    }
+    case ActivationRecomputeMethod::kUniform: {
+        size_t layer_idx = 0;
+        while (layer_idx < num_layers) {
+            auto chunk_end = std::min(layer_idx + recompute_num_layers, num_layers);
+            layer_inputs[0] = x1;
+            x1 = CheckpointTransformerLayers(layers, layer_idx, chunk_end, layer_inputs)[0];
+            layer_idx = chunk_end;
+        }
+        break;
+    }
+    case ActivationRecomputeMethod::kBlock: {
+        const size_t checkpoint_layers = std::min(recompute_num_layers, num_layers);
+        for (size_t layer_idx = 0; layer_idx < num_layers; ++layer_idx) {
+            layer_inputs[0] = x1;
+            if (layer_idx < checkpoint_layers) {
+                x1 = CheckpointTransformerLayers(layers, layer_idx, layer_idx + 1, layer_inputs)[0];
+            } else {
+                x1 = RunTransformerLayers(layers, layer_idx, layer_idx + 1, layer_inputs)[0];
+            }
+        }
+        break;
+    }
+    default:
+        LOG(FATAL) << "Unsupported activation recompute method.";
+    }
+
     return {x1};
 }
 
diff --git a/infini_train/src/nn/modules/transformer/transformer_config.cc b/infini_train/src/nn/modules/transformer/transformer_config.cc
index b8947d4b6..e992f9474 100644
--- a/infini_train/src/nn/modules/transformer/transformer_config.cc
+++ b/infini_train/src/nn/modules/transformer/transformer_config.cc
@@ -1,9 +1,47 @@
 #include "infini_train/include/nn/modules/transformer/transformer_config.h"
 
+#include <array>
+#include <utility>
+
+#include "glog/logging.h"
+
 #include "infini_train/include/nn/parallel/global.h"
 #include "infini_train/include/nn/parallel/pp/pipeline_parallel.h"
 
 namespace infini_train::nn {
+namespace {
+template <typename Enum, size_t N>
+Enum ParseEnum(std::string_view value, const std::array<std::pair<std::string_view, Enum>, N> &entries,
+               std::string_view name, std::string_view expected) {
+    for (const auto &[key, parsed] : entries) {
+        if (value == key) {
+            return parsed;
+        }
+    }
+    LOG(FATAL) << "Unknown " << name << ": " << value << ". Expected " << expected << ".";
+    return entries[0].second;
+}
+
+ActivationRecomputeGranularity ParseActivationRecomputeGranularity(std::string_view value) {
+    static constexpr std::array kEntries = {
+        std::pair<std::string_view, ActivationRecomputeGranularity>{"none", ActivationRecomputeGranularity::kNone},
+        std::pair<std::string_view, ActivationRecomputeGranularity>{"full", ActivationRecomputeGranularity::kFull},
+        std::pair<std::string_view, ActivationRecomputeGranularity>{"selective",
+                                                                    ActivationRecomputeGranularity::kSelective},
+    };
+    return ParseEnum(value, kEntries, "recompute_granularity", "none|full|selective");
+}
+
+ActivationRecomputeMethod ParseActivationRecomputeMethod(std::string_view value) {
+    static constexpr std::array kEntries = {
+        std::pair<std::string_view, ActivationRecomputeMethod>{"none", ActivationRecomputeMethod::kNone},
+        std::pair<std::string_view, ActivationRecomputeMethod>{"uniform", ActivationRecomputeMethod::kUniform},
+        std::pair<std::string_view, ActivationRecomputeMethod>{"block", ActivationRecomputeMethod::kBlock},
+    };
+    return ParseEnum(value, kEntries, "recompute_method", "none|uniform|block");
+}
+} // namespace
+
 bool TransformerConfig::UseGQA() const { return n_kv_head < n_head; }
 
 int TransformerConfig::GetChunkSize() const {
@@ -12,4 +50,35 @@ int TransformerConfig::GetChunkSize() const {
                                                                parallel::global::GetVirtualPipelineParallelSize());
     return stage_info.layer_ranges_per_chunk.size();
 }
+
+bool TransformerConfig::RecomputeEnabled() const {
+    return recompute_granularity != ActivationRecomputeGranularity::kNone;
+}
+
+ActivationRecomputeOptions ParseActivationRecomputeOptions(bool enabled, std::string_view granularity,
+                                                           std::string_view method, int64_t num_layers) {
+    ActivationRecomputeOptions options;
+    if (!enabled) {
+        return options;
+    }
+    options.granularity = ParseActivationRecomputeGranularity(granularity);
+    if (options.granularity == ActivationRecomputeGranularity::kNone
+        || options.granularity == ActivationRecomputeGranularity::kSelective) {
+        return options;
+    }
+    options.method = ParseActivationRecomputeMethod(method);
+    if (options.method == ActivationRecomputeMethod::kNone) {
+        return options;
+    }
+    CHECK_GE(num_layers, 1) << "recompute_num_layers must be >= 1 when recompute_method is uniform or block.";
+    options.num_layers = num_layers;
+    return options;
+}
+
+void ApplyActivationRecomputeOptions(TransformerConfig *config, const ActivationRecomputeOptions &options) {
+    CHECK(config);
+    config->recompute_granularity = options.granularity;
+    config->recompute_method = options.method;
+    config->recompute_num_layers = options.num_layers;
+}
 } // namespace infini_train::nn
diff --git a/infini_train/src/tensor.cc b/infini_train/src/tensor.cc
index 70d9591af..710be8f6e 100644
--- a/infini_train/src/tensor.cc
+++ b/infini_train/src/tensor.cc
@@ -3,6 +3,7 @@
 #include <cstring>
 #include <fstream>
 #include <functional>
+#include <iomanip>
 #include <memory>
 #include <numeric>
 
diff --git a/infini_train/src/utils/checkpoint.cc b/infini_train/src/utils/checkpoint.cc
new file mode 100644
index 000000000..dd67bbdb2
--- /dev/null
+++ b/infini_train/src/utils/checkpoint.cc
@@ -0,0 +1,354 @@
+#include "infini_train/include/utils/checkpoint.h"
+
+#include <algorithm>
+#include <atomic>
+#include <exception>
+#include <unordered_map>
+
+#include "glog/logging.h"
+
+#include "infini_train/include/autocast.h"
+#include "infini_train/include/autograd/function.h"
+#include "infini_train/include/autograd/grad_mode.h"
+#include "infini_train/include/tensor.h"
+
+namespace infini_train::utils::checkpoint {
+namespace {
+constexpr char kCheckpointType[] = "Checkpoint";
+std::atomic<int64_t> g_checkpoint_gid{0};
+
+int64_t NextCheckpointGid() { return g_checkpoint_gid.fetch_add(1) + 1; }
+
+class CheckpointFunction : public autograd::Function {
+public:
+    explicit CheckpointFunction(CheckpointForwardFn forward_fn);
+
+    std::vector<std::shared_ptr<Tensor>> Forward(const std::vector<std::shared_ptr<Tensor>> &input_tensors) override;
+    void SetupContext(const std::vector<std::shared_ptr<Tensor>> &input_tensors,
+                      const std::vector<std::shared_ptr<Tensor>> &output_tensors) override;
+    std::vector<std::shared_ptr<Tensor>> Backward(const std::vector<std::shared_ptr<Tensor>> &grad_outputs) override;
+
+private:
+    CheckpointForwardFn forward_fn_;
+    std::vector<std::shared_ptr<Tensor>> saved_inputs_;
+    std::vector<bool> saved_inputs_requires_grad_;
+};
+
+struct SavedTensorMeta {
+    std::vector<int64_t> dims;
+    DataType dtype = DataType::kFLOAT32;
+    Device::DeviceType device_type = Device::DeviceType::kCPU;
+};
+
+bool MetaEquals(const SavedTensorMeta &a, const SavedTensorMeta &b) {
+    return a.dtype == b.dtype && a.device_type == b.device_type && a.dims == b.dims;
+}
+
+struct CheckpointFrame;
+
+struct SavedTensorHolder {
+    std::shared_ptr<CheckpointFrame> frame;
+    size_t index = 0;
+    std::unordered_map<int64_t, bool> handles;
+    std::shared_ptr<Tensor> tensor;
+};
+
+struct CheckpointFrame {
+    CheckpointForwardFn forward_fn;
+    std::vector<std::shared_ptr<Tensor>> inputs;
+    std::vector<bool> inputs_requires_grad;
+    AutocastContext autocast_context;
+    std::vector<SavedTensorMeta> forward_metas;
+    bool early_stop = true;
+    bool determinism_check = true;
+    std::unordered_map<int64_t, bool> recomputed;
+    std::vector<std::weak_ptr<SavedTensorHolder>> weak_holders;
+    int64_t current_gid = -1;
+
+    struct StopRecomputeError : public std::exception {};
+
+    bool HasActiveHandles(int64_t gid) const {
+        for (const auto &weak_holder : weak_holders) {
+            auto holder = weak_holder.lock();
+            if (!holder) {
+                continue;
+            }
+            auto it = holder->handles.find(gid);
+            if (it != holder->handles.end() && it->second) {
+                return true;
+            }
+        }
+        return false;
+    }
+
+    size_t CountAliveHolders() const {
+        size_t alive = 0;
+        for (const auto &weak_holder : weak_holders) {
+            if (weak_holder.lock()) {
+                ++alive;
+            }
+        }
+        return alive;
+    }
+
+    int64_t GetOrCreateGid() {
+        if (current_gid < 0) {
+            current_gid = NextCheckpointGid();
+            return current_gid;
+        }
+        if (recomputed[current_gid] && !HasActiveHandles(current_gid)) {
+            current_gid = NextCheckpointGid();
+        }
+        return current_gid;
+    }
+
+    void Recompute(int64_t gid) {
+        if (recomputed[gid]) {
+            return;
+        }
+        const size_t alive_needed = CountAliveHolders();
+        size_t filled = 0;
+        size_t recompute_index = 0;
+
+        std::vector<std::shared_ptr<Tensor>> detached_inputs;
+        detached_inputs.reserve(inputs.size());
+        for (size_t i = 0; i < inputs.size(); ++i) {
+            if (!inputs[i]) {
+                detached_inputs.push_back(nullptr);
+                continue;
+            }
+            auto detached = inputs[i]->Detach();
+            detached->set_requires_grad(inputs_requires_grad[i]);
+            detached->set_is_leaf(true);
+            detached->set_grad_fn(nullptr);
+            detached_inputs.push_back(detached);
+        }
+
+        AutocastGuard autocast_guard(autocast_context);
+        // Unlike PyTorch's engine, this autograd implementation mutates
+        // dependency counters eagerly while building a graph. The recompute
+        // graph is not traversed by non-reentrant checkpoint here, so building
+        // it would pollute parameter/input dependency counts and break
+        // accumulation. Recompute under no-grad, while Function::Apply still
+        // propagates requires_grad metadata and populates needs_input_grad_ so
+        // SetupContext saves the same tensors as the original forward.
+        autograd::NoGradGuard no_grad;
+        autograd::PropagateRequiresGradGuard propagate_requires_grad;
+
+        autograd::FunctionCtx::SavedTensorHooks recompute_hooks;
+        recompute_hooks.pack = [this, gid, alive_needed, &filled,
+                                &recompute_index](const std::shared_ptr<Tensor> &tensor) -> std::shared_ptr<void> {
+            size_t idx = recompute_index++;
+            if (idx >= weak_holders.size()) {
+                LOG(FATAL) << "Checkpoint: recomputed more tensors than saved during forward.";
+            }
+            auto holder = weak_holders[idx].lock();
+            if (tensor) {
+                if (determinism_check) {
+                    SavedTensorMeta meta;
+                    meta.dims = tensor->Dims();
+                    meta.dtype = tensor->Dtype();
+                    meta.device_type = tensor->GetDevice().type();
+                    if (!MetaEquals(meta, forward_metas[idx])) {
+                        LOG(FATAL) << "Checkpoint: recomputed tensor metadata mismatch at index " << idx << ".";
+                    }
+                }
+                if (holder) {
+                    holder->handles[gid] = true;
+                    holder->tensor = tensor;
+                    ++filled;
+                }
+            } else {
+                if (holder) {
+                    holder->handles[gid] = true;
+                    holder->tensor.reset();
+                    ++filled;
+                }
+            }
+            if (early_stop && filled >= alive_needed) {
+                throw StopRecomputeError();
+            }
+            return tensor;
+        };
+        recompute_hooks.unpack = [](const std::shared_ptr<void> &state) -> std::shared_ptr<Tensor> {
+            return std::static_pointer_cast<Tensor>(state);
+        };
+        autograd::FunctionCtx::SavedTensorHooksGuard guard(std::move(recompute_hooks));
+
+        try {
+            forward_fn(detached_inputs);
+        } catch (const StopRecomputeError &) {
+            // Early-stop: expected when all needed tensors are recomputed.
+        }
+        if (filled < alive_needed) {
+            LOG(FATAL) << "Checkpoint: recomputed fewer tensors (" << filled << ") than required (" << alive_needed
+                       << ").";
+        }
+
+        recomputed[gid] = true;
+
+        // Break potential reference cycles once recomputation is done.
+        inputs.clear();
+        inputs_requires_grad.clear();
+        forward_fn = nullptr;
+    }
+};
+
+CheckpointFunction::CheckpointFunction(CheckpointForwardFn forward_fn)
+    : autograd::Function(kCheckpointType), forward_fn_(std::move(forward_fn)) {}
+
+std::vector<std::shared_ptr<Tensor>>
+CheckpointFunction::Forward(const std::vector<std::shared_ptr<Tensor>> &input_tensors) {
+    saved_inputs_.clear();
+    saved_inputs_requires_grad_.clear();
+    saved_inputs_.reserve(input_tensors.size());
+    saved_inputs_requires_grad_.reserve(input_tensors.size());
+    for (const auto &input : input_tensors) {
+        if (!input) {
+            saved_inputs_requires_grad_.push_back(false);
+            saved_inputs_.push_back(nullptr);
+            continue;
+        }
+        saved_inputs_requires_grad_.push_back(input->requires_grad());
+        saved_inputs_.push_back(input->Detach());
+    }
+
+    // TODO(zbl): RNG state is not captured yet. Dropout or random ops are not supported.
+    return forward_fn_(input_tensors);
+}
+
+void CheckpointFunction::SetupContext(const std::vector<std::shared_ptr<Tensor>> &,
+                                      const std::vector<std::shared_ptr<Tensor>> &) {
+    // Intentionally empty: checkpoint avoids saving intermediate tensors.
+}
+
+std::vector<std::shared_ptr<Tensor>>
+CheckpointFunction::Backward(const std::vector<std::shared_ptr<Tensor>> &grad_outputs) {
+    // TODO(zbl): RNG state is not captured yet. Dropout or random ops are not supported.
+
+    CHECK(!saved_inputs_.empty());
+    CHECK_EQ(grad_outputs.size(), 1) << "Checkpoint currently supports single-output forward only.";
+
+    AutocastGuard autocast_guard(ctx_.forward_autocast_context());
+    autograd::EnableGradGuard enable_grad;
+
+    std::vector<std::shared_ptr<Tensor>> detached_inputs;
+    detached_inputs.reserve(saved_inputs_.size());
+    for (size_t i = 0; i < saved_inputs_.size(); ++i) {
+        if (!saved_inputs_[i]) {
+            detached_inputs.push_back(nullptr);
+            continue;
+        }
+        auto detached = saved_inputs_[i]->Detach();
+        detached->set_requires_grad(saved_inputs_requires_grad_[i]);
+        detached->set_is_leaf(true);
+        detached->set_grad_fn(nullptr);
+        detached_inputs.push_back(detached);
+    }
+
+    auto outputs = forward_fn_(detached_inputs);
+    // TODO(zbl): Support multiple-output forward.
+    CHECK_EQ(outputs.size(), 1) << "Checkpoint currently supports single-output forward only.";
+
+    if (grad_outputs[0]) {
+        outputs[0]->Backward(grad_outputs[0]);
+    }
+
+    std::vector<std::shared_ptr<Tensor>> grad_inputs;
+    grad_inputs.reserve(detached_inputs.size());
+    for (const auto &detached : detached_inputs) {
+        if (detached && detached->requires_grad()) {
+            grad_inputs.push_back(detached->grad());
+        } else {
+            grad_inputs.push_back(nullptr);
+        }
+    }
+
+    saved_inputs_.clear();
+    saved_inputs_requires_grad_.clear();
+    return grad_inputs;
+}
+
+} // namespace
+
+std::vector<std::shared_ptr<Tensor>> Checkpoint(const CheckpointForwardFn &forward_fn,
+                                                const std::vector<std::shared_ptr<Tensor>> &inputs, bool use_reentrant,
+                                                bool preserve_rng_state, bool determinism_check, bool early_stop) {
+    if (preserve_rng_state) {
+        // TODO(zbl): Preserve and restore RNG state for CPU/CUDA.
+        LOG(FATAL) << "Checkpoint RNG state preservation is not implemented. Pass preserve_rng_state=false.";
+    }
+    if (!autograd::GradMode::IsEnabled()) {
+        return forward_fn(inputs);
+    }
+
+    if (use_reentrant) {
+        const bool any_requires_grad = std::any_of(
+            inputs.begin(), inputs.end(), [](const std::shared_ptr<Tensor> &t) { return t && t->requires_grad(); });
+        if (!any_requires_grad) {
+            return forward_fn(inputs);
+        }
+        auto func = std::make_shared<CheckpointFunction>(forward_fn);
+        return func->Apply(inputs);
+    }
+
+    auto frame = std::make_shared<CheckpointFrame>();
+    frame->forward_fn = forward_fn;
+    frame->early_stop = early_stop;
+    frame->determinism_check = determinism_check;
+    frame->inputs.reserve(inputs.size());
+    frame->inputs_requires_grad.reserve(inputs.size());
+    for (const auto &input : inputs) {
+        if (input) {
+            frame->inputs.push_back(input->Detach());
+            frame->inputs_requires_grad.push_back(input->requires_grad());
+        } else {
+            frame->inputs.push_back(nullptr);
+            frame->inputs_requires_grad.push_back(false);
+        }
+    }
+    frame->autocast_context = GetCurrentAutocastContext();
+
+    autograd::FunctionCtx::SavedTensorHooks placeholder_hooks;
+    placeholder_hooks.pack = [frame](const std::shared_ptr<Tensor> &tensor) -> std::shared_ptr<void> {
+        auto holder = std::make_shared<SavedTensorHolder>();
+        holder->frame = frame;
+        holder->index = frame->forward_metas.size();
+        frame->weak_holders.push_back(holder);
+        if (tensor) {
+            SavedTensorMeta meta;
+            meta.dims = tensor->Dims();
+            meta.dtype = tensor->Dtype();
+            meta.device_type = tensor->GetDevice().type();
+            frame->forward_metas.push_back(std::move(meta));
+        } else {
+            frame->forward_metas.push_back({});
+        }
+        return holder;
+    };
+    placeholder_hooks.unpack = [](const std::shared_ptr<void> &state) -> std::shared_ptr<Tensor> {
+        auto holder = std::static_pointer_cast<SavedTensorHolder>(state);
+        auto frame = holder->frame;
+        const int64_t gid = frame->GetOrCreateGid();
+        if (!frame->recomputed[gid]) {
+            frame->Recompute(gid);
+        }
+        auto it = holder->handles.find(gid);
+        if (it == holder->handles.end() || !it->second) {
+            LOG(FATAL) << "Checkpoint: unpack called more than once for index " << holder->index << ".";
+        }
+        auto recomputed = holder->tensor;
+        // SaveVariables bypasses hooks for nullptr entries. Every holder here
+        // represents a real tensor, including zero-dimensional CPU FP32 tensors.
+        CHECK(recomputed) << "Checkpoint: recomputed null tensor for saved non-null entry.";
+        // Release recomputed tensor as soon as it's unpacked to reduce peak memory.
+        holder->tensor.reset();
+        it->second = false;
+        return recomputed;
+    };
+
+    autograd::FunctionCtx::SavedTensorHooksGuard guard(std::move(placeholder_hooks));
+    return forward_fn(inputs);
+}
+
+} // namespace infini_train::utils::checkpoint
diff --git a/scripts/test_config.json b/scripts/test_config.json
index aaf537135..df11bfd79 100644
--- a/scripts/test_config.json
+++ b/scripts/test_config.json
@@ -18,8 +18,8 @@
         "RUN_PROFILE_TEST": "true",
         "MIXTRAL_INPUT_BIN": "/data1/shared/InfiniTrain-dev/data/llmc/llama3/tinyshakespeare/tiny_shakespeare_train.bin",
         "MIXTRAL_LLMC_FILEPATH": "/data1/shared/InfiniTrain-dev/data/llmc/mixtral/mixtral_megatron_export.bin",
-        "GPT2_TEST_GROUPS": "basic,zero,lora,checkpoint,lr_scheduler,8_proc",
-        "LLAMA3_TEST_GROUPS": "basic,zero,lora,checkpoint,lr_scheduler,8_proc",
+        "GPT2_TEST_GROUPS": "basic,zero,lora,checkpoint,lr_scheduler,8_proc,recompute",
+        "LLAMA3_TEST_GROUPS": "basic,zero,lora,checkpoint,lr_scheduler,8_proc,recompute",
         "QWEN3_TEST_GROUPS": "qwen3",
         "MIXTRAL_TEST_GROUPS": "moe"
     },
@@ -396,7 +396,7 @@
                         "batch_size": 10,
                         "total_batch_size": 5120,
                         "zero_stage": 2,
-                        "learning_rate": 0.00001,
+                        "learning_rate": 1e-05,
                         "lr_decay_style": "none"
                     }
                 },
@@ -409,8 +409,8 @@
                         "batch_size": 40,
                         "total_batch_size": 5120,
                         "tensor_parallel": 4,
-                        "learning_rate": 0.00001,
-                        "min_lr": 0.000001,
+                        "learning_rate": 1e-05,
+                        "min_lr": 1e-06,
                         "lr_decay_style": "constant",
                         "lr_warmup_iters": 0,
                         "lr_decay_iters": 0
@@ -427,8 +427,8 @@
                         "tensor_parallel": 4,
                         "sequence_parallel": true,
                         "zero_stage": 1,
-                        "learning_rate": 0.00001,
-                        "min_lr": 0.000001,
+                        "learning_rate": 1e-05,
+                        "min_lr": 1e-06,
                         "lr_decay_style": "linear",
                         "lr_warmup_iters": 2,
                         "lr_warmup_init": 0.0,
@@ -444,8 +444,8 @@
                         "batch_size": 10,
                         "total_batch_size": 5120,
                         "pipeline_parallel": 8,
-                        "learning_rate": 0.00001,
-                        "min_lr": 0.000001,
+                        "learning_rate": 1e-05,
+                        "min_lr": 1e-06,
                         "lr_decay_style": "cosine",
                         "lr_warmup_iters": 2,
                         "lr_warmup_init": 0.0,
@@ -462,8 +462,8 @@
                         "total_batch_size": 5120,
                         "pipeline_parallel": 4,
                         "virtual_pipeline_parallel": 2,
-                        "learning_rate": 0.00001,
-                        "min_lr": 0.000001,
+                        "learning_rate": 1e-05,
+                        "min_lr": 1e-06,
                         "lr_decay_style": "inverse-square-root",
                         "lr_warmup_iters": 2,
                         "lr_warmup_init": 0.0,
@@ -483,8 +483,8 @@
                         "pipeline_parallel": 2,
                         "virtual_pipeline_parallel": 2,
                         "zero_stage": 1,
-                        "learning_rate": 0.00001,
-                        "min_lr": 0.000001,
+                        "learning_rate": 1e-05,
+                        "min_lr": 1e-06,
                         "lr_decay_style": "cosine",
                         "lr_warmup_iters": 2,
                         "lr_warmup_init": 0.0,
@@ -499,8 +499,8 @@
                         "num_iteration": 10,
                         "batch_size": 10,
                         "total_batch_size": 5120,
-                        "learning_rate": 0.00001,
-                        "min_lr": 0.000001,
+                        "learning_rate": 1e-05,
+                        "min_lr": 1e-06,
                         "lr_decay_style": "linear",
                         "lr_warmup_iters": 2,
                         "lr_warmup_init": 0.0,
@@ -517,8 +517,8 @@
                         "total_batch_size": 5120,
                         "tensor_parallel": 4,
                         "zero_stage": 1,
-                        "learning_rate": 0.00001,
-                        "min_lr": 0.000001,
+                        "learning_rate": 1e-05,
+                        "min_lr": 1e-06,
                         "lr_decay_style": "inverse-square-root",
                         "lr_warmup_iters": 2,
                         "lr_warmup_init": 0.0,
@@ -535,8 +535,8 @@
                         "total_batch_size": 5120,
                         "tensor_parallel": 4,
                         "sequence_parallel": true,
-                        "learning_rate": 0.00001,
-                        "min_lr": 0.000001,
+                        "learning_rate": 1e-05,
+                        "min_lr": 1e-06,
                         "lr_decay_style": "constant",
                         "lr_warmup_iters": 0,
                         "lr_decay_iters": 10
@@ -554,12 +554,237 @@
                         "sequence_parallel": true,
                         "pipeline_parallel": 2,
                         "virtual_pipeline_parallel": 2,
-                        "learning_rate": 0.00001,
+                        "learning_rate": 1e-05,
                         "lr_decay_style": "none"
                     }
                 }
             ]
         },
+        {
+            "tag": "recompute",
+            "tests": [
+                {
+                    "id": "2_recompute_uniform2",
+                    "args": {
+                        "dtype": "float32",
+                        "num_iteration": 10,
+                        "batch_size": 80,
+                        "total_batch_size": 5120,
+                        "activation_recompute": true,
+                        "recompute_granularity": "full",
+                        "recompute_method": "uniform",
+                        "recompute_num_layers": 2
+                    }
+                },
+                {
+                    "id": "2_recompute_block2_bfloat16",
+                    "args": {
+                        "dtype": "bfloat16",
+                        "num_iteration": 10,
+                        "batch_size": 80,
+                        "total_batch_size": 5120,
+                        "activation_recompute": true,
+                        "recompute_granularity": "full",
+                        "recompute_method": "block",
+                        "recompute_num_layers": 2
+                    }
+                },
+                {
+                    "id": "3_recompute_uniform1",
+                    "args": {
+                        "dtype": "float32",
+                        "nthread_per_process": 8,
+                        "num_iteration": 10,
+                        "batch_size": 10,
+                        "total_batch_size": 5120,
+                        "activation_recompute": true,
+                        "recompute_granularity": "full",
+                        "recompute_method": "uniform",
+                        "recompute_num_layers": 1
+                    }
+                },
+                {
+                    "id": "3_recompute_block2_bfloat16",
+                    "args": {
+                        "dtype": "bfloat16",
+                        "nthread_per_process": 8,
+                        "num_iteration": 10,
+                        "batch_size": 10,
+                        "total_batch_size": 5120,
+                        "activation_recompute": true,
+                        "recompute_granularity": "full",
+                        "recompute_method": "block",
+                        "recompute_num_layers": 2
+                    }
+                },
+                {
+                    "id": "6_recompute_uniform1",
+                    "args": {
+                        "dtype": "float32",
+                        "nthread_per_process": 8,
+                        "num_iteration": 10,
+                        "batch_size": 10,
+                        "total_batch_size": 5120,
+                        "pipeline_parallel": 8,
+                        "activation_recompute": true,
+                        "recompute_granularity": "full",
+                        "recompute_method": "uniform",
+                        "recompute_num_layers": 1
+                    }
+                },
+                {
+                    "id": "6_recompute_block1_bfloat16",
+                    "args": {
+                        "dtype": "bfloat16",
+                        "nthread_per_process": 8,
+                        "num_iteration": 10,
+                        "batch_size": 10,
+                        "total_batch_size": 5120,
+                        "pipeline_parallel": 8,
+                        "activation_recompute": true,
+                        "recompute_granularity": "full",
+                        "recompute_method": "block",
+                        "recompute_num_layers": 1
+                    }
+                },
+                {
+                    "id": "7_recompute_uniform2",
+                    "args": {
+                        "dtype": "float32",
+                        "nthread_per_process": 4,
+                        "num_iteration": 10,
+                        "batch_size": 10,
+                        "total_batch_size": 5120,
+                        "pipeline_parallel": 4,
+                        "virtual_pipeline_parallel": 2,
+                        "activation_recompute": true,
+                        "recompute_granularity": "full",
+                        "recompute_method": "uniform",
+                        "recompute_num_layers": 2
+                    }
+                },
+                {
+                    "id": "7_recompute_block2_bfloat16",
+                    "args": {
+                        "dtype": "bfloat16",
+                        "nthread_per_process": 4,
+                        "num_iteration": 10,
+                        "batch_size": 10,
+                        "total_batch_size": 5120,
+                        "pipeline_parallel": 4,
+                        "virtual_pipeline_parallel": 2,
+                        "activation_recompute": true,
+                        "recompute_granularity": "full",
+                        "recompute_method": "block",
+                        "recompute_num_layers": 2
+                    }
+                },
+                {
+                    "id": "8_recompute_uniform3",
+                    "args": {
+                        "dtype": "float32",
+                        "nthread_per_process": 8,
+                        "num_iteration": 10,
+                        "batch_size": 40,
+                        "total_batch_size": 5120,
+                        "tensor_parallel": 2,
+                        "sequence_parallel": true,
+                        "pipeline_parallel": 2,
+                        "virtual_pipeline_parallel": 2,
+                        "activation_recompute": true,
+                        "recompute_granularity": "full",
+                        "recompute_method": "uniform",
+                        "recompute_num_layers": 3
+                    }
+                },
+                {
+                    "id": "8_recompute_block2_bfloat16",
+                    "args": {
+                        "dtype": "bfloat16",
+                        "nthread_per_process": 8,
+                        "num_iteration": 10,
+                        "batch_size": 40,
+                        "total_batch_size": 5120,
+                        "tensor_parallel": 2,
+                        "sequence_parallel": true,
+                        "pipeline_parallel": 2,
+                        "virtual_pipeline_parallel": 2,
+                        "activation_recompute": true,
+                        "recompute_granularity": "full",
+                        "recompute_method": "block",
+                        "recompute_num_layers": 2
+                    }
+                },
+                {
+                    "id": "3_recompute_uniform1_distopt",
+                    "args": {
+                        "dtype": "float32",
+                        "nthread_per_process": 8,
+                        "num_iteration": 10,
+                        "batch_size": 10,
+                        "total_batch_size": 5120,
+                        "zero_stage": 1,
+                        "activation_recompute": true,
+                        "recompute_granularity": "full",
+                        "recompute_method": "uniform",
+                        "recompute_num_layers": 1
+                    }
+                },
+                {
+                    "id": "3_recompute_block2_distopt_bfloat16",
+                    "args": {
+                        "dtype": "bfloat16",
+                        "nthread_per_process": 8,
+                        "num_iteration": 10,
+                        "batch_size": 10,
+                        "total_batch_size": 5120,
+                        "zero_stage": 1,
+                        "activation_recompute": true,
+                        "recompute_granularity": "full",
+                        "recompute_method": "block",
+                        "recompute_num_layers": 2
+                    }
+                },
+                {
+                    "id": "8_recompute_uniform3_distopt",
+                    "args": {
+                        "dtype": "float32",
+                        "nthread_per_process": 8,
+                        "num_iteration": 10,
+                        "batch_size": 40,
+                        "total_batch_size": 5120,
+                        "tensor_parallel": 2,
+                        "sequence_parallel": true,
+                        "pipeline_parallel": 2,
+                        "virtual_pipeline_parallel": 2,
+                        "zero_stage": 1,
+                        "activation_recompute": true,
+                        "recompute_granularity": "full",
+                        "recompute_method": "uniform",
+                        "recompute_num_layers": 3
+                    }
+                },
+                {
+                    "id": "8_recompute_block2_distopt_bfloat16",
+                    "args": {
+                        "dtype": "bfloat16",
+                        "nthread_per_process": 8,
+                        "num_iteration": 10,
+                        "batch_size": 40,
+                        "total_batch_size": 5120,
+                        "tensor_parallel": 2,
+                        "sequence_parallel": true,
+                        "pipeline_parallel": 2,
+                        "virtual_pipeline_parallel": 2,
+                        "zero_stage": 1,
+                        "activation_recompute": true,
+                        "recompute_granularity": "full",
+                        "recompute_method": "block",
+                        "recompute_num_layers": 2
+                    }
+                }
+            ]
+        },
         {
             "tag": "lora",
             "tests": [
diff --git a/tests/autograd/test_activation_checkpoint.cc b/tests/autograd/test_activation_checkpoint.cc
new file mode 100644
index 000000000..abb52a7ca
--- /dev/null
+++ b/tests/autograd/test_activation_checkpoint.cc
@@ -0,0 +1,144 @@
+#include <memory>
+#include <stdexcept>
+#include <vector>
+
+#include "gtest/gtest.h"
+
+#include "infini_train/include/autograd/function.h"
+#include "infini_train/include/autograd/grad_mode.h"
+#include "infini_train/include/tensor.h"
+#include "infini_train/include/utils/checkpoint.h"
+
+#include "tests/common/test_utils.h"
+
+using namespace infini_train;
+
+namespace {
+// A custom Function may implement its forward using other autograd Functions.
+// Both saved tensors are needed to detect a shifted pack/unpack sequence.
+class NestedSquare : public autograd::Function {
+public:
+    explicit NestedSquare(bool nested_setup = false) : Function("NestedSquare"), nested_setup_(nested_setup) {}
+
+    std::vector<std::shared_ptr<Tensor>> Forward(const std::vector<std::shared_ptr<Tensor>> &inputs) override {
+        return {inputs[0] * inputs[0]};
+    }
+
+    void SetupContext(const std::vector<std::shared_ptr<Tensor>> &inputs,
+                      const std::vector<std::shared_ptr<Tensor>> &outputs) override {
+        auto square = nested_setup_ ? inputs[0] * inputs[0] : outputs[0];
+        ctx_.SaveForBackward({inputs[0], square});
+    }
+
+    std::vector<std::shared_ptr<Tensor>> Backward(const std::vector<std::shared_ptr<Tensor>> &grads) override {
+        auto saved = ctx_.GetSavedTensors();
+        // Test inputs are nonzero. This equals 2*x and uses both saved tensors.
+        return {grads[0] * (saved[1] / saved[0]) * 2.0f};
+    }
+
+private:
+    bool nested_setup_;
+};
+
+class SaveScalarAndNull : public autograd::Function {
+public:
+    SaveScalarAndNull() : Function("SaveScalarAndNull") {}
+
+    std::vector<std::shared_ptr<Tensor>> Forward(const std::vector<std::shared_ptr<Tensor>> &inputs) override {
+        auto out = std::make_shared<Tensor>(inputs[0]->Dims(), inputs[0]->Dtype(), inputs[0]->GetDevice());
+        out->CopyFrom(inputs[0]);
+        return {out};
+    }
+
+    void SetupContext(const std::vector<std::shared_ptr<Tensor>> &inputs,
+                      const std::vector<std::shared_ptr<Tensor>> &) override {
+        ctx_.SaveForBackward({nullptr, inputs[0], nullptr});
+    }
+
+    std::vector<std::shared_ptr<Tensor>> Backward(const std::vector<std::shared_ptr<Tensor>> &grads) override {
+        auto saved = ctx_.GetSavedTensors();
+        EXPECT_EQ(saved.size(), 3);
+        EXPECT_EQ(saved.at(0), nullptr);
+        EXPECT_EQ(saved.at(2), nullptr);
+        EXPECT_TRUE(saved.at(1)->Dims().empty());
+        test::ExpectTensorFloatEqual(saved.at(1), 3.0f);
+        return {grads[0]};
+    }
+};
+} // namespace
+
+class ActivationCheckpointTest : public test::InfiniTrainTest {};
+
+TEST_P(ActivationCheckpointTest, NestedFunctionsPreserveSavedTensorsAndGradients) {
+    for (bool nested_setup : {false, true}) {
+        for (bool early_stop : {false, true}) {
+            SCOPED_TRACE(::testing::Message() << "nested_setup=" << nested_setup << " early_stop=" << early_stop);
+            const float data[] = {2.0f, 3.0f};
+            auto input = std::make_shared<Tensor>(data, std::vector<int64_t>{2}, DataType::kFLOAT32, GetDevice());
+            input->RequiresGrad();
+            auto forward = [nested_setup](const auto &inputs) {
+                return std::make_shared<NestedSquare>(nested_setup)->Apply(inputs);
+            };
+            auto outputs = utils::checkpoint::Checkpoint(forward, {input}, false, false, true, early_stop);
+            test::ExpectTensorFloatEqual(outputs[0], {4.0f, 9.0f});
+            auto gradient = std::make_shared<Tensor>(input->Dims(), DataType::kFLOAT32, GetDevice());
+            gradient->Fill(1.0f);
+            outputs[0]->Backward(gradient);
+            test::ExpectTensorNear(input->grad(), {4.0f, 6.0f}, 1e-6f);
+            EXPECT_TRUE(autograd::GradMode::IsEnabled());
+            EXPECT_FALSE(autograd::GradMode::PropagateRequiresGrad());
+        }
+    }
+}
+
+TEST_P(ActivationCheckpointTest, ScalarAndNullSavedEntriesRemainDistinct) {
+    for (bool determinism_check : {false, true}) {
+        auto input = std::make_shared<Tensor>(std::vector<int64_t>{}, DataType::kFLOAT32, GetDevice(), true);
+        input->Fill(3.0f);
+        auto forward = [](const auto &inputs) { return std::make_shared<SaveScalarAndNull>()->Apply(inputs); };
+        auto outputs = utils::checkpoint::Checkpoint(forward, {input}, false, false, determinism_check);
+        outputs[0]->Backward();
+        test::ExpectTensorFloatEqual(input->grad(), 1.0f);
+    }
+}
+
+TEST_P(ActivationCheckpointTest, PropagationDoesNotLeakIntoNestedFunctions) {
+    for (bool grad_enabled : {false, true}) {
+        for (bool propagate : {false, true}) {
+            SCOPED_TRACE(::testing::Message() << "grad_enabled=" << grad_enabled << " propagate=" << propagate);
+            auto input = std::make_shared<Tensor>(std::vector<int64_t>{2}, DataType::kFLOAT32, GetDevice(), true);
+            input->Fill(2.0f);
+            size_t packed = 0;
+            autograd::FunctionCtx::SavedTensorHooks hooks;
+            hooks.pack = [&packed](const auto &tensor) -> std::shared_ptr<void> {
+                ++packed;
+                return tensor;
+            };
+            hooks.unpack = [](const auto &state) { return std::static_pointer_cast<Tensor>(state); };
+            autograd::FunctionCtx::SavedTensorHooksGuard hooks_guard(std::move(hooks));
+            autograd::NoGradGuard no_grad;
+            autograd::GradMode::SetEnabled(grad_enabled);
+            autograd::PropagateRequiresGradGuard propagation(propagate);
+            auto outputs = std::make_shared<NestedSquare>(true)->Apply({input});
+            EXPECT_EQ(packed, grad_enabled || propagate ? 2 : 0);
+            EXPECT_EQ(outputs[0]->requires_grad(), grad_enabled || propagate);
+            EXPECT_EQ(outputs[0]->grad_fn() != nullptr, grad_enabled);
+        }
+    }
+}
+
+TEST(ActivationCheckpointGuardTest, RestoresPropagationAfterException) {
+    EXPECT_FALSE(autograd::GradMode::PropagateRequiresGrad());
+    {
+        autograd::PropagateRequiresGradGuard enabled;
+        try {
+            autograd::PropagateRequiresGradGuard disabled(false);
+            EXPECT_FALSE(autograd::GradMode::PropagateRequiresGrad());
+            throw std::runtime_error("test");
+        } catch (const std::runtime_error &) {}
+        EXPECT_TRUE(autograd::GradMode::PropagateRequiresGrad());
+    }
+    EXPECT_FALSE(autograd::GradMode::PropagateRequiresGrad());
+}
+
+INFINI_TRAIN_REGISTER_TEST(ActivationCheckpointTest);