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

@@ -79,6 +79,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");
@@ -212,22 +218,25 @@ void Train(const nn::parallel::Rank &rank) {
     nn::TransformerConfig model_config = gpt2::GPT2Config();
     std::shared_ptr<nn::Module> model = nullptr;
 
-    if (!FLAGS_llmc_filepath.empty()) {
-        model = gpt2::LoadFromLLMC(FLAGS_llmc_filepath);
-    } else if (kModelToConfigs.count(FLAGS_model)) {
+    if (FLAGS_llmc_filepath.empty() && kModelToConfigs.count(FLAGS_model)) {
         model_config = kModelToConfigs.at(FLAGS_model);
         gpt2::SanitizeGPT2Config(model_config);
+    }
+    nn::SetActivationRecomputeConfig(&model_config, 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, nn::GetActivationRecomputeOptions(model_config));
+    } else {
         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;
     }
@@ -82,6 +83,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

@@ -77,6 +77,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");
@@ -198,13 +204,17 @@ void Train(const nn::parallel::Rank &rank) {
 
     nn::TransformerConfig model_config = llama3::LLaMA3Config();
     std::shared_ptr<nn::Module> model = nullptr;
+    nn::SetActivationRecomputeConfig(&model_config, 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, nn::GetActivationRecomputeOptions(model_config));
     } else {
         llama3::SanitizeLLaMA3Config(model_config);
         model = std::make_shared<nn::TransformerModel>(model_config);
     }
 
+    CHECK(model) << "LLaMA3 example expects LLaMA3 model.";
+
     model->To(device);
 
     utils::PrecisionChecker::BuildNameMap(model.get());

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});
@@ -101,7 +101,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);

infini_train/include/autocast.h

@@ -143,6 +143,38 @@ struct AutocastContext {
 // Global thread-local storage for autocast context
 inline thread_local AutocastContext tls_autocast_context;
 
+// Lightweight snapshot for autocast state
+struct AutocastState {
+    bool enabled = false;
+    Device::DeviceType device_type = Device::DeviceType::kCPU;
+    DataType autocast_dtype = DataType::kBFLOAT16;
+};
+
+inline AutocastState GetAutocastState() {
+    return AutocastState{tls_autocast_context.enabled, tls_autocast_context.device_type,
+                         tls_autocast_context.autocast_dtype};
+}
+
+inline void SetAutocastState(const AutocastState &state) {
+    tls_autocast_context.enabled = state.enabled;
+    tls_autocast_context.device_type = state.device_type;
+    tls_autocast_context.autocast_dtype = state.autocast_dtype;
+}
+
+class AutocastStateGuard {
+public:
+    explicit AutocastStateGuard(const AutocastState &state) : saved_state_(GetAutocastState()) {
+        SetAutocastState(state);
+    }
+    ~AutocastStateGuard() { SetAutocastState(saved_state_); }
+
+    AutocastStateGuard(const AutocastStateGuard &) = delete;
+    AutocastStateGuard &operator=(const AutocastStateGuard &) = delete;
+
+private:
+    AutocastState saved_state_;
+};
+
 // RAII guard to enable/disable autocast in a scope
 class AutocastGuard {
 public:

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,19 @@ 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.
+class PropagateRequiresGradGuard {
+public:
+    PropagateRequiresGradGuard() : prev_(GradMode::PropagateRequiresGrad()) {
+        GradMode::SetPropagateRequiresGrad(true);
+    }
+    ~PropagateRequiresGradGuard() { GradMode::SetPropagateRequiresGrad(prev_); }
+
+private:
+    bool prev_;
+};
+
 } // namespace infini_train::autograd

infini_train/include/nn/modules/module.h

@@ -65,7 +65,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
@@ -108,6 +108,9 @@ class Module : public std::enable_shared_from_this<Module> {
     std::vector<ModulePreHook> backward_pre_hooks_;
     std::vector<ModulePostHook> backward_post_hooks_;
 
+    // Full module call protocol: hooks followed by Forward.
+    std::vector<std::shared_ptr<Tensor>> Call(const std::vector<std::shared_ptr<Tensor>> &input_tensors);
+
 private:
     friend std::vector<std::shared_ptr<Module>> parallel::function::Replicate(const std::shared_ptr<Module> &network,
                                                                               const std::vector<Device> &devices);

infini_train/include/nn/modules/transformer/transformer_config.h

@@ -2,6 +2,7 @@
 
 #include <cstdint>
 #include <optional>
+#include <string_view>
 
 namespace infini_train::nn {
 
@@ -20,6 +21,24 @@ enum class NormType {
     kRMSNorm    // RMSNorm
 };
 
+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
@@ -51,12 +70,23 @@ 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
     int64_t max_gen_batch_size = 4; // max batch size during inference
 
     bool UseGQA() const;
     int GetChunkSize() const;
+    bool RecomputeEnabled() const;
 };
+
+void SetActivationRecomputeConfig(TransformerConfig *config, bool enabled, std::string_view granularity,
+                                  std::string_view method, int64_t num_layers);
+ActivationRecomputeOptions GetActivationRecomputeOptions(const TransformerConfig &config);
+void ApplyActivationRecomputeOptions(TransformerConfig *config, const ActivationRecomputeOptions &options);
 } // namespace infini_train::nn

infini_train/include/utils/checkpoint.h

@@ -0,0 +1,21 @@
+#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>> &)>;
+
+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