diff --git a/docs/lora_usage_guide.md b/docs/lora_usage_guide.md
index f7512b4c..2a3470db 100644
--- a/docs/lora_usage_guide.md
+++ b/docs/lora_usage_guide.md
@@ -30,7 +30,7 @@ model = GetLoRAModel(model, config);
 PrintLoRASummary(model);
 
 auto params = GetLoRAParameters(model);
-auto optimizer = infini_train::optimizers::Adam::Create(/*learning_rate=*/1e-4)(params);
+auto optimizer = infini_train::optimizers::Adam::Create(/*learning_rate=*/1e-4)(params, model->NamedParameters());
 
 for (int step = 0; step < num_steps; ++step) {
     optimizer->ZeroGrad();
diff --git a/example/gpt2/main.cc b/example/gpt2/main.cc
index 068d1e7c..5a5cfc65 100644
--- a/example/gpt2/main.cc
+++ b/example/gpt2/main.cc
@@ -327,17 +327,28 @@ void Train(const nn::parallel::Rank &rank) {
 
     // TODO(dcj): support more complex optimizer later
     // auto optimizer = optimizers::SGD(model->Parameters(), FLAGS_learning_rate);
-    auto optimizer_creator = optimizers::SGD::Create(FLAGS_learning_rate);
+    auto optimizer_creator = optimizers::SGD::CreateNamed(FLAGS_learning_rate);
     std::shared_ptr<Optimizer> optimizer = nullptr;
+    std::unordered_set<const Tensor *> params_to_optimize_set;
+    params_to_optimize_set.reserve(params_to_optimize.size());
+    for (const auto &param : params_to_optimize) { params_to_optimize_set.insert(param.get()); }
+
+    NamedParameterList named_parameters;
+    for (const auto &[name, param] : model->NamedParameters()) {
+        if (params_to_optimize_set.contains(param.get())) {
+            named_parameters.emplace_back(name, param);
+        }
+    }
+    CHECK_EQ(named_parameters.size(), params_to_optimize.size());
 
     if (FLAGS_zero_stage >= 1) {
         auto model_chunks = (pp_world_size > 1)
                               ? *(dynamic_cast<nn::parallel::PipelineParallel *>(model.get())->mutable_chunks())
                               : std::vector<std::shared_ptr<nn::Module>>{model};
-        optimizer = std::make_shared<nn::parallel::DistributedOptimizer>(optimizer_creator, params_to_optimize,
+        optimizer = std::make_shared<nn::parallel::DistributedOptimizer>(optimizer_creator, named_parameters,
                                                                          model_chunks, ddp_world_size, ddp_rank);
     } else {
-        optimizer = optimizer_creator(params_to_optimize);
+        optimizer = optimizer_creator(named_parameters);
     }
 
     const int64_t lr_decay_iters = FLAGS_lr_decay_iters > 0 ? FLAGS_lr_decay_iters : FLAGS_num_iteration;
diff --git a/example/llama3/main.cc b/example/llama3/main.cc
index 9ce5a7e7..c4642cc2 100644
--- a/example/llama3/main.cc
+++ b/example/llama3/main.cc
@@ -300,7 +300,7 @@ void Train(const nn::parallel::Rank &rank) {
 
     // TODO(dcj): support more complex optimizer later
     // auto optimizer = optimizers::Adam(model->Parameters(), FLAGS_learning_rate);
-    auto optimizer_creator = optimizers::Adam::Create(FLAGS_learning_rate);
+    auto optimizer_creator = optimizers::Adam::CreateNamed(FLAGS_learning_rate);
     std::shared_ptr<Optimizer> optimizer = nullptr;
 
     std::vector<std::shared_ptr<Tensor>> params_to_optimize;
@@ -311,15 +311,26 @@ void Train(const nn::parallel::Rank &rank) {
         params_to_optimize = model->Parameters();
         LOG(INFO) << "Optimizing " << params_to_optimize.size() << " model parameters";
     }
+    std::unordered_set<const Tensor *> params_to_optimize_set;
+    params_to_optimize_set.reserve(params_to_optimize.size());
+    for (const auto &param : params_to_optimize) { params_to_optimize_set.insert(param.get()); }
+
+    NamedParameterList named_parameters;
+    for (const auto &[name, param] : model->NamedParameters()) {
+        if (params_to_optimize_set.contains(param.get())) {
+            named_parameters.emplace_back(name, param);
+        }
+    }
+    CHECK_EQ(named_parameters.size(), params_to_optimize.size());
 
     if (FLAGS_zero_stage >= 1) {
         auto model_chunks = (pp_world_size > 1)
                               ? *(dynamic_cast<nn::parallel::PipelineParallel *>(model.get())->mutable_chunks())
                               : std::vector<std::shared_ptr<nn::Module>>{model};
-        optimizer = std::make_shared<nn::parallel::DistributedOptimizer>(optimizer_creator, params_to_optimize,
+        optimizer = std::make_shared<nn::parallel::DistributedOptimizer>(optimizer_creator, named_parameters,
                                                                          model_chunks, ddp_world_size, ddp_rank);
     } else {
-        optimizer = optimizer_creator(params_to_optimize);
+        optimizer = optimizer_creator(named_parameters);
     }
 
     const int64_t lr_decay_iters = FLAGS_lr_decay_iters > 0 ? FLAGS_lr_decay_iters : FLAGS_num_iteration;
diff --git a/example/mixtral/main.cc b/example/mixtral/main.cc
index ceabfad9..553d7a62 100644
--- a/example/mixtral/main.cc
+++ b/example/mixtral/main.cc
@@ -104,8 +104,8 @@ int main(int argc, char *argv[]) {
     }
 
     auto loss_fn = std::make_shared<infini_train::nn::CrossEntropyLoss>();
-    auto optimizer
-        = infini_train::optimizers::Adam::Create(static_cast<float>(FLAGS_learning_rate))(model->Parameters());
+    auto optimizer = infini_train::optimizers::Adam::CreateNamed(static_cast<float>(FLAGS_learning_rate))(
+        model->NamedParameters());
 
     auto device_impl = infini_train::core::GetDeviceGuardImpl(train_device.type());
     std::vector<double> step_duration_ms;
diff --git a/infini_train/include/nn/modules/module.h b/infini_train/include/nn/modules/module.h
index 2b8cd6dc..1d42b2ac 100644
--- a/infini_train/include/nn/modules/module.h
+++ b/infini_train/include/nn/modules/module.h
@@ -47,8 +47,12 @@ class Module : public std::enable_shared_from_this<Module> {
 
     const std::string &type() const;
 
-    // TODO: Change return type to filterable iterator (like PyTorch's named_parameters with prefix matching)
     virtual std::vector<std::shared_ptr<Tensor>> Parameters() const;
+
+    // InfiniTrain's NamedParameters returns results ordered by full parameter name.
+    // TODO: Align with PyTorch's ordering in the future.
+    std::vector<std::pair<std::string, std::shared_ptr<Tensor>>>
+    NamedParameters(const std::string &prefix = "", bool recurse = true, bool remove_duplicate = true) const;
     bool has_parameter(const std::string &name) const;
     std::shared_ptr<Tensor> *mutable_parameter(const std::string &name);
     const std::shared_ptr<Tensor> &parameter(const std::string &name) const;
diff --git a/infini_train/include/nn/parallel/ddp/distributed_optimizer.h b/infini_train/include/nn/parallel/ddp/distributed_optimizer.h
index d694ab2a..fa9592ae 100644
--- a/infini_train/include/nn/parallel/ddp/distributed_optimizer.h
+++ b/infini_train/include/nn/parallel/ddp/distributed_optimizer.h
@@ -1,6 +1,7 @@
 #pragma once
 
 #include <cstdint>
+#include <functional>
 #include <memory>
 #include <unordered_map>
 #include <vector>
@@ -24,6 +25,10 @@ class DistributedOptimizer final : public infini_train::Optimizer {
                          const std::vector<std::shared_ptr<Module>> &model_chunks, size_t ddp_world_size,
                          size_t ddp_rank);
 
+    DistributedOptimizer(OptimizerCreatorNamed base_optimizer_creator, const NamedParameterList &named_parameters,
+                         const std::vector<std::shared_ptr<Module>> &model_chunks, size_t ddp_world_size,
+                         size_t ddp_rank);
+
     void Step() override;
 
     void ZeroGrad(bool set_to_none = true) override;
@@ -42,7 +47,10 @@ class DistributedOptimizer final : public infini_train::Optimizer {
     virtual float learning_rate() const override;
 
 private:
-    void BuildShardParamsAndBindGrads();
+    using AddShard = std::function<void(const std::shared_ptr<Tensor> &, const std::shared_ptr<Tensor> &)>;
+
+    void InitializeModelChunks(const std::vector<std::shared_ptr<Module>> &model_chunks);
+    void BuildShardParamsAndBindGrads(const AddShard &add_shard);
 
 private:
     // Inherit from DDP model
@@ -53,9 +61,6 @@ class DistributedOptimizer final : public infini_train::Optimizer {
     size_t ddp_world_size_;
     size_t ddp_rank_;
 
-    // shard params
-    std::vector<std::shared_ptr<Tensor>> shard_params_;
-
     // Base optimizer (SGD, Adam and etc.)
     std::shared_ptr<Optimizer> base_optimizer_;
 };
diff --git a/infini_train/include/optimizer.h b/infini_train/include/optimizer.h
index 0a33f0b8..d85b1ace 100644
--- a/infini_train/include/optimizer.h
+++ b/infini_train/include/optimizer.h
@@ -5,6 +5,7 @@
 #include <memory>
 #include <string>
 #include <unordered_map>
+#include <utility>
 #include <vector>
 
 namespace infini_train {
@@ -13,11 +14,16 @@ class Tensor;
 namespace infini_train {
 class Optimizer;
 
+using NamedParameter = std::pair<std::string, std::shared_ptr<Tensor>>;
+using NamedParameterList = std::vector<NamedParameter>;
 using OptimizerCreator = std::function<std::shared_ptr<Optimizer>(const std::vector<std::shared_ptr<Tensor>> &params)>;
+using OptimizerCreatorNamed = std::function<std::shared_ptr<Optimizer>(const NamedParameterList &named_params)>;
 
 class Optimizer {
 public:
-    explicit Optimizer(const std::vector<std::shared_ptr<Tensor>> &params, float learning_rate = 0.0f);
+    explicit Optimizer(const std::vector<std::shared_ptr<Tensor>> &params, float learning_rate);
+
+    Optimizer(const NamedParameterList &named_params, float learning_rate);
 
     virtual void ZeroGrad(bool set_to_none = true);
 
@@ -39,6 +45,7 @@ class Optimizer {
 
 protected:
     std::vector<std::shared_ptr<Tensor>> params_;
+    std::vector<std::string> parameter_names_;
     float learning_rate_ = 0.0f;
     float initial_learning_rate_ = 0.0f;
     bool initial_lr_set_ = false;
@@ -48,16 +55,20 @@ namespace optimizers {
 class SGD : public Optimizer {
 public:
     SGD(const std::vector<std::shared_ptr<Tensor>> &params, float learning_rate);
+    SGD(const NamedParameterList &named_params, float learning_rate);
 
     void Step() override;
 
     static OptimizerCreator Create(float learning_rate);
+    static OptimizerCreatorNamed CreateNamed(float learning_rate);
 };
 
 class Adam : public Optimizer {
 public:
     Adam(const std::vector<std::shared_ptr<Tensor>> &params, float learning_rate = 1e-3, float beta1 = 0.9,
          float beta2 = 0.999, float eps = 1e-8);
+    Adam(const NamedParameterList &named_params, float learning_rate = 1e-3, float beta1 = 0.9, float beta2 = 0.999,
+         float eps = 1e-8);
 
     void Step() override;
 
@@ -66,6 +77,8 @@ class Adam : public Optimizer {
     void LoadStateDict(const std::unordered_map<std::string, std::shared_ptr<Tensor>> &state_dict) override;
     static OptimizerCreator Create(float learning_rate = 1e-3, float beta1 = 0.9, float beta2 = 0.999,
                                    float eps = 1e-8);
+    static OptimizerCreatorNamed CreateNamed(float learning_rate = 1e-3, float beta1 = 0.9, float beta2 = 0.999,
+                                             float eps = 1e-8);
 
 private:
     int64_t t_;
diff --git a/infini_train/src/nn/modules/module.cc b/infini_train/src/nn/modules/module.cc
index 81068fe8..9475d49f 100644
--- a/infini_train/src/nn/modules/module.cc
+++ b/infini_train/src/nn/modules/module.cc
@@ -28,24 +28,55 @@ Module::Module(const std::string &type) : type_(type), device_(Device()) {}
 const std::string &Module::type() const { return type_; }
 
 std::vector<std::shared_ptr<Tensor>> Module::Parameters() const {
+    const auto &named_parameters = NamedParameters();
+
     std::vector<std::shared_ptr<Tensor>> params;
-    std::unordered_set<const Tensor *> visited;
+    params.reserve(named_parameters.size());
+
+    for (const auto &[_, param] : named_parameters) { params.emplace_back(param); }
+
+    return params;
+}
 
-    auto AddIfUnvisited = [&](const std::shared_ptr<Tensor> &param) {
-        if (visited.insert(param.get()).second) {
-            params.push_back(param);
+std::vector<std::pair<std::string, std::shared_ptr<Tensor>>>
+Module::NamedParameters(const std::string &prefix, bool recurse, bool remove_duplicate) const {
+    std::vector<std::pair<std::string, std::shared_ptr<Tensor>>> named_parameters;
+    std::unordered_set<const Tensor *> visited_parameters;
+
+    std::vector<std::pair<std::string, std::shared_ptr<Module>>> named_modules;
+
+    if (recurse) {
+        named_modules = const_cast<Module *>(this)->NamedModules(
+            /*memory=*/nullptr, prefix, remove_duplicate);
+    } else {
+        named_modules.emplace_back(prefix, std::const_pointer_cast<Module>(shared_from_this()));
+    }
+
+    for (const auto &[module_prefix, module] : named_modules) {
+        std::vector<std::pair<std::string, std::shared_ptr<Tensor>>> local_parameters;
+        local_parameters.reserve(module->parameters_.size());
+
+        for (const auto &[name, parameter] : module->parameters_) {
+            if (parameter != nullptr) {
+                local_parameters.emplace_back(name, parameter);
+            }
         }
-    };
 
-    // Add parameters of this module
-    for (const auto &[_, param] : parameters_) { AddIfUnvisited(param); }
+        std::sort(local_parameters.begin(), local_parameters.end(),
+                  [](const auto &lhs, const auto &rhs) { return lhs.first < rhs.first; });
+
+        for (const auto &[name, parameter] : local_parameters) {
+            if (remove_duplicate && !visited_parameters.insert(parameter.get()).second) {
+                continue;
+            }
+
+            const std::string full_name = module_prefix.empty() ? name : module_prefix + "." + name;
 
-    // Recursively add parameters of submodules
-    for (const auto &[_, module] : modules_) {
-        for (const auto &param : module->Parameters()) { AddIfUnvisited(param); }
+            named_parameters.emplace_back(full_name, parameter);
+        }
     }
 
-    return params;
+    return named_parameters;
 }
 
 bool Module::has_parameter(const std::string &name) const { return parameters_.find(name) != parameters_.end(); }
diff --git a/infini_train/src/nn/parallel/ddp/distributed_optimizer.cc b/infini_train/src/nn/parallel/ddp/distributed_optimizer.cc
index 022a4758..02d0425b 100644
--- a/infini_train/src/nn/parallel/ddp/distributed_optimizer.cc
+++ b/infini_train/src/nn/parallel/ddp/distributed_optimizer.cc
@@ -10,8 +10,47 @@ DistributedOptimizer::DistributedOptimizer(OptimizerCreator creator,
                                            const std::vector<std::shared_ptr<Tensor>> &full_params,
                                            const std::vector<std::shared_ptr<Module>> &model_chunks,
                                            size_t ddp_world_size, size_t ddp_rank)
-    : Optimizer(full_params), ddp_world_size_(ddp_world_size), ddp_rank_(ddp_rank) {
+    : Optimizer(full_params, /*learning_rate=*/0.0f), ddp_world_size_(ddp_world_size), ddp_rank_(ddp_rank) {
+    InitializeModelChunks(model_chunks);
 
+    std::vector<std::shared_ptr<Tensor>> shard_params;
+    BuildShardParamsAndBindGrads(
+        [&shard_params](const std::shared_ptr<Tensor> &, const std::shared_ptr<Tensor> &param_piece) {
+            shard_params.push_back(param_piece);
+        });
+
+    base_optimizer_ = creator(shard_params);
+    CHECK(base_optimizer_) << "DistributedOptimizer: failed to create base optimizer.";
+}
+
+DistributedOptimizer::DistributedOptimizer(OptimizerCreatorNamed creator, const NamedParameterList &named_parameters,
+                                           const std::vector<std::shared_ptr<Module>> &model_chunks,
+                                           size_t ddp_world_size, size_t ddp_rank)
+    : Optimizer(named_parameters, /*learning_rate=*/0.0f), ddp_world_size_(ddp_world_size), ddp_rank_(ddp_rank) {
+    InitializeModelChunks(model_chunks);
+
+    std::unordered_map<const Tensor *, std::string> parameter_name_by_tensor;
+    parameter_name_by_tensor.reserve(named_parameters.size());
+    for (const auto &[name, parameter] : named_parameters) {
+        CHECK(parameter);
+        parameter_name_by_tensor.emplace(parameter.get(), name);
+    }
+
+    NamedParameterList shard_named_parameters;
+    BuildShardParamsAndBindGrads(
+        [&parameter_name_by_tensor, &shard_named_parameters](const std::shared_ptr<Tensor> &parameter,
+                                                             const std::shared_ptr<Tensor> &param_piece) {
+            const auto name_it = parameter_name_by_tensor.find(parameter.get());
+            CHECK(name_it != parameter_name_by_tensor.end())
+                << "DistributedOptimizer parameter is not registered in the model";
+            shard_named_parameters.emplace_back(name_it->second, param_piece);
+        });
+
+    base_optimizer_ = creator(shard_named_parameters);
+    CHECK(base_optimizer_) << "DistributedOptimizer: failed to create base optimizer.";
+}
+
+void DistributedOptimizer::InitializeModelChunks(const std::vector<std::shared_ptr<Module>> &model_chunks) {
     CHECK(ddp_world_size_ > 1) << "DistributedOptimizer: ddp_world_size must be greater than 1.";
 
     for (size_t i = 0; i < model_chunks.size(); ++i) {
@@ -23,16 +62,10 @@ DistributedOptimizer::DistributedOptimizer(OptimizerCreator creator,
         bucket_groups_.insert(bucket_groups_.end(), ddp_chunk->bucket_groups().begin(),
                               ddp_chunk->bucket_groups().end());
     }
-
-    BuildShardParamsAndBindGrads();
-
-    // Build base optimizer
-    base_optimizer_ = creator(shard_params_);
-    CHECK(base_optimizer_) << "DistributedOptimizer: failed to create base optimizer.";
 }
 
-void DistributedOptimizer::BuildShardParamsAndBindGrads() {
-    shard_params_.clear();
+void DistributedOptimizer::BuildShardParamsAndBindGrads(const AddShard &add_shard) {
+    size_t num_shard_params = 0;
 
     for (const auto &group : bucket_groups_) {
         const bool use_grad_shard = group->config().zero_stage >= 2;
@@ -82,12 +115,13 @@ void DistributedOptimizer::BuildShardParamsAndBindGrads() {
                 // NOTE(zbl): Do not call `param->set_grad(grad_piece);` under ZeRO-2.
                 //            The base optimizer updates param_piece views only; original param->grad()
                 //            would be a partial flattened shard and does not represent the full parameter grad.
-                shard_params_.push_back(param_piece);
+                add_shard(param, param_piece);
+                ++num_shard_params;
             }
         }
     }
 
-    CHECK(!shard_params_.empty()) << "DistributedOptimizer: this DP rank owns no param pieces. "
+    CHECK_GT(num_shard_params, 0) << "DistributedOptimizer: this DP rank owns no param pieces. "
                                   << "Check bucket padding/divisibility and param bucketing order.";
 }
 
diff --git a/infini_train/src/optimizer.cc b/infini_train/src/optimizer.cc
index 1e97bfe0..39b999c7 100644
--- a/infini_train/src/optimizer.cc
+++ b/infini_train/src/optimizer.cc
@@ -1,6 +1,6 @@
 #include "infini_train/include/optimizer.h"
 
-#include <format>
+#include <unordered_map>
 #include <vector>
 
 #include "infini_train/include/core/runtime/device_guard.h"
@@ -12,6 +12,19 @@ namespace infini_train {
 Optimizer::Optimizer(const std::vector<std::shared_ptr<Tensor>> &params, float learning_rate)
     : params_(params), learning_rate_(learning_rate) {}
 
+Optimizer::Optimizer(const NamedParameterList &named_params, float learning_rate) : learning_rate_(learning_rate) {
+    if (named_params.empty()) {
+        return;
+    }
+
+    params_.reserve(named_params.size());
+    parameter_names_.reserve(named_params.size());
+    for (const auto &[name, parameter] : named_params) {
+        params_.push_back(parameter);
+        parameter_names_.push_back(name);
+    }
+}
+
 void Optimizer::ZeroGrad(bool set_to_none) {
     for (auto param : params_) { param->ZeroGrad(set_to_none); }
 }
@@ -33,10 +46,13 @@ void Optimizer::set_initial_learning_rate(float lr) {
     initial_learning_rate_ = lr;
     initial_lr_set_ = true;
 }
+
 namespace optimizers {
 
 SGD::SGD(const std::vector<std::shared_ptr<Tensor>> &params, float learning_rate) : Optimizer(params, learning_rate) {}
 
+SGD::SGD(const NamedParameterList &named_params, float learning_rate) : Optimizer(named_params, learning_rate) {}
+
 void SGD::Step() {
     for (auto param : params_) {
         if (!param->grad()) {
@@ -56,6 +72,12 @@ OptimizerCreator SGD::Create(float learning_rate) {
     };
 }
 
+OptimizerCreatorNamed SGD::CreateNamed(float learning_rate) {
+    return [learning_rate](const NamedParameterList &named_params) {
+        return std::make_shared<SGD>(named_params, learning_rate);
+    };
+}
+
 Adam::Adam(const std::vector<std::shared_ptr<Tensor>> &params, float learning_rate, float beta1, float beta2, float eps)
     : Optimizer(params, learning_rate), t_(0), beta1_(beta1), beta2_(beta2), eps_(eps) {
 
@@ -67,6 +89,16 @@ Adam::Adam(const std::vector<std::shared_ptr<Tensor>> &params, float learning_ra
     }
 }
 
+Adam::Adam(const NamedParameterList &named_params, float learning_rate, float beta1, float beta2, float eps)
+    : Optimizer(named_params, learning_rate), t_(0), beta1_(beta1), beta2_(beta2), eps_(eps) {
+    for (const auto &[name, param] : named_params) {
+        m_.emplace_back(std::make_shared<Tensor>(param->Dims(), param->Dtype(), param->GetDevice()));
+        v_.emplace_back(std::make_shared<Tensor>(param->Dims(), param->Dtype(), param->GetDevice()));
+        m_.back()->Fill(0.0);
+        v_.back()->Fill(0.0);
+    }
+}
+
 void Adam::Step() {
     ++t_;
 
@@ -93,11 +125,18 @@ OptimizerCreator Adam::Create(float learning_rate, float beta1, float beta2, flo
     };
 }
 
+OptimizerCreatorNamed Adam::CreateNamed(float learning_rate, float beta1, float beta2, float eps) {
+    return [=](const NamedParameterList &named_params) {
+        return std::make_shared<Adam>(named_params, learning_rate, beta1, beta2, eps);
+    };
+}
+
 std::unordered_map<std::string, std::shared_ptr<Tensor>> Adam::StateDict() const {
     std::unordered_map<std::string, std::shared_ptr<Tensor>> state;
     for (size_t i = 0; i < m_.size(); ++i) {
-        state.emplace(std::format("adam.m.{}", i), m_[i]);
-        state.emplace(std::format("adam.v.{}", i), v_[i]);
+        const auto suffix = parameter_names_.empty() ? std::to_string(i) : parameter_names_[i];
+        state.emplace("adam.m." + suffix, m_[i]);
+        state.emplace("adam.v." + suffix, v_[i]);
     }
 
     auto t_tensor = std::make_shared<Tensor>(std::vector<int64_t>{}, DataType::kINT64, Device());
@@ -108,8 +147,9 @@ std::unordered_map<std::string, std::shared_ptr<Tensor>> Adam::StateDict() const
 
 void Adam::LoadStateDict(const std::unordered_map<std::string, std::shared_ptr<Tensor>> &state_dict) {
     for (size_t i = 0; i < m_.size(); ++i) {
-        const auto m_key = std::format("adam.m.{}", i);
-        const auto v_key = std::format("adam.v.{}", i);
+        const auto suffix = parameter_names_.empty() ? std::to_string(i) : parameter_names_[i];
+        const auto m_key = "adam.m." + suffix;
+        const auto v_key = "adam.v." + suffix;
         CHECK(state_dict.contains(m_key)) << "Missing optimizer state: " << m_key;
         CHECK(state_dict.contains(v_key)) << "Missing optimizer state: " << v_key;
         m_[i]->CopyFrom(state_dict.at(m_key));
diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt
index f7ced03d..3bfaa548 100644
--- a/tests/CMakeLists.txt
+++ b/tests/CMakeLists.txt
@@ -9,6 +9,8 @@ add_subdirectory(common)
 
 # Distributed tests
 add_subdirectory(distributed)
+# Module tests
+add_subdirectory(module)
 
 # Tensor tests
 add_subdirectory(tensor)
diff --git a/tests/module/CMakeLists.txt b/tests/module/CMakeLists.txt
new file mode 100644
index 00000000..84b78098
--- /dev/null
+++ b/tests/module/CMakeLists.txt
@@ -0,0 +1,5 @@
+file(GLOB MODULE_SOURCES ${CMAKE_CURRENT_SOURCE_DIR}/test_*.cc)
+
+infini_train_add_test_suite(test_module
+  SOURCES ${MODULE_SOURCES}
+)
diff --git a/tests/module/test_named_parameters.cc b/tests/module/test_named_parameters.cc
new file mode 100644
index 00000000..c1a4a621
--- /dev/null
+++ b/tests/module/test_named_parameters.cc
@@ -0,0 +1,74 @@
+#include <memory>
+#include <string>
+#include <unordered_map>
+#include <unordered_set>
+#include <vector>
+
+#include "gtest/gtest.h"
+
+#include "infini_train/include/nn/modules/container.h"
+#include "infini_train/include/nn/modules/linear.h"
+#include "infini_train/include/tensor.h"
+
+#include "tests/common/test_utils.h"
+
+using namespace infini_train;
+
+class ModuleNamedParametersTest : public test::InfiniTrainTest {};
+
+TEST_P(ModuleNamedParametersTest, SupportsPrefixAndNonRecursiveLookup) {
+    auto linear = std::make_shared<nn::Linear>(2, 3, /*bias=*/true, GetDevice());
+
+    const auto parameters = linear->NamedParameters("model", false);
+    const std::unordered_map<std::string, std::shared_ptr<Tensor>> by_name(parameters.begin(), parameters.end());
+
+    ASSERT_EQ(by_name.size(), 2);
+    EXPECT_EQ(by_name.at("model.weight"), linear->parameter(nn::Linear::kParamWeightName));
+    EXPECT_EQ(by_name.at("model.bias"), linear->parameter(nn::Linear::kParamBiasName));
+}
+
+TEST_P(ModuleNamedParametersTest, SupportsRecursionAndSharedParameterDeduplication) {
+    auto shared = std::make_shared<nn::Linear>(2, 3, /*bias=*/true, GetDevice());
+    auto root = std::make_shared<nn::Sequential>(std::vector<std::shared_ptr<nn::Module>>{shared, shared});
+
+    const auto deduplicated = root->NamedParameters("model");
+    ASSERT_EQ(deduplicated.size(), 2);
+    EXPECT_EQ(deduplicated[0].first, "model.0.bias");
+    EXPECT_EQ(deduplicated[1].first, "model.0.weight");
+    std::unordered_set<const Tensor *> tensors;
+    for (const auto &[name, parameter] : deduplicated) { tensors.insert(parameter.get()); }
+    EXPECT_TRUE(tensors.contains(shared->parameter(nn::Linear::kParamWeightName).get()));
+    EXPECT_TRUE(tensors.contains(shared->parameter(nn::Linear::kParamBiasName).get()));
+
+    const auto aliases = root->NamedParameters("model", true, false);
+    const std::unordered_map<std::string, std::shared_ptr<Tensor>> by_name(aliases.begin(), aliases.end());
+    ASSERT_EQ(by_name.size(), 4);
+    EXPECT_EQ(by_name.at("model.0.weight"), by_name.at("model.1.weight"));
+    EXPECT_EQ(by_name.at("model.0.bias"), by_name.at("model.1.bias"));
+}
+
+TEST_P(ModuleNamedParametersTest, ReturnsNestedParametersInStableNameOrder) {
+    auto first = std::make_shared<nn::Linear>(2, 3, /*bias=*/false, GetDevice());
+    auto second = std::make_shared<nn::Linear>(3, 4, /*bias=*/false, GetDevice());
+    auto nested = std::make_shared<nn::Sequential>(std::vector<std::shared_ptr<nn::Module>>{first, second});
+    auto root_linear = std::make_shared<nn::Linear>(2, 2, false, GetDevice());
+    auto root = std::make_shared<nn::Sequential>(std::vector<std::shared_ptr<nn::Module>>{root_linear, nested});
+
+    const auto parameters = root->NamedParameters();
+
+    ASSERT_EQ(parameters.size(), 3);
+    EXPECT_EQ(parameters[0].first, "0.weight");
+    EXPECT_EQ(parameters[1].first, "1.0.weight");
+    EXPECT_EQ(parameters[2].first, "1.1.weight");
+    EXPECT_EQ(parameters[0].second, root_linear->parameter(nn::Linear::kParamWeightName));
+    EXPECT_EQ(parameters[1].second, first->parameter(nn::Linear::kParamWeightName));
+    EXPECT_EQ(parameters[2].second, second->parameter(nn::Linear::kParamWeightName));
+}
+
+TEST_P(ModuleNamedParametersTest, SkipsNullSubmodules) {
+    auto root = std::make_shared<nn::Sequential>(std::vector<std::shared_ptr<nn::Module>>{nullptr});
+
+    EXPECT_TRUE(root->NamedParameters().empty());
+}
+
+INFINI_TRAIN_REGISTER_TEST(ModuleNamedParametersTest);
diff --git a/tests/optimizer/CMakeLists.txt b/tests/optimizer/CMakeLists.txt
index c0bfbd50..bce88694 100644
--- a/tests/optimizer/CMakeLists.txt
+++ b/tests/optimizer/CMakeLists.txt
@@ -7,3 +7,17 @@ file(GLOB OPTIMIZER_SOURCES ${CMAKE_CURRENT_SOURCE_DIR}/test_*.cc)
 infini_train_add_test_suite(test_optimizer
   SOURCES ${OPTIMIZER_SOURCES}
 )
+
+add_test(
+  NAME OptimizerParameterNamesTest.DistributedOptimizerPropagatesNamesToShardOptimizer
+  COMMAND ${CMAKE_COMMAND} -E env
+    "PROC_WORLD_SIZE=2"
+    $<TARGET_FILE:test_optimizer_cuda>
+    --gtest_filter=CUDA/OptimizerParameterNamesTest.DistributedOptimizerPropagatesNamesToShardOptimizer/*
+)
+set_tests_properties(
+  OptimizerParameterNamesTest.DistributedOptimizerPropagatesNamesToShardOptimizer
+  PROPERTIES
+    LABELS "cuda;distributed"
+    TIMEOUT 30
+)
diff --git a/tests/optimizer/test_optimizer_parameter_names.cc b/tests/optimizer/test_optimizer_parameter_names.cc
new file mode 100644
index 00000000..31943454
--- /dev/null
+++ b/tests/optimizer/test_optimizer_parameter_names.cc
@@ -0,0 +1,96 @@
+#include <memory>
+#include <vector>
+
+#include "gtest/gtest.h"
+
+#include "infini_train/include/nn/modules/linear.h"
+#include "infini_train/include/nn/parallel/ddp/distributed_data_parallel.h"
+#include "infini_train/include/nn/parallel/ddp/distributed_data_parallel_config.h"
+#include "infini_train/include/nn/parallel/ddp/distributed_optimizer.h"
+#include "infini_train/include/nn/parallel/global.h"
+#include "infini_train/include/nn/parallel/process_group.h"
+#include "infini_train/include/nn/parallel/rank.h"
+#include "infini_train/include/nn/parallel/utils.h"
+#include "infini_train/include/optimizer.h"
+#include "infini_train/include/tensor.h"
+
+#include "tests/common/test_utils.h"
+
+using namespace infini_train;
+
+class OptimizerParameterNamesTest : public test::InfiniTrainTest {};
+
+TEST_P(OptimizerParameterNamesTest, AdamStateDictUsesStableParameterNames) {
+    auto first = std::make_shared<Tensor>(std::vector<int64_t>{2, 2}, DataType::kFLOAT32, GetDevice());
+    auto second = std::make_shared<Tensor>(std::vector<int64_t>{3}, DataType::kFLOAT32, GetDevice());
+    const NamedParameterList named_parameters{{"transformer.h.0.weight", first}, {"transformer.h.0.bias", second}};
+    auto adam = optimizers::Adam::CreateNamed(0.001)(named_parameters);
+
+    const auto state = adam->StateDict();
+    EXPECT_TRUE(state.contains("adam.m.transformer.h.0.weight"));
+    EXPECT_TRUE(state.contains("adam.v.transformer.h.0.weight"));
+    EXPECT_TRUE(state.contains("adam.m.transformer.h.0.bias"));
+    EXPECT_TRUE(state.contains("adam.v.transformer.h.0.bias"));
+    EXPECT_TRUE(state.contains("adam.t"));
+
+    auto restored = optimizers::Adam::CreateNamed(0.001)(named_parameters);
+    restored->LoadStateDict(state);
+    EXPECT_EQ(restored->StateDict().size(), state.size());
+}
+
+TEST_P(OptimizerParameterNamesTest, ConstructorMatchesNamesToOptimizerParameterOrder) {
+    auto first = std::make_shared<Tensor>(std::vector<int64_t>{2, 2}, DataType::kFLOAT32, GetDevice());
+    auto second = std::make_shared<Tensor>(std::vector<int64_t>{3}, DataType::kFLOAT32, GetDevice());
+    const NamedParameterList named_parameters{{"second", second}, {"first", first}};
+
+    auto adam = optimizers::Adam::CreateNamed(0.001)(named_parameters);
+    const auto state = adam->StateDict();
+
+    EXPECT_TRUE(state.contains("adam.m.second"));
+    EXPECT_TRUE(state.contains("adam.v.second"));
+    EXPECT_TRUE(state.contains("adam.m.first"));
+    EXPECT_TRUE(state.contains("adam.v.first"));
+}
+
+TEST_P(OptimizerParameterNamesTest, DistributedOptimizerPropagatesNamesToShardOptimizer) {
+    ONLY_CUDA();
+    REQUIRE_MIN_DEVICES(2);
+    if (nn::parallel::global::GetDataParallelSize() != 2) {
+        GTEST_SKIP() << "requires PROC_WORLD_SIZE=2";
+    }
+
+    const nn::parallel::Rank rank(/*process_rank=*/0, /*thread_rank=*/0, /*process_size=*/1, /*thread_size=*/2);
+    auto *pg_factory = nn::parallel::ProcessGroupFactory::Instance(Device::DeviceType::kCUDA);
+    pg_factory->GetOrCreate(nn::parallel::GetDataParallelProcessGroupName(rank.GlobalRank()),
+                            nn::parallel::GetDataParallelGroupRanks(rank.GlobalRank()));
+
+    auto model = std::make_shared<nn::Linear>(4, 4, /*bias=*/false, GetDevice());
+    const auto named_parameters = model->NamedParameters();
+
+    nn::parallel::DistributedDataParallelConfig ddp_config;
+    ddp_config.zero_stage = 1;
+    ddp_config.overlap_grad_reduce = false;
+    ddp_config.overlap_param_gather = false;
+    auto ddp_model = std::make_shared<nn::parallel::DistributedDataParallel>(model, rank, ddp_config);
+
+    nn::parallel::DistributedOptimizer optimizer(optimizers::Adam::CreateNamed(0.001), named_parameters,
+                                                 std::vector<std::shared_ptr<nn::Module>>{ddp_model},
+                                                 /*ddp_world_size=*/2, /*ddp_rank=*/0);
+    const auto state = optimizer.StateDict();
+
+    EXPECT_TRUE(state.contains("adam.m.weight"));
+    EXPECT_TRUE(state.contains("adam.v.weight"));
+    EXPECT_FALSE(state.contains("adam.m.0"));
+    EXPECT_FALSE(state.contains("adam.v.0"));
+}
+
+TEST_P(OptimizerParameterNamesTest, PreservesNumericKeysWhenNamesAreNotSet) {
+    auto parameter = std::make_shared<Tensor>(std::vector<int64_t>{2, 2}, DataType::kFLOAT32, GetDevice());
+    auto adam = std::make_shared<optimizers::Adam>(std::vector<std::shared_ptr<Tensor>>{parameter}, 0.001);
+
+    const auto state = adam->StateDict();
+    EXPECT_TRUE(state.contains("adam.m.0"));
+    EXPECT_TRUE(state.contains("adam.v.0"));
+}
+
+INFINI_TRAIN_REGISTER_TEST(OptimizerParameterNamesTest);