diff --git a/.gitmodules b/.gitmodules
index 470cf466..d33d0a64 100644
--- a/.gitmodules
+++ b/.gitmodules
@@ -7,3 +7,6 @@
 [submodule "third_party/eigen"]
 	path = third_party/eigen
 	url = git@github.com:InfiniTensor/eigen-mirror.git
+[submodule "third_party/googletest"]
+	path = third_party/googletest
+	url = git@github.com:google/googletest.git
diff --git a/CMakeLists.txt b/CMakeLists.txt
index 9ea3529c..7cc7ea84 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -4,6 +4,7 @@ option(USE_CUDA "Support NVIDIA CUDA" OFF)
 option(PROFILE_MODE "ENABLE PROFILE MODE" OFF)
 option(USE_OMP "Use OpenMP as backend for Eigen" ON)
 option(USE_NCCL "Build project for distributed running" ON)
+option(BUILD_TEST "Build InfiniTrain tests" OFF)
 
 project(infini_train VERSION 0.5.0 LANGUAGES CXX)
 
@@ -14,6 +15,19 @@ set(CMAKE_CXX_EXTENSIONS OFF)
 # Generate compile_commands.json
 set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
 
+# ------------------------------------------------------------------------------
+# GoogleTest (submodule)
+# ------------------------------------------------------------------------------
+if(BUILD_TEST)
+  if(NOT EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/third_party/googletest/CMakeLists.txt)
+    message(FATAL_ERROR "googletest submodule not found at third_party/googletest. "
+      "Run: git submodule update --init third_party/googletest")
+  endif()
+  set(gtest_force_shared_crt ON CACHE BOOL "" FORCE)
+  add_subdirectory(third_party/googletest)
+  enable_testing()
+endif()
+
 # ------------------------------------------------------------------------------
 # Third-party deps
 # ------------------------------------------------------------------------------
