Feat: Activation Checkpointing

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

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

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) {

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 ==========

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

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

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);

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 ==========

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

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);
     }

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

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

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

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