@@ -48,6 +62,10 @@ endif()
 # Framework core sources (*.cc), excluding cpu kernels (they are built separately)
 file(GLOB_RECURSE SRC ${PROJECT_SOURCE_DIR}/infini_train/src/*.cc)
 list(FILTER SRC EXCLUDE REGEX ".*kernels/cpu/.*")
+if(NOT USE_CUDA)
+  list(FILTER SRC EXCLUDE REGEX ".*runtime/cuda/.*")
+  list(FILTER SRC EXCLUDE REGEX ".*ccl/cuda/.*")
+endif()
 if(NOT USE_NCCL)
   list(FILTER SRC EXCLUDE REGEX ".*infini_train/src/core/ccl/cuda/.*")
 endif()
@@ -195,34 +213,6 @@ add_subdirectory(tools/infini_run)
 set_target_properties(infini_run PROPERTIES RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR})
 
 # Tests
-add_executable(test_hook test/hook/test_hook.cc)
-link_infini_train_exe(test_hook)
-
-add_executable(test_precision_check test/hook/test_precision_check.cc)
-link_infini_train_exe(test_precision_check)
-
-add_executable(test_lora test/lora/test_lora.cc)
-link_infini_train_exe(test_lora)
-
-add_executable(test_transformer_architecture test/transformer/test_transformer_architecture.cc)
-link_infini_train_exe(test_transformer_architecture)
-
-add_executable(test_scalar test/dtype/test_scalar.cc)
-link_infini_train_exe(test_scalar)
-
-add_executable(test_dtype_dispatch test/dtype/test_dtype_dispatch.cc)
-link_infini_train_exe(test_dtype_dispatch)
-
-# Negative compile test: missing dtype registration must fail at compile time.
-set(DTYPE_DISPATCH_COMPILE_FAIL_SOURCE
-  ${PROJECT_SOURCE_DIR}/test/dtype/test_dtype_dispatch_compile_fail.cc)
-
-add_executable(test_dtype_dispatch_compile_fail EXCLUDE_FROM_ALL
-  ${DTYPE_DISPATCH_COMPILE_FAIL_SOURCE}
-)
-
-target_include_directories(test_dtype_dispatch_compile_fail PRIVATE
-  ${PROJECT_SOURCE_DIR}
-)
-
-link_infini_train_exe(test_dtype_dispatch_compile_fail)
+if(BUILD_TEST)
+  add_subdirectory(tests)
+endif()
diff --git a/cmake/test_macros.cmake b/cmake/test_macros.cmake
new file mode 100644
index 00000000..500103b9
--- /dev/null
+++ b/cmake/test_macros.cmake
@@ -0,0 +1,128 @@
+# ============================================================================
+# InfiniTrain Test Macros
+# ============================================================================
+# Unified test configuration interface to reduce boilerplate.
+#
+# Usage:
+#   1. Include this file in tests/CMakeLists.txt
+#   2. Use infini_train_add_test macro to register tests
+#
+# Examples:
+#   infini_train_add_test(
+#     test_tensor_create
+#     SOURCES test_tensor_create.cc
+#     LABELS cpu cuda
+#   )
+# ============================================================================
+
+include_guard(GLOBAL)
+
+# -----------------------------------------------------------------------------
+# Load GoogleTest module (provides gtest_discover_tests)
+# -----------------------------------------------------------------------------
+include(GoogleTest)
+
+# -----------------------------------------------------------------------------
+# infini_train_add_test - Test registration macro
+# -----------------------------------------------------------------------------
+# Features:
+#   1. Create executable target
+#   2. Configure compile options, link libraries, and include paths
+#   3. Use gtest_discover_tests to auto-discover test cases
+#   4. Set test labels
+#
+# Arguments:
+#   SOURCES:    Source file list (required)
+#   LABELS:     Test labels, e.g. "cpu" "cuda" "distributed" (optional, default "cpu")
+#   TEST_FILTER: gtest test filter pattern (optional)
+#
+# Examples:
+#   # Single-label test (one liner)
+#   infini_train_add_test(test_example SOURCES test_example.cc LABELS cpu)
+#
+#   # Filter same binary by label suffix (one call per label)
+#   infini_train_add_test(test_example SOURCES test_example.cc LABELS cpu TEST_FILTER "-*CUDA*")
+#   infini_train_add_test(test_example_cuda SOURCES test_example.cc LABELS cuda TEST_FILTER "*CUDA*")
+# -----------------------------------------------------------------------------
+macro(infini_train_add_test)
+  cmake_parse_arguments(ARG "" "TEST_NAME;TEST_FILTER" "SOURCES;LABELS" ${ARGN})
+
+  if(NOT ARG_TEST_NAME)
+    set(ARG_TEST_NAME ${ARG_UNPARSED_ARGUMENTS})
+  endif()
+
+  if(NOT ARG_SOURCES)
+    message(FATAL_ERROR "infini_train_add_test: TEST_NAME and SOURCES are required")
+  endif()
+
+  # 1. Create executable target
+  add_executable(${ARG_TEST_NAME} ${ARG_SOURCES} $<TARGET_OBJECTS:test_main>)
+
+  # 2. Disable -Werror so tests can run under relaxed warning levels
+  target_compile_options(${ARG_TEST_NAME} PRIVATE -Wno-error)
+
+  # 3. Link Google Test (uses custom main from test_main that initializes GlobalEnv)
+  target_link_libraries(${ARG_TEST_NAME} PRIVATE GTest::gtest)
+
+  # 4. Add include paths
+  target_include_directories(${ARG_TEST_NAME} PRIVATE
+    ${glog_SOURCE_DIR}/src
+  )
+
+  # 5. Link project library (reuses framework linking strategy)
+  link_infini_train_exe(${ARG_TEST_NAME})
+
+  # 6. Auto-discover gtest cases and register as ctest tests
+  set(labels "cpu")
+  if(ARG_LABELS)
+    set(labels "${ARG_LABELS}")
+  endif()
+
+  if(ARG_TEST_FILTER)
+    gtest_discover_tests(${ARG_TEST_NAME}
+      EXTRA_ARGS --gtest_output=xml:%T.xml
+      TEST_FILTER "${ARG_TEST_FILTER}"
+      DISCOVERY_TIMEOUT 10
+      PROPERTIES LABELS "${labels}" TIMEOUT 10
+    )
+  else()
+    gtest_discover_tests(${ARG_TEST_NAME}
+      EXTRA_ARGS --gtest_output=xml:%T.xml
+      PROPERTIES LABELS "${labels}" TIMEOUT 10
+    )
+  endif()
+endmacro()
+
+# -----------------------------------------------------------------------------
+# infini_train_add_test_suite - Register cpu/cuda/distributed targets in one call
+# -----------------------------------------------------------------------------
+# Calls infini_train_add_test three times (or fewer) with the correct
+# TEST_FILTER and LABELS derived from the label list.
+#
+# Arguments:
+#   <name>   Base name; each target is named <name>_<label>
+#   SOURCES  Source file list (required)
+#   LABELS   Subset of {cpu cuda} (optional, default: both)
+#
+# Examples:
+#   infini_train_add_test_suite(test_tensor SOURCES ${TENSOR_TEST_SOURCES})
+#   infini_train_add_test_suite(test_lora   SOURCES test_lora.cc LABELS cpu)
+# -----------------------------------------------------------------------------
+macro(infini_train_add_test_suite)
+  cmake_parse_arguments(SUITE "" "" "SOURCES;LABELS" ${ARGN})
+  set(_suite_name ${SUITE_UNPARSED_ARGUMENTS})
+
+  if(NOT SUITE_LABELS)
+    set(SUITE_LABELS cpu cuda)
+  endif()
+
+  foreach(_label IN LISTS SUITE_LABELS)
+    string(TOUPPER ${_label} _label_upper)
+    set(_filter "${_label_upper}/*")
+    infini_train_add_test(${_suite_name}_${_label}
+      SOURCES ${SUITE_SOURCES}
+      LABELS ${_label}
+      TEST_FILTER "${_filter}"
+    )
+  endforeach()
+endmacro()
diff --git a/docs/hook_mechanism.md b/docs/hook_mechanism_design.md
similarity index 100%
rename from docs/hook_mechanism.md
rename to docs/hook_mechanism_design.md
diff --git a/docs/lora_usage.md b/docs/lora_usage_guide.md
similarity index 100%
rename from docs/lora_usage.md
rename to docs/lora_usage_guide.md
diff --git a/docs/test_infrastructure_design.md b/docs/test_infrastructure_design.md
new file mode 100644
index 00000000..8869e3b2
--- /dev/null
+++ b/docs/test_infrastructure_design.md
@@ -0,0 +1,259 @@
+# 测试体系设计
+
+**核心思路：** 测试屏蔽平台差异，底层实例化不同平台，特化需单独处理的测试。
+
+## 1. 架构
+
+```
+tests/
+├── CMakeLists.txt              # 顶层：include 宏 + add_subdirectory
+├── common/
+│   ├── CMakeLists.txt          # header-only interface library
+│   └── test_utils.h            # C++ 基类、skip 宏、填充工具函数
+├── tensor/                     # Tensor 创建 / 拷贝 / 销毁 / 算子
+├── optimizer/                  # Optimizer 创建 / step
+├── autograd/                   # 各 autograd op 的 forward / backward
+├── hook/                       # Module hook + precision check
+├── lora/                       # LoRA 相关
+├── dtype/                      # Scalar / dtype dispatch + 编译期负面测试
+└── transformer/                # Transformer 架构测试
+
+cmake/
+└── test_macros.cmake           # CMake 宏：infini_train_add_test / infini_train_add_test_suite
+```
+
+### 核心设计：设备参数化
+
+测试不区分 CPU / CUDA 平台。一个测试定义通过 GTest 参数化自动在所有可用设备上运行：
+
+- `INFINI_TRAIN_REGISTER_TEST(TestName)` — 注册 CPU + CUDA 两个实例
+
+无 GPU 时 CUDA 实例在注册阶段直接跳过（不会出现在测试列表里），并打印一条 `LOG(INFO)` 提示。
+
+### 基类层次
+
+| 基类 | 用途 | 提供的能力 |
+|------|------|-----------|
+| `InfiniTrainTest` | 通用参数化测试 | `GetDevice()`（当前参数化的 `Device`） |
+
+测试中的张量直接通过 `Tensor` 构造接口创建：
+
+```cpp
+auto t = std::make_shared<Tensor>(shape, DataType::kFLOAT32, GetDevice());
+auto g = std::make_shared<Tensor>(shape, DataType::kFLOAT32, GetDevice(), /*requires_grad=*/true);
+t->Fill(1.0f);                   // 常量填充（framework 内置 API）
+```
+
+### 跳过特定平台
+
+这些宏函数涉及到了具体平台，用来针对性检验或跳过某些测试样例。
+
+在个别测试内部按需跳过：
+
+```cpp
+// 跳过 CPU 实例（用于硬编码加速器设备的测试，未来新平台仍会运行）
+SKIP_CPU();
+
+// 只在 CPU 实例运行（用于硬编码 CPU 设备的测试）
+ONLY_CPU();
+
+// 只在 CUDA 实例运行（用于硬编码 CUDA 设备的测试）
+ONLY_CUDA();
+
+// 需要 ≥n 个加速器设备
+REQUIRE_MIN_DEVICES(n);
+```
+
+### CMake 宏
+
+`test_macros.cmake` 提供两个宏减少 CMakeLists 样板：
+
+- `infini_train_add_test(name SOURCES ... LABELS ...)` — 创建可执行文件、链接 GTest + 框架库、用 `gtest_discover_tests` 自动发现用例
+- `infini_train_add_test_suite(name SOURCES ... LABELS ...)` — 按 label（cpu/cuda）拆分为多个 CTest target，通过 `TEST_FILTER` 路由到对应的参数化前缀（`CPU/*`, `CUDA/*`）
+
+## 2. 构建与运行
+
+```bash
+# 构建（从 build 目录）
+cmake -DBUILD_TEST=ON -DUSE_CUDA=ON ..
+make -j$(nproc)
+
+# 运行全部测试
+ctest --output-on-failure
+
+# 只运行 CPU 测试
+ctest -L cpu --output-on-failure
+
+# 只运行 CUDA 测试
+ctest -L cuda --output-on-failure
+
+# 运行单个测试二进制（看完整 GTest 输出）
+./test_tensor_cpu
+./test_autograd_cuda
+
+# GTest filter 过滤特定用例
+./test_tensor_cpu --gtest_filter="CPU/TensorCreateTest.*"
+```
+
+无 GPU 机器上 `cmake -DBUILD_TEST=ON -DUSE_CUDA=OFF ..` 即可，CUDA 测试实例不会注册。
+
+## 3. 新增测试
+
+### 3.1 新增 GTest 参数化测试（推荐）
+
+以新增 `tests/foo/` 为例，完整流程：
+
+**Step 1: 创建目录和测试文件**
+
+```bash
+mkdir tests/foo
+```
+
+```cpp
+// tests/foo/test_foo_basic.cc
+#include <gtest/gtest.h>
+#include "infini_train/include/tensor.h"
+#include "tests/common/test_utils.h"
+
+using namespace infini_train;
+
+class FooBasicTest : public infini_train::test::InfiniTrainTest {};
+
+TEST_P(FooBasicTest, CreateTensor) {
+    auto tensor = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice());
+    EXPECT_NE(tensor, nullptr);
+}
+
+TEST_P(FooBasicTest, CUDAOnlyFeature) {
+    SKIP_CPU();
+    // CUDA-specific logic ...
+}
+
+INFINI_TRAIN_REGISTER_TEST(FooBasicTest);
+```
+
+**基类选择：**
+
+所有测试类都继承 `InfiniTrainTest`。需要梯度时，给 `Tensor` 构造传 `requires_grad=true`；需要填充数据时用 `Tensor::Fill`。
+
+**Step 2: 写 CMakeLists.txt**
+
+```cmake
+# tests/foo/CMakeLists.txt
+file(GLOB FOO_SOURCES ${CMAKE_CURRENT_SOURCE_DIR}/test_*.cc)
+
+infini_train_add_test_suite(test_foo
+  SOURCES ${FOO_SOURCES}
+)
+```
+
+`file(GLOB test_*.cc)` 会自动拾取目录下所有测试文件。以后新增 `test_foo_advanced.cc` 只需放入目录，无需改 CMakeLists.txt（重新 cmake 即可）。
+
+**Step 3: 注册到顶层**
+
+在 `tests/CMakeLists.txt` 中添加一行：
+
+```cmake
+add_subdirectory(foo)
+```
+
+**生成的 CTest target：** `test_foo_cpu`、`test_foo_cuda`，可通过 `ctest -L cpu` 等按标签筛选。
+
+### 3.2 在已有目录新增测试文件
+
+所有使用 `file(GLOB ...)` 的目录（autograd、tensor、optimizer、hook、lora）：
+
+1. 把新文件放入对应目录，命名为 `test_*.cc`
+2. 重新 `cmake ..`（glob 在 configure 时求值）
+3. 完成
+
+无需修改任何 CMakeLists.txt。
+
+### 3.3 工具函数速查
+
+`test_utils.h` 提供的常用工具：
+
+| 函数 / 宏 | 用途 |
+|-----------|------|
+| `GetDevice()` | 返回当前参数化的 `Device`（基类方法） |
+| `SKIP_CPU()` | 跳过 CPU 实例 |
+| `ONLY_CPU()` | 只在 CPU 实例运行 |
+| `ONLY_CUDA()` | 只在 CUDA 实例运行 |
+| `REQUIRE_MIN_DEVICES(n)` | 加速器设备不足时 skip |
+
+## 4. 扩展新设备平台（以沐曦 MACA 为例）
+
+当前测试体系围绕 CPU / CUDA 两种设备参数化。如果需要支持新平台（以沐曦 MACA 为例），需要改动以下几处：
+
+### 4.1 框架层：注册新设备类型
+
+在 `infini_train/include/device.h` 的 `DeviceType` 枚举中新增：
+
+```cpp
+enum class DeviceType : int8_t {
+    kCPU = 0,
+    kCUDA = 1,
+    kMACA = 2,  // 新增
+};
+```
+
+### 4.2 测试工具层：`test_utils.h`
+
+1. 新增 MACA 头文件的编译期引入（和 CUDA 对称）：
+
+```cpp
+#if defined(USE_MACA)
+#include <maca_runtime_api.h>
+#endif
+```
+
+2. 新增 `ONLY_MACA()` 宏：
+
+```cpp
+#define ONLY_MACA() \
+    do { if (GetParam() != infini_train::Device::DeviceType::kMACA) { GTEST_SKIP() << "MACA-only test"; } } while (0)
+```
+
+如果希望有类似 `REQUIRE_MIN_DEVICES(n)` 但针对 MACA 的语义，可以按 `USE_CUDA` 分支的写法增加一个新的宏；同理 `USE_MACA` 不开时该宏直接 skip 即可。
+
+### 4.3 注册宏：新增 MACA 实例
+
+沿用 `USE_CUDA` 的做法，未开启编译开关时不注册对应实例：
+
+```cpp
+#if defined(USE_CUDA) && defined(USE_MACA)
+#define INFINI_TRAIN_REGISTER_TEST(TestName)                                    \
+    INSTANTIATE_TEST_SUITE_P(CPU, TestName,                                     \
+        ::testing::Values(infini_train::Device::DeviceType::kCPU));             \
+    INSTANTIATE_TEST_SUITE_P(CUDA, TestName,                                    \
+        ::testing::Values(infini_train::Device::DeviceType::kCUDA));            \
+    INSTANTIATE_TEST_SUITE_P(MACA, TestName,                                    \
+        ::testing::Values(infini_train::Device::DeviceType::kMACA))
+#elif defined(USE_CUDA)
+#define INFINI_TRAIN_REGISTER_TEST(TestName) /* CPU + CUDA, 同现状 */
+#elif defined(USE_MACA)
+#define INFINI_TRAIN_REGISTER_TEST(TestName)                                    \
+    INSTANTIATE_TEST_SUITE_P(CPU, TestName,                                     \
+        ::testing::Values(infini_train::Device::DeviceType::kCPU));             \
+    INSTANTIATE_TEST_SUITE_P(MACA, TestName,                                    \
+        ::testing::Values(infini_train::Device::DeviceType::kMACA))
+#else
+#define INFINI_TRAIN_REGISTER_TEST(TestName) /* 仅 CPU */
+#endif
+```
+
+运行时如果机器上没有对应设备（例如 `USE_MACA` 编译但无 MACA 硬件），让测试直接报错而不是静默跳过。
+
+### 4.4 CMake 层：`test_macros.cmake`
+
+将默认 label 列表从 `cpu cuda` 扩展为 `cpu cuda maca`
+
+### 4.5 检查清单
+
+| 步骤 | 文件 | 改动 |
+|------|------|------|
+| 1 | `device.h` | `DeviceType` 枚举新增 `kMACA` |
+| 2 | `test_utils.h` | 新增 `USE_MACA` 下的 `<maca_runtime_api.h>` 引入、`ONLY_MACA()` 宏 |
+| 3 | `test_utils.h` | `INFINI_TRAIN_REGISTER_TEST` 按 `USE_MACA` 条件新增 MACA 实例 |
+| 4 | `test_macros.cmake` | 将默认 label 列表扩展为 `cpu cuda maca` |
+| 5 | `CMakeLists.txt`（根） | 新增 `USE_MACA` option + MACA SDK 查找 + kernel 编译 |
diff --git a/docs/test_usage_guide.md b/docs/test_usage_guide.md
new file mode 100644
index 00000000..108f995a
--- /dev/null
+++ b/docs/test_usage_guide.md
@@ -0,0 +1,148 @@
+# 测试使用指南
+
+本指南介绍如何在 InfiniTrain 中构建、运行和编写测试。底层设计详见 [test_infrastructure_design.md](test_infrastructure_design.md) 和 [test_lifecycle.md](test_lifecycle.md)。
+
+---
+
+## 构建测试
+
+```bash
+mkdir build && cd build
+cmake -DBUILD_TEST=ON -DUSE_CUDA=ON ..
+make -j$(nproc)
+```
+
+仅构建 CPU（不含 CUDA）：
+
+```bash
+cmake -DBUILD_TEST=ON -DUSE_CUDA=OFF ..
+```
+
+---
+
+## 运行测试
+
+```bash
+# 运行所有测试
+ctest --output-on-failure
+
+# 仅运行 CPU 测试
+ctest -L cpu --output-on-failure
+
+# 仅运行 CUDA 测试
+ctest -L cuda --output-on-failure
+
+# 运行指定测试套件
+ctest -R tensor --output-on-failure
+
+# 直接运行测试二进制，使用 GTest 过滤器
+./tests/tensor/test_tensor_create_cpu --gtest_filter="CPU/TensorCreateTest.*"
+./tests/tensor/test_tensor_create_cuda --gtest_filter="CUDA/TensorCreateTest.*"
+```
+
+---
+
+## 编写新测试
+
+### 1. 创建测试文件
+
+在 `tests/` 下对应子目录中新建文件，例如 `tests/tensor/test_tensor_copy.cc`：
+
+```cpp
+#include "tests/common/test_utils.h"
+
+class TensorCopyTest : public infini_train::test::InfiniTrainTest {};
+
+TEST_P(TensorCopyTest, CopiesDataCorrectly) {
+    auto src = std::make_shared<Tensor>(std::vector<int64_t>{4}, DataType::kFLOAT32, GetDevice());
+    src->Fill(1.0f);
+
+    auto dst = std::make_shared<Tensor>(std::vector<int64_t>{4}, DataType::kFLOAT32, GetDevice());
+    // ... 执行拷贝并断言 ...
+    EXPECT_EQ(dst->Dims(), src->Dims());
+}
+
+INFINI_TRAIN_REGISTER_TEST(TensorCopyTest);
+```
+
+注意事项：
+- 继承 `InfiniTrainTest`。
+- 使用 `TEST_P`，设备参数由框架自动注入。
+- 文件末尾调用 `INFINI_TRAIN_REGISTER_TEST`，自动实例化 CPU 和 CUDA 两个变体。
+
+### 2. 在 CMakeLists.txt 中注册
+
+在子目录的 `CMakeLists.txt`（例如 `tests/tensor/CMakeLists.txt`）中添加：
+
+```cmake
+infini_train_add_test_suite(test_tensor_copy test_tensor_copy.cc)
+```
+
+这会生成两个 CTest 目标：`test_tensor_copy_cpu`（标签 `cpu`）和 `test_tensor_copy_cuda`（标签 `cuda`）。
+
+---
+
+## 基类辅助方法
+
+以下方法在 `TEST_P` 体内可直接调用（通过 `this->` 隐式访问）：
+
+| 方法 | 说明 |
+|---|---|
+| `GetDevice()` | 返回当前测试实例的设备（CPU 或 CUDA） |
+| `tensor->Fill(value)` | 用常量填充张量所有元素（`Tensor` 内置方法） |
+
+张量创建直接使用 `std::make_shared<Tensor>(shape, dtype, GetDevice(), requires_grad)`，`requires_grad` 参数默认 `false`，需要梯度的测试传 `true` 即可。
+
+---
+
+## 条件执行宏
+
+在 `TEST_P` 体内使用，跳过不适用的场景：
+
+```cpp
+TEST_P(MyTest, 仅CUDA) {
+    ONLY_CUDA();   // 在 CPU 上跳过
+    // ...
+}
+
+TEST_P(MyTest, 仅CPU) {
+    ONLY_CPU();    // 在 CUDA 上跳过
+    // ...
+}
+
+TEST_P(MyTest, 需要多卡) {
+    REQUIRE_MIN_DEVICES(2);  // GPU 数量不足时跳过
+    // ...
+}
+```
+
+---
+
+## 添加新测试模块
+
+1. 在 `tests/` 下新建子目录，例如 `tests/mymodule/`。
+2. 在其中创建 `CMakeLists.txt`，对每个测试文件调用 `infini_train_add_test_suite()`。
+3. 在 `tests/CMakeLists.txt` 中添加 `add_subdirectory(mymodule)`。
+
+---
+
+## Autograd 测试
+
+创建启用自动微分的张量时，给 `Tensor` 构造的第四个参数传 `true`：
+
+```cpp
+#include "tests/common/test_utils.h"
+
+class MyOpTest : public infini_train::test::InfiniTrainTest {};
+
+TEST_P(MyOpTest, 前向传播) {
+    auto input = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    input->Fill(1.0f);
+    auto weight = std::make_shared<Tensor>(std::vector<int64_t>{4, 3}, DataType::kFLOAT32, GetDevice(), true);
+    weight->Fill(0.5f);
+    auto output = MyOp(input, weight);
+    EXPECT_NE(output, nullptr);
+}
+
+INFINI_TRAIN_REGISTER_TEST(MyOpTest);
+```
diff --git a/example/gpt2/checkpoint_loader.cc b/example/gpt2/checkpoint_loader.cc
index 57064423..4a7789e9 100644
--- a/example/gpt2/checkpoint_loader.cc
+++ b/example/gpt2/checkpoint_loader.cc
@@ -87,6 +87,7 @@ std::shared_ptr<nn::TransformerModel> LoadFromLLMC(const std::string &filepath)
     gpt2_config.n_layer = n_layer;
     gpt2_config.n_head = n_head;
     gpt2_config.n_embd = n_embd;
+    gpt2::SanitizeGPT2Config(gpt2_config);
     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/config.h b/example/gpt2/config.h
index 978c0ca6..078f9fd5 100644
--- a/example/gpt2/config.h
+++ b/example/gpt2/config.h
@@ -1,5 +1,7 @@
 #pragma once
 
+#include "glog/logging.h"
+
 #include "infini_train/include/nn/modules/transformer/transformer_config.h"
 
 namespace nn = infini_train::nn;
@@ -23,4 +25,18 @@ inline nn::TransformerConfig GPT2Config() {
             .multiple_of = 1};
 }
 
+inline void SanitizeGPT2Config(const nn::TransformerConfig &c) {
+    CHECK_GT(c.block_size, 0);
+    CHECK_GT(c.vocab_size, 0);
+    CHECK_GE(c.vocab_size, c.original_vocab_size);
+    CHECK_GT(c.n_layer, 0);
+    CHECK_GT(c.n_head, 0);
+    CHECK_GT(c.n_embd, 0);
+    CHECK_EQ(c.n_embd % c.n_head, 0) << "n_embd must be divisible by n_head";
+    CHECK_EQ(c.n_kv_head, c.n_head) << "GPT-2 does not use GQA; n_kv_head must equal n_head";
+    CHECK(c.attention_type == nn::AttentionType::kStandard) << "GPT-2 requires standard attention";
+    CHECK(c.activation_type == nn::MLPType::kGELU) << "GPT-2 requires GELU activation";
+    CHECK(c.norm_type == nn::NormType::kLayerNorm) << "GPT-2 requires LayerNorm";
+}
+
 } // namespace gpt2
diff --git a/example/gpt2/main.cc b/example/gpt2/main.cc
index f69736f5..c12b5a28 100644
--- a/example/gpt2/main.cc
+++ b/example/gpt2/main.cc
@@ -190,6 +190,7 @@ void Train(const nn::parallel::Rank &rank) {
         model = gpt2::LoadFromLLMC(FLAGS_llmc_filepath);
     } else if (kModelToConfigs.count(FLAGS_model)) {
         model_config = kModelToConfigs.at(FLAGS_model);
+        gpt2::SanitizeGPT2Config(model_config);
         model = std::make_shared<nn::TransformerModel>(model_config);
     }
 
diff --git a/example/llama3/checkpoint_loader.cc b/example/llama3/checkpoint_loader.cc
index a31d1748..f29bc540 100644
--- a/example/llama3/checkpoint_loader.cc
+++ b/example/llama3/checkpoint_loader.cc
@@ -80,6 +80,7 @@ std::shared_ptr<nn::TransformerModel> LoadFromLLMC(const std::string &filepath)
     llama3_config.use_scaled_rope = static_cast<bool>(use_scaled_rope);
     llama3_config.norm_eps = norm_eps;
     llama3_config.max_gen_batch_size = max_gen_bs;
+    llama3::SanitizeLLaMA3Config(llama3_config);
     auto llama3 = std::make_shared<nn::TransformerModel>(llama3_config);
 
     // ========== pp_size：num_stages; vpp_size: num_chunks_per_stage ==========
diff --git a/example/llama3/config.h b/example/llama3/config.h
index a9eef863..6bc9124d 100644
--- a/example/llama3/config.h
+++ b/example/llama3/config.h
@@ -1,5 +1,7 @@
 #pragma once
 
+#include "glog/logging.h"
+
 #include "infini_train/include/nn/modules/transformer/transformer_config.h"
 
 namespace nn = infini_train::nn;
@@ -22,4 +24,24 @@ inline nn::TransformerConfig LLaMA3Config() {
             .ffn_dim_multiplier = 1.5f,
             .multiple_of = 256};
 }
+
+inline void SanitizeLLaMA3Config(const nn::TransformerConfig &c) {
+    CHECK_GT(c.block_size, 0);
+    CHECK_GT(c.vocab_size, 0);
+    CHECK_GE(c.vocab_size, c.original_vocab_size);
+    CHECK_GT(c.n_layer, 0);
+    CHECK_GT(c.n_head, 0);
+    CHECK_GT(c.n_kv_head, 0);
+    CHECK_LE(c.n_kv_head, c.n_head);
+    CHECK_EQ(c.n_head % c.n_kv_head, 0) << "n_head must be divisible by n_kv_head for GQA";
+    CHECK_GT(c.n_embd, 0);
+    CHECK_EQ(c.n_embd % c.n_head, 0) << "n_embd must be divisible by n_head";
+    CHECK(c.attention_type == nn::AttentionType::kRoPE) << "LLaMA-3 requires RoPE attention";
+    CHECK(c.activation_type == nn::MLPType::kSwiGLU) << "LLaMA-3 requires SwiGLU activation";
+    CHECK(c.norm_type == nn::NormType::kRMSNorm) << "LLaMA-3 requires RMSNorm";
+    CHECK(!c.add_bias_linear) << "LLaMA-3 has no bias in linear layers";
+    CHECK(!c.tie_weights) << "LLaMA-3 does not tie embedding and lm_head weights";
+    CHECK(c.ffn_dim_multiplier.has_value()) << "LLaMA-3 requires ffn_dim_multiplier";
+    CHECK_GT(c.multiple_of, 0);
+}
 } // namespace llama3
diff --git a/example/llama3/main.cc b/example/llama3/main.cc
index da9a1027..117551d5 100644
--- a/example/llama3/main.cc
+++ b/example/llama3/main.cc
@@ -174,6 +174,7 @@ void Train(const nn::parallel::Rank &rank) {
     if (!FLAGS_llmc_filepath.empty()) {
         model = llama3::LoadFromLLMC(FLAGS_llmc_filepath);
     } else {
+        llama3::SanitizeLLaMA3Config(model_config);
         model = std::make_shared<nn::TransformerModel>(model_config);
     }
 
diff --git a/infini_train/include/tensor.h b/infini_train/include/tensor.h
index 4f4ed94b..12f45f57 100644
--- a/infini_train/include/tensor.h
+++ b/infini_train/include/tensor.h
@@ -58,7 +58,7 @@ class Tensor : public std::enable_shared_from_this<Tensor> {
 public:
     Tensor() = default;
 
-    Tensor(const std::vector<int64_t> &dims, DataType dtype, Device device);
+    Tensor(const std::vector<int64_t> &dims, DataType dtype, Device device, bool requires_grad = false);
     Tensor(const std::vector<int64_t> &dims, DataType dtype) : Tensor(dims, dtype, Device()) {}
 
     Tensor(const Tensor &tensor, size_t offset, const std::vector<int64_t> &dims);
diff --git a/infini_train/src/tensor.cc b/infini_train/src/tensor.cc
index c41aa974..f7947030 100644
--- a/infini_train/src/tensor.cc
+++ b/infini_train/src/tensor.cc
@@ -46,7 +46,8 @@ Device TensorBuffer::GetDevice() const { return device_; }
 size_t TensorBuffer::Size() const { return size_; }
 
 // Tensor implementation
-Tensor::Tensor(const std::vector<int64_t> &dims, DataType dtype, Device device) : dims_(dims), dtype_(dtype) {
+Tensor::Tensor(const std::vector<int64_t> &dims, DataType dtype, Device device, bool requires_grad)
+    : dims_(dims), dtype_(dtype), requires_grad_(requires_grad) {
     num_elements_ = std::accumulate(dims.begin(), dims.end(), 1, std::multiplies<int64_t>());
     buffer_ = std::make_shared<TensorBuffer>(device, kDataTypeToSize.at(dtype) * num_elements_);
 }
diff --git a/test/dtype/test_dtype_dispatch.cc b/test/dtype/test_dtype_dispatch.cc
deleted file mode 100644
index a79a14fc..00000000
--- a/test/dtype/test_dtype_dispatch.cc
+++ /dev/null
@@ -1,120 +0,0 @@
-#include <cstdlib>
-#include <iostream>
-#include <string>
-#include <type_traits>
-
-#include "glog/logging.h"
-
-#include "infini_train/include/datatype.h"
-#include "infini_train/include/dtype_dispatch.h"
-
-#include "infini_train/src/core/runtime/cpu/cpu_dispatch.h"
-
-using namespace infini_train;
-
-// ============================================================================
-// Test 1: HasMappedType_v intercepts backends missing FP16 / BF16
-// ============================================================================
-
-// A backend TypeMap that only registers kFLOAT32 — FP16 / BF16 are absent.
-template <DataType DType> struct LowPrecisionAbsentTypeMap;
-
-template <> struct LowPrecisionAbsentTypeMap<DataType::kFLOAT32> {
-    using type = float;
-};
-
-static_assert(HasMappedType_v<LowPrecisionAbsentTypeMap, DataType::kFLOAT32>,
-              "sanity: registered dtype must be detected as present");
-static_assert(!HasMappedType_v<LowPrecisionAbsentTypeMap, DataType::kFLOAT16>,
-              "unregistered kFLOAT16 must be intercepted by HasMappedType_v");
-static_assert(!HasMappedType_v<LowPrecisionAbsentTypeMap, DataType::kBFLOAT16>,
-              "unregistered kBFLOAT16 must be intercepted by HasMappedType_v");
-
-// ============================================================================
-// Test 2: CpuTypeMap resolves FP16 / BF16 to framework scalar types
-// ============================================================================
-
-static_assert(std::is_same_v<MappedType_t<core::cpu::CpuTypeMap, DataType::kFLOAT16>, FP16>,
-              "CpuTypeMap<kFLOAT16> must resolve to framework FP16");
-static_assert(std::is_same_v<MappedType_t<core::cpu::CpuTypeMap, DataType::kBFLOAT16>, BF16>,
-              "CpuTypeMap<kBFLOAT16> must resolve to framework BF16");
-
-// ============================================================================
-// Test 3: Runtime dispatch of kFLOAT16 / kBFLOAT16
-// ============================================================================
-
-void TestRuntimeDispatchLowPrecision() {
-    std::cout << "\n=== Test 3: Runtime dispatch of kFLOAT16 / kBFLOAT16 ===" << std::endl;
-
-    // kFLOAT16 must dispatch to framework FP16
-    bool called_fp16 = false;
-    core::cpu::DispatchCpuFunc<DataType::kFLOAT16, DataType::kBFLOAT16>(
-        DataType::kFLOAT16,
-        [&called_fp16]<typename T>() {
-            if constexpr (std::is_same_v<T, FP16>) {
-                called_fp16 = true;
-            }
-        },
-        "dispatch kFLOAT16");
-    CHECK(called_fp16) << "DispatchCpuFunc did not invoke functor for kFLOAT16";
-
-    // kBFLOAT16 must dispatch to framework BF16
-    bool called_bf16 = false;
-    core::cpu::DispatchCpuFunc<DataType::kFLOAT16, DataType::kBFLOAT16>(
-        DataType::kBFLOAT16,
-        [&called_bf16]<typename T>() {
-            if constexpr (std::is_same_v<T, BF16>) {
-                called_bf16 = true;
-            }
-        },
-        "dispatch kBFLOAT16");
-    CHECK(called_bf16) << "DispatchCpuFunc did not invoke functor for kBFLOAT16";
-
-    std::cout << "Low-precision dispatch OK." << std::endl;
-}
-
-// ============================================================================
-// Test 4: Runtime dispatch of a low-precision dtype outside AllowedDTypes
-//         must fatal
-// ============================================================================
-
-// Sub-process entry: tries to dispatch kFLOAT16 with only kFLOAT32 allowed.
-void TriggerRuntimeUnsupportedLowPrecisionFatal() {
-    core::cpu::DispatchCpuFunc<DataType::kFLOAT32>(
-        DataType::kFLOAT16,
-        []<typename T>() { (void)sizeof(T); },
-        "intercept kFLOAT16 when only kFLOAT32 is allowed");
-}
-
-void TestRuntimeInterceptLowPrecision(const char *argv0) {
-    std::cout << "\n=== Test 4: Runtime intercept of kFLOAT16 outside AllowedDTypes ===" << std::endl;
-    const std::string cmd = std::string(argv0) + " --expect-runtime-fatal > /dev/null 2>&1";
-    const int status = std::system(cmd.c_str());
-    CHECK_NE(status, 0) << "Expected non-zero exit when dispatching an unallowed low-precision dtype";
-    std::cout << "Low-precision runtime intercept OK." << std::endl;
-}
-
-// ============================================================================
-// Main
-// ============================================================================
-
-int main(int argc, char *argv[]) {
-    google::InitGoogleLogging(argv[0]);
-
-    if (argc > 1 && std::string(argv[1]) == "--expect-runtime-fatal") {
-        TriggerRuntimeUnsupportedLowPrecisionFatal();
-        return 0;
-    }
-
-    std::cout << "========================================" << std::endl;
-    std::cout << "  Low-precision Dtype Dispatch Test Suite" << std::endl;
-    std::cout << "========================================" << std::endl;
-
-    std::cout << "Compile-time checks: PASSED" << std::endl;
-
-    TestRuntimeDispatchLowPrecision();
-    TestRuntimeInterceptLowPrecision(argv[0]);
-
-    std::cout << "\nAll low-precision dtype dispatch tests passed." << std::endl;
-    return 0;
-}
diff --git a/test/dtype/test_scalar.cc b/test/dtype/test_scalar.cc
deleted file mode 100644
index bd97f473..00000000
--- a/test/dtype/test_scalar.cc
+++ /dev/null
@@ -1,324 +0,0 @@
-#include <cmath>
-#include <cstdint>
-#include <iostream>
-#include <limits>
-
-#include "glog/logging.h"
-
-#include "infini_train/include/scalar.h"
-
-using namespace infini_train;
-
-// ============================================================================
-// Test 1: Default Constructor
-// ============================================================================
-void TestDefaultConstructor() {
-    std::cout << "\n=== Test 1: Default Constructor ===" << std::endl;
-
-    Scalar default_scalar;
-    CHECK_EQ(static_cast<int>(default_scalar.kind), static_cast<int>(Scalar::Kind::kInt64));
-    CHECK_EQ(default_scalar.i, 0);
-
-    std::cout << "Default constructor test passed!" << std::endl;
-}
-
-// ============================================================================
-// Test 2: Bool Constructor
-// ============================================================================
-void TestBoolConstructor() {
-    std::cout << "\n=== Test 2: Bool Constructor ===" << std::endl;
-
-    Scalar scalar_true(true);
-    CHECK_EQ(static_cast<int>(scalar_true.kind), static_cast<int>(Scalar::Kind::kBool));
-    CHECK_EQ(scalar_true.u, static_cast<uint64_t>(1));
-
-    Scalar scalar_false(false);
-    CHECK_EQ(static_cast<int>(scalar_false.kind), static_cast<int>(Scalar::Kind::kBool));
-    CHECK_EQ(scalar_false.u, static_cast<uint64_t>(0));
-
-    std::cout << "Bool constructor test passed!" << std::endl;
-}
-
-// ============================================================================
-// Test 3: Signed Integer Constructor
-// ============================================================================
-void TestSignedIntConstructor() {
-    std::cout << "\n=== Test 3: Signed Integer Constructor ===" << std::endl;
-
-    Scalar scalar_positive(42);
-    CHECK_EQ(static_cast<int>(scalar_positive.kind), static_cast<int>(Scalar::Kind::kInt64));
-    CHECK_EQ(scalar_positive.i, static_cast<int64_t>(42));
-
-    Scalar scalar_negative(-7);
-    CHECK_EQ(scalar_negative.i, static_cast<int64_t>(-7));
-
-    int16_t short_val = -100;
-    Scalar scalar_short(short_val);
-    CHECK_EQ(static_cast<int>(scalar_short.kind), static_cast<int>(Scalar::Kind::kInt64));
-    CHECK_EQ(scalar_short.i, static_cast<int64_t>(-100));
-
-    std::cout << "Signed integer constructor test passed!" << std::endl;
-}
-
-// ============================================================================
-// Test 4: Unsigned Integer Constructor
-// ============================================================================
-void TestUnsignedIntConstructor() {
-    std::cout << "\n=== Test 4: Unsigned Integer Constructor ===" << std::endl;
-
-    unsigned int uint_val = 99u;
-    Scalar scalar_uint(uint_val);
-    CHECK_EQ(static_cast<int>(scalar_uint.kind), static_cast<int>(Scalar::Kind::kUInt64));
-    CHECK_EQ(scalar_uint.u, static_cast<uint64_t>(99));
-
-    uint64_t uint64_max = std::numeric_limits<uint64_t>::max();
-    Scalar scalar_max(uint64_max);
-    CHECK_EQ(scalar_max.u, uint64_max);
-
-    std::cout << "Unsigned integer constructor test passed!" << std::endl;
-}
-
-// ============================================================================
-// Test 5: Float Constructor
-// ============================================================================
-void TestFloatConstructor() {
-    std::cout << "\n=== Test 5: Float Constructor ===" << std::endl;
-
-    Scalar scalar_float(3.14f);
-    CHECK_EQ(static_cast<int>(scalar_float.kind), static_cast<int>(Scalar::Kind::kDouble));
-    CHECK(std::abs(scalar_float.d - 3.14) < 1e-5) << "Expected ~3.14, got " << scalar_float.d;
-
-    std::cout << "Float constructor test passed!" << std::endl;
-}
-
-// ============================================================================
-// Test 6: Double Constructor
-// ============================================================================
-void TestDoubleConstructor() {
-    std::cout << "\n=== Test 6: Double Constructor ===" << std::endl;
-
-    Scalar scalar_double(2.718281828);
-    CHECK_EQ(static_cast<int>(scalar_double.kind), static_cast<int>(Scalar::Kind::kDouble));
-    CHECK(std::abs(scalar_double.d - 2.718281828) < 1e-12) << "Expected ~2.718281828, got " << scalar_double.d;
-
-    std::cout << "Double constructor test passed!" << std::endl;
-}
-
-// ============================================================================
-// Test 7: FP16 / BF16 Constructor
-// ============================================================================
-void TestHalfPrecisionConstructor() {
-    std::cout << "\n=== Test 7: FP16 / BF16 Constructor ===" << std::endl;
-
-    FP16 fp16_val(1.5f);
-    Scalar scalar_from_fp16(fp16_val);
-    CHECK_EQ(static_cast<int>(scalar_from_fp16.kind), static_cast<int>(Scalar::Kind::kDouble));
-    CHECK(std::abs(scalar_from_fp16.d - 1.5) < 1e-3) << "FP16 constructor: expected ~1.5, got " << scalar_from_fp16.d;
-
-    BF16 bf16_val(2.0f);
-    Scalar scalar_from_bf16(bf16_val);
-    CHECK_EQ(static_cast<int>(scalar_from_bf16.kind), static_cast<int>(Scalar::Kind::kDouble));
-    CHECK(std::abs(scalar_from_bf16.d - 2.0) < 1e-2) << "BF16 constructor: expected ~2.0, got " << scalar_from_bf16.d;
-
-    std::cout << "FP16 / BF16 constructor test passed!" << std::endl;
-}
-
-// ============================================================================
-// Test 8: to<T>() Same-Type and Numeric Conversions
-// ============================================================================
-void TestToNumericConversions() {
-    std::cout << "\n=== Test 8: to<T>() Same-Type and Numeric Conversions ===" << std::endl;
-
-    // Same type
-    Scalar scalar_int(10);
-    CHECK_EQ(scalar_int.to<int64_t>(), static_cast<int64_t>(10));
-
-    Scalar scalar_double(3.14);
-    CHECK(std::abs(scalar_double.to<double>() - 3.14) < 1e-12) << "to<double> failed";
-
-    // Int -> float
-    Scalar scalar_positive(42);
-    CHECK(std::abs(scalar_positive.to<float>() - 42.0f) < 1e-6) << "to<float> failed";
-    CHECK(std::abs(scalar_positive.to<double>() - 42.0) < 1e-12) << "to<double> failed";
-
-    // Float -> int (truncation)
-    Scalar scalar_fractional(7.9);
-    CHECK_EQ(scalar_fractional.to<int64_t>(), static_cast<int64_t>(7));
-    CHECK_EQ(scalar_fractional.to<int32_t>(), static_cast<int32_t>(7));
-
-    // Negative int -> float
-    Scalar scalar_negative(-42);
-    CHECK(std::abs(scalar_negative.to<float>() - (-42.0f)) < 1e-6) << "negative to<float> failed";
-
-    std::cout << "Numeric conversion test passed!" << std::endl;
-}
-
-// ============================================================================
-// Test 9: to<T>() Bool Conversions
-// ============================================================================
-void TestToBoolConversions() {
-    std::cout << "\n=== Test 9: to<T>() Bool Conversions ===" << std::endl;
-
-    // Bool -> int
-    Scalar scalar_true(true);
-    CHECK_EQ(scalar_true.to<int64_t>(), static_cast<int64_t>(1));
-
-    Scalar scalar_false(false);
-    CHECK_EQ(scalar_false.to<int64_t>(), static_cast<int64_t>(0));
-
-    // Bool -> double
-    CHECK(std::abs(scalar_true.to<double>() - 1.0) < 1e-12) << "true to<double> failed";
-    CHECK(std::abs(scalar_false.to<double>() - 0.0) < 1e-12) << "false to<double> failed";
-
-    // Int -> bool
-    Scalar scalar_zero(0);
-    CHECK_EQ(scalar_zero.to<bool>(), false);
-
-    Scalar scalar_nonzero(5);
-    CHECK_EQ(scalar_nonzero.to<bool>(), true);
-
-    std::cout << "Bool conversion test passed!" << std::endl;
-}
-
-// ============================================================================
-// Test 10: to<T>() Unsigned Conversions
-// ============================================================================
-void TestToUnsignedConversions() {
-    std::cout << "\n=== Test 10: to<T>() Unsigned Conversions ===" << std::endl;
-
-    // uint -> double
-    uint64_t uint_val = 12345;
-    Scalar scalar_uint(uint_val);
-    CHECK(std::abs(scalar_uint.to<double>() - 12345.0) < 1e-6) << "uint to<double> failed";
-
-    // double -> uint
-    Scalar scalar_double(100.0);
-    CHECK_EQ(scalar_double.to<uint64_t>(), static_cast<uint64_t>(100));
-
-    std::cout << "Unsigned conversion test passed!" << std::endl;
-}
-
-// ============================================================================
-// Test 11: to<T>() FP16 / BF16 Conversions
-// ============================================================================
-void TestToHalfPrecisionConversions() {
-    std::cout << "\n=== Test 11: to<T>() FP16 / BF16 Conversions ===" << std::endl;
-
-    // NOTE(dcj): These tests exercise scalar.to<FP16/BF16>(), which goes through
-    // common::cpu::Cast and follows a double->float->bf16/fp16 two-step path.
-    // This differs from what happens in CUDA kernels where dispatch resolves T
-    // to __nv_bfloat16/__half and Cast falls through to a one-step static_cast
-    // (double->bf16 directly). The two paths may produce different rounding
-    // results. See the TODO in scalar.h for the planned fix.
-
-    // Double -> FP16
-    Scalar scalar_double(1.5);
-    auto to_fp16 = scalar_double.to<FP16>();
-    CHECK(std::abs(static_cast<float>(to_fp16) - 1.5f) < 1e-3) << "to<FP16> from double failed";
-
-    // Int -> FP16
-    Scalar scalar_int(3);
-    auto int_to_fp16 = scalar_int.to<FP16>();
-    CHECK(std::abs(static_cast<float>(int_to_fp16) - 3.0f) < 1e-3) << "to<FP16> from int failed";
-
-    // Double -> BF16
-    Scalar scalar_double2(2.0);
-    auto to_bf16 = scalar_double2.to<BF16>();
-    CHECK(std::abs(static_cast<float>(to_bf16) - 2.0f) < 1e-2) << "to<BF16> from double failed";
-
-    // Int -> BF16
-    Scalar scalar_int2(5);
-    auto int_to_bf16 = scalar_int2.to<BF16>();
-    CHECK(std::abs(static_cast<float>(int_to_bf16) - 5.0f) < 1e-1) << "to<BF16> from int failed";
-
-    std::cout << "FP16 / BF16 conversion test passed!" << std::endl;
-}
-
-// ============================================================================
-// Test 12: Edge Cases — Numeric Limits
-// ============================================================================
-void TestNumericLimits() {
-    std::cout << "\n=== Test 12: Edge Cases — Numeric Limits ===" << std::endl;
-
-    int64_t int64_max = std::numeric_limits<int64_t>::max();
-    Scalar scalar_int64_max(int64_max);
-    CHECK_EQ(scalar_int64_max.to<int64_t>(), int64_max);
-
-    int64_t int64_min = std::numeric_limits<int64_t>::min();
-    Scalar scalar_int64_min(int64_min);
-    CHECK_EQ(scalar_int64_min.to<int64_t>(), int64_min);
-
-    uint64_t uint64_max = std::numeric_limits<uint64_t>::max();
-    Scalar scalar_uint64_max(uint64_max);
-    CHECK_EQ(scalar_uint64_max.to<uint64_t>(), uint64_max);
-
-    std::cout << "Numeric limits test passed!" << std::endl;
-}
-
-// ============================================================================
-// Test 13: Edge Cases — Zero Values
-// ============================================================================
-void TestZeroValues() {
-    std::cout << "\n=== Test 13: Edge Cases — Zero Values ===" << std::endl;
-
-    Scalar scalar_int_zero(0);
-    CHECK_EQ(scalar_int_zero.to<int64_t>(), static_cast<int64_t>(0));
-    CHECK(std::abs(scalar_int_zero.to<double>()) < 1e-12) << "int zero to<double> failed";
-
-    Scalar scalar_double_zero(0.0);
-    CHECK(std::abs(scalar_double_zero.to<double>()) < 1e-12) << "double zero to<double> failed";
-    CHECK_EQ(scalar_double_zero.to<int64_t>(), static_cast<int64_t>(0));
-
-    std::cout << "Zero values test passed!" << std::endl;
-}
-
-// ============================================================================
-// Test 14: FP16 / BF16 Roundtrip
-// ============================================================================
-void TestHalfPrecisionRoundtrip() {
-    std::cout << "\n=== Test 14: FP16 / BF16 Roundtrip ===" << std::endl;
-
-    FP16 original_fp16(0.5f);
-    Scalar scalar_from_fp16(original_fp16);
-    auto roundtrip_fp16 = scalar_from_fp16.to<FP16>();
-    CHECK(std::abs(static_cast<float>(roundtrip_fp16) - 0.5f) < 1e-3) << "FP16 roundtrip failed";
-
-    BF16 original_bf16(4.0f);
-    Scalar scalar_from_bf16(original_bf16);
-    auto roundtrip_bf16 = scalar_from_bf16.to<BF16>();
-    CHECK(std::abs(static_cast<float>(roundtrip_bf16) - 4.0f) < 1e-2) << "BF16 roundtrip failed";
-
-    std::cout << "FP16 / BF16 roundtrip test passed!" << std::endl;
-}
-
-// ============================================================================
-// Main
-// ============================================================================
-int main(int argc, char *argv[]) {
-    google::InitGoogleLogging(argv[0]);
-
-    std::cout << "========================================" << std::endl;
-    std::cout << "       Scalar Test Suite" << std::endl;
-    std::cout << "========================================" << std::endl;
-
-    TestDefaultConstructor();
-    TestBoolConstructor();
-    TestSignedIntConstructor();
-    TestUnsignedIntConstructor();
-    TestFloatConstructor();
-    TestDoubleConstructor();
-    TestHalfPrecisionConstructor();
-    TestToNumericConversions();
-    TestToBoolConversions();
-    TestToUnsignedConversions();
-    TestToHalfPrecisionConversions();
-    TestNumericLimits();
-    TestZeroValues();
-    TestHalfPrecisionRoundtrip();
-
-    std::cout << "\n========================================" << std::endl;
-    std::cout << "    All Tests Completed Successfully" << std::endl;
-    std::cout << "========================================" << std::endl;
-
-    return 0;
-}
diff --git a/test/hook/test_hook.cc b/test/hook/test_hook.cc
deleted file mode 100644
index 32c7e097..00000000
--- a/test/hook/test_hook.cc
+++ /dev/null
@@ -1,179 +0,0 @@
-#include <iostream>
-#include <memory>
-
-#include "glog/logging.h"
-
-#include "infini_train/include/autograd/elementwise.h"
-#include "infini_train/include/autograd/function.h"
-#include "infini_train/include/autograd/function_hook.h"
-#include "infini_train/include/common/hook.h"
-#include "infini_train/include/nn/modules/module.h"
-#include "infini_train/include/nn/parallel/global.h"
-#include "infini_train/include/tensor.h"
-
-using namespace infini_train;
-
-// ============================================================================
-// Test 1: Basic Module Hooks
-// ============================================================================
-void test_basic_hooks() {
-    std::cout << "\n=== Test 1: Basic Module Hooks ===" << std::endl;
-
-    auto x = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32);
-    x->set_requires_grad(true);
-
-    // Module hook example
-    class MyModule : public nn::Module {
-    public:
-        MyModule() : Module("MyModule") {}
-
-        std::vector<std::shared_ptr<Tensor>> Forward(const std::vector<std::shared_ptr<Tensor>> &inputs) override {
-            std::cout << "Forward pass executing..." << std::endl;
-            return inputs;
-        }
-    };
-
-    auto module = std::make_shared<MyModule>();
-
-    // Register forward pre-hook
-    auto pre_hook
-        = module->RegisterForwardPreHook([](nn::Module *mod, const std::vector<std::shared_ptr<Tensor>> &inputs) {
-              std::cout << "Forward pre-hook: Module type = " << mod->type() << std::endl;
-          });
-
-    // Register forward post-hook
-    auto fwd_hook
-        = module->RegisterForwardPostHook([](nn::Module *mod, const std::vector<std::shared_ptr<Tensor>> &inputs,
-                                             const std::vector<std::shared_ptr<Tensor>> &outputs) {
-              std::cout << "Forward post-hook: Got " << outputs.size() << " outputs" << std::endl;
-          });
-
-    // Register backward pre-hook
-    auto bwd_pre_hook = module->RegisterBackwardPreHook(
-        [](nn::Module *mod, const std::vector<std::shared_ptr<Tensor>> &grad_outputs) {
-            std::cout << "Backward pre-hook called!" << std::endl;
-        });
-
-    // Register backward post-hook
-    auto bwd_post_hook
-        = module->RegisterBackwardPostHook([](nn::Module *mod, const std::vector<std::shared_ptr<Tensor>> &grad_inputs,
-                                              const std::vector<std::shared_ptr<Tensor>> &grad_outputs) {
-              std::cout << "Backward post-hook called!" << std::endl;
-          });
-
-    // Test forward pass
-    std::vector<std::shared_ptr<Tensor>> inputs = {x};
-    auto outputs = (*module)(inputs);
-
-    std::cout << "Module hook test completed!" << std::endl;
-}
-
-// ============================================================================
-// Test 2: Hook Remove() Functionality Test
-// ============================================================================
-void test_hook_remove() {
-    std::cout << "\n=== Test 2: Hook Remove() Functionality Test ===" << std::endl;
-
-    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 2}, DataType::kFLOAT32);
-    auto b = std::make_shared<Tensor>(std::vector<int64_t>{2, 2}, DataType::kFLOAT32);
-    a->set_requires_grad(true);
-    b->set_requires_grad(true);
-
-    int hook1_count = 0;
-    int hook2_count = 0;
-    int hook3_count = 0;
-
-    auto add_fn = std::make_shared<autograd::Add>();
-
-    // Register three forward pre-hooks
-    auto handle1 = add_fn->RegisterForwardPreHook(
-        [&hook1_count](autograd::Function *, const std::vector<std::shared_ptr<Tensor>> &) {
-            hook1_count++;
-            std::cout << "Hook 1 called (count: " << hook1_count << ")" << std::endl;
-        });
-
-    auto handle2 = add_fn->RegisterForwardPreHook(
-        [&hook2_count](autograd::Function *, const std::vector<std::shared_ptr<Tensor>> &) {
-            hook2_count++;
-            std::cout << "Hook 2 called (count: " << hook2_count << ")" << std::endl;
-        });
-
-    auto handle3 = add_fn->RegisterForwardPreHook(
-        [&hook3_count](autograd::Function *, const std::vector<std::shared_ptr<Tensor>> &) {
-            hook3_count++;
-            std::cout << "Hook 3 called (count: " << hook3_count << ")" << std::endl;
-        });
-
-    // First call - all hooks should fire
-    std::cout << "\n--- First Apply (all hooks active) ---" << std::endl;
-    std::vector<std::shared_ptr<Tensor>> inputs;
-    inputs.push_back(a);
-    inputs.push_back(b);
-    auto result1 = add_fn->Apply(inputs);
-    std::cout << "Hook counts: " << hook1_count << ", " << hook2_count << ", " << hook3_count << std::endl;
-
-    // Remove hook 2
-    std::cout << "\n--- Removing Hook 2 ---" << std::endl;
-    handle2->Remove();
-
-    // Second call - hook 2 should not fire
-    std::cout << "\n--- Second Apply (hook 2 removed) ---" << std::endl;
-    auto result2 = add_fn->Apply(inputs);
-    std::cout << "Hook counts: " << hook1_count << ", " << hook2_count << ", " << hook3_count << std::endl;
-
-    // Remove hook 1
-    std::cout << "\n--- Removing Hook 1 ---" << std::endl;
-    handle1->Remove();
-
-    // Third call - only hook 3 should fire
-    std::cout << "\n--- Third Apply (hooks 1 and 2 removed) ---" << std::endl;
-    auto result3 = add_fn->Apply(inputs);
-    std::cout << "Hook counts: " << hook1_count << ", " << hook2_count << ", " << hook3_count << std::endl;
-
-    // Verify results
-    std::cout << "\n=== Test Results ===" << std::endl;
-    bool test_passed = true;
-
-    if (hook1_count != 2) {
-        std::cout << "FAIL: Hook 1 should be called 2 times, got " << hook1_count << std::endl;
-        test_passed = false;
-    }
-
-    if (hook2_count != 1) {
-        std::cout << "FAIL: Hook 2 should be called 1 time, got " << hook2_count << std::endl;
-        test_passed = false;
-    }
-
-    if (hook3_count != 3) {
-        std::cout << "FAIL: Hook 3 should be called 3 times, got " << hook3_count << std::endl;
-        test_passed = false;
-    }
-
-    if (test_passed) {
-        std::cout << "SUCCESS: All hooks behaved correctly!" << std::endl;
-        std::cout << "  - Hook 1: called 2 times (before removal)" << std::endl;
-        std::cout << "  - Hook 2: called 1 time (removed after first call)" << std::endl;
-        std::cout << "  - Hook 3: called 3 times (never removed)" << std::endl;
-    }
-}
-
-// ============================================================================
-// Main
-// ============================================================================
-int main(int argc, char *argv[]) {
-    google::InitGoogleLogging(argv[0]);
-    nn::parallel::global::GlobalEnv::Instance().Init(1, 1, false, 1, 1);
-
-    std::cout << "========================================" << std::endl;
-    std::cout << "    Hook Mechanism Tests" << std::endl;
-    std::cout << "========================================" << std::endl;
-
-    test_basic_hooks();
-    test_hook_remove();
-
-    std::cout << "\n========================================" << std::endl;
-    std::cout << "    All Tests Completed Successfully" << std::endl;
-    std::cout << "========================================" << std::endl;
-
-    return 0;
-}
diff --git a/test/hook/test_precision_check.cc b/test/hook/test_precision_check.cc
deleted file mode 100644
index eae3aaf3..00000000
--- a/test/hook/test_precision_check.cc
+++ /dev/null
@@ -1,241 +0,0 @@
-#include <filesystem>
-#include <iostream>
-#include <memory>
-
-#include "glog/logging.h"
-
-#include "infini_train/include/nn/modules/module.h"
-#include "infini_train/include/nn/parallel/global.h"
-#include "infini_train/include/tensor.h"
-#include "infini_train/include/utils/global_module_hook_registry.h"
-#include "infini_train/include/utils/precision_check_config.h"
-#include "infini_train/include/utils/precision_checker.h"
-
-using namespace infini_train;
-
-class MyModel : public nn::Module {
-public:
-    MyModel() : Module("MyModel") {}
-
-    std::vector<std::shared_ptr<Tensor>> Forward(const std::vector<std::shared_ptr<Tensor>> &inputs) override {
-        auto x = inputs[0];
-        x->RequiresGrad();
-        auto y = x->Mul(x);
-        return {y};
-    }
-};
-
-// Simple model for multi-iteration test
-class SimpleModel : public nn::Module {
-public:
-    SimpleModel() : Module("SimpleModel") {}
-
-    std::vector<std::shared_ptr<Tensor>> Forward(const std::vector<std::shared_ptr<Tensor>> &inputs) override {
-        auto x = inputs[0];
-        x->RequiresGrad();
-        auto y = x->Mul(x)->Mul(x); // x^3
-        return {y};
-    }
-};
-
-void RunModelForwardBackward(const std::shared_ptr<nn::Module> &model) {
-    auto x = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32);
-    x->Fill(2.0f);
-    x->RequiresGrad();
-
-    std::vector<std::shared_ptr<Tensor>> inputs = {x};
-    auto outputs = (*model)(inputs);
-    auto loss = outputs[0]->Sum(0, false)->Sum(0, false);
-    loss->Backward();
-}
-
-void TestFunctionLevel(const std::string &config_str) {
-    std::cout << "\n========================================" << std::endl;
-    std::cout << "  Function-Level Test: " << config_str << std::endl;
-    std::cout << "========================================" << std::endl;
-
-    auto x = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32);
-    x->Fill(2.0f);
-    x->RequiresGrad();
-
-    auto y = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32);
-    y->Fill(3.0f);
-    y->RequiresGrad();
-
-    auto z = x->Mul(y);
-    auto loss = z->Sum(0, false)->Sum(0, false);
-    loss->Backward();
-
-    std::cout << "Test completed." << std::endl;
-}
-
-void TestModuleLevel(const std::string &config_str) {
-    std::cout << "\n========================================" << std::endl;
-    std::cout << "  Module-Level Test: " << config_str << std::endl;
-    std::cout << "========================================" << std::endl;
-
-    auto model = std::make_shared<MyModel>();
-    RunModelForwardBackward(model);
-
-    std::cout << "Test completed." << std::endl;
-}
-
-// Test: Simple format output (level=2, format=simple)
-void TestSimpleFormat() {
-    std::cout << "\n========================================" << std::endl;
-    std::cout << "  Test: Simple Format (level=2, format=simple)" << std::endl;
-    std::cout << "========================================" << std::endl;
-
-    auto x = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32);
-    x->Fill(2.0f);
-    x->RequiresGrad();
-
-    auto y = x->Mul(x);
-    auto loss = y->Sum(0, false)->Sum(0, false); // Two Sum ops to produce scalar
-    loss->Backward();
-
-    std::cout << "Simple format test completed - check output for min/max/mean values." << std::endl;
-}
-
-// Test: MD5 format output (level=2, format=md5)
-void TestMd5Format() {
-    std::cout << "\n========================================" << std::endl;
-    std::cout << "  Test: MD5 Format (level=2, format=md5)" << std::endl;
-    std::cout << "========================================" << std::endl;
-
-    auto x = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32);
-    x->Fill(2.0f);
-    x->RequiresGrad();
-
-    auto y = x->Mul(x);
-    auto loss = y->Sum(0, false)->Sum(0, false); // Two Sum ops to produce scalar
-    loss->Backward();
-
-    std::cout << "MD5 format test completed - check output for md5 hashes." << std::endl;
-}
-
-// Test: Save tensors to NPY files (level=1, save_tensors=true)
-void TestSaveTensors() {
-    std::cout << "\n========================================" << std::endl;
-    std::cout << "  Test: Save Tensors (level=1, save_tensors=true)" << std::endl;
-    std::cout << "========================================" << std::endl;
-
-    std::string output_path = "/tmp/precision_check_npy";
-
-    auto model = std::make_shared<MyModel>();
-    RunModelForwardBackward(model);
-
-    // Verify NPY files were created
-    namespace fs = std::filesystem;
-    bool found_npy = false;
-    if (fs::exists(output_path)) {
-        for (const auto &entry : fs::recursive_directory_iterator(output_path)) {
-            if (entry.path().extension() == ".npy") {
-                found_npy = true;
-                std::cout << "Found NPY file: " << entry.path() << std::endl;
-            }
-        }
-    }
-
-    if (found_npy) {
-        std::cout << "Save tensors test PASSED - NPY files created successfully." << std::endl;
-    } else {
-        std::cout << "Save tensors test completed - check output directory for NPY files." << std::endl;
-    }
-}
-
-// Test: Multi-iteration file overwrite (level=1, save_tensors=true, iter=3)
-void TestMultiIterOverwrite() {
-    std::cout << "\n========================================" << std::endl;
-    std::cout << "  Test: Multi-Iteration File Overwrite" << std::endl;
-    std::cout << "========================================" << std::endl;
-
-    std::string output_path = "/tmp/precision_check_overwrite";
-
-    auto model = std::make_shared<SimpleModel>();
-    int num_iters = 3;
-
-    // Run multiple iterations - files should be overwritten
-    for (int i = 0; i < num_iters; ++i) {
-        std::cout << "Iteration " << (i + 1) << "/" << num_iters << std::endl;
-        utils::PrecisionCheckEnv::ResetCounters(); // Reset counters each iteration
-        RunModelForwardBackward(model);
-    }
-
-    namespace fs = std::filesystem;
-    int npy_count = 0;
-    if (fs::exists(output_path)) {
-        for (const auto &entry : fs::recursive_directory_iterator(output_path)) {
-            if (entry.path().extension() == ".npy") {
-                ++npy_count;
-            }
-        }
-    }
-
-    std::cout << "Multi-iteration test completed - found " << npy_count << " NPY files after " << num_iters
-              << " iterations." << std::endl;
-    std::cout << "(Files should be overwritten each iteration, count should be consistent with 1 iter)" << std::endl;
-}
-
-int main(int argc, char *argv[]) {
-    google::InitGoogleLogging(argv[0]);
-
-    std::string config_str = argc > 1 ? argv[1] : "";
-
-    std::cout << "========================================" << std::endl;
-    std::cout << "  Precision Check Test Suite" << std::endl;
-    std::cout << "========================================" << std::endl;
-
-    nn::parallel::global::InitAllEnv(1, 1, false, 1, 1);
-
-    // If no config argument, run all format tests
-    if (config_str.empty()) {
-        auto config = utils::PrecisionCheckConfig::Parse("level=2,format=simple");
-        utils::PrecisionCheckEnv::Instance().Init(config);
-
-        std::cout << "\nRunning all precision check format tests..." << std::endl;
-
-        // Test 1: Simple format
-        TestSimpleFormat();
-
-        // Test 2: MD5 format
-        auto md5_config = utils::PrecisionCheckConfig::Parse("level=2,format=md5");
-        utils::PrecisionCheckEnv::Instance().Init(md5_config);
-        TestMd5Format();
-
-        // Test 3: Save tensors
-        auto npy_config = utils::PrecisionCheckConfig::Parse("level=1,save_tensors=true");
-        utils::PrecisionCheckEnv::Instance().Init(npy_config);
-        TestSaveTensors();
-
-        // Test 4: Multi-iteration overwrite
-        auto iter_config = utils::PrecisionCheckConfig::Parse("level=1,save_tensors=true");
-        utils::PrecisionCheckEnv::Instance().Init(iter_config);
-        TestMultiIterOverwrite();
-
-        std::cout << "\n========================================" << std::endl;
-        std::cout << "  All Tests Completed Successfully" << std::endl;
-        std::cout << "========================================" << std::endl;
-        return 0;
-    }
-
-    // If config provided, run single test (original behavior)
-    auto config = utils::PrecisionCheckConfig::Parse(config_str);
-    utils::PrecisionCheckEnv::Instance().Init(config);
-
-    std::cout << "Config: " << config_str << std::endl;
-
-    if (config.level == utils::PrecisionCheckLevel::MODULE) {
-        TestModuleLevel(config_str);
-    } else if (config.level == utils::PrecisionCheckLevel::FUNCTION) {
-        TestFunctionLevel(config_str);
-    } else {
-        std::cout << "No tests to run (level=0)" << std::endl;
-    }
-
-    std::cout << "\n========================================" << std::endl;
-    std::cout << "  Test Completed" << std::endl;
-    std::cout << "========================================" << std::endl;
-
-    return 0;
-}
diff --git a/test/lora/test_lora.cc b/test/lora/test_lora.cc
deleted file mode 100644
index 06966809..00000000
--- a/test/lora/test_lora.cc
+++ /dev/null
@@ -1,860 +0,0 @@
-#include <cmath>
-#include <iostream>
-#include <memory>
-
-#include "glog/logging.h"
-
-#include "infini_train/include/nn/lora/lora_config.h"
-#include "infini_train/include/nn/lora/lora_linear.h"
-#include "infini_train/include/nn/lora/lora_utils.h"
-#include "infini_train/include/nn/modules/container.h"
-#include "infini_train/include/nn/modules/linear.h"
-#include "infini_train/include/nn/modules/module.h"
-#include "infini_train/include/nn/parallel/global.h"
-#include "infini_train/include/tensor.h"
-
-using namespace infini_train;
-using namespace infini_train::nn::lora;
-
-// ============================================================================
-// Test 1: LoRAConfig
-// ============================================================================
-void test_lora_config() {
-    std::cout << "\n=== Test 1: LoRAConfig ===" << std::endl;
-
-    LoRAConfig config;
-    config.rank = 8;
-    config.alpha = 16.0f;
-
-    // Test scaling calculation
-    float expected_scaling = 16.0f / 8.0f;
-    CHECK_EQ(config.Scaling(), expected_scaling) << "Scaling calculation failed";
-    std::cout << "Scaling: " << config.Scaling() << " (expected: " << expected_scaling << ")" << std::endl;
-
-    // Test ShouldApplyLoRA
-    CHECK(config.ShouldApplyLoRA("c_attn")) << "Should match c_attn";
-    CHECK(config.ShouldApplyLoRA("transformer.h.0.attn.c_attn")) << "Should match nested c_attn";
-    CHECK(config.ShouldApplyLoRA("c_proj")) << "Should match c_proj";
-    CHECK(!config.ShouldApplyLoRA("c_fc")) << "Should not match c_fc (not in default targets)";
-    CHECK(!config.ShouldApplyLoRA("random_layer")) << "Should not match random_layer";
-
-    std::cout << "LoRAConfig tests passed!" << std::endl;
-}
-
-// ============================================================================
-// Test 2: LoRALinear Initialization
-// ============================================================================
-void test_lora_linear_init() {
-    std::cout << "\n=== Test 2: LoRALinear Initialization ===" << std::endl;
-
-    LoRAConfig config;
-    config.rank = 4;
-    config.alpha = 8.0f;
-
-    int64_t in_features = 64;
-    int64_t out_features = 128;
-
-    auto lora_linear
-        = std::shared_ptr<LoRALinear>(new LoRALinear(in_features, out_features, config, /*bias=*/true, nullptr));
-
-    // Check parameter shapes
-    auto weight = lora_linear->parameter(nn::Linear::kParamWeightName);
-    auto bias = lora_linear->parameter(nn::Linear::kParamBiasName);
-    auto lora_A = lora_linear->parameter(LoRALinear::kParamLoraAName);
-    auto lora_B = lora_linear->parameter(LoRALinear::kParamLoraBName);
-
-    CHECK_EQ(weight->Dims().size(), 2);
-    CHECK_EQ(weight->Dims()[0], out_features);
-    CHECK_EQ(weight->Dims()[1], in_features);
-    std::cout << "Weight shape: [" << weight->Dims()[0] << ", " << weight->Dims()[1] << "]" << std::endl;
-
-    CHECK_EQ(bias->Dims().size(), 1);
-    CHECK_EQ(bias->Dims()[0], out_features);
-    std::cout << "Bias shape: [" << bias->Dims()[0] << "]" << std::endl;
-
-    CHECK_EQ(lora_A->Dims().size(), 2);
-    CHECK_EQ(lora_A->Dims()[0], config.rank);
-    CHECK_EQ(lora_A->Dims()[1], in_features);
-    std::cout << "LoRA A shape: [" << lora_A->Dims()[0] << ", " << lora_A->Dims()[1] << "]" << std::endl;
-
-    CHECK_EQ(lora_B->Dims().size(), 2);
-    CHECK_EQ(lora_B->Dims()[0], out_features);
-    CHECK_EQ(lora_B->Dims()[1], config.rank);
-    std::cout << "LoRA B shape: [" << lora_B->Dims()[0] << ", " << lora_B->Dims()[1] << "]" << std::endl;
-
-    // Check requires_grad
-    CHECK(!weight->requires_grad()) << "Base weight should be frozen";
-    CHECK(!bias->requires_grad()) << "Base bias should be frozen";
-    CHECK(lora_A->requires_grad()) << "LoRA A should be trainable";
-    CHECK(lora_B->requires_grad()) << "LoRA B should be trainable";
-    std::cout << "requires_grad check passed!" << std::endl;
-
-    // Check LoRAParameters() returns only LoRA params
-    auto params = lora_linear->LoRAParameters();
-    CHECK_EQ(params.size(), 2) << "LoRAParameters() should return only LoRA params";
-    std::cout << "LoRAParameters() returns " << params.size() << " tensors (LoRA A and B)" << std::endl;
-
-    std::cout << "LoRALinear initialization tests passed!" << std::endl;
-}
-
-// ============================================================================
-// Test 3: LoRALinear Forward Pass
-// ============================================================================
-void test_lora_linear_forward() {
-    std::cout << "\n=== Test 3: LoRALinear Forward Pass ===" << std::endl;
-
-    LoRAConfig config;
-    config.rank = 4;
-    config.alpha = 8.0f;
-
-    int64_t in_features = 64;
-    int64_t out_features = 128;
-    int64_t batch_size = 2;
-    int64_t seq_len = 10;
-
-    auto lora_linear
-        = std::shared_ptr<LoRALinear>(new LoRALinear(in_features, out_features, config, /*bias=*/true, nullptr));
-
-    // Create input tensor
-    auto input = std::make_shared<Tensor>(std::vector<int64_t>{batch_size, seq_len, in_features}, DataType::kFLOAT32);
-
-    // Forward pass
-    auto output = (*lora_linear)({input})[0];
-
-    // Check output shape
-    CHECK_EQ(output->Dims().size(), 3);
-    CHECK_EQ(output->Dims()[0], batch_size);
-    CHECK_EQ(output->Dims()[1], seq_len);
-    CHECK_EQ(output->Dims()[2], out_features);
-    std::cout << "Output shape: [" << output->Dims()[0] << ", " << output->Dims()[1] << ", " << output->Dims()[2] << "]"
-              << std::endl;
-
-    std::cout << "LoRALinear forward pass tests passed!" << std::endl;
-}
-
-// ============================================================================
-// Test 4: LoRALinear Weight Merging
-// ============================================================================
-void test_lora_linear_merge() {
-    std::cout << "\n=== Test 4: LoRALinear Weight Merging ===" << std::endl;
-
-    LoRAConfig config;
-    config.rank = 4;
-    config.alpha = 8.0f;
-
-    int64_t in_features = 32;
-    int64_t out_features = 64;
-
-    auto lora_linear
-        = std::shared_ptr<LoRALinear>(new LoRALinear(in_features, out_features, config, /*bias=*/false, nullptr));
-
-    // Print weight sum before merge
-    auto weight_before = lora_linear->parameter(nn::Linear::kParamWeightName);
-    auto lora_A = lora_linear->parameter(LoRALinear::kParamLoraAName);
-    auto lora_B = lora_linear->parameter(LoRALinear::kParamLoraBName);
-
-    float weight_before_sum = weight_before->EigenMatrix().sum();
-    float lora_A_sum = lora_A->EigenMatrix().sum();
-    float lora_B_sum = lora_B->EigenMatrix().sum();
-
-    std::cout << "\n--- Before Merge ---" << std::endl;
-    std::cout << "Base weight sum: " << weight_before_sum << std::endl;
-    std::cout << "LoRA A sum: " << lora_A_sum << std::endl;
-    std::cout << "LoRA B sum: " << lora_B_sum << std::endl;
-    std::cout << "Scaling (alpha/r): " << config.Scaling() << std::endl;
-
-    // Create input
-    auto input = std::make_shared<Tensor>(std::vector<int64_t>{2, 5, in_features}, DataType::kFLOAT32);
-    input->EigenMatrix().setRandom();
-
-    // Get output before merge
-    auto output_before = (*lora_linear)({input})[0];
-    float output_before_sum = output_before->EigenMatrix().sum();
-    std::cout << "Output sum before merge: " << output_before_sum << std::endl;
-
-    // Merge weights
-    CHECK(!lora_linear->IsMerged()) << "Should not be merged initially";
-    lora_linear->MergeWeights();
-    CHECK(lora_linear->IsMerged()) << "Should be merged after MergeWeights()";
-
-    // Verify LoRA params are frozen after merge
-    CHECK(!lora_A->requires_grad()) << "lora_A should be frozen after merge";
-    CHECK(!lora_B->requires_grad()) << "lora_B should be frozen after merge";
-    std::cout << "\nWeights merged successfully, LoRA params frozen" << std::endl;
-
-    // Print weight sum after merge
-    auto weight_after = lora_linear->parameter(nn::Linear::kParamWeightName);
-    float weight_after_sum = weight_after->EigenMatrix().sum();
-    std::cout << "\n--- After Merge ---" << std::endl;
-    std::cout << "Base weight sum after merge: " << weight_after_sum << std::endl;
-    std::cout << "Weight change (should be ~LoRA contribution): " << (weight_after_sum - weight_before_sum)
-              << std::endl;
-
-    // Get output after merge
-    auto output_merged = (*lora_linear)({input})[0];
-    float output_merged_sum = output_merged->EigenMatrix().sum();
-    std::cout << "Output sum after merge: " << output_merged_sum << std::endl;
-
-    // Verify: output_after should equal output_before (numerically)
-    std::cout << "\nVerification: output_before == output_after? " << std::endl;
-    std::cout << "  Before: " << output_before_sum << std::endl;
-    std::cout << "  After:  " << output_merged_sum << std::endl;
-    std::cout << "  Diff:   " << std::abs(output_before_sum - output_merged_sum) << std::endl;
-    CHECK(std::abs(output_before_sum - output_merged_sum) < 1e-3) << "Outputs should be numerically identical!";
-
-    // Shape comparison (always same)
-    std::cout << "\nOutput shape: [" << output_before->Dims()[0] << ", " << output_before->Dims()[1] << ", "
-              << output_before->Dims()[2] << "] (unchanged)" << std::endl;
-
-    // Unmerge weights
-    lora_linear->UnmergeWeights();
-    CHECK(!lora_linear->IsMerged()) << "Should not be merged after UnmergeWeights()";
-
-    // Verify LoRA params are trainable again after unmerge
-    CHECK(lora_A->requires_grad()) << "lora_A should be trainable after unmerge";
-    CHECK(lora_B->requires_grad()) << "lora_B should be trainable after unmerge";
-
-    // Print weight sum after unmerge
-    auto weight_unmerged = lora_linear->parameter(nn::Linear::kParamWeightName);
-    float weight_unmerged_sum = weight_unmerged->EigenMatrix().sum();
-    std::cout << "\n--- After Unmerge ---" << std::endl;
-    std::cout << "Base weight sum after unmerge: " << weight_unmerged_sum << std::endl;
-
-    // Verify: weight should be restored to original value
-    std::cout << "\nVerification: weight restored after unmerge? " << std::endl;
-    std::cout << "  Original: " << weight_before_sum << std::endl;
-    std::cout << "  Unmerged: " << weight_unmerged_sum << std::endl;
-    std::cout << "  Diff:     " << std::abs(weight_before_sum - weight_unmerged_sum) << std::endl;
-    CHECK(std::abs(weight_before_sum - weight_unmerged_sum) < 1e-4) << "Weight should be restored!";
-
-    // Get output after unmerge
-    auto output_unmerged = (*lora_linear)({input})[0];
-    float output_unmerged_sum = output_unmerged->EigenMatrix().sum();
-    std::cout << "Output sum after unmerge: " << output_unmerged_sum << std::endl;
-
-    // Shape comparison: merge doesn't change shape, only weights
-    CHECK(output_before->Dims() == output_merged->Dims()) << "Shape should be identical after merge";
-    CHECK(output_merged->Dims() == output_unmerged->Dims()) << "Shape should be identical after unmerge";
-
-    std::cout << "\nLoRALinear weight merging tests passed!" << std::endl;
-}
-
-// ============================================================================
-// Test 5: LoRA Utility Functions
-// ============================================================================
-void test_lora_utils() {
-    std::cout << "\n=== Test 5: LoRA Utility Functions ===" << std::endl;
-
-    LoRAConfig config;
-    config.rank = 4;
-    config.alpha = 8.0f;
-
-    auto lora_linear = std::shared_ptr<LoRALinear>(new LoRALinear(32, 64, config, /*bias=*/true, nullptr));
-
-    // Test GetLoRAParameters
-    auto lora_params = GetLoRAParameters(lora_linear);
-    CHECK_EQ(lora_params.size(), 2) << "Should have 2 LoRA parameters";
-    std::cout << "GetLoRAParameters returned " << lora_params.size() << " parameters" << std::endl;
-
-    // Test CountTrainableParameters
-    int64_t trainable = CountTrainableParameters(lora_linear);
-    int64_t expected_trainable = config.rank * 32 + 64 * config.rank; // A: [4, 32], B: [64, 4]
-    CHECK_EQ(trainable, expected_trainable) << "Trainable parameter count mismatch";
-    std::cout << "Trainable parameters: " << trainable << " (expected: " << expected_trainable << ")" << std::endl;
-
-    // Test CountTotalParameters
-    int64_t total = CountTotalParameters(lora_linear);
-    int64_t expected_total = 64 * 32 + 64 + config.rank * 32 + 64 * config.rank; // weight + bias + A + B
-    CHECK_EQ(total, expected_total) << "Total parameter count mismatch";
-    std::cout << "Total parameters: " << total << " (expected: " << expected_total << ")" << std::endl;
-
-    // Test PrintLoRASummary
-    std::cout << "\nLoRA Summary:" << std::endl;
-    PrintLoRASummary(lora_linear);
-
-    std::cout << "LoRA utility function tests passed!" << std::endl;
-}
-
-// ============================================================================
-// Test 6: LoRALinear from existing Linear
-// ============================================================================
-void test_lora_from_linear() {
-    std::cout << "\n=== Test 6: LoRALinear from existing Linear ===" << std::endl;
-
-    // Create a standard Linear layer
-    auto linear = std::make_shared<nn::Linear>(64, 128, /*bias=*/true);
-
-    // Wrap it with LoRA
-    LoRAConfig config;
-    config.rank = 8;
-    config.alpha = 16.0f;
-
-    auto lora_linear = std::make_shared<LoRALinear>(linear, config);
-
-    // Check dimensions
-    CHECK_EQ(lora_linear->in_features(), 64);
-    CHECK_EQ(lora_linear->out_features(), 128);
-    CHECK_EQ(lora_linear->rank(), 8);
-    std::cout << "LoRALinear created from Linear: in=" << lora_linear->in_features()
-              << ", out=" << lora_linear->out_features() << ", rank=" << lora_linear->rank() << std::endl;
-
-    // Test forward pass
-    auto input = std::make_shared<Tensor>(std::vector<int64_t>{2, 10, 64}, DataType::kFLOAT32);
-    auto output = (*lora_linear)({input})[0];
-
-    CHECK_EQ(output->Dims()[0], 2);
-    CHECK_EQ(output->Dims()[1], 10);
-    CHECK_EQ(output->Dims()[2], 128);
-    std::cout << "Forward pass successful, output shape: [" << output->Dims()[0] << ", " << output->Dims()[1] << ", "
-              << output->Dims()[2] << "]" << std::endl;
-
-    std::cout << "LoRALinear from existing Linear tests passed!" << std::endl;
-}
-
-// ============================================================================
-// Test 7: LoRALinear from existing Linear (tests LoRA utilities)
-// ============================================================================
-void test_lora_model_wrapper() {
-    std::cout << "\n=== Test 7: LoRALinear from existing Linear ===" << std::endl;
-
-    // Create LoRA config
-    LoRAConfig lora_config;
-    lora_config.rank = 8;
-    lora_config.alpha = 16.0f;
-
-    // Create base Linear module (simple test without InjectLoRALayers)
-    auto base_linear = std::make_shared<nn::Linear>(64, 128, /*bias=*/true);
-
-    // Create a minimal wrapper test by manually testing what LoRAModel does
-    // Apply LoRA directly to the Linear layer
-    auto lora_linear = std::make_shared<LoRALinear>(base_linear, lora_config);
-
-    // Replace the base_linear in its container
-    // Note: In a real use case, you would use InjectLoRALayers on a transformer model
-
-    // Test GetLoRAParameters on the LoRA Linear
-    auto lora_params = GetLoRAParameters(lora_linear);
-    CHECK_GT(lora_params.size(), 0) << "Should have trainable parameters";
-    std::cout << "LoRA parameters extracted: " << lora_params.size() << std::endl;
-
-    // Test CountTrainableParameters
-    int64_t trainable = CountTrainableParameters(lora_linear);
-    CHECK_EQ(trainable, lora_config.rank * 64 + 128 * lora_config.rank);
-    std::cout << "Trainable parameters: " << trainable << std::endl;
-
-    // Test PrintSummary
-    std::cout << "\nLoRA Summary for Linear wrapper:" << std::endl;
-    PrintLoRASummary(lora_linear);
-
-    // Test Save/Load LoRA on the LoRA Linear
-    const std::string test_path = "/tmp/test_lora_linear.bin";
-    SaveLoRAWeights(lora_linear, test_path);
-    std::cout << "SaveLoRAWeights completed" << std::endl;
-
-    LoadLoRAWeights(lora_linear, test_path);
-    std::cout << "LoadLoRAWeights completed" << std::endl;
-
-    // Test Merge/Unmerge on LoRA Linear
-    CHECK(!lora_linear->IsMerged()) << "Should not be merged initially";
-    lora_linear->MergeWeights();
-    CHECK(lora_linear->IsMerged()) << "Should be merged after MergeWeights()";
-    std::cout << "MergeWeights completed" << std::endl;
-
-    lora_linear->UnmergeWeights();
-    CHECK(!lora_linear->IsMerged()) << "Should be unmerged after UnmergeWeights()";
-    std::cout << "UnmergeWeights completed" << std::endl;
-
-    std::cout << "LoRALinear utility tests passed!" << std::endl;
-}
-
-// ============================================================================
-// Test 8: Save/Load LoRA Weights
-// ============================================================================
-void test_lora_save_load_weights() {
-    std::cout << "\n=== Test 8: Save/Load LoRA Weights ===" << std::endl;
-
-    // Create a LoRALinear
-    LoRAConfig config;
-    config.rank = 4;
-    config.alpha = 8.0f;
-
-    int64_t in_features = 32;
-    int64_t out_features = 64;
-
-    auto linear = std::make_shared<nn::Linear>(in_features, out_features, /*bias=*/true);
-    auto lora_linear = std::make_shared<LoRALinear>(linear, config);
-
-    // Get references to lora_A and lora_B
-    auto lora_A = lora_linear->parameter(LoRALinear::kParamLoraAName);
-    auto lora_B = lora_linear->parameter(LoRALinear::kParamLoraBName);
-
-    // Set specific values to lora_A and lora_B
-    // lora_A: [rank, in_features] = [4, 32]
-    // lora_B: [out_features, rank] = [64, 4]
-    lora_A->EigenMatrix().setZero();
-    lora_B->EigenMatrix().setZero();
-
-    // Set lora_A to all 1s
-    for (int64_t i = 0; i < lora_A->Dims()[0]; ++i) {
-        for (int64_t j = 0; j < lora_A->Dims()[1]; ++j) { lora_A->EigenMatrix()(i, j) = 1.0f; }
-    }
-
-    // Set lora_B to all 2s
-    for (int64_t i = 0; i < lora_B->Dims()[0]; ++i) {
-        for (int64_t j = 0; j < lora_B->Dims()[1]; ++j) { lora_B->EigenMatrix()(i, j) = 2.0f; }
-    }
-
-    // Record original sums
-    float lora_A_sum_orig = lora_A->EigenMatrix().sum();
-    float lora_B_sum_orig = lora_B->EigenMatrix().sum();
-    // lora_A: all 1.0f, shape [rank, in_features] = [4, 32]
-    // lora_B: all 2.0f, shape [out_features, rank] = [64, 4]
-    float expected_lora_A_sum = config.rank * in_features * 1.0f;  // 4 * 32 * 1 = 128
-    float expected_lora_B_sum = out_features * config.rank * 2.0f; // 64 * 4 * 2 = 512
-    std::cout << "Original lora_A sum: " << lora_A_sum_orig << " (expected: " << expected_lora_A_sum << ")"
-              << std::endl;
-    std::cout << "Original lora_B sum: " << lora_B_sum_orig << " (expected: " << expected_lora_B_sum << ")"
-              << std::endl;
-
-    CHECK_EQ(lora_A_sum_orig, expected_lora_A_sum);
-    CHECK_EQ(lora_B_sum_orig, expected_lora_B_sum);
-
-    // Save to file
-    const std::string test_path = "/tmp/test_lora_save_load.bin";
-    SaveLoRAWeights(lora_linear, test_path);
-    std::cout << "Saved LoRA weights to: " << test_path << std::endl;
-
-    // Modify weights to different values
-    lora_A->EigenMatrix().setConstant(9.0f);
-    lora_B->EigenMatrix().setConstant(9.0f);
-
-    float lora_A_sum_modified = lora_A->EigenMatrix().sum();
-    float lora_B_sum_modified = lora_B->EigenMatrix().sum();
-    std::cout << "Modified lora_A sum: " << lora_A_sum_modified << std::endl;
-    std::cout << "Modified lora_B sum: " << lora_B_sum_modified << std::endl;
-
-    CHECK_NE(lora_A_sum_modified, lora_A_sum_orig);
-    CHECK_NE(lora_B_sum_modified, lora_B_sum_orig);
-
-    // Load from file
-    LoadLoRAWeights(lora_linear, test_path);
-    std::cout << "Loaded LoRA weights from: " << test_path << std::endl;
-
-    // Verify weights are restored
-    float lora_A_sum_loaded = lora_A->EigenMatrix().sum();
-    float lora_B_sum_loaded = lora_B->EigenMatrix().sum();
-    std::cout << "Loaded lora_A sum: " << lora_A_sum_loaded << std::endl;
-    std::cout << "Loaded lora_B sum: " << lora_B_sum_loaded << std::endl;
-
-    CHECK_EQ(lora_A_sum_loaded, lora_A_sum_orig) << "lora_A should be restored to original values";
-    CHECK_EQ(lora_B_sum_loaded, lora_B_sum_orig) << "lora_B should be restored to original values";
-
-    // Also verify individual elements
-    for (int64_t i = 0; i < lora_A->Dims()[0]; ++i) {
-        for (int64_t j = 0; j < lora_A->Dims()[1]; ++j) {
-            CHECK_EQ(lora_A->EigenMatrix()(i, j), 1.0f) << "lora_A element mismatch at (" << i << "," << j << ")";
-        }
-    }
-
-    for (int64_t i = 0; i < lora_B->Dims()[0]; ++i) {
-        for (int64_t j = 0; j < lora_B->Dims()[1]; ++j) {
-            CHECK_EQ(lora_B->EigenMatrix()(i, j), 2.0f) << "lora_B element mismatch at (" << i << "," << j << ")";
-        }
-    }
-
-    std::cout << "All elements verified correctly!" << std::endl;
-
-    // Cleanup
-    std::remove(test_path.c_str());
-    std::cout << "Test 8: Save/Load LoRA Weights passed!" << std::endl;
-}
-
-// ============================================================================
-// Test 8: ParseLoRATargetModules parsing
-// ============================================================================
-void test_set_target_modules() {
-    std::cout << "\n=== Test 8: ParseLoRATargetModules Parsing ===" << std::endl;
-
-    // Test single target
-    auto modules = ParseLoRATargetModules("c_attn");
-    CHECK_EQ(modules.size(), 1);
-    CHECK(modules.count("c_attn"));
-    std::cout << "Single target: OK" << std::endl;
-
-    // Test multiple targets
-    modules = ParseLoRATargetModules("c_attn,c_proj,c_fc");
-    CHECK_EQ(modules.size(), 3);
-    CHECK(modules.count("c_attn"));
-    CHECK(modules.count("c_proj"));
-    CHECK(modules.count("c_fc"));
-    std::cout << "Multiple targets: OK" << std::endl;
-
-    // Test with spaces
-    modules = ParseLoRATargetModules("c_attn, c_proj , c_fc");
-    CHECK_EQ(modules.size(), 3);
-    std::cout << "Targets with spaces: OK" << std::endl;
-
-    // Test empty/whitespace
-    modules = ParseLoRATargetModules("c_attn,,c_proj");
-    CHECK_EQ(modules.size(), 2);
-    std::cout << "Empty entries ignored: OK" << std::endl;
-
-    std::cout << "ParseLoRATargetModules tests passed!" << std::endl;
-}
-
-// ============================================================================
-// Test 9: ShouldApplyLoRA edge cases (attn.c_proj vs mlp.c_proj)
-// ============================================================================
-void test_should_apply_lora_edge_cases() {
-    std::cout << "\n=== Test 9: ShouldApplyLoRA Edge Cases ===" << std::endl;
-
-    // Test: Only attn.c_proj in target_modules
-    {
-        LoRAConfig config{8, 16.0f, 0.0f, ParseLoRATargetModules("c_attn,attn.c_proj")};
-
-        // Should match attention paths
-        CHECK(config.ShouldApplyLoRA("attn.c_proj"));
-        CHECK(config.ShouldApplyLoRA("transformer.h.0.attn.c_proj"));
-        CHECK(config.ShouldApplyLoRA("transformer.h.1.attn.c_proj"));
-
-        // Should NOT match mlp paths
-        CHECK(!config.ShouldApplyLoRA("mlp.c_proj"));
-        CHECK(!config.ShouldApplyLoRA("transformer.h.0.mlp.c_proj"));
-        std::cout << "attn.c_proj only: OK" << std::endl;
-    }
-
-    // Test: Only mlp.c_proj in target_modules
-    {
-        LoRAConfig config{8, 16.0f, 0.0f, ParseLoRATargetModules("c_attn,mlp.c_proj")};
-
-        // Should NOT match attention paths
-        CHECK(!config.ShouldApplyLoRA("attn.c_proj"));
-        CHECK(!config.ShouldApplyLoRA("transformer.h.0.attn.c_proj"));
-
-        // Should match mlp paths
-        CHECK(config.ShouldApplyLoRA("mlp.c_proj"));
-        CHECK(config.ShouldApplyLoRA("transformer.h.0.mlp.c_proj"));
-        std::cout << "mlp.c_proj only: OK" << std::endl;
-    }
-
-    // Test: Generic c_proj in target_modules (matches both)
-    {
-        LoRAConfig config{8, 16.0f, 0.0f, ParseLoRATargetModules("c_attn,c_proj")};
-
-        // Should match both attention and mlp
-        CHECK(config.ShouldApplyLoRA("transformer.h.0.attn.c_proj"));
-        CHECK(config.ShouldApplyLoRA("transformer.h.0.mlp.c_proj"));
-        std::cout << "Generic c_proj (matches both): OK" << std::endl;
-    }
-
-    // Test: All targets
-    {
-        LoRAConfig config{8, 16.0f, 0.0f, ParseLoRATargetModules("c_attn,attn.c_proj,c_fc,c_fc2,mlp.c_proj")};
-
-        CHECK(config.ShouldApplyLoRA("transformer.h.0.attn.c_attn"));
-        CHECK(config.ShouldApplyLoRA("transformer.h.0.attn.c_proj"));
-        CHECK(config.ShouldApplyLoRA("transformer.h.0.mlp.c_fc"));
-        CHECK(config.ShouldApplyLoRA("transformer.h.0.mlp.c_fc2"));
-        CHECK(config.ShouldApplyLoRA("transformer.h.0.mlp.c_proj"));
-        std::cout << "All targets: OK" << std::endl;
-    }
-
-    std::cout << "ShouldApplyLoRA edge cases tests passed!" << std::endl;
-}
-
-// ============================================================================
-// Test 10: ReplaceModuleByPath
-// ============================================================================
-void test_replace_module_by_path() {
-    std::cout << "\n=== Test 10: ReplaceModuleByPath ===" << std::endl;
-
-    // Test ReplaceModuleByPath by wrapping a Linear with LoRA directly
-    // This tests the core functionality that ReplaceModuleByPath provides
-
-    // Create base Linear
-    auto base_linear = std::make_shared<nn::Linear>(64, 128, /*bias=*/true);
-
-    // Configure LoRA
-    LoRAConfig lora_config;
-    lora_config.rank = 4;
-    lora_config.alpha = 8.0f;
-
-    // Wrap with LoRA - this is what ReplaceModuleByPath does internally
-    auto lora_linear = std::make_shared<LoRALinear>(base_linear, lora_config);
-
-    // Verify LoRA was applied correctly
-    auto params = lora_linear->LoRAParameters();
-    CHECK_EQ(params.size(), 2) << "LoRALinear should have 2 trainable parameters (lora_A and lora_B)";
-    std::cout << "LoRALinear has " << params.size() << " trainable parameters" << std::endl;
-
-    // Verify parameter shapes
-    auto lora_a = params[0];
-    auto lora_b = params[1];
-    CHECK_EQ(lora_a->Dims()[0], lora_config.rank); // rank x in_features
-    CHECK_EQ(lora_a->Dims()[1], 64);
-    CHECK_EQ(lora_b->Dims()[0], 128); // out_features x rank
-    CHECK_EQ(lora_b->Dims()[1], lora_config.rank);
-    std::cout << "LoRA parameter shapes: OK" << std::endl;
-
-    // Verify base parameters are frozen (use named parameters instead of index)
-    auto weight = lora_linear->parameter(nn::Linear::kParamWeightName);
-    auto lora_a_param = lora_linear->parameter(LoRALinear::kParamLoraAName);
-    auto lora_b_param = lora_linear->parameter(LoRALinear::kParamLoraBName);
-    CHECK(weight != nullptr);
-    CHECK(lora_a_param != nullptr);
-    CHECK(lora_b_param != nullptr);
-    CHECK(!weight->requires_grad());      // weight is frozen
-    CHECK(lora_a_param->requires_grad()); // lora_A is trainable
-    CHECK(lora_b_param->requires_grad()); // lora_B is trainable
-    std::cout << "Base weight frozen, LoRA params trainable: OK" << std::endl;
-
-    std::cout << "ReplaceModuleByPath tests passed!" << std::endl;
-}
-
-// ============================================================================
-// Test 11: FreezeBaseModel / UnfreezeModel
-// ============================================================================
-void test_freeze_unfreeze() {
-    std::cout << "\n=== Test 11: FreezeBaseModel / UnfreezeModel ===" << std::endl;
-
-    // Test with LoRALinear directly - it has both base and LoRA params
-    LoRAConfig lora_config;
-    lora_config.rank = 4;
-    lora_config.alpha = 8.0f;
-
-    auto linear = std::make_shared<nn::Linear>(64, 128, /*bias=*/true);
-    auto lora_linear = std::make_shared<LoRALinear>(linear, lora_config);
-
-    // Get all parameters from LoRALinear (includes base + LoRA)
-    auto all_params = lora_linear->Parameters();
-
-    // Initially only LoRA params should be trainable (base weights are frozen by constructor)
-    int64_t total_params = 0;
-    for (const auto &p : all_params) {
-        if (p->requires_grad()) {
-            total_params += p->NumElements();
-        }
-    }
-    // Expected: only LoRA params (lora_A + lora_B) = 4*64 + 128*4 = 256 + 512 = 768
-    // Note: LoRALinear freezes base weights in constructor by design
-    int64_t expected_total = lora_config.rank * 64 + 128 * lora_config.rank;
-    CHECK_EQ(total_params, expected_total);
-    std::cout << "Initial trainable params: " << total_params << " (expected: " << expected_total << ")" << std::endl;
-
-    // FreezeBaseModel on LoRALinear
-    FreezeBaseModel(lora_linear);
-
-    // After freeze, only LoRA params should be trainable
-    int64_t after_freeze = 0;
-    for (const auto &p : all_params) {
-        if (p->requires_grad()) {
-            after_freeze += p->NumElements();
-        }
-    }
-    // LoRA params: A (rank x in) + B (out x rank) = 4*64 + 128*4 = 256 + 512 = 768
-    int64_t expected_lora = lora_config.rank * 64 + 128 * lora_config.rank;
-    CHECK_EQ(after_freeze, expected_lora);
-    std::cout << "After freeze trainable: " << after_freeze << " (expected: " << expected_lora << ")" << std::endl;
-
-    // Unfreeze all
-    UnfreezeModel(lora_linear);
-    int64_t after_unfreeze = 0;
-    for (const auto &p : all_params) {
-        if (p->requires_grad()) {
-            after_unfreeze += p->NumElements();
-        }
-    }
-    // Should be back to all params trainable (base + LoRA)
-    int64_t expected_after_unfreeze = 64 * 128 + 128 + lora_config.rank * 64 + 128 * lora_config.rank;
-    CHECK_EQ(after_unfreeze, expected_after_unfreeze);
-    std::cout << "After unfreeze trainable: " << after_unfreeze << std::endl;
-
-    std::cout << "FreezeBaseModel / UnfreezeModel tests passed!" << std::endl;
-}
-
-// ============================================================================
-// Test 12: LoRAStateDict
-// ============================================================================
-void test_lora_state_dict() {
-    std::cout << "\n=== Test 12: LoRAStateDict ===" << std::endl;
-
-    // Test with a single LoRALinear
-    LoRAConfig lora_config;
-    lora_config.rank = 4;
-    lora_config.alpha = 8.0f;
-
-    auto linear = std::make_shared<nn::Linear>(64, 128, /*bias=*/true);
-    auto lora_linear = std::make_shared<LoRALinear>(linear, lora_config);
-
-    // Get state dict - it contains all parameters with their names
-    auto state_dict = lora_linear->StateDict();
-
-    // Check that we have all expected parameters
-    CHECK(state_dict.count("weight")) << "Should have weight parameter";
-    CHECK(state_dict.count("bias")) << "Should have bias parameter";
-    CHECK(state_dict.count("lora_A")) << "Should have lora_A parameter";
-    CHECK(state_dict.count("lora_B")) << "Should have lora_B parameter";
-    std::cout << "State dict contains: weight, bias, lora_A, lora_B" << std::endl;
-
-    // Verify LoRA parameters exist and are trainable
-    CHECK(state_dict.at("lora_A")->requires_grad()) << "lora_A should be trainable";
-    CHECK(state_dict.at("lora_B")->requires_grad()) << "lora_B should be trainable";
-    CHECK(!state_dict.at("weight")->requires_grad()) << "weight should be frozen";
-    std::cout << "LoRA parameters are trainable, base weight is frozen: OK" << std::endl;
-
-    // Verify shapes
-    CHECK_EQ(state_dict.at("lora_A")->Dims()[0], lora_config.rank);
-    CHECK_EQ(state_dict.at("lora_A")->Dims()[1], 64);
-    CHECK_EQ(state_dict.at("lora_B")->Dims()[0], 128);
-    CHECK_EQ(state_dict.at("lora_B")->Dims()[1], lora_config.rank);
-    std::cout << "LoRA parameter shapes: OK" << std::endl;
-
-    std::cout << "LoRAStateDict tests passed!" << std::endl;
-}
-
-// ============================================================================
-// Test 13: GetLoRAModel simplified API
-// ============================================================================
-void test_get_lora_model() {
-    std::cout << "\n=== Test 13: GetLoRAModel Simplified API ===" << std::endl;
-
-    // Test GetLoRAModel with a simple Linear layer
-    // We'll wrap it with LoRA directly and verify the wrapper works
-
-    // Create base Linear
-    auto base_linear = std::make_shared<nn::Linear>(64, 128, /*bias=*/true);
-
-    // Configure LoRA
-    LoRAConfig config{4, 8.0f, 0.0f, ParseLoRATargetModules("Linear")};
-
-    // Use GetLoRAModel with the linear as the "model"
-    // Note: GetLoRAModel returns the modified model (in-place injection)
-    auto model = GetLoRAModel(base_linear, config);
-
-    CHECK(model != nullptr);
-    std::cout << "GetLoRAModel returned valid pointer" << std::endl;
-
-    // Test that LoRA was applied - check trainable parameters
-    auto lora_params = GetLoRAParameters(model);
-    // GetLoRAParameters returns vector<shared_ptr<Tensor>>, size() is the count of tensors
-    // LoRALinear has 2 trainable tensors: lora_A (rank x in) and lora_B (out x rank)
-    CHECK_EQ(lora_params.size(), 2);
-    std::cout << "Trainable parameter tensors: " << lora_params.size() << " (expected: 2)" << std::endl;
-
-    // Also verify total element count
-    int64_t total_elements = 0;
-    for (const auto &t : lora_params) { total_elements += t->NumElements(); }
-    int64_t expected_elements = config.rank * 64 + 128 * config.rank; // 768
-    CHECK_EQ(total_elements, expected_elements);
-    std::cout << "Total trainable elements: " << total_elements << " (expected: " << expected_elements << ")"
-              << std::endl;
-
-    // Test PrintSummary
-    std::cout << "\nLoRA Model Summary:" << std::endl;
-    PrintLoRASummary(model);
-
-    // Test Merge/Unmerge using utility functions
-    MergeLoRAWeights(model);
-    // Verify LoRA params frozen after merge
-    auto *lora_mod = dynamic_cast<LoRALinear *>(model.get());
-    CHECK(lora_mod != nullptr);
-    CHECK(!lora_mod->LoRAParameters()[0]->requires_grad()) << "lora_A should be frozen after merge";
-    CHECK(!lora_mod->LoRAParameters()[1]->requires_grad()) << "lora_B should be frozen after merge";
-    std::cout << "Merge: OK (LoRA params frozen)" << std::endl;
-
-    UnmergeLoRAWeights(model);
-    CHECK(lora_mod->LoRAParameters()[0]->requires_grad()) << "lora_A should be trainable after unmerge";
-    CHECK(lora_mod->LoRAParameters()[1]->requires_grad()) << "lora_B should be trainable after unmerge";
-    std::cout << "Unmerge: OK (LoRA params trainable)" << std::endl;
-
-    std::cout << "GetLoRAModel in-place injection tests passed!" << std::endl;
-}
-
-// ============================================================================
-// Test 14: MergeAndUnload
-// ============================================================================
-void test_merge_and_unload() {
-    std::cout << "\n=== Test 14: MergeAndUnload ===" << std::endl;
-
-    // Create base Linear and apply LoRA
-    auto base_linear = std::make_shared<nn::Linear>(64, 128, /*bias=*/true);
-    LoRAConfig config{4, 8.0f, 0.0f, ParseLoRATargetModules("Linear")};
-    auto model = GetLoRAModel(base_linear, config);
-
-    // Verify it's a LoRA module
-    CHECK(dynamic_cast<LoRALinear *>(model.get()) != nullptr) << "Should be LoRALinear";
-
-    // Create input and get output before merge_and_unload
-    auto input = std::make_shared<Tensor>(std::vector<int64_t>{2, 5, 64}, DataType::kFLOAT32);
-    input->EigenMatrix().setRandom();
-    auto output_before = (*model)({input})[0];
-    float output_before_sum = output_before->EigenMatrix().sum();
-    std::cout << "Output sum before MergeAndUnload: " << output_before_sum << std::endl;
-
-    // MergeAndUnload
-    auto unloaded_model = MergeAndUnload(model);
-    CHECK(unloaded_model != nullptr) << "MergeAndUnload should return valid model";
-
-    // Verify it's no longer a LoRA module
-    CHECK(dynamic_cast<LoRALinear *>(unloaded_model.get()) == nullptr) << "Should be plain Linear after MergeAndUnload";
-    std::cout << "Model is no longer LoRALinear: OK" << std::endl;
-
-    // Verify no LoRA parameters exist (check state dict)
-    auto state_dict = unloaded_model->StateDict();
-    for (const auto &[name, param] : state_dict) {
-        CHECK(name.find("lora_A") == std::string::npos && name.find("lora_B") == std::string::npos)
-            << "Should not have LoRA parameters after MergeAndUnload, found: " << name;
-    }
-    std::cout << "No LoRA parameters in state dict: OK" << std::endl;
-
-    // Verify forward output matches (merged output should equal unmerged LoRA output)
-    auto output_after = (*unloaded_model)({input})[0];
-    float output_after_sum = output_after->EigenMatrix().sum();
-    std::cout << "Output sum after MergeAndUnload: " << output_after_sum << std::endl;
-    std::cout << "Diff: " << std::abs(output_before_sum - output_after_sum) << std::endl;
-    CHECK(std::abs(output_before_sum - output_after_sum) < 1e-3) << "Output should match after MergeAndUnload";
-
-    // Verify all parameters have requires_grad = true (unfrozen)
-    for (const auto &param : unloaded_model->Parameters()) {
-        CHECK(param->requires_grad()) << "All parameters should be trainable after MergeAndUnload";
-    }
-    std::cout << "All parameters trainable: OK" << std::endl;
-
-    std::cout << "MergeAndUnload tests passed!" << std::endl;
-}
-
-int main(int argc, char **argv) {
-    google::InitGoogleLogging(argv[0]);
-    FLAGS_logtostderr = 1;
-
-    // Initialize parallel settings (required for some tensor operations)
-    // Parameters: nthread_per_process, tensor_parallel_size, sequence_parallel_enabled,
-    //             pipeline_parallel_size, virtual_pipeline_parallel_size
-    nn::parallel::global::InitAllEnv(1, 1, false, 1, 1);
-
-    std::cout << "========================================" << std::endl;
-    std::cout << "       LoRA Module Unit Tests          " << std::endl;
-    std::cout << "========================================" << std::endl;
-
-    test_lora_config();
-    test_lora_linear_init();
-    test_lora_linear_forward();
-    test_lora_linear_merge();
-    test_lora_utils();
-    test_lora_from_linear();
-    test_lora_model_wrapper();
-    test_lora_save_load_weights();
-    test_set_target_modules();
-    test_should_apply_lora_edge_cases();
-    test_replace_module_by_path();
-    test_freeze_unfreeze();
-    test_lora_state_dict();
-    test_get_lora_model();
-    test_merge_and_unload();
-
-    std::cout << "\n========================================" << std::endl;
-    std::cout << "       All LoRA Tests Passed!          " << std::endl;
-    std::cout << "========================================" << std::endl;
-
-    return 0;
-}
diff --git a/test/transformer/test_transformer_architecture.cc b/test/transformer/test_transformer_architecture.cc
deleted file mode 100644
index 1c6ed28f..00000000
--- a/test/transformer/test_transformer_architecture.cc
+++ /dev/null
@@ -1,556 +0,0 @@
-#include <cmath>
-#include <iostream>
-#include <memory>
-
-#include "glog/logging.h"
-
-#include "example/gpt2/config.h"
-#include "example/llama3/config.h"
-#include "infini_train/include/nn/modules/activations.h"
-#include "infini_train/include/nn/modules/normalization.h"
-#include "infini_train/include/nn/modules/sparse.h"
-#include "infini_train/include/nn/modules/transformer/causal_self_attention.h"
-#include "infini_train/include/nn/modules/transformer/mlp.h"
-#include "infini_train/include/nn/modules/transformer/transformer.h"
-#include "infini_train/include/nn/modules/transformer/transformer_config.h"
-#include "infini_train/include/nn/modules/transformer/utils.h"
-#include "infini_train/include/nn/parallel/global.h"
-#include "infini_train/include/tensor.h"
-
-using namespace infini_train;
-namespace nn = infini_train::nn;
-
-// ============================================================================
-// Test 1: TransformerConfig Validation
-// ============================================================================
-void TestConfigValidation() {
-    std::cout << "\n=== Test 1: TransformerConfig Validation ===" << std::endl;
-
-    bool all_passed = true;
-
-    // Test GPT2 config
-    auto gpt2_config = gpt2::GPT2Config();
-    if (gpt2_config.attention_type != nn::AttentionType::kStandard) {
-        std::cout << "FAIL: GPT2 config should use Standard attention" << std::endl;
-        all_passed = false;
-    }
-    if (gpt2_config.activation_type != nn::MLPType::kGELU) {
-        std::cout << "FAIL: GPT2 config should use GELU activation" << std::endl;
-        all_passed = false;
-    }
-    if (gpt2_config.norm_type != nn::NormType::kLayerNorm) {
-        std::cout << "FAIL: GPT2 config should use LayerNorm" << std::endl;
-        all_passed = false;
-    }
-    if (!gpt2_config.add_bias_linear) {
-        std::cout << "FAIL: GPT2 config should have bias enabled" << std::endl;
-        all_passed = false;
-    }
-    if (!gpt2_config.tie_weights) {
-        std::cout << "FAIL: GPT2 config should have tied weights" << std::endl;
-        all_passed = false;
-    }
-
-    // Test LLaMA3 config
-    auto llama3_config = llama3::LLaMA3Config();
-    if (llama3_config.attention_type != nn::AttentionType::kRoPE) {
-        std::cout << "FAIL: LLaMA3 config should use RoPE attention" << std::endl;
-        all_passed = false;
-    }
-    if (llama3_config.activation_type != nn::MLPType::kSwiGLU) {
-        std::cout << "FAIL: LLaMA3 config should use SwiGLU activation" << std::endl;
-        all_passed = false;
-    }
-    if (llama3_config.norm_type != nn::NormType::kRMSNorm) {
-        std::cout << "FAIL: LLaMA3 config should use RMSNorm" << std::endl;
-        all_passed = false;
-    }
-    if (llama3_config.add_bias_linear) {
-        std::cout << "FAIL: LLaMA3 config should have bias disabled" << std::endl;
-        all_passed = false;
-    }
-    if (llama3_config.tie_weights) {
-        std::cout << "FAIL: LLaMA3 config should not have tied weights" << std::endl;
-        all_passed = false;
-    }
-
-    // Test GQA detection
-    if (!llama3_config.UseGQA()) {
-        std::cout << "FAIL: LLaMA3 config should detect GQA (n_kv_head < n_head)" << std::endl;
-        all_passed = false;
-    }
-    if (gpt2_config.UseGQA()) {
-        std::cout << "FAIL: GPT2 config should not detect GQA (n_kv_head == n_head)" << std::endl;
-        all_passed = false;
-    }
-
-    if (all_passed) {
-        std::cout << "SUCCESS: All config validations passed!" << std::endl;
-    }
-}
-
-// ============================================================================
-// Test 2: Embedding Layer
-// ============================================================================
-void TestEmbedding() {
-    std::cout << "\n=== Test 2: Embedding Layer ===" << std::endl;
-
-    const int64_t vocab_size = 1000;
-    const int64_t embedding_dim = 128;
-    const int64_t batch_size = 2;
-    const int64_t seq_len = 16;
-
-    try {
-        auto embedding = std::make_shared<nn::Embedding>(vocab_size, embedding_dim);
-
-        // Check parameters
-        auto params = embedding->Parameters();
-        if (params.size() != 1) {
-            std::cout << "FAIL: Embedding should have 1 parameter, got " << params.size() << std::endl;
-            return;
-        }
-
-        // Check weight shape
-        auto weight = embedding->parameter(nn::Embedding::kParamWeightName);
-        if (weight->Dims() != std::vector<int64_t>{vocab_size, embedding_dim}) {
-            std::cout << "FAIL: Embedding weight shape mismatch" << std::endl;
-            return;
-        }
-
-        // Forward pass
-        auto input = std::make_shared<Tensor>(std::vector<int64_t>{batch_size, seq_len}, DataType::kINT64);
-        auto output = (*embedding)({input});
-
-        if (output.size() != 1) {
-            std::cout << "FAIL: Embedding forward should return 1 tensor" << std::endl;
-            return;
-        }
-
-        const auto &out_dims = output[0]->Dims();
-        if (out_dims != std::vector<int64_t>{batch_size, seq_len, embedding_dim}) {
-            std::cout << "FAIL: Embedding output shape mismatch. Expected [" << batch_size << ", " << seq_len << ", "
-                      << embedding_dim << "], got [" << out_dims[0] << ", " << out_dims[1] << ", " << out_dims[2] << "]"
-                      << std::endl;
-            return;
-        }
-
-        std::cout << "SUCCESS: Embedding layer works correctly!" << std::endl;
-
-    } catch (const std::exception &e) { std::cout << "FAIL: Exception: " << e.what() << std::endl; }
-}
-
-// ============================================================================
-// Test 3: Normalization Layers (LayerNorm vs RMSNorm)
-// ============================================================================
-void TestNormalization() {
-    std::cout << "\n=== Test 3: Normalization Layers ===" << std::endl;
-
-    const int64_t hidden_size = 64;
-    const int64_t batch_size = 2;
-    const int64_t seq_len = 8;
-
-    try {
-        // Test LayerNorm
-        auto layernorm = std::make_shared<nn::LayerNorm>(std::vector<int64_t>{hidden_size});
-        auto ln_params = layernorm->Parameters();
-        if (ln_params.size() != 2) {
-            std::cout << "FAIL: LayerNorm should have 2 parameters (weight, bias), got " << ln_params.size()
-                      << std::endl;
-            return;
-        }
-
-        // Test RMSNorm
-        auto rmsnorm = std::make_shared<nn::RMSNorm>(hidden_size);
-        auto rms_params = rmsnorm->Parameters();
-        if (rms_params.size() != 1) {
-            std::cout << "FAIL: RMSNorm should have 1 parameter (weight), got " << rms_params.size() << std::endl;
-            return;
-        }
-
-        // Forward pass for both
-        auto input
-            = std::make_shared<Tensor>(std::vector<int64_t>{batch_size, seq_len, hidden_size}, DataType::kFLOAT32);
-
-        auto ln_output = (*layernorm)({input});
-        auto rms_output = (*rmsnorm)({input});
-
-        if (ln_output[0]->Dims() != input->Dims()) {
-            std::cout << "FAIL: LayerNorm output shape mismatch" << std::endl;
-            return;
-        }
-
-        if (rms_output[0]->Dims() != input->Dims()) {
-            std::cout << "FAIL: RMSNorm output shape mismatch" << std::endl;
-            return;
-        }
-
-        std::cout << "SUCCESS: Normalization layers work correctly!" << std::endl;
-
-    } catch (const std::exception &e) { std::cout << "FAIL: Exception: " << e.what() << std::endl; }
-}
-
-// ============================================================================
-// Test 4: MLP Layer (GELU vs SwiGLU)
-// ============================================================================
-void TestMlp() {
-    std::cout << "\n=== Test 4: MLP Layer ===" << std::endl;
-
-    const int64_t hidden_size = 64;
-    const int64_t batch_size = 2;
-    const int64_t seq_len = 8;
-
-    try {
-        // Test GPT2-style MLP (GELU)
-        nn::TransformerConfig gpt2_mlp_config;
-        gpt2_mlp_config.n_embd = hidden_size;
-        gpt2_mlp_config.activation_type = nn::MLPType::kGELU;
-        gpt2_mlp_config.ffn_expansion_ratio = 4.0f;
-        gpt2_mlp_config.add_bias_linear = true;
-
-        auto gpt2_mlp = std::make_shared<nn::MLP>(gpt2_mlp_config);
-        auto gpt2_params = gpt2_mlp->Parameters();
-
-        // GPT2 MLP should have: c_fc.weight, c_fc.bias, c_proj.weight, c_proj.bias
-        if (gpt2_params.size() != 4) {
-            std::cout << "FAIL: GPT2 MLP should have 4 parameters, got " << gpt2_params.size() << std::endl;
-            return;
-        }
-
-        // Test LLaMA3-style MLP (SwiGLU)
-        nn::TransformerConfig llama3_mlp_config;
-        llama3_mlp_config.n_embd = hidden_size;
-        llama3_mlp_config.activation_type = nn::MLPType::kSwiGLU;
-        llama3_mlp_config.ffn_expansion_ratio = 4.0f;
-        llama3_mlp_config.add_bias_linear = false;
-        llama3_mlp_config.ffn_dim_multiplier = 1.5f;
-        llama3_mlp_config.multiple_of = 256;
-
-        auto llama3_mlp = std::make_shared<nn::MLP>(llama3_mlp_config);
-        auto llama3_params = llama3_mlp->Parameters();
-
-        // LLaMA3 MLP should have: c_fc.weight, c_fc2.weight, c_proj.weight (no bias)
-        if (llama3_params.size() != 3) {
-            std::cout << "FAIL: LLaMA3 MLP should have 3 parameters, got " << llama3_params.size() << std::endl;
-            return;
-        }
-
-        // Forward pass
-        auto input
-            = std::make_shared<Tensor>(std::vector<int64_t>{batch_size, seq_len, hidden_size}, DataType::kFLOAT32);
-
-        auto gpt2_output = (*gpt2_mlp)({input});
-        auto llama3_output = (*llama3_mlp)({input});
-
-        // Output should have same hidden dimension
-        if (gpt2_output[0]->Dims()[2] != hidden_size) {
-            std::cout << "FAIL: GPT2 MLP output hidden dim mismatch" << std::endl;
-            return;
-        }
-
-        if (llama3_output[0]->Dims()[2] != hidden_size) {
-            std::cout << "FAIL: LLaMA3 MLP output hidden dim mismatch" << std::endl;
-            return;
-        }
-
-        std::cout << "SUCCESS: MLP layers work correctly!" << std::endl;
-
-    } catch (const std::exception &e) { std::cout << "FAIL: Exception: " << e.what() << std::endl; }
-}
-
-// ============================================================================
-// Test 5: CausalSelfAttention
-// ============================================================================
-void TestAttention() {
-    std::cout << "\n=== Test 5: CausalSelfAttention ===" << std::endl;
-
-    const int64_t hidden_size = 64;
-    const int64_t batch_size = 2;
-    const int64_t seq_len = 8;
-    const int64_t n_head = 4;
-
-    try {
-        // Test standard attention (GPT2-style)
-        nn::TransformerConfig standard_config;
-        standard_config.n_embd = hidden_size;
-        standard_config.n_head = n_head;
-        standard_config.n_kv_head = n_head;
-        standard_config.attention_type = nn::AttentionType::kStandard;
-        standard_config.add_bias_linear = true;
-
-        auto standard_attn = std::make_shared<nn::CausalSelfAttention>(standard_config);
-        auto standard_params = standard_attn->Parameters();
-
-        // Should have c_attn (QKV combined) and c_proj with biases
-        if (standard_params.size() != 4) {
-            std::cout << "FAIL: Standard attention should have 4 parameters, got " << standard_params.size()
-                      << std::endl;
-            return;
-        }
-
-        // Test RoPE attention with GQA (LLaMA3-style)
-        nn::TransformerConfig rope_config;
-        rope_config.n_embd = hidden_size;
-        rope_config.n_head = n_head;
-        rope_config.n_kv_head = 2; // GQA: fewer KV heads
-        rope_config.attention_type = nn::AttentionType::kRoPE;
-        rope_config.add_bias_linear = false;
-
-        auto rope_attn = std::make_shared<nn::CausalSelfAttention>(rope_config);
-        auto rope_params = rope_attn->Parameters();
-
-        // RoPE attention without bias should have fewer params
-        if (rope_params.empty()) {
-            std::cout << "FAIL: RoPE attention should have parameters" << std::endl;
-            return;
-        }
-
-        // Forward pass
-        auto input
-            = std::make_shared<Tensor>(std::vector<int64_t>{batch_size, seq_len, hidden_size}, DataType::kFLOAT32);
-
-        auto standard_output = (*standard_attn)({input});
-        if (standard_output[0]->Dims() != input->Dims()) {
-            std::cout << "FAIL: Standard attention output shape mismatch" << std::endl;
-            return;
-        }
-
-        std::cout << "SUCCESS: CausalSelfAttention works correctly!" << std::endl;
-
-    } catch (const std::exception &e) { std::cout << "FAIL: Exception: " << e.what() << std::endl; }
-}
-
-// ============================================================================
-// Test 6: TransformerLayer
-// ============================================================================
-void TestTransformerLayer() {
-    std::cout << "\n=== Test 6: TransformerLayer ===" << std::endl;
-
-    const int64_t hidden_size = 64;
-    const int64_t batch_size = 2;
-    const int64_t seq_len = 8;
-
-    try {
-        // Test GPT2-style layer
-        auto gpt2_config = gpt2::GPT2Config();
-        gpt2_config.n_embd = hidden_size;
-        gpt2_config.n_head = 4;
-        gpt2_config.n_layer = 1;
-
-        auto gpt2_layer = std::make_shared<nn::TransformerLayer>(gpt2_config);
-        auto gpt2_params = gpt2_layer->Parameters();
-
-        if (gpt2_params.empty()) {
-            std::cout << "FAIL: GPT2 TransformerLayer should have parameters" << std::endl;
-            return;
-        }
-
-        // Forward pass
-        auto input
-            = std::make_shared<Tensor>(std::vector<int64_t>{batch_size, seq_len, hidden_size}, DataType::kFLOAT32);
-
-        auto output = (*gpt2_layer)({input});
-        if (output[0]->Dims() != input->Dims()) {
-            std::cout << "FAIL: TransformerLayer output shape mismatch" << std::endl;
-            return;
-        }
-
-        // Test LLaMA3-style layer
-        auto llama3_config = llama3::LLaMA3Config();
-        llama3_config.n_embd = hidden_size;
-        llama3_config.n_head = 4;
-        llama3_config.n_kv_head = 2;
-        llama3_config.n_layer = 1;
-
-        auto llama3_layer = std::make_shared<nn::TransformerLayer>(llama3_config);
-        auto llama3_params = llama3_layer->Parameters();
-
-        if (llama3_params.empty()) {
-            std::cout << "FAIL: LLaMA3 TransformerLayer should have parameters" << std::endl;
-            return;
-        }
-
-        std::cout << "SUCCESS: TransformerLayer works correctly!" << std::endl;
-
-    } catch (const std::exception &e) { std::cout << "FAIL: Exception: " << e.what() << std::endl; }
-}
-
-// ============================================================================
-// Test 7: TransformerModel Instantiation (GPT2)
-// ============================================================================
-void TestGpt2Model() {
-    std::cout << "\n=== Test 7: GPT2 Model Instantiation ===" << std::endl;
-
-    auto config = gpt2::GPT2Config();
-    // Use smaller config for faster testing
-    config.n_layer = 2;
-    config.n_head = 4;
-    config.n_embd = 64;
-
-    try {
-        auto model = std::make_shared<nn::TransformerModel>(config);
-
-        if (model == nullptr) {
-            std::cout << "FAIL: Failed to create GPT2 model" << std::endl;
-            return;
-        }
-
-        auto params = model->Parameters();
-        if (params.empty()) {
-            std::cout << "FAIL: GPT2 model has no parameters" << std::endl;
-            return;
-        }
-
-        std::cout << "SUCCESS: GPT2 model created with " << params.size() << " parameters!" << std::endl;
-
-    } catch (const std::exception &e) { std::cout << "FAIL: Exception: " << e.what() << std::endl; }
-}
-
-// ============================================================================
-// Test 8: TransformerModel Instantiation (LLaMA3)
-// ============================================================================
-void TestLlama3Model() {
-    std::cout << "\n=== Test 8: LLaMA3 Model Instantiation ===" << std::endl;
-
-    auto config = llama3::LLaMA3Config();
-    // Use smaller config for faster testing
-    config.n_layer = 2;
-    config.n_head = 4;
-    config.n_kv_head = 2;
-    config.n_embd = 64;
-
-    try {
-        auto model = std::make_shared<nn::TransformerModel>(config);
-
-        if (model == nullptr) {
-            std::cout << "FAIL: Failed to create LLaMA3 model" << std::endl;
-            return;
-        }
-
-        auto params = model->Parameters();
-        if (params.empty()) {
-            std::cout << "FAIL: LLaMA3 model has no parameters" << std::endl;
-            return;
-        }
-
-        std::cout << "SUCCESS: LLaMA3 model created with " << params.size() << " parameters!" << std::endl;
-
-    } catch (const std::exception &e) { std::cout << "FAIL: Exception: " << e.what() << std::endl; }
-}
-
-// ============================================================================
-// Test 9: RoPE Utilities
-// ============================================================================
-void TestRopeUtils() {
-    std::cout << "\n=== Test 9: RoPE Utilities ===" << std::endl;
-
-    const int64_t head_dim = 64;
-    const int64_t seq_len = 128;
-
-    try {
-        // Test precompute freqs_cis
-        auto freqs_cis = PrecomputeFreqsCis(head_dim, seq_len);
-
-        // freqs_cis shape: [seq_len, head_dim/2, 2] (cos and sin stacked on last dim)
-        const auto &dims = freqs_cis->Dims();
-        if (dims.size() != 3) {
-            std::cout << "FAIL: freqs_cis should be 3D, got " << dims.size() << "D" << std::endl;
-            return;
-        }
-        if (dims[0] != seq_len) {
-            std::cout << "FAIL: freqs_cis seq_len mismatch. Expected " << seq_len << ", got " << dims[0] << std::endl;
-            return;
-        }
-        if (dims[1] != head_dim / 2) {
-            std::cout << "FAIL: freqs_cis head_dim/2 mismatch. Expected " << head_dim / 2 << ", got " << dims[1]
-                      << std::endl;
-            return;
-        }
-        if (dims[2] != 2) {
-            std::cout << "FAIL: freqs_cis last dim should be 2 (cos, sin), got " << dims[2] << std::endl;
-            return;
-        }
-
-        std::cout << "SUCCESS: RoPE utilities work correctly!" << std::endl;
-
-    } catch (const std::exception &e) { std::cout << "FAIL: Exception: " << e.what() << std::endl; }
-}
-
-// ============================================================================
-// Test 10: Model StateDict
-// ============================================================================
-void TestStateDict() {
-    std::cout << "\n=== Test 10: Model StateDict ===" << std::endl;
-
-    nn::TransformerConfig config;
-    config.n_layer = 1;
-    config.n_head = 2;
-    config.n_kv_head = 2; // Must set explicitly
-    config.n_embd = 32;
-    config.vocab_size = 1000;
-    config.attention_type = nn::AttentionType::kStandard;
-    config.activation_type = nn::MLPType::kGELU;
-    config.norm_type = nn::NormType::kLayerNorm;
-    config.add_bias_linear = true;
-
-    try {
-        auto model = std::make_shared<nn::TransformerModel>(config);
-        auto state_dict = model->StateDict();
-
-        if (state_dict.empty()) {
-            std::cout << "FAIL: StateDict should not be empty" << std::endl;
-            return;
-        }
-
-        // StateDict includes both parameters and buffers, so it should have >= parameters count
-        auto params = model->Parameters();
-        auto buffers = model->Buffers();
-
-        if (state_dict.size() < params.size()) {
-            std::cout << "FAIL: StateDict size (" << state_dict.size() << ") should be >= parameter count ("
-                      << params.size() << ")" << std::endl;
-            return;
-        }
-
-        // Expected: state_dict.size() == params.size() + buffers.size()
-        size_t expected_size = params.size() + buffers.size();
-        if (state_dict.size() != expected_size) {
-            std::cout << "FAIL: StateDict size (" << state_dict.size() << ") should equal params (" << params.size()
-                      << ") + buffers (" << buffers.size() << ") = " << expected_size << std::endl;
-            return;
-        }
-
-        std::cout << "SUCCESS: StateDict works correctly with " << state_dict.size() << " entries (" << params.size()
-                  << " params + " << buffers.size() << " buffers)!" << std::endl;
-
-    } catch (const std::exception &e) { std::cout << "FAIL: Exception: " << e.what() << std::endl; }
-}
-
-// ============================================================================
-// Main
-// ============================================================================
-int main(int argc, char *argv[]) {
-    google::InitGoogleLogging(argv[0]);
-
-    nn::parallel::global::GlobalEnv::Instance().Init(1, 1, false, 1, 1);
-
-    std::cout << "========================================" << std::endl;
-    std::cout << "    Transformer architecture Tests" << std::endl;
-    std::cout << "========================================" << std::endl;
-
-    TestConfigValidation();
-    TestEmbedding();
-    TestNormalization();
-    TestMlp();
-    TestAttention();
-    TestTransformerLayer();
-    TestGpt2Model();
-    TestLlama3Model();
-    TestRopeUtils();
-    TestStateDict();
-
-    std::cout << "\n========================================" << std::endl;
-    std::cout << "    All Tests Completed" << std::endl;
-    std::cout << "========================================" << std::endl;
-
-    return 0;
-}
diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt
new file mode 100644
index 00000000..b2ed5009
--- /dev/null
+++ b/tests/CMakeLists.txt
@@ -0,0 +1,29 @@
+# Tests CMakeLists.txt
+# This file manages the test infrastructure for InfiniTrain
+
+# Include shared test macros (must be before any test subdirectory)
+include(${CMAKE_SOURCE_DIR}/cmake/test_macros.cmake)
+
+# Common test utilities
+add_subdirectory(common)
+
+# Tensor tests
+add_subdirectory(tensor)
+
+# Optimizer tests
+add_subdirectory(optimizer)
+
+# Autograd operator tests
+add_subdirectory(autograd)
+
+# LoRA tests
+add_subdirectory(lora)
+
+# Hook tests
+add_subdirectory(hook)
+
+# Dtype dispatch tests
+add_subdirectory(dtype)
+
+# Transformer architecture tests
+add_subdirectory(transformer)
diff --git a/tests/autograd/CMakeLists.txt b/tests/autograd/CMakeLists.txt
new file mode 100644
index 00000000..50bd6096
--- /dev/null
+++ b/tests/autograd/CMakeLists.txt
@@ -0,0 +1,9 @@
+# ============================================================================
+# Autograd tests
+# ============================================================================
+
+file(GLOB AUTOGRAD_SOURCES ${CMAKE_CURRENT_SOURCE_DIR}/test_*.cc)
+
+infini_train_add_test_suite(test_autograd
+  SOURCES ${AUTOGRAD_SOURCES}
+)
diff --git a/tests/autograd/test_autograd.cc b/tests/autograd/test_autograd.cc
new file mode 100644
index 00000000..c0e1546b
--- /dev/null
+++ b/tests/autograd/test_autograd.cc
@@ -0,0 +1,385 @@
+#include <cmath>
+#include <vector>
+
+#include "infini_train/include/autograd/activations.h"
+#include "infini_train/include/autograd/elementwise.h"
+#include "infini_train/include/autograd/function.h"
+#include "infini_train/include/autograd/linear.h"
+#include "infini_train/include/autograd/matmul.h"
+#include "infini_train/include/autograd/misc.h"
+#include "infini_train/include/autograd/normalization.h"
+#include "infini_train/include/autograd/outer.h"
+#include "infini_train/include/autograd/reduction.h"
+#include "infini_train/include/autograd/softmax.h"
+#include "infini_train/include/autograd/transform.h"
+#include "infini_train/include/tensor.h"
+#include "tests/common/test_utils.h"
+#include "gtest/gtest.h"
+
+using namespace infini_train;
+
+// ============================================================================
+// Forward / Backward — CPU + CUDA
+// ============================================================================
+
+class AutogradForwardTest : public infini_train::test::InfiniTrainTest {};
+class AutogradBackwardTest : public infini_train::test::InfiniTrainTest {};
+
+TEST_P(AutogradForwardTest, AddForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(1.0f);
+    auto b = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    b->Fill(2.0f);
+    auto result = std::make_shared<autograd::Add>()->Apply({a, b});
+    EXPECT_EQ(result.size(), 1);
+    EXPECT_EQ(result[0]->Dims(), (std::vector<int64_t>{2, 3}));
+}
+
+TEST_P(AutogradForwardTest, SubForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(5.0f);
+    auto b = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    b->Fill(3.0f);
+    auto result = std::make_shared<autograd::Sub>()->Apply({a, b});
+    EXPECT_EQ(result.size(), 1);
+}
+
+TEST_P(AutogradForwardTest, MulForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(2.0f);
+    auto b = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    b->Fill(3.0f);
+    auto result = std::make_shared<autograd::Mul>()->Apply({a, b});
+    EXPECT_EQ(result.size(), 1);
+}
+
+TEST_P(AutogradForwardTest, DivForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(6.0f);
+    auto b = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    b->Fill(2.0f);
+    auto result = std::make_shared<autograd::Div>()->Apply({a, b});
+    EXPECT_EQ(result.size(), 1);
+}
+
+TEST_P(AutogradForwardTest, NegForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(5.0f);
+    auto result = std::make_shared<autograd::Neg>()->Apply({a});
+    EXPECT_EQ(result.size(), 1);
+}
+
+TEST_P(AutogradForwardTest, SinForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(0.0f);
+    auto result = std::make_shared<autograd::Sin>()->Apply({a});
+    EXPECT_EQ(result.size(), 1);
+}
+
+TEST_P(AutogradForwardTest, CosForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(0.0f);
+    auto result = std::make_shared<autograd::Cos>()->Apply({a});
+    EXPECT_EQ(result.size(), 1);
+}
+
+TEST_P(AutogradForwardTest, TanhForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(0.0f);
+    auto result = std::make_shared<autograd::Tanh>()->Apply({a});
+    EXPECT_EQ(result.size(), 1);
+}
+
+TEST_P(AutogradForwardTest, ExpForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(1.0f);
+    auto result = std::make_shared<autograd::Exp>()->Apply({a});
+    EXPECT_EQ(result.size(), 1);
+}
+
+TEST_P(AutogradForwardTest, LogForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(2.0f);
+    auto result = std::make_shared<autograd::Log>()->Apply({a});
+    EXPECT_EQ(result.size(), 1);
+}
+
+TEST_P(AutogradForwardTest, ReciprocalForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(2.0f);
+    auto result = std::make_shared<autograd::Reciprocal>()->Apply({a});
+    EXPECT_EQ(result.size(), 1);
+}
+
+TEST_P(AutogradForwardTest, PowForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(2.0f);
+    auto result = std::make_shared<autograd::Pow>(2.0f)->Apply({a});
+    EXPECT_EQ(result.size(), 1);
+}
+
+TEST_P(AutogradForwardTest, RsqrtForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(4.0f);
+    auto result = std::make_shared<autograd::Rsqrt>()->Apply({a});
+    EXPECT_EQ(result.size(), 1);
+}
+
+TEST_P(AutogradForwardTest, SigmoidForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(0.0f);
+    auto result = std::make_shared<autograd::Sigmoid>()->Apply({a});
+    EXPECT_EQ(result.size(), 1);
+}
+
+TEST_P(AutogradForwardTest, MatmulForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(1.0f);
+    auto b = std::make_shared<Tensor>(std::vector<int64_t>{3, 4}, DataType::kFLOAT32, GetDevice(), true);
+    b->Fill(1.0f);
+    auto result = std::make_shared<autograd::Matmul>()->Apply({a, b});
+    EXPECT_EQ(result.size(), 1);
+    EXPECT_EQ(result[0]->Dims(), (std::vector<int64_t>{2, 4}));
+}
+
+TEST_P(AutogradForwardTest, SumForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(1.0f);
+    auto result = std::make_shared<autograd::Sum>(1, false)->Apply({a});
+    EXPECT_EQ(result.size(), 1);
+}
+
+TEST_P(AutogradForwardTest, MeanForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(1.0f);
+    auto result = std::make_shared<autograd::Mean>(1, false)->Apply({a});
+    EXPECT_EQ(result.size(), 1);
+}
+
+TEST_P(AutogradForwardTest, MaxForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(1.0f);
+    auto result = std::make_shared<autograd::Max>(1, false)->Apply({a});
+    EXPECT_EQ(result.size(), 1);
+}
+
+TEST_P(AutogradForwardTest, MinForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(1.0f);
+    auto result = std::make_shared<autograd::Min>(1, false)->Apply({a});
+    EXPECT_EQ(result.size(), 1);
+}
+
+TEST_P(AutogradForwardTest, SoftmaxForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(1.0f);
+    auto result = std::make_shared<autograd::Softmax>(1)->Apply({a});
+    EXPECT_EQ(result.size(), 1);
+    EXPECT_EQ(result[0]->Dims(), (std::vector<int64_t>{2, 3}));
+}
+
+TEST_P(AutogradForwardTest, LayerNormForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3, 4}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(1.0f);
+    auto weight = std::make_shared<Tensor>(std::vector<int64_t>{4}, DataType::kFLOAT32, GetDevice(), true);
+    weight->Fill(1.0f);
+    auto bias = std::make_shared<Tensor>(std::vector<int64_t>{4}, DataType::kFLOAT32, GetDevice(), true);
+    bias->Fill(0.0f);
+    auto result = std::make_shared<autograd::LayerNorm>(1e-5f)->Apply({a, weight, bias});
+    EXPECT_EQ(result.size(), 1);
+}
+
+TEST_P(AutogradForwardTest, LinearForward) {
+    auto input = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    input->Fill(1.0f);
+    auto weight = std::make_shared<Tensor>(std::vector<int64_t>{4, 3}, DataType::kFLOAT32, GetDevice(), true);
+    weight->Fill(1.0f);
+    auto bias = std::make_shared<Tensor>(std::vector<int64_t>{4}, DataType::kFLOAT32, GetDevice(), true);
+    bias->Fill(0.0f);
+    auto result = std::make_shared<autograd::Linear>()->Apply({input, weight, bias});
+    EXPECT_EQ(result.size(), 1);
+    EXPECT_EQ(result[0]->Dims(), (std::vector<int64_t>{2, 4}));
+}
+
+TEST_P(AutogradForwardTest, TransposeForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(1.0f);
+    auto result = std::make_shared<autograd::Transpose>(0, 1)->Apply({a});
+    EXPECT_EQ(result.size(), 1);
+    EXPECT_EQ(result[0]->Dims(), (std::vector<int64_t>{3, 2}));
+}
+
+TEST_P(AutogradForwardTest, SliceForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{4, 4}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(1.0f);
+    auto result = std::make_shared<autograd::Slice>(std::vector<int64_t>{1, 1}, std::vector<int64_t>{3, 3},
+                                                    std::vector<int64_t>{1, 1})
+                      ->Apply({a});
+    EXPECT_EQ(result.size(), 1);
+}
+
+TEST_P(AutogradForwardTest, SplitForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{4, 4}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(1.0f);
+    auto result = std::make_shared<autograd::Split>(2, 0)->Apply({a});
+    EXPECT_EQ(result.size(), 2);
+}
+
+TEST_P(AutogradForwardTest, ConcatForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 2}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(1.0f);
+    auto b = std::make_shared<Tensor>(std::vector<int64_t>{2, 2}, DataType::kFLOAT32, GetDevice(), true);
+    b->Fill(2.0f);
+    auto result = std::make_shared<autograd::Concat>(0)->Apply({a, b});
+    EXPECT_EQ(result.size(), 1);
+    EXPECT_EQ(result[0]->Dims(), (std::vector<int64_t>{4, 2}));
+}
+
+TEST_P(AutogradForwardTest, StackForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(1.0f);
+    auto b = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    b->Fill(2.0f);
+    auto result = std::make_shared<autograd::Stack>(0)->Apply({a, b});
+    EXPECT_EQ(result.size(), 1);
+    EXPECT_EQ(result[0]->Dims(), (std::vector<int64_t>{2, 2, 3}));
+}
+
+TEST_P(AutogradForwardTest, TrilForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{3, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(1.0f);
+    auto result = std::make_shared<autograd::Tril>(0)->Apply({a});
+    EXPECT_EQ(result.size(), 1);
+}
+
+TEST_P(AutogradForwardTest, TriuForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{3, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(1.0f);
+    auto result = std::make_shared<autograd::Triu>(0)->Apply({a});
+    EXPECT_EQ(result.size(), 1);
+}
+
+TEST_P(AutogradForwardTest, OuterForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(1.0f);
+    auto b = std::make_shared<Tensor>(std::vector<int64_t>{4}, DataType::kFLOAT32, GetDevice(), true);
+    b->Fill(1.0f);
+    auto result = std::make_shared<autograd::Outer>()->Apply({a, b});
+    EXPECT_EQ(result.size(), 1);
+    EXPECT_EQ(result[0]->Dims(), (std::vector<int64_t>{3, 4}));
+}
+
+TEST_P(AutogradForwardTest, AddScalarForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(1.0f);
+    auto result = std::make_shared<autograd::AddScalar>(2.0f)->Apply({a});
+    EXPECT_EQ(result.size(), 1);
+}
+
+TEST_P(AutogradForwardTest, MulScalarForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(2.0f);
+    auto result = std::make_shared<autograd::MulScalar>(3.0f)->Apply({a});
+    EXPECT_EQ(result.size(), 1);
+}
+
+TEST_P(AutogradForwardTest, LtForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(5.0f);
+    auto b = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    b->Fill(3.0f);
+    auto result = std::make_shared<autograd::Lt>()->Apply({a, b});
+    EXPECT_EQ(result.size(), 1);
+}
+
+TEST_P(AutogradForwardTest, LeForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(3.0f);
+    auto b = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    b->Fill(3.0f);
+    auto result = std::make_shared<autograd::Le>()->Apply({a, b});
+    EXPECT_EQ(result.size(), 1);
+}
+
+TEST_P(AutogradForwardTest, GtForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(5.0f);
+    auto b = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    b->Fill(3.0f);
+    auto result = std::make_shared<autograd::Gt>()->Apply({a, b});
+    EXPECT_EQ(result.size(), 1);
+}
+
+TEST_P(AutogradForwardTest, GeForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(3.0f);
+    auto b = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    b->Fill(3.0f);
+    auto result = std::make_shared<autograd::Ge>()->Apply({a, b});
+    EXPECT_EQ(result.size(), 1);
+}
+
+TEST_P(AutogradForwardTest, EqualsForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(3.0f);
+    auto b = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    b->Fill(3.0f);
+    auto result = std::make_shared<autograd::Equals>()->Apply({a, b});
+    EXPECT_EQ(result.size(), 1);
+}
+
+TEST_P(AutogradForwardTest, AndForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(1.0f);
+    auto b = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    b->Fill(1.0f);
+    auto result = std::make_shared<autograd::And>()->Apply({a, b});
+    EXPECT_EQ(result.size(), 1);
+}
+
+TEST_P(AutogradForwardTest, OrForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(0.0f);
+    auto b = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    b->Fill(1.0f);
+    auto result = std::make_shared<autograd::Or>()->Apply({a, b});
+    EXPECT_EQ(result.size(), 1);
+}
+
+TEST_P(AutogradForwardTest, NoOpForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(1.0f);
+    auto result = std::make_shared<autograd::NoOp>(std::vector<int64_t>{2, 3})->Apply({a});
+    EXPECT_EQ(result.size(), 1);
+    EXPECT_EQ(result[0]->Dims(), (std::vector<int64_t>{2, 3}));
+}
+
+TEST_P(AutogradBackwardTest, AddBackward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(1.0f);
+    auto b = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    b->Fill(2.0f);
+    auto add_fn = std::make_shared<autograd::Add>();
+    auto result = add_fn->Apply({a, b});
+    auto grad = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    grad->Fill(1.0f);
+    auto grad_inputs = add_fn->Backward({grad});
+    EXPECT_EQ(grad_inputs.size(), 2);
+}
+
+TEST_P(AutogradBackwardTest, MulBackward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(2.0f);
+    auto b = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    b->Fill(3.0f);
+    auto mul_fn = std::make_shared<autograd::Mul>();
+    auto result = mul_fn->Apply({a, b});
+    auto grad = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    grad->Fill(1.0f);
+    auto grad_inputs = mul_fn->Backward({grad});
+    EXPECT_EQ(grad_inputs.size(), 2);
+}
+
+INFINI_TRAIN_REGISTER_TEST(AutogradForwardTest);
+
+INFINI_TRAIN_REGISTER_TEST(AutogradBackwardTest);
diff --git a/tests/autograd/test_autograd_elementwise_backward.cc b/tests/autograd/test_autograd_elementwise_backward.cc
new file mode 100644
index 00000000..8954a943
--- /dev/null
+++ b/tests/autograd/test_autograd_elementwise_backward.cc
@@ -0,0 +1,165 @@
+#include <cmath>
+#include <vector>
+
+#include "infini_train/include/autograd/elementwise.h"
+#include "infini_train/include/nn/parallel/global.h"
+#include "infini_train/include/tensor.h"
+#include "tests/common/test_utils.h"
+#include "gtest/gtest.h"
+
+using namespace infini_train;
+
+class AutogradElementwiseBackwardTest : public infini_train::test::InfiniTrainTest {};
+
+TEST_P(AutogradElementwiseBackwardTest, AddBackward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(1.0f);
+    auto b = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    b->Fill(2.0f);
+    auto add_fn = std::make_shared<autograd::Add>();
+    auto result = add_fn->Apply({a, b});
+    auto grad = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    grad->Fill(1.0f);
+    auto grad_inputs = add_fn->Backward({grad});
+    EXPECT_EQ(grad_inputs.size(), 2);
+}
+
+TEST_P(AutogradElementwiseBackwardTest, SubBackward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(5.0f);
+    auto b = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    b->Fill(3.0f);
+    auto sub_fn = std::make_shared<autograd::Sub>();
+    auto result = sub_fn->Apply({a, b});
+    auto grad = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    grad->Fill(1.0f);
+    auto grad_inputs = sub_fn->Backward({grad});
+    EXPECT_EQ(grad_inputs.size(), 2);
+}
+
+TEST_P(AutogradElementwiseBackwardTest, MulBackward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(2.0f);
+    auto b = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    b->Fill(3.0f);
+    auto mul_fn = std::make_shared<autograd::Mul>();
+    auto result = mul_fn->Apply({a, b});
+    auto grad = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    grad->Fill(1.0f);
+    auto grad_inputs = mul_fn->Backward({grad});
+    EXPECT_EQ(grad_inputs.size(), 2);
+}
+
+TEST_P(AutogradElementwiseBackwardTest, DivBackward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(6.0f);
+    auto b = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    b->Fill(2.0f);
+    auto div_fn = std::make_shared<autograd::Div>();
+    auto result = div_fn->Apply({a, b});
+    auto grad = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    grad->Fill(1.0f);
+    auto grad_inputs = div_fn->Backward({grad});
+    EXPECT_EQ(grad_inputs.size(), 2);
+}
+
+TEST_P(AutogradElementwiseBackwardTest, NegBackward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(5.0f);
+    auto neg_fn = std::make_shared<autograd::Neg>();
+    auto result = neg_fn->Apply({a});
+    auto grad = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    grad->Fill(1.0f);
+    auto grad_inputs = neg_fn->Backward({grad});
+    EXPECT_EQ(grad_inputs.size(), 1);
+}
+
+TEST_P(AutogradElementwiseBackwardTest, SinBackward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(0.0f);
+    auto sin_fn = std::make_shared<autograd::Sin>();
+    auto result = sin_fn->Apply({a});
+    auto grad = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    grad->Fill(1.0f);
+    auto grad_inputs = sin_fn->Backward({grad});
+    EXPECT_EQ(grad_inputs.size(), 1);
+}
+
+TEST_P(AutogradElementwiseBackwardTest, CosBackward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(0.0f);
+    auto cos_fn = std::make_shared<autograd::Cos>();
+    auto result = cos_fn->Apply({a});
+    auto grad = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    grad->Fill(1.0f);
+    auto grad_inputs = cos_fn->Backward({grad});
+    EXPECT_EQ(grad_inputs.size(), 1);
+}
+
+TEST_P(AutogradElementwiseBackwardTest, TanhBackward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(0.0f);
+    auto tanh_fn = std::make_shared<autograd::Tanh>();
+    auto result = tanh_fn->Apply({a});
+    auto grad = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    grad->Fill(1.0f);
+    auto grad_inputs = tanh_fn->Backward({grad});
+    EXPECT_EQ(grad_inputs.size(), 1);
+}
+
+TEST_P(AutogradElementwiseBackwardTest, ExpBackward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(1.0f);
+    auto exp_fn = std::make_shared<autograd::Exp>();
+    auto result = exp_fn->Apply({a});
+    auto grad = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    grad->Fill(1.0f);
+    auto grad_inputs = exp_fn->Backward({grad});
+    EXPECT_EQ(grad_inputs.size(), 1);
+}
+
+TEST_P(AutogradElementwiseBackwardTest, LogBackward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(2.0f);
+    auto log_fn = std::make_shared<autograd::Log>();
+    auto result = log_fn->Apply({a});
+    auto grad = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    grad->Fill(1.0f);
+    auto grad_inputs = log_fn->Backward({grad});
+    EXPECT_EQ(grad_inputs.size(), 1);
+}
+
+TEST_P(AutogradElementwiseBackwardTest, ReciprocalBackward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(2.0f);
+    auto reciprocal_fn = std::make_shared<autograd::Reciprocal>();
+    auto result = reciprocal_fn->Apply({a});
+    auto grad = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    grad->Fill(1.0f);
+    auto grad_inputs = reciprocal_fn->Backward({grad});
+    EXPECT_EQ(grad_inputs.size(), 1);
+}
+
+TEST_P(AutogradElementwiseBackwardTest, PowBackward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(2.0f);
+    auto pow_fn = std::make_shared<autograd::Pow>(2.0f);
+    auto result = pow_fn->Apply({a});
+    auto grad = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    grad->Fill(1.0f);
+    auto grad_inputs = pow_fn->Backward({grad});
+    EXPECT_EQ(grad_inputs.size(), 1);
+}
+
+TEST_P(AutogradElementwiseBackwardTest, RsqrtBackward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(4.0f);
+    auto rsqrt_fn = std::make_shared<autograd::Rsqrt>();
+    auto result = rsqrt_fn->Apply({a});
+    auto grad = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    grad->Fill(1.0f);
+    auto grad_inputs = rsqrt_fn->Backward({grad});
+    EXPECT_EQ(grad_inputs.size(), 1);
+}
+
+INFINI_TRAIN_REGISTER_TEST(AutogradElementwiseBackwardTest);
diff --git a/tests/autograd/test_autograd_elementwise_forward.cc b/tests/autograd/test_autograd_elementwise_forward.cc
new file mode 100644
index 00000000..b2c8d1a3
--- /dev/null
+++ b/tests/autograd/test_autograd_elementwise_forward.cc
@@ -0,0 +1,222 @@
+#include <cmath>
+#include <vector>
+
+#include "infini_train/include/autograd/activations.h"
+#include "infini_train/include/autograd/elementwise.h"
+#include "infini_train/include/nn/parallel/global.h"
+#include "infini_train/include/tensor.h"
+#include "tests/common/test_utils.h"
+#include "gtest/gtest.h"
+
+using namespace infini_train;
+
+class AutogradElementwiseForwardTest : public infini_train::test::InfiniTrainTest {};
+
+TEST_P(AutogradElementwiseForwardTest, AddForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(1.0f);
+    auto b = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    b->Fill(2.0f);
+    auto add_fn = std::make_shared<autograd::Add>();
+    auto result = add_fn->Apply({a, b});
+    EXPECT_EQ(result.size(), 1);
+    EXPECT_EQ(result[0]->Dims(), (std::vector<int64_t>{2, 3}));
+}
+
+TEST_P(AutogradElementwiseForwardTest, SubForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(5.0f);
+    auto b = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    b->Fill(3.0f);
+    auto sub_fn = std::make_shared<autograd::Sub>();
+    auto result = sub_fn->Apply({a, b});
+    EXPECT_EQ(result.size(), 1);
+}
+
+TEST_P(AutogradElementwiseForwardTest, MulForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(2.0f);
+    auto b = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    b->Fill(3.0f);
+    auto mul_fn = std::make_shared<autograd::Mul>();
+    auto result = mul_fn->Apply({a, b});
+    EXPECT_EQ(result.size(), 1);
+}
+
+TEST_P(AutogradElementwiseForwardTest, DivForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(6.0f);
+    auto b = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    b->Fill(2.0f);
+    auto div_fn = std::make_shared<autograd::Div>();
+    auto result = div_fn->Apply({a, b});
+    EXPECT_EQ(result.size(), 1);
+}
+
+TEST_P(AutogradElementwiseForwardTest, NegForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(5.0f);
+    auto neg_fn = std::make_shared<autograd::Neg>();
+    auto result = neg_fn->Apply({a});
+    EXPECT_EQ(result.size(), 1);
+}
+
+TEST_P(AutogradElementwiseForwardTest, SinForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(0.0f);
+    auto sin_fn = std::make_shared<autograd::Sin>();
+    auto result = sin_fn->Apply({a});
+    EXPECT_EQ(result.size(), 1);
+}
+
+TEST_P(AutogradElementwiseForwardTest, CosForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(0.0f);
+    auto cos_fn = std::make_shared<autograd::Cos>();
+    auto result = cos_fn->Apply({a});
+    EXPECT_EQ(result.size(), 1);
+}
+
+TEST_P(AutogradElementwiseForwardTest, TanhForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(0.0f);
+    auto tanh_fn = std::make_shared<autograd::Tanh>();
+    auto result = tanh_fn->Apply({a});
+    EXPECT_EQ(result.size(), 1);
+}
+
+TEST_P(AutogradElementwiseForwardTest, ExpForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(1.0f);
+    auto exp_fn = std::make_shared<autograd::Exp>();
+    auto result = exp_fn->Apply({a});
+    EXPECT_EQ(result.size(), 1);
+}
+
+TEST_P(AutogradElementwiseForwardTest, LogForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(2.0f);
+    auto log_fn = std::make_shared<autograd::Log>();
+    auto result = log_fn->Apply({a});
+    EXPECT_EQ(result.size(), 1);
+}
+
+TEST_P(AutogradElementwiseForwardTest, ReciprocalForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(2.0f);
+    auto reciprocal_fn = std::make_shared<autograd::Reciprocal>();
+    auto result = reciprocal_fn->Apply({a});
+    EXPECT_EQ(result.size(), 1);
+}
+
+TEST_P(AutogradElementwiseForwardTest, PowForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(2.0f);
+    auto pow_fn = std::make_shared<autograd::Pow>(2.0f);
+    auto result = pow_fn->Apply({a});
+    EXPECT_EQ(result.size(), 1);
+}
+
+TEST_P(AutogradElementwiseForwardTest, RsqrtForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(4.0f);
+    auto rsqrt_fn = std::make_shared<autograd::Rsqrt>();
+    auto result = rsqrt_fn->Apply({a});
+    EXPECT_EQ(result.size(), 1);
+}
+
+TEST_P(AutogradElementwiseForwardTest, SigmoidForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(0.0f);
+    auto sigmoid_fn = std::make_shared<autograd::Sigmoid>();
+    auto result = sigmoid_fn->Apply({a});
+    EXPECT_EQ(result.size(), 1);
+}
+
+TEST_P(AutogradElementwiseForwardTest, AddScalarForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(1.0f);
+    auto add_scalar_fn = std::make_shared<autograd::AddScalar>(2.0f);
+    auto result = add_scalar_fn->Apply({a});
+    EXPECT_EQ(result.size(), 1);
+}
+
+TEST_P(AutogradElementwiseForwardTest, MulScalarForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(2.0f);
+    auto mul_scalar_fn = std::make_shared<autograd::MulScalar>(3.0f);
+    auto result = mul_scalar_fn->Apply({a});
+    EXPECT_EQ(result.size(), 1);
+}
+
+TEST_P(AutogradElementwiseForwardTest, LtForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(5.0f);
+    auto b = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    b->Fill(3.0f);
+    auto lt_fn = std::make_shared<autograd::Lt>();
+    auto result = lt_fn->Apply({a, b});
+    EXPECT_EQ(result.size(), 1);
+}
+
+TEST_P(AutogradElementwiseForwardTest, LeForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(3.0f);
+    auto b = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    b->Fill(3.0f);
+    auto le_fn = std::make_shared<autograd::Le>();
+    auto result = le_fn->Apply({a, b});
+    EXPECT_EQ(result.size(), 1);
+}
+
+TEST_P(AutogradElementwiseForwardTest, GtForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(5.0f);
+    auto b = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    b->Fill(3.0f);
+    auto gt_fn = std::make_shared<autograd::Gt>();
+    auto result = gt_fn->Apply({a, b});
+    EXPECT_EQ(result.size(), 1);
+}
+
+TEST_P(AutogradElementwiseForwardTest, GeForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(3.0f);
+    auto b = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    b->Fill(3.0f);
+    auto ge_fn = std::make_shared<autograd::Ge>();
+    auto result = ge_fn->Apply({a, b});
+    EXPECT_EQ(result.size(), 1);
+}
+
+TEST_P(AutogradElementwiseForwardTest, EqualsForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(3.0f);
+    auto b = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    b->Fill(3.0f);
+    auto eq_fn = std::make_shared<autograd::Equals>();
+    auto result = eq_fn->Apply({a, b});
+    EXPECT_EQ(result.size(), 1);
+}
+
+TEST_P(AutogradElementwiseForwardTest, AndForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(1.0f);
+    auto b = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    b->Fill(1.0f);
+    auto and_fn = std::make_shared<autograd::And>();
+    auto result = and_fn->Apply({a, b});
+    EXPECT_EQ(result.size(), 1);
+}
+
+TEST_P(AutogradElementwiseForwardTest, OrForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(0.0f);
+    auto b = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    b->Fill(1.0f);
+    auto or_fn = std::make_shared<autograd::Or>();
+    auto result = or_fn->Apply({a, b});
+    EXPECT_EQ(result.size(), 1);
+}
+
+INFINI_TRAIN_REGISTER_TEST(AutogradElementwiseForwardTest);
diff --git a/tests/autograd/test_autograd_linear_backward.cc b/tests/autograd/test_autograd_linear_backward.cc
new file mode 100644
index 00000000..e5ca6ebf
--- /dev/null
+++ b/tests/autograd/test_autograd_linear_backward.cc
@@ -0,0 +1,41 @@
+#include <vector>
+
+#include "infini_train/include/autograd/linear.h"
+#include "infini_train/include/nn/parallel/global.h"
+#include "infini_train/include/tensor.h"
+#include "tests/common/test_utils.h"
+#include "gtest/gtest.h"
+
+using namespace infini_train;
+
+class AutogradLinearBackwardTest : public infini_train::test::InfiniTrainTest {};
+
+TEST_P(AutogradLinearBackwardTest, LinearBackward) {
+    auto input = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    input->Fill(1.0f);
+    auto weight = std::make_shared<Tensor>(std::vector<int64_t>{4, 3}, DataType::kFLOAT32, GetDevice(), true);
+    weight->Fill(1.0f);
+    auto bias = std::make_shared<Tensor>(std::vector<int64_t>{4}, DataType::kFLOAT32, GetDevice(), true);
+    bias->Fill(0.0f);
+    auto linear_fn = std::make_shared<autograd::Linear>();
+    auto result = linear_fn->Apply({input, weight, bias});
+    auto grad = std::make_shared<Tensor>(std::vector<int64_t>{2, 4}, DataType::kFLOAT32, GetDevice(), true);
+    grad->Fill(1.0f);
+    auto grad_inputs = linear_fn->Backward({grad});
+    EXPECT_EQ(grad_inputs.size(), 3);
+}
+
+TEST_P(AutogradLinearBackwardTest, LinearBackwardNoBias) {
+    auto input = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    input->Fill(1.0f);
+    auto weight = std::make_shared<Tensor>(std::vector<int64_t>{4, 3}, DataType::kFLOAT32, GetDevice(), true);
+    weight->Fill(1.0f);
+    auto linear_fn = std::make_shared<autograd::Linear>();
+    auto result = linear_fn->Apply({input, weight});
+    auto grad = std::make_shared<Tensor>(std::vector<int64_t>{2, 4}, DataType::kFLOAT32, GetDevice(), true);
+    grad->Fill(1.0f);
+    auto grad_inputs = linear_fn->Backward({grad});
+    EXPECT_EQ(grad_inputs.size(), 2);
+}
+
+INFINI_TRAIN_REGISTER_TEST(AutogradLinearBackwardTest);
diff --git a/tests/autograd/test_autograd_linear_forward.cc b/tests/autograd/test_autograd_linear_forward.cc
new file mode 100644
index 00000000..79cfa165
--- /dev/null
+++ b/tests/autograd/test_autograd_linear_forward.cc
@@ -0,0 +1,50 @@
+#include <vector>
+
+#include "infini_train/include/autograd/linear.h"
+#include "infini_train/include/nn/parallel/global.h"
+#include "infini_train/include/tensor.h"
+#include "tests/common/test_utils.h"
+#include "gtest/gtest.h"
+
+using namespace infini_train;
+
+class AutogradLinearForwardTest : public infini_train::test::InfiniTrainTest {};
+
+TEST_P(AutogradLinearForwardTest, LinearForward) {
+    auto input = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    input->Fill(1.0f);
+    auto weight = std::make_shared<Tensor>(std::vector<int64_t>{4, 3}, DataType::kFLOAT32, GetDevice(), true);
+    weight->Fill(1.0f);
+    auto bias = std::make_shared<Tensor>(std::vector<int64_t>{4}, DataType::kFLOAT32, GetDevice(), true);
+    bias->Fill(0.0f);
+    auto linear_fn = std::make_shared<autograd::Linear>();
+    auto result = linear_fn->Apply({input, weight, bias});
+    EXPECT_EQ(result.size(), 1);
+    EXPECT_EQ(result[0]->Dims(), (std::vector<int64_t>{2, 4}));
+}
+
+TEST_P(AutogradLinearForwardTest, LinearNoBias) {
+    auto input = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    input->Fill(1.0f);
+    auto weight = std::make_shared<Tensor>(std::vector<int64_t>{4, 3}, DataType::kFLOAT32, GetDevice(), true);
+    weight->Fill(1.0f);
+    auto linear_fn = std::make_shared<autograd::Linear>();
+    auto result = linear_fn->Apply({input, weight});
+    EXPECT_EQ(result.size(), 1);
+    EXPECT_EQ(result[0]->Dims(), (std::vector<int64_t>{2, 4}));
+}
+
+TEST_P(AutogradLinearForwardTest, LinearBatch) {
+    auto input = std::make_shared<Tensor>(std::vector<int64_t>{32, 128}, DataType::kFLOAT32, GetDevice(), true);
+    input->Fill(1.0f);
+    auto weight = std::make_shared<Tensor>(std::vector<int64_t>{64, 128}, DataType::kFLOAT32, GetDevice(), true);
+    weight->Fill(1.0f);
+    auto bias = std::make_shared<Tensor>(std::vector<int64_t>{64}, DataType::kFLOAT32, GetDevice(), true);
+    bias->Fill(0.0f);
+    auto linear_fn = std::make_shared<autograd::Linear>();
+    auto result = linear_fn->Apply({input, weight, bias});
+    EXPECT_EQ(result.size(), 1);
+    EXPECT_EQ(result[0]->Dims(), (std::vector<int64_t>{32, 64}));
+}
+
+INFINI_TRAIN_REGISTER_TEST(AutogradLinearForwardTest);
diff --git a/tests/autograd/test_autograd_matmul_backward.cc b/tests/autograd/test_autograd_matmul_backward.cc
new file mode 100644
index 00000000..06f7770f
--- /dev/null
+++ b/tests/autograd/test_autograd_matmul_backward.cc
@@ -0,0 +1,52 @@
+#include <vector>
+
+#include "infini_train/include/autograd/matmul.h"
+#include "infini_train/include/nn/parallel/global.h"
+#include "infini_train/include/tensor.h"
+#include "tests/common/test_utils.h"
+#include "gtest/gtest.h"
+
+using namespace infini_train;
+
+class AutogradMatmulBackwardTest : public infini_train::test::InfiniTrainTest {};
+
+TEST_P(AutogradMatmulBackwardTest, MatmulBackward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(1.0f);
+    auto b = std::make_shared<Tensor>(std::vector<int64_t>{3, 4}, DataType::kFLOAT32, GetDevice(), true);
+    b->Fill(1.0f);
+    auto matmul_fn = std::make_shared<autograd::Matmul>();
+    auto result = matmul_fn->Apply({a, b});
+    auto grad = std::make_shared<Tensor>(std::vector<int64_t>{2, 4}, DataType::kFLOAT32, GetDevice(), true);
+    grad->Fill(1.0f);
+    auto grad_inputs = matmul_fn->Backward({grad});
+    EXPECT_EQ(grad_inputs.size(), 2);
+}
+
+TEST_P(AutogradMatmulBackwardTest, MatmulBackwardSquare) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{3, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(2.0f);
+    auto b = std::make_shared<Tensor>(std::vector<int64_t>{3, 3}, DataType::kFLOAT32, GetDevice(), true);
+    b->Fill(3.0f);
+    auto matmul_fn = std::make_shared<autograd::Matmul>();
+    auto result = matmul_fn->Apply({a, b});
+    auto grad = std::make_shared<Tensor>(std::vector<int64_t>{3, 3}, DataType::kFLOAT32, GetDevice(), true);
+    grad->Fill(1.0f);
+    auto grad_inputs = matmul_fn->Backward({grad});
+    EXPECT_EQ(grad_inputs.size(), 2);
+}
+
+TEST_P(AutogradMatmulBackwardTest, MatmulBackwardDifferentShapes) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{3, 4}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(1.5f);
+    auto b = std::make_shared<Tensor>(std::vector<int64_t>{4, 2}, DataType::kFLOAT32, GetDevice(), true);
+    b->Fill(2.5f);
+    auto matmul_fn = std::make_shared<autograd::Matmul>();
+    auto result = matmul_fn->Apply({a, b});
+    auto grad = std::make_shared<Tensor>(std::vector<int64_t>{3, 2}, DataType::kFLOAT32, GetDevice(), true);
+    grad->Fill(1.0f);
+    auto grad_inputs = matmul_fn->Backward({grad});
+    EXPECT_EQ(grad_inputs.size(), 2);
+}
+
+INFINI_TRAIN_REGISTER_TEST(AutogradMatmulBackwardTest);
diff --git a/tests/autograd/test_autograd_matmul_forward.cc b/tests/autograd/test_autograd_matmul_forward.cc
new file mode 100644
index 00000000..bbb52a54
--- /dev/null
+++ b/tests/autograd/test_autograd_matmul_forward.cc
@@ -0,0 +1,57 @@
+#include <vector>
+
+#include "infini_train/include/autograd/matmul.h"
+#include "infini_train/include/nn/parallel/global.h"
+#include "infini_train/include/tensor.h"
+#include "tests/common/test_utils.h"
+#include "gtest/gtest.h"
+
+using namespace infini_train;
+
+class AutogradMatmulForwardTest : public infini_train::test::InfiniTrainTest {};
+
+TEST_P(AutogradMatmulForwardTest, MatmulForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(1.0f);
+    auto b = std::make_shared<Tensor>(std::vector<int64_t>{3, 4}, DataType::kFLOAT32, GetDevice(), true);
+    b->Fill(1.0f);
+    auto matmul_fn = std::make_shared<autograd::Matmul>();
+    auto result = matmul_fn->Apply({a, b});
+    EXPECT_EQ(result.size(), 1);
+    EXPECT_EQ(result[0]->Dims(), (std::vector<int64_t>{2, 4}));
+}
+
+TEST_P(AutogradMatmulForwardTest, MatmulDifferentShapes) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{3, 4}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(1.0f);
+    auto b = std::make_shared<Tensor>(std::vector<int64_t>{4, 2}, DataType::kFLOAT32, GetDevice(), true);
+    b->Fill(1.0f);
+    auto matmul_fn = std::make_shared<autograd::Matmul>();
+    auto result = matmul_fn->Apply({a, b});
+    EXPECT_EQ(result.size(), 1);
+    EXPECT_EQ(result[0]->Dims(), (std::vector<int64_t>{3, 2}));
+}
+
+TEST_P(AutogradMatmulForwardTest, MatmulBatch) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3, 4}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(1.0f);
+    auto b = std::make_shared<Tensor>(std::vector<int64_t>{2, 4, 5}, DataType::kFLOAT32, GetDevice(), true);
+    b->Fill(1.0f);
+    auto matmul_fn = std::make_shared<autograd::Matmul>();
+    auto result = matmul_fn->Apply({a, b});
+    EXPECT_EQ(result.size(), 1);
+    EXPECT_EQ(result[0]->Dims(), (std::vector<int64_t>{2, 3, 5}));
+}
+
+TEST_P(AutogradMatmulForwardTest, MatmulSquare) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{3, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(1.0f);
+    auto b = std::make_shared<Tensor>(std::vector<int64_t>{3, 3}, DataType::kFLOAT32, GetDevice(), true);
+    b->Fill(1.0f);
+    auto matmul_fn = std::make_shared<autograd::Matmul>();
+    auto result = matmul_fn->Apply({a, b});
+    EXPECT_EQ(result.size(), 1);
+    EXPECT_EQ(result[0]->Dims(), (std::vector<int64_t>{3, 3}));
+}
+
+INFINI_TRAIN_REGISTER_TEST(AutogradMatmulForwardTest);
diff --git a/tests/autograd/test_autograd_normalization_backward.cc b/tests/autograd/test_autograd_normalization_backward.cc
new file mode 100644
index 00000000..cfbaaf98
--- /dev/null
+++ b/tests/autograd/test_autograd_normalization_backward.cc
@@ -0,0 +1,43 @@
+#include <vector>
+
+#include "infini_train/include/autograd/normalization.h"
+#include "infini_train/include/nn/parallel/global.h"
+#include "infini_train/include/tensor.h"
+#include "tests/common/test_utils.h"
+#include "gtest/gtest.h"
+
+using namespace infini_train;
+
+class AutogradNormalizationBackwardTest : public infini_train::test::InfiniTrainTest {};
+
+TEST_P(AutogradNormalizationBackwardTest, LayerNormBackward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3, 4}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(1.0f);
+    auto weight = std::make_shared<Tensor>(std::vector<int64_t>{4}, DataType::kFLOAT32, GetDevice(), true);
+    weight->Fill(1.0f);
+    auto bias = std::make_shared<Tensor>(std::vector<int64_t>{4}, DataType::kFLOAT32, GetDevice(), true);
+    bias->Fill(0.0f);
+    auto layernorm_fn = std::make_shared<autograd::LayerNorm>(1e-5f);
+    auto result = layernorm_fn->Apply({a, weight, bias});
+    auto grad = std::make_shared<Tensor>(std::vector<int64_t>{2, 3, 4}, DataType::kFLOAT32, GetDevice(), true);
+    grad->Fill(1.0f);
+    auto grad_inputs = layernorm_fn->Backward({grad});
+    EXPECT_EQ(grad_inputs.size(), 3);
+}
+
+TEST_P(AutogradNormalizationBackwardTest, LayerNormBackwardZeroBias) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3, 4}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(1.0f);
+    auto weight = std::make_shared<Tensor>(std::vector<int64_t>{4}, DataType::kFLOAT32, GetDevice(), true);
+    weight->Fill(1.0f);
+    auto bias = std::make_shared<Tensor>(std::vector<int64_t>{4}, DataType::kFLOAT32, GetDevice(), true);
+    bias->Fill(0.0f);
+    auto layernorm_fn = std::make_shared<autograd::LayerNorm>(1e-5f);
+    auto result = layernorm_fn->Apply({a, weight, bias});
+    auto grad = std::make_shared<Tensor>(std::vector<int64_t>{2, 3, 4}, DataType::kFLOAT32, GetDevice(), true);
+    grad->Fill(1.0f);
+    auto grad_inputs = layernorm_fn->Backward({grad});
+    EXPECT_EQ(grad_inputs.size(), 3);
+}
+
+INFINI_TRAIN_REGISTER_TEST(AutogradNormalizationBackwardTest);
diff --git a/tests/autograd/test_autograd_normalization_forward.cc b/tests/autograd/test_autograd_normalization_forward.cc
new file mode 100644
index 00000000..755f7d43
--- /dev/null
+++ b/tests/autograd/test_autograd_normalization_forward.cc
@@ -0,0 +1,50 @@
+#include <vector>
+
+#include "infini_train/include/autograd/normalization.h"
+#include "infini_train/include/nn/parallel/global.h"
+#include "infini_train/include/tensor.h"
+#include "tests/common/test_utils.h"
+#include "gtest/gtest.h"
+
+using namespace infini_train;
+
+class AutogradNormalizationForwardTest : public infini_train::test::InfiniTrainTest {};
+
+TEST_P(AutogradNormalizationForwardTest, LayerNormForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3, 4}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(1.0f);
+    auto weight = std::make_shared<Tensor>(std::vector<int64_t>{4}, DataType::kFLOAT32, GetDevice(), true);
+    weight->Fill(1.0f);
+    auto bias = std::make_shared<Tensor>(std::vector<int64_t>{4}, DataType::kFLOAT32, GetDevice(), true);
+    bias->Fill(0.0f);
+    auto layernorm_fn = std::make_shared<autograd::LayerNorm>(1e-5f);
+    auto result = layernorm_fn->Apply({a, weight, bias});
+    EXPECT_EQ(result.size(), 1);
+}
+
+TEST_P(AutogradNormalizationForwardTest, LayerNormZeroBias) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3, 4}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(1.0f);
+    auto weight = std::make_shared<Tensor>(std::vector<int64_t>{4}, DataType::kFLOAT32, GetDevice(), true);
+    weight->Fill(1.0f);
+    auto bias = std::make_shared<Tensor>(std::vector<int64_t>{4}, DataType::kFLOAT32, GetDevice(), true);
+    bias->Fill(0.0f);
+    auto layernorm_fn = std::make_shared<autograd::LayerNorm>(1e-5f);
+    auto result = layernorm_fn->Apply({a, weight, bias});
+    EXPECT_EQ(result.size(), 1);
+}
+
+TEST_P(AutogradNormalizationForwardTest, LayerNormThreeDim) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 1, 4}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(1.0f);
+    auto weight = std::make_shared<Tensor>(std::vector<int64_t>{4}, DataType::kFLOAT32, GetDevice(), true);
+    weight->Fill(1.0f);
+    auto bias = std::make_shared<Tensor>(std::vector<int64_t>{4}, DataType::kFLOAT32, GetDevice(), true);
+    bias->Fill(0.0f);
+    auto layernorm_fn = std::make_shared<autograd::LayerNorm>(1e-5f);
+    auto result = layernorm_fn->Apply({a, weight, bias});
+    EXPECT_EQ(result.size(), 1);
+    EXPECT_EQ(result[0]->Dims(), (std::vector<int64_t>{2, 1, 4}));
+}
+
+INFINI_TRAIN_REGISTER_TEST(AutogradNormalizationForwardTest);
diff --git a/tests/autograd/test_autograd_reduction_backward.cc b/tests/autograd/test_autograd_reduction_backward.cc
new file mode 100644
index 00000000..66de9d16
--- /dev/null
+++ b/tests/autograd/test_autograd_reduction_backward.cc
@@ -0,0 +1,79 @@
+#include <vector>
+
+#include "infini_train/include/autograd/reduction.h"
+#include "infini_train/include/nn/parallel/global.h"
+#include "infini_train/include/tensor.h"
+#include "tests/common/test_utils.h"
+#include "gtest/gtest.h"
+
+using namespace infini_train;
+
+class AutogradReductionBackwardTest : public infini_train::test::InfiniTrainTest {};
+
+TEST_P(AutogradReductionBackwardTest, SumBackward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(1.0f);
+    auto sum_fn = std::make_shared<autograd::Sum>(1, false);
+    auto result = sum_fn->Apply({a});
+    auto grad = std::make_shared<Tensor>(std::vector<int64_t>{2}, DataType::kFLOAT32, GetDevice(), true);
+    grad->Fill(1.0f);
+    auto grad_inputs = sum_fn->Backward({grad});
+    EXPECT_EQ(grad_inputs.size(), 1);
+}
+
+TEST_P(AutogradReductionBackwardTest, MeanBackward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(1.0f);
+    auto mean_fn = std::make_shared<autograd::Mean>(1, false);
+    auto result = mean_fn->Apply({a});
+    auto grad = std::make_shared<Tensor>(std::vector<int64_t>{2}, DataType::kFLOAT32, GetDevice(), true);
+    grad->Fill(1.0f);
+    auto grad_inputs = mean_fn->Backward({grad});
+    EXPECT_EQ(grad_inputs.size(), 1);
+}
+
+TEST_P(AutogradReductionBackwardTest, MaxBackward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(1.0f);
+    auto max_fn = std::make_shared<autograd::Max>(1, false);
+    auto result = max_fn->Apply({a});
+    auto grad = std::make_shared<Tensor>(std::vector<int64_t>{2}, DataType::kFLOAT32, GetDevice(), true);
+    grad->Fill(1.0f);
+    auto grad_inputs = max_fn->Backward({grad});
+    EXPECT_EQ(grad_inputs.size(), 1);
+}
+
+TEST_P(AutogradReductionBackwardTest, MinBackward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(1.0f);
+    auto min_fn = std::make_shared<autograd::Min>(1, false);
+    auto result = min_fn->Apply({a});
+    auto grad = std::make_shared<Tensor>(std::vector<int64_t>{2}, DataType::kFLOAT32, GetDevice(), true);
+    grad->Fill(1.0f);
+    auto grad_inputs = min_fn->Backward({grad});
+    EXPECT_EQ(grad_inputs.size(), 1);
+}
+
+TEST_P(AutogradReductionBackwardTest, SumBackwardKeepDim) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(1.0f);
+    auto sum_fn = std::make_shared<autograd::Sum>(1, true);
+    auto result = sum_fn->Apply({a});
+    auto grad = std::make_shared<Tensor>(std::vector<int64_t>{2, 1}, DataType::kFLOAT32, GetDevice(), true);
+    grad->Fill(1.0f);
+    auto grad_inputs = sum_fn->Backward({grad});
+    EXPECT_EQ(grad_inputs.size(), 1);
+}
+
+TEST_P(AutogradReductionBackwardTest, MeanBackwardKeepDim) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(1.0f);
+    auto mean_fn = std::make_shared<autograd::Mean>(1, true);
+    auto result = mean_fn->Apply({a});
+    auto grad = std::make_shared<Tensor>(std::vector<int64_t>{2, 1}, DataType::kFLOAT32, GetDevice(), true);
+    grad->Fill(1.0f);
+    auto grad_inputs = mean_fn->Backward({grad});
+    EXPECT_EQ(grad_inputs.size(), 1);
+}
+
+INFINI_TRAIN_REGISTER_TEST(AutogradReductionBackwardTest);
diff --git a/tests/autograd/test_autograd_reduction_forward.cc b/tests/autograd/test_autograd_reduction_forward.cc
new file mode 100644
index 00000000..6d394e15
--- /dev/null
+++ b/tests/autograd/test_autograd_reduction_forward.cc
@@ -0,0 +1,61 @@
+#include <vector>
+
+#include "infini_train/include/autograd/reduction.h"
+#include "infini_train/include/nn/parallel/global.h"
+#include "infini_train/include/tensor.h"
+#include "tests/common/test_utils.h"
+#include "gtest/gtest.h"
+
+using namespace infini_train;
+
+class AutogradReductionForwardTest : public infini_train::test::InfiniTrainTest {};
+
+TEST_P(AutogradReductionForwardTest, SumForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(1.0f);
+    auto sum_fn = std::make_shared<autograd::Sum>(1, false);
+    auto result = sum_fn->Apply({a});
+    EXPECT_EQ(result.size(), 1);
+}
+
+TEST_P(AutogradReductionForwardTest, MeanForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(1.0f);
+    auto mean_fn = std::make_shared<autograd::Mean>(1, false);
+    auto result = mean_fn->Apply({a});
+    EXPECT_EQ(result.size(), 1);
+}
+
+TEST_P(AutogradReductionForwardTest, MaxForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(1.0f);
+    auto max_fn = std::make_shared<autograd::Max>(1, false);
+    auto result = max_fn->Apply({a});
+    EXPECT_EQ(result.size(), 1);
+}
+
+TEST_P(AutogradReductionForwardTest, MinForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(1.0f);
+    auto min_fn = std::make_shared<autograd::Min>(1, false);
+    auto result = min_fn->Apply({a});
+    EXPECT_EQ(result.size(), 1);
+}
+
+TEST_P(AutogradReductionForwardTest, SumKeepDim) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(1.0f);
+    auto sum_fn = std::make_shared<autograd::Sum>(1, true);
+    auto result = sum_fn->Apply({a});
+    EXPECT_EQ(result.size(), 1);
+}
+
+TEST_P(AutogradReductionForwardTest, MeanKeepDim) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(1.0f);
+    auto mean_fn = std::make_shared<autograd::Mean>(1, true);
+    auto result = mean_fn->Apply({a});
+    EXPECT_EQ(result.size(), 1);
+}
+
+INFINI_TRAIN_REGISTER_TEST(AutogradReductionForwardTest);
diff --git a/tests/autograd/test_autograd_softmax_backward.cc b/tests/autograd/test_autograd_softmax_backward.cc
new file mode 100644
index 00000000..7e15c4b7
--- /dev/null
+++ b/tests/autograd/test_autograd_softmax_backward.cc
@@ -0,0 +1,35 @@
+#include <vector>
+
+#include "infini_train/include/autograd/softmax.h"
+#include "infini_train/include/nn/parallel/global.h"
+#include "infini_train/include/tensor.h"
+#include "tests/common/test_utils.h"
+#include "gtest/gtest.h"
+
+using namespace infini_train;
+
+class AutogradSoftmaxBackwardTest : public infini_train::test::InfiniTrainTest {};
+
+TEST_P(AutogradSoftmaxBackwardTest, SoftmaxBackward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(1.0f);
+    auto softmax_fn = std::make_shared<autograd::Softmax>(1);
+    auto result = softmax_fn->Apply({a});
+    auto grad = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    grad->Fill(1.0f);
+    auto grad_inputs = softmax_fn->Backward({grad});
+    EXPECT_EQ(grad_inputs.size(), 1);
+}
+
+TEST_P(AutogradSoftmaxBackwardTest, SoftmaxBackwardDim0) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{4, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(1.0f);
+    auto softmax_fn = std::make_shared<autograd::Softmax>(0);
+    auto result = softmax_fn->Apply({a});
+    auto grad = std::make_shared<Tensor>(std::vector<int64_t>{4, 3}, DataType::kFLOAT32, GetDevice(), true);
+    grad->Fill(1.0f);
+    auto grad_inputs = softmax_fn->Backward({grad});
+    EXPECT_EQ(grad_inputs.size(), 1);
+}
+
+INFINI_TRAIN_REGISTER_TEST(AutogradSoftmaxBackwardTest);
diff --git a/tests/autograd/test_autograd_softmax_forward.cc b/tests/autograd/test_autograd_softmax_forward.cc
new file mode 100644
index 00000000..39bde108
--- /dev/null
+++ b/tests/autograd/test_autograd_softmax_forward.cc
@@ -0,0 +1,40 @@
+#include <vector>
+
+#include "infini_train/include/autograd/softmax.h"
+#include "infini_train/include/nn/parallel/global.h"
+#include "infini_train/include/tensor.h"
+#include "tests/common/test_utils.h"
+#include "gtest/gtest.h"
+
+using namespace infini_train;
+
+class AutogradSoftmaxForwardTest : public infini_train::test::InfiniTrainTest {};
+
+TEST_P(AutogradSoftmaxForwardTest, SoftmaxForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(1.0f);
+    auto softmax_fn = std::make_shared<autograd::Softmax>(1);
+    auto result = softmax_fn->Apply({a});
+    EXPECT_EQ(result.size(), 1);
+    EXPECT_EQ(result[0]->Dims(), (std::vector<int64_t>{2, 3}));
+}
+
+TEST_P(AutogradSoftmaxForwardTest, SoftmaxDim0) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{4, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(1.0f);
+    auto softmax_fn = std::make_shared<autograd::Softmax>(0);
+    auto result = softmax_fn->Apply({a});
+    EXPECT_EQ(result.size(), 1);
+    EXPECT_EQ(result[0]->Dims(), (std::vector<int64_t>{4, 3}));
+}
+
+TEST_P(AutogradSoftmaxForwardTest, SoftmaxLastDim) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3, 4}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(1.0f);
+    auto softmax_fn = std::make_shared<autograd::Softmax>(2);
+    auto result = softmax_fn->Apply({a});
+    EXPECT_EQ(result.size(), 1);
+    EXPECT_EQ(result[0]->Dims(), (std::vector<int64_t>{2, 3, 4}));
+}
+
+INFINI_TRAIN_REGISTER_TEST(AutogradSoftmaxForwardTest);
diff --git a/tests/autograd/test_autograd_transform_backward.cc b/tests/autograd/test_autograd_transform_backward.cc
new file mode 100644
index 00000000..90461ad4
--- /dev/null
+++ b/tests/autograd/test_autograd_transform_backward.cc
@@ -0,0 +1,24 @@
+#include <vector>
+
+#include "infini_train/include/autograd/transform.h"
+#include "infini_train/include/nn/parallel/global.h"
+#include "infini_train/include/tensor.h"
+#include "tests/common/test_utils.h"
+#include "gtest/gtest.h"
+
+using namespace infini_train;
+
+class AutogradTransformBackwardTest : public infini_train::test::InfiniTrainTest {};
+
+TEST_P(AutogradTransformBackwardTest, TransposeBackward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(1.0f);
+    auto transpose_fn = std::make_shared<autograd::Transpose>(0, 1);
+    auto result = transpose_fn->Apply({a});
+    auto grad = std::make_shared<Tensor>(std::vector<int64_t>{3, 2}, DataType::kFLOAT32, GetDevice(), true);
+    grad->Fill(1.0f);
+    auto grad_inputs = transpose_fn->Backward({grad});
+    EXPECT_EQ(grad_inputs.size(), 1);
+}
+
+INFINI_TRAIN_REGISTER_TEST(AutogradTransformBackwardTest);
diff --git a/tests/autograd/test_autograd_transform_forward.cc b/tests/autograd/test_autograd_transform_forward.cc
new file mode 100644
index 00000000..acca3a61
--- /dev/null
+++ b/tests/autograd/test_autograd_transform_forward.cc
@@ -0,0 +1,78 @@
+#include <vector>
+
+#include "infini_train/include/autograd/misc.h"
+#include "infini_train/include/autograd/transform.h"
+#include "infini_train/include/nn/parallel/global.h"
+#include "infini_train/include/tensor.h"
+#include "tests/common/test_utils.h"
+#include "gtest/gtest.h"
+
+using namespace infini_train;
+
+class AutogradTransformForwardTest : public infini_train::test::InfiniTrainTest {};
+
+TEST_P(AutogradTransformForwardTest, TransposeForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(1.0f);
+    auto transpose_fn = std::make_shared<autograd::Transpose>(0, 1);
+    auto result = transpose_fn->Apply({a});
+    EXPECT_EQ(result.size(), 1);
+    EXPECT_EQ(result[0]->Dims(), (std::vector<int64_t>{3, 2}));
+}
+
+TEST_P(AutogradTransformForwardTest, SliceForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{4, 4}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(1.0f);
+    auto slice_fn = std::make_shared<autograd::Slice>(std::vector<int64_t>{1, 1}, std::vector<int64_t>{3, 3},
+                                                      std::vector<int64_t>{1, 1});
+    auto result = slice_fn->Apply({a});
+    EXPECT_EQ(result.size(), 1);
+}
+
+TEST_P(AutogradTransformForwardTest, SplitForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{4, 4}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(1.0f);
+    auto split_fn = std::make_shared<autograd::Split>(2, 0);
+    auto result = split_fn->Apply({a});
+    EXPECT_EQ(result.size(), 2);
+}
+
+TEST_P(AutogradTransformForwardTest, ConcatForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 2}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(1.0f);
+    auto b = std::make_shared<Tensor>(std::vector<int64_t>{2, 2}, DataType::kFLOAT32, GetDevice(), true);
+    b->Fill(2.0f);
+    auto concat_fn = std::make_shared<autograd::Concat>(0);
+    auto result = concat_fn->Apply({a, b});
+    EXPECT_EQ(result.size(), 1);
+    EXPECT_EQ(result[0]->Dims(), (std::vector<int64_t>{4, 2}));
+}
+
+TEST_P(AutogradTransformForwardTest, StackForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(1.0f);
+    auto b = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    b->Fill(2.0f);
+    auto stack_fn = std::make_shared<autograd::Stack>(0);
+    auto result = stack_fn->Apply({a, b});
+    EXPECT_EQ(result.size(), 1);
+    EXPECT_EQ(result[0]->Dims(), (std::vector<int64_t>{2, 2, 3}));
+}
+
+TEST_P(AutogradTransformForwardTest, TrilForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{3, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(1.0f);
+    auto tril_fn = std::make_shared<autograd::Tril>(0);
+    auto result = tril_fn->Apply({a});
+    EXPECT_EQ(result.size(), 1);
+}
+
+TEST_P(AutogradTransformForwardTest, TriuForward) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{3, 3}, DataType::kFLOAT32, GetDevice(), true);
+    a->Fill(1.0f);
+    auto triu_fn = std::make_shared<autograd::Triu>(0);
+    auto result = triu_fn->Apply({a});
+    EXPECT_EQ(result.size(), 1);
+}
+
+INFINI_TRAIN_REGISTER_TEST(AutogradTransformForwardTest);
diff --git a/tests/common/CMakeLists.txt b/tests/common/CMakeLists.txt
new file mode 100644
index 00000000..84c92e51
--- /dev/null
+++ b/tests/common/CMakeLists.txt
@@ -0,0 +1,7 @@
+# Common test utilities + custom test main (initializes GlobalEnv before RUN_ALL_TESTS).
+add_library(test_main OBJECT test_main.cc)
+target_include_directories(test_main PUBLIC
+  ${CMAKE_SOURCE_DIR}
+  ${glog_SOURCE_DIR}/src
+)
+target_link_libraries(test_main PUBLIC GTest::gtest)
diff --git a/tests/common/test_main.cc b/tests/common/test_main.cc
new file mode 100644
index 00000000..23b6643e
--- /dev/null
+++ b/tests/common/test_main.cc
@@ -0,0 +1,8 @@
+#include "infini_train/include/nn/parallel/global.h"
+#include "gtest/gtest.h"
+
+int main(int argc, char **argv) {
+    ::testing::InitGoogleTest(&argc, argv);
+    infini_train::nn::parallel::global::GlobalEnv::Instance().Init(1, 1, false, 1, 1);
+    return RUN_ALL_TESTS();
+}
diff --git a/tests/common/test_utils.h b/tests/common/test_utils.h
new file mode 100644
index 00000000..06a8bf04
--- /dev/null
+++ b/tests/common/test_utils.h
@@ -0,0 +1,63 @@
+#pragma once
+
+#if defined(USE_CUDA)
+#include <cuda_runtime_api.h>
+#endif
+
+#include "infini_train/include/device.h"
+#include "gtest/gtest.h"
+
+namespace infini_train {
+namespace test {
+
+#if defined(USE_CUDA)
+#define REQUIRE_MIN_DEVICES(n)                                                                                         \
+    do {                                                                                                               \
+        int available_gpus = 0;                                                                                        \
+        cudaGetDeviceCount(&available_gpus);                                                                           \
+        if (available_gpus < (n)) {                                                                                    \
+            GTEST_SKIP() << "requires at least " << (n) << " GPUs (found " << available_gpus << ")";                   \
+        }                                                                                                              \
+    } while (0)
+#else
+#define REQUIRE_MIN_DEVICES(n)                                                                                         \
+    do { GTEST_SKIP() << "requires at least " << (n) << " GPUs (CUDA disabled)"; } while (0)
+#endif
+
+#define SKIP_CPU()                                                                                                     \
+    do {                                                                                                               \
+        if (GetParam() == infini_train::Device::DeviceType::kCPU) {                                                    \
+            GTEST_SKIP() << "skipped on CPU";                                                                          \
+        }                                                                                                              \
+    } while (0)
+
+#define ONLY_CPU()                                                                                                     \
+    do {                                                                                                               \
+        if (GetParam() != infini_train::Device::DeviceType::kCPU) {                                                    \
+            GTEST_SKIP() << "CPU-only test";                                                                           \
+        }                                                                                                              \
+    } while (0)
+
+#define ONLY_CUDA()                                                                                                    \
+    do {                                                                                                               \
+        if (GetParam() != infini_train::Device::DeviceType::kCUDA) {                                                   \
+            GTEST_SKIP() << "CUDA-only test";                                                                          \
+        }                                                                                                              \
+    } while (0)
+
+class InfiniTrainTest : public ::testing::TestWithParam<Device::DeviceType> {
+protected:
+    Device GetDevice() const { return Device(GetParam(), 0); }
+};
+
+} // namespace test
+} // namespace infini_train
+
+#if defined(USE_CUDA)
+#define INFINI_TRAIN_REGISTER_TEST(TestName)                                                                           \
+    INSTANTIATE_TEST_SUITE_P(CPU, TestName, ::testing::Values(infini_train::Device::DeviceType::kCPU));                \
+    INSTANTIATE_TEST_SUITE_P(CUDA, TestName, ::testing::Values(infini_train::Device::DeviceType::kCUDA))
+#else
+#define INFINI_TRAIN_REGISTER_TEST(TestName)                                                                           \
+    INSTANTIATE_TEST_SUITE_P(CPU, TestName, ::testing::Values(infini_train::Device::DeviceType::kCPU))
+#endif
diff --git a/tests/dtype/CMakeLists.txt b/tests/dtype/CMakeLists.txt
new file mode 100644
index 00000000..031bb797
--- /dev/null
+++ b/tests/dtype/CMakeLists.txt
@@ -0,0 +1,40 @@
+# ==========================================================================
+# Dtype dispatch tests
+# ==========================================================================
+
+file(GLOB DTYPE_SOURCES ${CMAKE_CURRENT_SOURCE_DIR}/test_*.cc)
+list(FILTER DTYPE_SOURCES EXCLUDE REGEX ".*compile_fail.*")
+
+infini_train_add_test_suite(test_dtype
+  SOURCES ${DTYPE_SOURCES}
+)
+
+# Negative compile test: missing dtype registration must fail at compile time.
+set(DTYPE_DISPATCH_COMPILE_FAIL_SOURCE
+  ${CMAKE_CURRENT_SOURCE_DIR}/test_dtype_dispatch_compile_fail.cc)
+
+try_compile(DTYPE_DISPATCH_COMPILE_UNEXPECTEDLY_SUCCEEDED
+  ${CMAKE_BINARY_DIR}/CMakeFiles/try_compile_dtype_dispatch_missing_map
+  SOURCES ${DTYPE_DISPATCH_COMPILE_FAIL_SOURCE}
+  CMAKE_FLAGS
+    "-DCMAKE_CXX_STANDARD=${CMAKE_CXX_STANDARD}"
+    "-DCMAKE_CXX_STANDARD_REQUIRED=ON"
+    "-DCMAKE_CXX_EXTENSIONS=OFF"
+    "-DCMAKE_CXX_FLAGS=-I${PROJECT_SOURCE_DIR}"
+  OUTPUT_VARIABLE DTYPE_DISPATCH_TRY_COMPILE_OUTPUT
+)
+
+if(DTYPE_DISPATCH_COMPILE_UNEXPECTEDLY_SUCCEEDED)
+  message(FATAL_ERROR
+    "dtype dispatch compile-fail test unexpectedly succeeded.\n"
+    "Source: ${DTYPE_DISPATCH_COMPILE_FAIL_SOURCE}\n"
+    "Output:\n${DTYPE_DISPATCH_TRY_COMPILE_OUTPUT}")
+else()
+  message(STATUS "compile-fail output:\n${DTYPE_DISPATCH_TRY_COMPILE_OUTPUT}")
+endif()
+
+add_custom_target(test_dtype_dispatch_compile_fail
+  COMMAND ${CMAKE_COMMAND} -E echo
+    "dtype dispatch compile-fail check passed (missing dtype registration correctly fails to compile)."
+  VERBATIM
+)
diff --git a/tests/dtype/test_dtype_dispatch.cc b/tests/dtype/test_dtype_dispatch.cc
new file mode 100644
index 00000000..10fc2122
--- /dev/null
+++ b/tests/dtype/test_dtype_dispatch.cc
@@ -0,0 +1,71 @@
+#include <type_traits>
+
+#include "infini_train/include/datatype.h"
+#include "infini_train/include/dtype_dispatch.h"
+#include "infini_train/src/core/runtime/cpu/cpu_dispatch.h"
+#include "tests/common/test_utils.h"
+#include "gtest/gtest.h"
+
+using namespace infini_train;
+
+// ============================================================================
+// Compile-time checks (static_assert, always verified at build time)
+// ============================================================================
+
+template <DataType DType> struct LowPrecisionAbsentTypeMap;
+template <> struct LowPrecisionAbsentTypeMap<DataType::kFLOAT32> {
+    using type = float;
+};
+
+static_assert(HasMappedType_v<LowPrecisionAbsentTypeMap, DataType::kFLOAT32>,
+              "sanity: registered dtype must be detected as present");
+static_assert(!HasMappedType_v<LowPrecisionAbsentTypeMap, DataType::kFLOAT16>,
+              "unregistered kFLOAT16 must be intercepted by HasMappedType_v");
+static_assert(!HasMappedType_v<LowPrecisionAbsentTypeMap, DataType::kBFLOAT16>,
+              "unregistered kBFLOAT16 must be intercepted by HasMappedType_v");
+
+static_assert(std::is_same_v<MappedType_t<core::cpu::CpuTypeMap, DataType::kFLOAT16>, FP16>,
+              "CpuTypeMap<kFLOAT16> must resolve to framework FP16");
+static_assert(std::is_same_v<MappedType_t<core::cpu::CpuTypeMap, DataType::kBFLOAT16>, BF16>,
+              "CpuTypeMap<kBFLOAT16> must resolve to framework BF16");
+
+// ============================================================================
+// Runtime tests
+// ============================================================================
+
+class DtypeDispatchTest : public infini_train::test::InfiniTrainTest {};
+
+TEST_P(DtypeDispatchTest, DispatchFP16) {
+    bool called = false;
+    core::cpu::DispatchCpuFunc<DataType::kFLOAT16, DataType::kBFLOAT16>(
+        DataType::kFLOAT16,
+        [&called]<typename T>() {
+            if constexpr (std::is_same_v<T, FP16>) {
+                called = true;
+            }
+        },
+        "dispatch kFLOAT16");
+    EXPECT_TRUE(called);
+}
+
+TEST_P(DtypeDispatchTest, DispatchBF16) {
+    bool called = false;
+    core::cpu::DispatchCpuFunc<DataType::kFLOAT16, DataType::kBFLOAT16>(
+        DataType::kBFLOAT16,
+        [&called]<typename T>() {
+            if constexpr (std::is_same_v<T, BF16>) {
+                called = true;
+            }
+        },
+        "dispatch kBFLOAT16");
+    EXPECT_TRUE(called);
+}
+
+TEST_P(DtypeDispatchTest, UnallowedDtypeFatals) {
+    EXPECT_DEATH(core::cpu::DispatchCpuFunc<DataType::kFLOAT32>(
+                     DataType::kFLOAT16, []<typename T>() { (void)sizeof(T); },
+                     "intercept kFLOAT16 when only kFLOAT32 is allowed"),
+                 "");
+}
+
+INFINI_TRAIN_REGISTER_TEST(DtypeDispatchTest);
diff --git a/test/dtype/test_dtype_dispatch_compile_fail.cc b/tests/dtype/test_dtype_dispatch_compile_fail.cc
similarity index 88%
rename from test/dtype/test_dtype_dispatch_compile_fail.cc
rename to tests/dtype/test_dtype_dispatch_compile_fail.cc
index 57c706a3..5beddd85 100644
--- a/test/dtype/test_dtype_dispatch_compile_fail.cc
+++ b/tests/dtype/test_dtype_dispatch_compile_fail.cc
@@ -21,8 +21,6 @@ int main() {
     // static_assert inside DispatchByTypeMap, failing this translation unit
     // before MappedType_t<TypeMap, kFLOAT16> is ever instantiated.
     DispatchByTypeMap<LowPrecisionMissingTypeMap, DataType::kFLOAT16>(
-        DataType::kFLOAT16,
-        []<typename T>() { (void)sizeof(T); },
-        "compile-fail: unregistered low-precision dtype");
+        DataType::kFLOAT16, []<typename T>() { (void)sizeof(T); }, "compile-fail: unregistered low-precision dtype");
     return 0;
 }
diff --git a/tests/dtype/test_scalar.cc b/tests/dtype/test_scalar.cc
new file mode 100644
index 00000000..79570ea9
--- /dev/null
+++ b/tests/dtype/test_scalar.cc
@@ -0,0 +1,145 @@
+#include <cmath>
+#include <cstdint>
+#include <limits>
+
+#include "infini_train/include/scalar.h"
+#include "tests/common/test_utils.h"
+#include "gtest/gtest.h"
+
+using namespace infini_train;
+
+class ScalarTest : public infini_train::test::InfiniTrainTest {};
+
+TEST_P(ScalarTest, DefaultConstructor) {
+    Scalar s;
+    EXPECT_EQ(static_cast<int>(s.kind), static_cast<int>(Scalar::Kind::kInt64));
+    EXPECT_EQ(s.i, 0);
+}
+
+TEST_P(ScalarTest, BoolConstructor) {
+    Scalar t(true);
+    EXPECT_EQ(static_cast<int>(t.kind), static_cast<int>(Scalar::Kind::kBool));
+    EXPECT_EQ(t.u, static_cast<uint64_t>(1));
+
+    Scalar f(false);
+    EXPECT_EQ(static_cast<int>(f.kind), static_cast<int>(Scalar::Kind::kBool));
+    EXPECT_EQ(f.u, static_cast<uint64_t>(0));
+}
+
+TEST_P(ScalarTest, SignedIntConstructor) {
+    Scalar pos(42);
+    EXPECT_EQ(static_cast<int>(pos.kind), static_cast<int>(Scalar::Kind::kInt64));
+    EXPECT_EQ(pos.i, 42);
+
+    Scalar neg(-7);
+    EXPECT_EQ(neg.i, -7);
+
+    int16_t short_val = -100;
+    Scalar s(short_val);
+    EXPECT_EQ(static_cast<int>(s.kind), static_cast<int>(Scalar::Kind::kInt64));
+    EXPECT_EQ(s.i, -100);
+}
+
+TEST_P(ScalarTest, UnsignedIntConstructor) {
+    unsigned int uint_val = 99u;
+    Scalar s(uint_val);
+    EXPECT_EQ(static_cast<int>(s.kind), static_cast<int>(Scalar::Kind::kUInt64));
+    EXPECT_EQ(s.u, static_cast<uint64_t>(99));
+
+    uint64_t uint64_max = std::numeric_limits<uint64_t>::max();
+    Scalar smax(uint64_max);
+    EXPECT_EQ(smax.u, uint64_max);
+}
+
+TEST_P(ScalarTest, FloatConstructor) {
+    Scalar s(3.14f);
+    EXPECT_EQ(static_cast<int>(s.kind), static_cast<int>(Scalar::Kind::kDouble));
+    EXPECT_NEAR(s.d, 3.14, 1e-5);
+}
+
+TEST_P(ScalarTest, DoubleConstructor) {
+    Scalar s(2.718281828);
+    EXPECT_EQ(static_cast<int>(s.kind), static_cast<int>(Scalar::Kind::kDouble));
+    EXPECT_NEAR(s.d, 2.718281828, 1e-12);
+}
+
+TEST_P(ScalarTest, HalfPrecisionConstructor) {
+    FP16 fp16_val(1.5f);
+    Scalar sfp16(fp16_val);
+    EXPECT_EQ(static_cast<int>(sfp16.kind), static_cast<int>(Scalar::Kind::kDouble));
+    EXPECT_NEAR(sfp16.d, 1.5, 1e-3);
+
+    BF16 bf16_val(2.0f);
+    Scalar sbf16(bf16_val);
+    EXPECT_EQ(static_cast<int>(sbf16.kind), static_cast<int>(Scalar::Kind::kDouble));
+    EXPECT_NEAR(sbf16.d, 2.0, 1e-2);
+}
+
+TEST_P(ScalarTest, ToNumericConversions) {
+    EXPECT_EQ(Scalar(10).to<int64_t>(), 10);
+    EXPECT_NEAR(Scalar(3.14).to<double>(), 3.14, 1e-12);
+    EXPECT_NEAR(Scalar(42).to<float>(), 42.0f, 1e-6);
+    EXPECT_NEAR(Scalar(42).to<double>(), 42.0, 1e-12);
+    EXPECT_EQ(Scalar(7.9).to<int64_t>(), 7);
+    EXPECT_EQ(Scalar(7.9).to<int32_t>(), 7);
+    EXPECT_NEAR(Scalar(-42).to<float>(), -42.0f, 1e-6);
+}
+
+TEST_P(ScalarTest, ToBoolConversions) {
+    EXPECT_EQ(Scalar(true).to<int64_t>(), 1);
+    EXPECT_EQ(Scalar(false).to<int64_t>(), 0);
+    EXPECT_NEAR(Scalar(true).to<double>(), 1.0, 1e-12);
+    EXPECT_NEAR(Scalar(false).to<double>(), 0.0, 1e-12);
+    EXPECT_EQ(Scalar(0).to<bool>(), false);
+    EXPECT_EQ(Scalar(5).to<bool>(), true);
+}
+
+TEST_P(ScalarTest, ToUnsignedConversions) {
+    uint64_t val = 12345;
+    EXPECT_NEAR(Scalar(val).to<double>(), 12345.0, 1e-6);
+    EXPECT_EQ(Scalar(100.0).to<uint64_t>(), static_cast<uint64_t>(100));
+}
+
+TEST_P(ScalarTest, ToHalfPrecisionConversions) {
+    auto to_fp16 = Scalar(1.5).to<FP16>();
+    EXPECT_NEAR(static_cast<float>(to_fp16), 1.5f, 1e-3);
+
+    auto int_to_fp16 = Scalar(3).to<FP16>();
+    EXPECT_NEAR(static_cast<float>(int_to_fp16), 3.0f, 1e-3);
+
+    auto to_bf16 = Scalar(2.0).to<BF16>();
+    EXPECT_NEAR(static_cast<float>(to_bf16), 2.0f, 1e-2);
+
+    auto int_to_bf16 = Scalar(5).to<BF16>();
+    EXPECT_NEAR(static_cast<float>(int_to_bf16), 5.0f, 1e-1);
+}
+
+TEST_P(ScalarTest, NumericLimits) {
+    int64_t i64max = std::numeric_limits<int64_t>::max();
+    EXPECT_EQ(Scalar(i64max).to<int64_t>(), i64max);
+
+    int64_t i64min = std::numeric_limits<int64_t>::min();
+    EXPECT_EQ(Scalar(i64min).to<int64_t>(), i64min);
+
+    uint64_t u64max = std::numeric_limits<uint64_t>::max();
+    EXPECT_EQ(Scalar(u64max).to<uint64_t>(), u64max);
+}
+
+TEST_P(ScalarTest, ZeroValues) {
+    EXPECT_EQ(Scalar(0).to<int64_t>(), 0);
+    EXPECT_NEAR(Scalar(0).to<double>(), 0.0, 1e-12);
+    EXPECT_NEAR(Scalar(0.0).to<double>(), 0.0, 1e-12);
+    EXPECT_EQ(Scalar(0.0).to<int64_t>(), 0);
+}
+
+TEST_P(ScalarTest, HalfPrecisionRoundtrip) {
+    FP16 fp16(0.5f);
+    auto rt_fp16 = Scalar(fp16).to<FP16>();
+    EXPECT_NEAR(static_cast<float>(rt_fp16), 0.5f, 1e-3);
+
+    BF16 bf16(4.0f);
+    auto rt_bf16 = Scalar(bf16).to<BF16>();
+    EXPECT_NEAR(static_cast<float>(rt_bf16), 4.0f, 1e-2);
+}
+
+INFINI_TRAIN_REGISTER_TEST(ScalarTest);
diff --git a/tests/hook/CMakeLists.txt b/tests/hook/CMakeLists.txt
new file mode 100644
index 00000000..ed419e04
--- /dev/null
+++ b/tests/hook/CMakeLists.txt
@@ -0,0 +1,9 @@
+# ==========================================================================
+# Hook tests
+# ==========================================================================
+
+file(GLOB HOOK_SOURCES ${CMAKE_CURRENT_SOURCE_DIR}/test_*.cc)
+
+infini_train_add_test_suite(test_hook
+  SOURCES ${HOOK_SOURCES}
+)
diff --git a/tests/hook/test_hook.cc b/tests/hook/test_hook.cc
new file mode 100644
index 00000000..cd815e80
--- /dev/null
+++ b/tests/hook/test_hook.cc
@@ -0,0 +1,76 @@
+#include <memory>
+#include <vector>
+
+#include "infini_train/include/autograd/elementwise.h"
+#include "infini_train/include/autograd/function.h"
+#include "infini_train/include/autograd/function_hook.h"
+#include "infini_train/include/common/hook.h"
+#include "infini_train/include/nn/modules/module.h"
+#include "infini_train/include/tensor.h"
+#include "tests/common/test_utils.h"
+#include "gtest/gtest.h"
+
+using namespace infini_train;
+
+class HookTest : public infini_train::test::InfiniTrainTest {};
+
+class TestModule : public nn::Module {
+public:
+    TestModule() : Module("TestModule") {}
+    std::vector<std::shared_ptr<Tensor>> Forward(const std::vector<std::shared_ptr<Tensor>> &inputs) override {
+        return inputs;
+    }
+};
+
+TEST_P(HookTest, BasicModuleHooks) {
+    auto x = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice());
+    x->set_requires_grad(true);
+
+    auto module = std::make_shared<TestModule>();
+    auto pre_hook = module->RegisterForwardPreHook([](nn::Module *, const std::vector<std::shared_ptr<Tensor>> &) {});
+    auto fwd_hook = module->RegisterForwardPostHook([](nn::Module *, const std::vector<std::shared_ptr<Tensor>> &,
+                                                       const std::vector<std::shared_ptr<Tensor>> &) {});
+    auto bwd_pre_hook
+        = module->RegisterBackwardPreHook([](nn::Module *, const std::vector<std::shared_ptr<Tensor>> &) {});
+    auto bwd_post_hook = module->RegisterBackwardPostHook([](nn::Module *, const std::vector<std::shared_ptr<Tensor>> &,
+                                                             const std::vector<std::shared_ptr<Tensor>> &) {});
+
+    auto outputs = (*module)({x});
+    EXPECT_EQ(outputs.size(), 1);
+}
+
+TEST_P(HookTest, HookRemove) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 2}, DataType::kFLOAT32, GetDevice());
+    auto b = std::make_shared<Tensor>(std::vector<int64_t>{2, 2}, DataType::kFLOAT32, GetDevice());
+    a->set_requires_grad(true);
+    b->set_requires_grad(true);
+
+    int hook1_count = 0, hook2_count = 0, hook3_count = 0;
+    auto add_fn = std::make_shared<autograd::Add>();
+
+    auto handle1 = add_fn->RegisterForwardPreHook(
+        [&hook1_count](autograd::Function *, const std::vector<std::shared_ptr<Tensor>> &) { hook1_count++; });
+    auto handle2 = add_fn->RegisterForwardPreHook(
+        [&hook2_count](autograd::Function *, const std::vector<std::shared_ptr<Tensor>> &) { hook2_count++; });
+    auto handle3 = add_fn->RegisterForwardPreHook(
+        [&hook3_count](autograd::Function *, const std::vector<std::shared_ptr<Tensor>> &) { hook3_count++; });
+
+    add_fn->Apply({a, b});
+    EXPECT_EQ(hook1_count, 1);
+    EXPECT_EQ(hook2_count, 1);
+    EXPECT_EQ(hook3_count, 1);
+
+    handle2->Remove();
+    add_fn->Apply({a, b});
+    EXPECT_EQ(hook1_count, 2);
+    EXPECT_EQ(hook2_count, 1);
+    EXPECT_EQ(hook3_count, 2);
+
+    handle1->Remove();
+    add_fn->Apply({a, b});
+    EXPECT_EQ(hook1_count, 2);
+    EXPECT_EQ(hook2_count, 1);
+    EXPECT_EQ(hook3_count, 3);
+}
+
+INFINI_TRAIN_REGISTER_TEST(HookTest);
diff --git a/tests/hook/test_precision_check.cc b/tests/hook/test_precision_check.cc
new file mode 100644
index 00000000..c9e7fc29
--- /dev/null
+++ b/tests/hook/test_precision_check.cc
@@ -0,0 +1,72 @@
+#include <filesystem>
+#include <memory>
+#include <vector>
+
+#include "infini_train/include/nn/modules/module.h"
+#include "infini_train/include/tensor.h"
+#include "infini_train/include/utils/precision_check_config.h"
+#include "infini_train/include/utils/precision_checker.h"
+#include "tests/common/test_utils.h"
+#include "gtest/gtest.h"
+
+using namespace infini_train;
+
+class PrecisionCheckTest : public infini_train::test::InfiniTrainTest {};
+
+class SimpleModel : public nn::Module {
+public:
+    SimpleModel() : Module("SimpleModel") {}
+
+    std::vector<std::shared_ptr<Tensor>> Forward(const std::vector<std::shared_ptr<Tensor>> &inputs) override {
+        auto x = inputs[0];
+        x->RequiresGrad();
+        auto y = x->Mul(x)->Mul(x);
+        return {y};
+    }
+};
+
+TEST_P(PrecisionCheckTest, SimpleFormat) {
+    auto x = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice());
+    x->Fill(2.0f);
+    x->RequiresGrad();
+
+    auto y = x->Mul(x);
+    auto loss = y->Sum(0, false)->Sum(0, false);
+    loss->Backward();
+
+    EXPECT_NE(x->DataPtr(), nullptr);
+}
+
+TEST_P(PrecisionCheckTest, ModuleForwardBackward) {
+    auto model = std::make_shared<SimpleModel>();
+
+    auto x = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice());
+    x->Fill(2.0f);
+    x->RequiresGrad();
+
+    std::vector<std::shared_ptr<Tensor>> inputs = {x};
+    auto outputs = (*model)(inputs);
+    auto loss = outputs[0]->Sum(0, false)->Sum(0, false);
+    loss->Backward();
+
+    EXPECT_TRUE(x->requires_grad());
+}
+
+TEST_P(PrecisionCheckTest, MultiIteration) {
+    auto model = std::make_shared<SimpleModel>();
+
+    for (int i = 0; i < 3; ++i) {
+        auto x = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice());
+        x->Fill(2.0f);
+        x->RequiresGrad();
+
+        std::vector<std::shared_ptr<Tensor>> inputs = {x};
+        auto outputs = (*model)(inputs);
+        auto loss = outputs[0]->Sum(0, false)->Sum(0, false);
+        loss->Backward();
+    }
+
+    SUCCEED();
+}
+
+INFINI_TRAIN_REGISTER_TEST(PrecisionCheckTest);
diff --git a/tests/lora/CMakeLists.txt b/tests/lora/CMakeLists.txt
new file mode 100644
index 00000000..1d985373
--- /dev/null
+++ b/tests/lora/CMakeLists.txt
@@ -0,0 +1,9 @@
+# ==========================================================================
+# LoRA tests
+# ==========================================================================
+
+file(GLOB LORA_SOURCES ${CMAKE_CURRENT_SOURCE_DIR}/test_*.cc)
+
+infini_train_add_test_suite(test_lora
+  SOURCES ${LORA_SOURCES}
+)
diff --git a/tests/lora/test_lora.cc b/tests/lora/test_lora.cc
new file mode 100644
index 00000000..caa707ca
--- /dev/null
+++ b/tests/lora/test_lora.cc
@@ -0,0 +1,338 @@
+#include <cmath>
+#include <memory>
+#include <vector>
+
+#include "infini_train/include/nn/lora/lora_config.h"
+#include "infini_train/include/nn/lora/lora_linear.h"
+#include "infini_train/include/nn/lora/lora_utils.h"
+#include "infini_train/include/nn/modules/container.h"
+#include "infini_train/include/nn/modules/linear.h"
+#include "infini_train/include/nn/modules/module.h"
+#include "infini_train/include/nn/parallel/global.h"
+#include "infini_train/include/tensor.h"
+#include "tests/common/test_utils.h"
+#include "gtest/gtest.h"
+
+using namespace infini_train;
+using namespace infini_train::nn::lora;
+
+class LoRATest : public infini_train::test::InfiniTrainTest {};
+
+// Helper: sum tensor values on any device by copying to CPU if needed.
+static float TensorSum(const std::shared_ptr<Tensor> &t) {
+    if (t->GetDevice().IsCPU()) {
+        return t->EigenMatrix().sum();
+    }
+    auto cpu = std::make_shared<Tensor>(t->Dims(), t->Dtype(), Device(Device::DeviceType::kCPU, 0));
+    cpu->CopyFrom(t);
+    return cpu->EigenMatrix().sum();
+}
+
+TEST_P(LoRATest, LoRAConfigScaling) {
+    LoRAConfig config;
+    config.rank = 8;
+    config.alpha = 16.0f;
+
+    float expected_scaling = 16.0f / 8.0f;
+    EXPECT_EQ(config.Scaling(), expected_scaling);
+}
+
+TEST_P(LoRATest, LoRAConfigShouldApply) {
+    LoRAConfig config;
+    config.rank = 8;
+    config.alpha = 16.0f;
+
+    EXPECT_TRUE(config.ShouldApplyLoRA("c_attn"));
+    EXPECT_TRUE(config.ShouldApplyLoRA("transformer.h.0.attn.c_attn"));
+    EXPECT_TRUE(config.ShouldApplyLoRA("c_proj"));
+    EXPECT_FALSE(config.ShouldApplyLoRA("c_fc"));
+    EXPECT_FALSE(config.ShouldApplyLoRA("random_layer"));
+}
+
+TEST_P(LoRATest, LoRALinearFromModel) {
+    auto base_linear = std::make_shared<nn::Linear>(64, 128, /*bias=*/true, GetDevice());
+
+    LoRAConfig config;
+    config.rank = 4;
+    config.alpha = 8.0f;
+    config.target_modules = {"Linear"};
+
+    auto model = GetLoRAModel(base_linear, config);
+
+    auto *lora_linear = dynamic_cast<LoRALinear *>(model.get());
+    ASSERT_NE(lora_linear, nullptr);
+
+    EXPECT_EQ(lora_linear->in_features(), 64);
+    EXPECT_EQ(lora_linear->out_features(), 128);
+    EXPECT_EQ(lora_linear->rank(), 4);
+
+    auto lora_A = lora_linear->parameter(LoRALinear::kParamLoraAName);
+    auto lora_B = lora_linear->parameter(LoRALinear::kParamLoraBName);
+    auto weight = lora_linear->parameter(nn::Linear::kParamWeightName);
+
+    EXPECT_EQ(lora_A->Dims()[0], config.rank);
+    EXPECT_EQ(lora_A->Dims()[1], 64);
+    EXPECT_EQ(lora_B->Dims()[0], 128);
+    EXPECT_EQ(lora_B->Dims()[1], config.rank);
+
+    EXPECT_FALSE(weight->requires_grad());
+    EXPECT_TRUE(lora_A->requires_grad());
+    EXPECT_TRUE(lora_B->requires_grad());
+
+    auto params = lora_linear->LoRAParameters();
+    EXPECT_EQ(params.size(), 2);
+}
+
+TEST_P(LoRATest, LoRALinearForward) {
+    auto base_linear = std::make_shared<nn::Linear>(64, 128, /*bias=*/true, GetDevice());
+
+    LoRAConfig config;
+    config.rank = 4;
+    config.alpha = 8.0f;
+    config.target_modules = {"Linear"};
+
+    auto model = GetLoRAModel(base_linear, config);
+
+    auto input = std::make_shared<Tensor>(std::vector<int64_t>{2, 10, 64}, DataType::kFLOAT32, GetDevice());
+
+    auto output = (*model)({input})[0];
+
+    EXPECT_EQ(output->Dims().size(), 3);
+    EXPECT_EQ(output->Dims()[0], 2);
+    EXPECT_EQ(output->Dims()[1], 10);
+    EXPECT_EQ(output->Dims()[2], 128);
+}
+
+TEST_P(LoRATest, LoRALinearMerge) {
+    auto base_linear = std::make_shared<nn::Linear>(32, 64, /*bias=*/false, GetDevice());
+
+    LoRAConfig config;
+    config.rank = 4;
+    config.alpha = 8.0f;
+    config.target_modules = {"Linear"};
+
+    auto model = GetLoRAModel(base_linear, config);
+
+    auto *lora_linear = dynamic_cast<LoRALinear *>(model.get());
+    ASSERT_NE(lora_linear, nullptr);
+
+    auto input = std::make_shared<Tensor>(std::vector<int64_t>{2, 5, 32}, DataType::kFLOAT32, GetDevice());
+    input->Fill(1.0f);
+
+    auto output_before = (*model)({input})[0];
+    float output_before_sum = TensorSum(output_before);
+
+    EXPECT_FALSE(lora_linear->IsMerged());
+    MergeLoRAWeights(model);
+    EXPECT_TRUE(lora_linear->IsMerged());
+
+    auto lora_A = lora_linear->parameter(LoRALinear::kParamLoraAName);
+    auto lora_B = lora_linear->parameter(LoRALinear::kParamLoraBName);
+    EXPECT_FALSE(lora_A->requires_grad());
+    EXPECT_FALSE(lora_B->requires_grad());
+
+    auto output_merged = (*model)({input})[0];
+    float output_merged_sum = TensorSum(output_merged);
+    EXPECT_NEAR(std::abs(output_before_sum - output_merged_sum), 0.0f, 1e-3);
+
+    UnmergeLoRAWeights(model);
+    EXPECT_FALSE(lora_linear->IsMerged());
+    EXPECT_TRUE(lora_A->requires_grad());
+    EXPECT_TRUE(lora_B->requires_grad());
+
+    auto output_unmerged = (*model)({input})[0];
+    EXPECT_EQ(output_before->Dims(), output_unmerged->Dims());
+}
+
+TEST_P(LoRATest, LoRAUtils) {
+    auto base_linear = std::make_shared<nn::Linear>(32, 64, /*bias=*/true, GetDevice());
+
+    LoRAConfig config;
+    config.rank = 4;
+    config.alpha = 8.0f;
+    config.target_modules = {"Linear"};
+
+    auto model = GetLoRAModel(base_linear, config);
+
+    auto lora_params = GetLoRAParameters(model);
+    EXPECT_EQ(lora_params.size(), 2);
+
+    int64_t trainable = CountTrainableParameters(model);
+    int64_t expected_trainable = config.rank * 32 + 64 * config.rank;
+    EXPECT_EQ(trainable, expected_trainable);
+
+    int64_t total = CountTotalParameters(model);
+    int64_t expected_total = 64 * 32 + 64 + config.rank * 32 + 64 * config.rank;
+    EXPECT_EQ(total, expected_total);
+}
+
+TEST_P(LoRATest, ParseLoRATargetModules) {
+    auto modules = ParseLoRATargetModules("c_attn");
+    EXPECT_EQ(modules.size(), 1);
+    EXPECT_TRUE(modules.count("c_attn"));
+
+    modules = ParseLoRATargetModules("c_attn,c_proj,c_fc");
+    EXPECT_EQ(modules.size(), 3);
+    EXPECT_TRUE(modules.count("c_attn"));
+    EXPECT_TRUE(modules.count("c_proj"));
+    EXPECT_TRUE(modules.count("c_fc"));
+
+    modules = ParseLoRATargetModules("c_attn, c_proj , c_fc");
+    EXPECT_EQ(modules.size(), 3);
+
+    modules = ParseLoRATargetModules("c_attn,,c_proj");
+    EXPECT_EQ(modules.size(), 2);
+}
+
+TEST_P(LoRATest, ShouldApplyLoRAEdgeCases) {
+    {
+        LoRAConfig config{8, 16.0f, 0.0f, ParseLoRATargetModules("c_attn,attn.c_proj")};
+        EXPECT_TRUE(config.ShouldApplyLoRA("attn.c_proj"));
+        EXPECT_TRUE(config.ShouldApplyLoRA("transformer.h.0.attn.c_proj"));
+        EXPECT_FALSE(config.ShouldApplyLoRA("mlp.c_proj"));
+    }
+
+    {
+        LoRAConfig config{8, 16.0f, 0.0f, ParseLoRATargetModules("c_attn,mlp.c_proj")};
+        EXPECT_FALSE(config.ShouldApplyLoRA("attn.c_proj"));
+        EXPECT_TRUE(config.ShouldApplyLoRA("mlp.c_proj"));
+    }
+
+    {
+        LoRAConfig config{8, 16.0f, 0.0f, ParseLoRATargetModules("c_attn,c_proj")};
+        EXPECT_TRUE(config.ShouldApplyLoRA("transformer.h.0.attn.c_proj"));
+        EXPECT_TRUE(config.ShouldApplyLoRA("transformer.h.0.mlp.c_proj"));
+    }
+}
+
+TEST_P(LoRATest, FreezeUnfreeze) {
+    auto base_linear = std::make_shared<nn::Linear>(64, 128, /*bias=*/true, GetDevice());
+
+    LoRAConfig config;
+    config.rank = 4;
+    config.alpha = 8.0f;
+    config.target_modules = {"Linear"};
+
+    auto model = GetLoRAModel(base_linear, config);
+
+    auto all_params = model->Parameters();
+
+    int64_t total_trainable = 0;
+    for (const auto &p : all_params) {
+        if (p->requires_grad()) {
+            total_trainable += p->NumElements();
+        }
+    }
+    int64_t expected = config.rank * 64 + 128 * config.rank;
+    EXPECT_EQ(total_trainable, expected);
+
+    FreezeBaseModel(model);
+
+    int64_t after_freeze = 0;
+    for (const auto &p : all_params) {
+        if (p->requires_grad()) {
+            after_freeze += p->NumElements();
+        }
+    }
+    EXPECT_EQ(after_freeze, expected);
+
+    UnfreezeModel(model);
+    int64_t after_unfreeze = 0;
+    for (const auto &p : all_params) {
+        if (p->requires_grad()) {
+            after_unfreeze += p->NumElements();
+        }
+    }
+    int64_t expected_unfreeze = 64 * 128 + 128 + config.rank * 64 + 128 * config.rank;
+    EXPECT_EQ(after_unfreeze, expected_unfreeze);
+}
+
+TEST_P(LoRATest, LoRAStateDict) {
+    auto base_linear = std::make_shared<nn::Linear>(64, 128, /*bias=*/true, GetDevice());
+
+    LoRAConfig config;
+    config.rank = 4;
+    config.alpha = 8.0f;
+    config.target_modules = {"Linear"};
+
+    auto model = GetLoRAModel(base_linear, config);
+
+    auto state_dict = model->StateDict();
+
+    EXPECT_TRUE(state_dict.count("weight"));
+    EXPECT_TRUE(state_dict.count("bias"));
+    EXPECT_TRUE(state_dict.count("lora_A"));
+    EXPECT_TRUE(state_dict.count("lora_B"));
+
+    EXPECT_TRUE(state_dict.at("lora_A")->requires_grad());
+    EXPECT_TRUE(state_dict.at("lora_B")->requires_grad());
+    EXPECT_FALSE(state_dict.at("weight")->requires_grad());
+
+    EXPECT_EQ(state_dict.at("lora_A")->Dims()[0], config.rank);
+    EXPECT_EQ(state_dict.at("lora_A")->Dims()[1], 64);
+    EXPECT_EQ(state_dict.at("lora_B")->Dims()[0], 128);
+    EXPECT_EQ(state_dict.at("lora_B")->Dims()[1], config.rank);
+}
+
+TEST_P(LoRATest, GetLoRAModel) {
+    auto base_linear = std::make_shared<nn::Linear>(64, 128, /*bias=*/true, GetDevice());
+
+    LoRAConfig config;
+    config.rank = 4;
+    config.alpha = 8.0f;
+    config.target_modules = {"Linear"};
+
+    auto model = GetLoRAModel(base_linear, config);
+
+    EXPECT_NE(model, nullptr);
+
+    auto lora_params = GetLoRAParameters(model);
+    EXPECT_EQ(lora_params.size(), 2);
+
+    int64_t total_elements = 0;
+    for (const auto &t : lora_params) { total_elements += t->NumElements(); }
+    int64_t expected_elements = config.rank * 64 + 128 * config.rank;
+    EXPECT_EQ(total_elements, expected_elements);
+
+    MergeLoRAWeights(model);
+    auto *lora_mod = dynamic_cast<LoRALinear *>(model.get());
+    EXPECT_NE(lora_mod, nullptr);
+    EXPECT_FALSE(lora_mod->LoRAParameters()[0]->requires_grad());
+
+    UnmergeLoRAWeights(model);
+    EXPECT_TRUE(lora_mod->LoRAParameters()[0]->requires_grad());
+}
+
+TEST_P(LoRATest, MergeAndUnload) {
+    auto base_linear = std::make_shared<nn::Linear>(64, 128, /*bias=*/true, GetDevice());
+    LoRAConfig config;
+    config.rank = 4;
+    config.alpha = 8.0f;
+    config.target_modules = {"Linear"};
+    auto model = GetLoRAModel(base_linear, config);
+
+    EXPECT_NE(dynamic_cast<LoRALinear *>(model.get()), nullptr);
+
+    auto input = std::make_shared<Tensor>(std::vector<int64_t>{2, 5, 64}, DataType::kFLOAT32, GetDevice());
+    input->Fill(1.0f);
+    auto output_before = (*model)({input})[0];
+    float output_before_sum = TensorSum(output_before);
+
+    auto unloaded_model = MergeAndUnload(model);
+    EXPECT_NE(unloaded_model, nullptr);
+    EXPECT_EQ(dynamic_cast<LoRALinear *>(unloaded_model.get()), nullptr);
+
+    auto state_dict = unloaded_model->StateDict();
+    for (const auto &[name, param] : state_dict) {
+        EXPECT_EQ(name.find("lora_A"), std::string::npos);
+        EXPECT_EQ(name.find("lora_B"), std::string::npos);
+    }
+
+    auto output_after = (*unloaded_model)({input})[0];
+    float output_after_sum = TensorSum(output_after);
+    EXPECT_NEAR(std::abs(output_before_sum - output_after_sum), 0.0f, 1e-3);
+
+    for (const auto &param : unloaded_model->Parameters()) { EXPECT_TRUE(param->requires_grad()); }
+}
+
+INFINI_TRAIN_REGISTER_TEST(LoRATest);
diff --git a/tests/optimizer/CMakeLists.txt b/tests/optimizer/CMakeLists.txt
new file mode 100644
index 00000000..c0bfbd50
--- /dev/null
+++ b/tests/optimizer/CMakeLists.txt
@@ -0,0 +1,9 @@
+# ==========================================================================
+# Optimizer tests
+# ==========================================================================
+
+file(GLOB OPTIMIZER_SOURCES ${CMAKE_CURRENT_SOURCE_DIR}/test_*.cc)
+
+infini_train_add_test_suite(test_optimizer
+  SOURCES ${OPTIMIZER_SOURCES}
+)
diff --git a/tests/optimizer/test_optimizer_creation.cc b/tests/optimizer/test_optimizer_creation.cc
new file mode 100644
index 00000000..927b9a46
--- /dev/null
+++ b/tests/optimizer/test_optimizer_creation.cc
@@ -0,0 +1,49 @@
+#include <vector>
+
+#include "infini_train/include/nn/parallel/global.h"
+#include "infini_train/include/optimizer.h"
+#include "infini_train/include/tensor.h"
+#include "tests/common/test_utils.h"
+#include "gtest/gtest.h"
+
+using namespace infini_train;
+
+class OptimizerCreationTest : public infini_train::test::InfiniTrainTest {};
+
+TEST_P(OptimizerCreationTest, SGDCreation) {
+    auto param = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice());
+    param->set_requires_grad(true);
+    auto optimizer = std::make_shared<optimizers::SGD>(std::vector<std::shared_ptr<Tensor>>{param}, 0.01);
+    EXPECT_NE(optimizer, nullptr);
+}
+
+TEST_P(OptimizerCreationTest, AdamCreation) {
+    auto param = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice());
+    param->set_requires_grad(true);
+    auto optimizer = std::make_shared<optimizers::Adam>(std::vector<std::shared_ptr<Tensor>>{param}, 0.001);
+    EXPECT_NE(optimizer, nullptr);
+}
+
+TEST_P(OptimizerCreationTest, SGDMultiParams) {
+    std::vector<std::shared_ptr<Tensor>> params;
+    for (int i = 0; i < 3; ++i) {
+        auto param = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice());
+        param->set_requires_grad(true);
+        params.push_back(param);
+    }
+    auto optimizer = std::make_shared<optimizers::SGD>(params, 0.01);
+    EXPECT_NE(optimizer, nullptr);
+}
+
+TEST_P(OptimizerCreationTest, AdamMultiParams) {
+    std::vector<std::shared_ptr<Tensor>> params;
+    for (int i = 0; i < 3; ++i) {
+        auto param = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice());
+        param->set_requires_grad(true);
+        params.push_back(param);
+    }
+    auto optimizer = std::make_shared<optimizers::Adam>(params, 0.001);
+    EXPECT_NE(optimizer, nullptr);
+}
+
+INFINI_TRAIN_REGISTER_TEST(OptimizerCreationTest);
diff --git a/tests/optimizer/test_optimizer_step.cc b/tests/optimizer/test_optimizer_step.cc
new file mode 100644
index 00000000..d1d08f21
--- /dev/null
+++ b/tests/optimizer/test_optimizer_step.cc
@@ -0,0 +1,57 @@
+#include <vector>
+
+#include "infini_train/include/nn/parallel/global.h"
+#include "infini_train/include/optimizer.h"
+#include "infini_train/include/tensor.h"
+#include "tests/common/test_utils.h"
+#include "gtest/gtest.h"
+
+using namespace infini_train;
+
+class OptimizerStepTest : public infini_train::test::InfiniTrainTest {};
+
+TEST_P(OptimizerStepTest, SGDStep) {
+    auto param = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice());
+    param->set_requires_grad(true);
+    param->Fill(1.0f);
+    auto optimizer = std::make_shared<optimizers::SGD>(std::vector<std::shared_ptr<Tensor>>{param}, 0.01);
+    optimizer->ZeroGrad();
+    optimizer->Step();
+}
+
+TEST_P(OptimizerStepTest, AdamStep) {
+    auto param = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice());
+    param->set_requires_grad(true);
+    param->Fill(1.0f);
+    auto optimizer = std::make_shared<optimizers::Adam>(std::vector<std::shared_ptr<Tensor>>{param}, 0.001);
+    optimizer->ZeroGrad();
+    optimizer->Step();
+}
+
+TEST_P(OptimizerStepTest, ZeroGrad) {
+    auto param = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice());
+    param->set_requires_grad(true);
+    auto optimizer = std::make_shared<optimizers::SGD>(std::vector<std::shared_ptr<Tensor>>{param}, 0.01);
+    optimizer->ZeroGrad();
+}
+
+TEST_P(OptimizerStepTest, ZeroGradWithNone) {
+    auto param = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice());
+    param->set_requires_grad(true);
+    auto optimizer = std::make_shared<optimizers::SGD>(std::vector<std::shared_ptr<Tensor>>{param}, 0.01);
+    optimizer->ZeroGrad(false);
+}
+
+TEST_P(OptimizerStepTest, SGDMultiParams) {
+    std::vector<std::shared_ptr<Tensor>> params;
+    for (int i = 0; i < 3; ++i) {
+        auto param = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice());
+        param->set_requires_grad(true);
+        params.push_back(param);
+    }
+    auto optimizer = std::make_shared<optimizers::SGD>(params, 0.01);
+    EXPECT_NE(optimizer, nullptr);
+    optimizer->ZeroGrad();
+}
+
+INFINI_TRAIN_REGISTER_TEST(OptimizerStepTest);
diff --git a/tests/tensor/CMakeLists.txt b/tests/tensor/CMakeLists.txt
new file mode 100644
index 00000000..840cd8c0
--- /dev/null
+++ b/tests/tensor/CMakeLists.txt
@@ -0,0 +1,9 @@
+# ==========================================================================
+# Tensor tests
+# ==========================================================================
+
+file(GLOB TENSOR_SOURCES ${CMAKE_CURRENT_SOURCE_DIR}/test_*.cc)
+
+infini_train_add_test_suite(test_tensor
+  SOURCES ${TENSOR_SOURCES}
+)
diff --git a/tests/tensor/test_tensor.cc b/tests/tensor/test_tensor.cc
new file mode 100644
index 00000000..6ae3577d
--- /dev/null
+++ b/tests/tensor/test_tensor.cc
@@ -0,0 +1,25 @@
+#include <memory>
+#include <vector>
+
+#include "infini_train/include/tensor.h"
+#include "tests/common/test_utils.h"
+#include "gtest/gtest.h"
+
+using namespace infini_train;
+
+// ============================================================================
+// Op tests — CPU + CUDA
+// ============================================================================
+
+class TensorOpTest : public infini_train::test::InfiniTrainTest {};
+
+TEST_P(TensorOpTest, MatmulAllocatesOutputs) {
+    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice());
+    auto b = std::make_shared<Tensor>(std::vector<int64_t>{3, 4}, DataType::kFLOAT32, GetDevice());
+    auto c = std::make_shared<Tensor>(std::vector<int64_t>{2, 4}, DataType::kFLOAT32, GetDevice());
+    EXPECT_NE(a->DataPtr(), nullptr);
+    EXPECT_NE(b->DataPtr(), nullptr);
+    EXPECT_NE(c->DataPtr(), nullptr);
+}
+
+INFINI_TRAIN_REGISTER_TEST(TensorOpTest);
diff --git a/tests/tensor/test_tensor_copy.cc b/tests/tensor/test_tensor_copy.cc
new file mode 100644
index 00000000..9af973f8
--- /dev/null
+++ b/tests/tensor/test_tensor_copy.cc
@@ -0,0 +1,86 @@
+#include <memory>
+#include <vector>
+
+#include "infini_train/include/tensor.h"
+#include "tests/common/test_utils.h"
+#include "gtest/gtest.h"
+
+using namespace infini_train;
+
+class TensorCopyTest : public infini_train::test::InfiniTrainTest {};
+
+TEST_P(TensorCopyTest, CopiesCPUToCPU) {
+    ONLY_CPU();
+    auto source
+        = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, Device(Device::DeviceType::kCPU, 0));
+    auto target
+        = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, Device(Device::DeviceType::kCPU, 0));
+    source->Fill(1.0f);
+    target->CopyFrom(source);
+    auto *target_data = static_cast<float *>(target->DataPtr());
+    for (int i = 0; i < 6; ++i) { EXPECT_FLOAT_EQ(target_data[i], 1.0f); }
+}
+
+TEST_P(TensorCopyTest, CopiesBetweenSameShape) {
+    auto source = std::make_shared<Tensor>(std::vector<int64_t>{4, 5, 6}, DataType::kFLOAT32, GetDevice());
+    auto target = std::make_shared<Tensor>(std::vector<int64_t>{4, 5, 6}, DataType::kFLOAT32, GetDevice());
+    source->Fill(0.0f);
+    target->CopyFrom(source);
+    EXPECT_EQ(source->Dims(), target->Dims());
+}
+
+TEST_P(TensorCopyTest, CopiesPreservesDataType) {
+    auto source = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice());
+    auto target = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice());
+    EXPECT_EQ(source->Dtype(), target->Dtype());
+    target->CopyFrom(source);
+    EXPECT_EQ(target->Dtype(), DataType::kFLOAT32);
+}
+
+TEST_P(TensorCopyTest, CopiesCPUToCUDA) {
+    ONLY_CUDA();
+    auto cpu_tensor
+        = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, Device(Device::DeviceType::kCPU, 0));
+    auto cuda_tensor = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32,
+                                                Device(Device::DeviceType::kCUDA, 0));
+    cpu_tensor->Fill(0.0f);
+    cuda_tensor->CopyFrom(cpu_tensor);
+    EXPECT_TRUE(cuda_tensor->GetDevice().IsCUDA());
+}
+
+TEST_P(TensorCopyTest, CopiesCUDAtoCPU) {
+    ONLY_CUDA();
+    auto cuda_tensor = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32,
+                                                Device(Device::DeviceType::kCUDA, 0));
+    auto cpu_tensor
+        = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, Device(Device::DeviceType::kCPU, 0));
+    cuda_tensor->Fill(1.0f);
+    cpu_tensor->CopyFrom(cuda_tensor);
+    EXPECT_TRUE(cpu_tensor->GetDevice().IsCPU());
+}
+
+TEST_P(TensorCopyTest, CopiesCUDAtoCUDA) {
+    ONLY_CUDA();
+    auto source = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32,
+                                           Device(Device::DeviceType::kCUDA, 0));
+    auto target = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32,
+                                           Device(Device::DeviceType::kCUDA, 0));
+    source->Fill(2.0f);
+    target->CopyFrom(source);
+    EXPECT_TRUE(target->GetDevice().IsCUDA());
+}
+
+TEST_P(TensorCopyTest, CopiesWithDifferentDeviceId) {
+    ONLY_CUDA();
+    REQUIRE_MIN_DEVICES(2);
+    auto source = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32,
+                                           Device(Device::DeviceType::kCUDA, 0));
+    auto target = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32,
+                                           Device(Device::DeviceType::kCUDA, 1));
+    source->Fill(5.0f);
+    target->CopyFrom(source);
+    EXPECT_EQ(source->GetDevice().index(), 0);
+    EXPECT_EQ(target->GetDevice().index(), 1);
+}
+
+INFINI_TRAIN_REGISTER_TEST(TensorCopyTest);
diff --git a/tests/tensor/test_tensor_create.cc b/tests/tensor/test_tensor_create.cc
new file mode 100644
index 00000000..aaafd90c
--- /dev/null
+++ b/tests/tensor/test_tensor_create.cc
@@ -0,0 +1,45 @@
+#include <memory>
+#include <vector>
+
+#include "infini_train/include/tensor.h"
+#include "tests/common/test_utils.h"
+#include "gtest/gtest.h"
+
+using namespace infini_train;
+
+class TensorCreateTest : public infini_train::test::InfiniTrainTest {};
+
+TEST_P(TensorCreateTest, CreatesTensorWithShapeAndType) {
+    auto tensor = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice());
+    EXPECT_NE(tensor, nullptr);
+    EXPECT_EQ(tensor->Dims(), (std::vector<int64_t>{2, 3}));
+    EXPECT_EQ(tensor->Dtype(), DataType::kFLOAT32);
+}
+
+TEST_P(TensorCreateTest, TracksRequiresGrad) {
+    auto tensor = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice());
+    EXPECT_FALSE(tensor->requires_grad());
+    tensor->set_requires_grad(true);
+    EXPECT_TRUE(tensor->requires_grad());
+}
+
+TEST_P(TensorCreateTest, ProvidesDataPointer) {
+    auto tensor = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice());
+    EXPECT_NE(tensor->DataPtr(), nullptr);
+}
+
+TEST_P(TensorCreateTest, SupportsMultipleShapes) {
+    for (const auto &shape : std::vector<std::vector<int64_t>>{{2, 3}, {4, 5, 6}, {10}, {1, 1, 1, 1}}) {
+        auto tensor = std::make_shared<Tensor>(shape, DataType::kFLOAT32, GetDevice());
+        EXPECT_EQ(tensor->Dims(), shape);
+    }
+}
+
+TEST_P(TensorCreateTest, SupportsMultipleDtypes) {
+    for (auto dtype : {DataType::kFLOAT32, DataType::kBFLOAT16}) {
+        auto tensor = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, dtype, GetDevice());
+        EXPECT_EQ(tensor->Dtype(), dtype);
+    }
+}
+
+INFINI_TRAIN_REGISTER_TEST(TensorCreateTest);
diff --git a/tests/tensor/test_tensor_delete.cc b/tests/tensor/test_tensor_delete.cc
new file mode 100644
index 00000000..8caedeb8
--- /dev/null
+++ b/tests/tensor/test_tensor_delete.cc
@@ -0,0 +1,66 @@
+#include <memory>
+#include <vector>
+
+#include "infini_train/include/tensor.h"
+#include "tests/common/test_utils.h"
+#include "gtest/gtest.h"
+
+using namespace infini_train;
+
+class TensorDeleteTest : public infini_train::test::InfiniTrainTest {};
+
+TEST_P(TensorDeleteTest, ReleasesResourcesOnReset) {
+    std::weak_ptr<Tensor> weak_tensor;
+    {
+        auto tensor = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(),
+                                               /*requires_grad=*/true);
+        weak_tensor = tensor;
+    }
+    EXPECT_TRUE(weak_tensor.expired());
+}
+
+TEST_P(TensorDeleteTest, MoveTransferKeepsData) {
+    auto tensor = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice());
+    tensor->Fill(5.0f);
+
+    // Verify data via CPU copy for any device.
+    auto cpu_copy
+        = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, Device(Device::DeviceType::kCPU, 0));
+    cpu_copy->CopyFrom(tensor);
+
+    auto moved = std::move(tensor);
+    EXPECT_EQ(tensor, nullptr);
+    ASSERT_NE(moved, nullptr);
+
+    auto *data = static_cast<float *>(cpu_copy->DataPtr());
+    for (int i = 0; i < 6; ++i) { EXPECT_FLOAT_EQ(data[i], 5.0f); }
+}
+
+TEST_P(TensorDeleteTest, NullifiesPointerOnMove) {
+    auto tensor = std::make_shared<Tensor>(std::vector<int64_t>{3, 3}, DataType::kFLOAT32, GetDevice());
+    EXPECT_NE(tensor, nullptr);
+    auto moved_tensor = std::move(tensor);
+    EXPECT_EQ(tensor, nullptr);
+    EXPECT_NE(moved_tensor, nullptr);
+}
+
+TEST_P(TensorDeleteTest, SharedPtrRefCountOnCopy) {
+    auto tensor = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice());
+    auto copy1 = tensor;
+    auto copy2 = tensor;
+    EXPECT_EQ(tensor.use_count(), 3);
+    copy1.reset();
+    EXPECT_EQ(tensor.use_count(), 2);
+    copy2.reset();
+    EXPECT_EQ(tensor.use_count(), 1);
+    EXPECT_NE(tensor, nullptr);
+}
+
+TEST_P(TensorDeleteTest, TensorDestroyedAfterScope) {
+    {
+        auto tensor = std::make_shared<Tensor>(std::vector<int64_t>{2, 2}, DataType::kFLOAT32, GetDevice());
+        EXPECT_NE(tensor, nullptr);
+    }
+}
+
+INFINI_TRAIN_REGISTER_TEST(TensorDeleteTest);
diff --git a/tests/transformer/CMakeLists.txt b/tests/transformer/CMakeLists.txt
new file mode 100644
index 00000000..555de735
--- /dev/null
+++ b/tests/transformer/CMakeLists.txt
@@ -0,0 +1,9 @@
+# ==========================================================================
+# Transformer architecture tests
+# ==========================================================================
+
+file(GLOB TRANSFORMER_SOURCES ${CMAKE_CURRENT_SOURCE_DIR}/test_*.cc)
+
+infini_train_add_test_suite(test_transformer
+  SOURCES ${TRANSFORMER_SOURCES}
+)
diff --git a/tests/transformer/test_transformer_architecture.cc b/tests/transformer/test_transformer_architecture.cc
new file mode 100644
index 00000000..b2fbcb46
--- /dev/null
+++ b/tests/transformer/test_transformer_architecture.cc
@@ -0,0 +1,190 @@
+#include <cmath>
+#include <memory>
+#include <vector>
+
+#include "infini_train/include/nn/modules/normalization.h"
+#include "infini_train/include/nn/modules/sparse.h"
+#include "infini_train/include/nn/modules/transformer/causal_self_attention.h"
+#include "infini_train/include/nn/modules/transformer/mlp.h"
+#include "infini_train/include/nn/modules/transformer/transformer.h"
+#include "infini_train/include/nn/modules/transformer/transformer_config.h"
+#include "infini_train/include/nn/modules/transformer/utils.h"
+#include "infini_train/include/tensor.h"
+#include "tests/common/test_utils.h"
+#include "gtest/gtest.h"
+
+using namespace infini_train;
+namespace nn = infini_train::nn;
+
+class TransformerModuleTest : public infini_train::test::InfiniTrainTest {};
+
+TEST_P(TransformerModuleTest, Embedding) {
+    SKIP_CPU();
+    auto embedding = std::make_shared<nn::Embedding>(1000, 128, GetDevice());
+    EXPECT_EQ(embedding->Parameters().size(), 1);
+
+    auto weight = embedding->parameter(nn::Embedding::kParamWeightName);
+    EXPECT_EQ(weight->Dims(), (std::vector<int64_t>{1000, 128}));
+
+    auto input = std::make_shared<Tensor>(std::vector<int64_t>{2, 16}, DataType::kINT64, GetDevice());
+    auto output = (*embedding)({input});
+    EXPECT_EQ(output.size(), 1);
+    EXPECT_EQ(output[0]->Dims(), (std::vector<int64_t>{2, 16, 128}));
+}
+
+TEST_P(TransformerModuleTest, LayerNorm) {
+    SKIP_CPU();
+    auto layernorm = std::make_shared<nn::LayerNorm>(std::vector<int64_t>{64});
+    layernorm->To(GetDevice());
+    EXPECT_EQ(layernorm->Parameters().size(), 2);
+
+    auto input = std::make_shared<Tensor>(std::vector<int64_t>{2, 8, 64}, DataType::kFLOAT32, GetDevice());
+    auto output = (*layernorm)({input});
+    EXPECT_EQ(output[0]->Dims(), input->Dims());
+}
+
+TEST_P(TransformerModuleTest, RMSNorm) {
+    SKIP_CPU();
+    auto rmsnorm = std::make_shared<nn::RMSNorm>(64);
+    rmsnorm->To(GetDevice());
+    EXPECT_EQ(rmsnorm->Parameters().size(), 1);
+
+    auto input = std::make_shared<Tensor>(std::vector<int64_t>{2, 8, 64}, DataType::kFLOAT32, GetDevice());
+    auto output = (*rmsnorm)({input});
+    EXPECT_EQ(output[0]->Dims(), input->Dims());
+}
+
+TEST_P(TransformerModuleTest, GPT2MLP) {
+    SKIP_CPU();
+    nn::TransformerConfig config;
+    config.n_embd = 64;
+    config.activation_type = nn::MLPType::kGELU;
+    config.ffn_expansion_ratio = 4.0f;
+    config.add_bias_linear = true;
+
+    auto mlp = std::make_shared<nn::MLP>(config);
+    mlp->To(GetDevice());
+    EXPECT_EQ(mlp->Parameters().size(), 4);
+
+    auto input = std::make_shared<Tensor>(std::vector<int64_t>{2, 8, 64}, DataType::kFLOAT32, GetDevice());
+    auto output = (*mlp)({input});
+    EXPECT_EQ(output[0]->Dims()[2], 64);
+}
+
+TEST_P(TransformerModuleTest, SwiGLUMLP) {
+    SKIP_CPU();
+    nn::TransformerConfig config;
+    config.n_embd = 64;
+    config.activation_type = nn::MLPType::kSwiGLU;
+    config.ffn_expansion_ratio = 4.0f;
+    config.add_bias_linear = false;
+    config.ffn_dim_multiplier = 1.5f;
+    config.multiple_of = 256;
+
+    auto mlp = std::make_shared<nn::MLP>(config);
+    mlp->To(GetDevice());
+    EXPECT_EQ(mlp->Parameters().size(), 3);
+
+    auto input = std::make_shared<Tensor>(std::vector<int64_t>{2, 8, 64}, DataType::kFLOAT32, GetDevice());
+    auto output = (*mlp)({input});
+    EXPECT_EQ(output[0]->Dims()[2], 64);
+}
+
+TEST_P(TransformerModuleTest, StandardAttention) {
+    SKIP_CPU();
+    nn::TransformerConfig config;
+    config.n_embd = 64;
+    config.n_head = 4;
+    config.n_kv_head = 4;
+    config.attention_type = nn::AttentionType::kStandard;
+    config.add_bias_linear = true;
+
+    auto attn = std::make_shared<nn::CausalSelfAttention>(config);
+    attn->To(GetDevice());
+    EXPECT_EQ(attn->Parameters().size(), 4);
+
+    auto input = std::make_shared<Tensor>(std::vector<int64_t>{2, 8, 64}, DataType::kFLOAT32, GetDevice());
+    auto output = (*attn)({input});
+    EXPECT_EQ(output[0]->Dims(), input->Dims());
+}
+
+TEST_P(TransformerModuleTest, GPT2TransformerLayer) {
+    SKIP_CPU();
+    nn::TransformerConfig config;
+    config.n_embd = 64;
+    config.n_head = 4;
+    config.n_kv_head = 4;
+    config.n_layer = 1;
+
+    auto layer = std::make_shared<nn::TransformerLayer>(config);
+    layer->To(GetDevice());
+    EXPECT_FALSE(layer->Parameters().empty());
+
+    auto input = std::make_shared<Tensor>(std::vector<int64_t>{2, 8, 64}, DataType::kFLOAT32, GetDevice());
+    auto output = (*layer)({input});
+    EXPECT_EQ(output[0]->Dims(), input->Dims());
+}
+
+TEST_P(TransformerModuleTest, GPT2Model) {
+    SKIP_CPU();
+    nn::TransformerConfig config;
+    config.n_layer = 2;
+    config.n_head = 4;
+    config.n_kv_head = 4;
+    config.n_embd = 64;
+
+    auto model = std::make_shared<nn::TransformerModel>(config);
+    model->To(GetDevice());
+    EXPECT_FALSE(model->Parameters().empty());
+}
+
+TEST_P(TransformerModuleTest, LLaMA3Model) {
+    SKIP_CPU();
+    nn::TransformerConfig config;
+    config.n_layer = 2;
+    config.n_head = 4;
+    config.n_kv_head = 2;
+    config.n_embd = 64;
+    config.attention_type = nn::AttentionType::kRoPE;
+    config.activation_type = nn::MLPType::kSwiGLU;
+    config.norm_type = nn::NormType::kRMSNorm;
+    config.add_bias_linear = false;
+    config.tie_weights = false;
+
+    auto model = std::make_shared<nn::TransformerModel>(config);
+    model->To(GetDevice());
+    EXPECT_FALSE(model->Parameters().empty());
+}
+
+TEST_P(TransformerModuleTest, RoPEUtils) {
+    auto freqs_cis = PrecomputeFreqsCis(64, 128);
+    EXPECT_EQ(freqs_cis->Dims().size(), 3);
+    EXPECT_EQ(freqs_cis->Dims()[0], 128);
+    EXPECT_EQ(freqs_cis->Dims()[1], 32);
+    EXPECT_EQ(freqs_cis->Dims()[2], 2);
+}
+
+TEST_P(TransformerModuleTest, StateDict) {
+    SKIP_CPU();
+    nn::TransformerConfig config;
+    config.n_layer = 1;
+    config.n_head = 2;
+    config.n_kv_head = 2;
+    config.n_embd = 32;
+    config.vocab_size = 1000;
+    config.attention_type = nn::AttentionType::kStandard;
+    config.activation_type = nn::MLPType::kGELU;
+    config.norm_type = nn::NormType::kLayerNorm;
+    config.add_bias_linear = true;
+
+    auto model = std::make_shared<nn::TransformerModel>(config);
+    model->To(GetDevice());
+    auto state_dict = model->StateDict();
+    auto params = model->Parameters();
+    auto buffers = model->Buffers();
+
+    EXPECT_FALSE(state_dict.empty());
+    EXPECT_GE(state_dict.size(), params.size());
+}
+
+INFINI_TRAIN_REGISTER_TEST(TransformerModuleTest);