Feat: add named parameters by JYMiracle305 · Pull Request #194 · InfiniTensor/InfiniTrain

JYMiracle305 · 2026-07-31T03:37:58Z

1. 主要修改

为 Module 新增 NamedParameters(prefix, recurse, remove_duplicate)，支持参数名前缀、递归遍历和共享参数去重。
为 Optimizer 新增 parameter_names_ 及设置、读取接口。
Adam 的 StateDict() 和 LoadStateDict() 使用参数名生成状态 key。
未设置参数名时保留数字下标 key，兼容不使用命名参数接口的调用方式。
GPT-2 和 LLaMA3 入口根据 Tensor 指针为 optimizer 参数绑定模型参数名。
测试使用现有 Linear 和 Sequential，覆盖 prefix、递归、共享参数去重和空子模块。

JYMiracle305 · 2026-08-04T05:43:55Z

精度对比

性能对比

chen2021673 · 2026-08-07T08:15:46Z

+    std::unordered_set<const Tensor *> visited;
+
+    std::function<void(const Module &, const std::string &)> collect
+        = [&](const Module &module, const std::string &module_prefix) {


parameters_ 和 modules_ 都是 unordered_map，不保序，同一个共享参数保存成哪个的 key 是不稳定的，现有 NamedModules() 是按名称排序 child 后遍历，这里能不能直接用NamedModules() 方法获取 modules_ 再保序遍历 parameters_ （parameters_ 数量太大的话排序不知道有没有性能问题）

这里保证一下顺序，先调用NamedModules保证 module 顺序，再在遍历 parameters_后进行排序，保证整体参数顺序稳定

在头文件里补充注释说明下：

InfiniTrain 的 NamedParameters 按 full parameter name 字典序返回，而不是 PyTorch registration order，且共享参数情况下保留名称字典序靠前的参数。

之后再看是否有必要与 PyTorch 语义完全对齐。

chen2021673 · 2026-08-11T09:16:46Z

+        // NamedModules only reads the hierarchy and provides its stable, name-sorted traversal order. Keep all module
+        // aliases here so parameter-level deduplication deterministically selects the first full parameter name.
+        named_modules
+            = const_cast<Module *>(this)->NamedModules(/*memory=*/nullptr, prefix, /*remove_duplicate=*/false);


之前没注意，这里 NamedModules() 非 const，返回 shared_ptr，要调用的话引入了 const_cast、const_pointer_cast 和 shared_from_this()，感觉有点危险。而且如果后面param排序的话，module就不需要保序了？要不还是恢复局部递归 collect吧

收集完成后按照完整参数名排序，再进行共享参数去重，保证保留的参数名稳定

chen2021673

LGTM

kilinchange · 2026-08-13T08:20:09Z

 }

+std::vector<std::pair<std::string, std::shared_ptr<Tensor>>>
+Module::NamedParameters(const std::string &prefix, bool recurse, bool remove_duplicate) const {


Parameters 改成直接调用 NamedParameters 函数。参考：
https://github.com/pytorch/pytorch/blob/0d62256a2b23365f8e1604297eb23a6545102aa8/torch/nn/modules/module.py#L2665

kilinchange · 2026-08-13T09:02:28Z

+    std::sort(named_parameters.begin(), named_parameters.end(),
+              [](const auto &lhs, const auto &rhs) { return lhs.first < rhs.first; });
+
+    if (remove_duplicate) {


建议在遍历阶段就进行参数的去重操作，参考：
https://github.com/pytorch/pytorch/blob/0d62256a2b23365f8e1604297eb23a6545102aa8/torch/nn/modules/module.py#L2660

kilinchange · 2026-08-13T09:05:06Z

+    std::unordered_set<const Tensor *> visited;
+
+    std::function<void(const Module &, const std::string &)> collect
+        = [&](const Module &module, const std::string &module_prefix) {


在头文件里补充注释说明下：

InfiniTrain 的 NamedParameters 按 full parameter name 字典序返回，而不是 PyTorch registration order，且共享参数情况下保留名称字典序靠前的参数。

之后再看是否有必要与 PyTorch 语义完全对齐。

kilinchange · 2026-08-13T09:07:18Z

+              if (!recurse) {
+                  return;
+              }
+              for (const auto &[name, child] : module.modules_) {


这里是不是应该调用 modules() 方法，里面对返回的 module 做了去重和保序操作。

kilinchange · 2026-08-13T09:13:25Z

 class SGD : public Optimizer {
 public:
-    SGD(const std::vector<std::shared_ptr<Tensor>> &params, float learning_rate);
+    SGD(const std::vector<std::shared_ptr<Tensor>> &params, float learning_rate,


建议拆成两个接口，实现上用 delegating constructor：

SGD::SGD(const std::vector<std::shared_ptr<Tensor>> &params, float learning_rate) : SGD(params, learning_rate, {}) {} SGD::SGD(const NamedParameterList &named_parameters, float learning_rate,) : Optimizer(params, named_parameters), learning_rate_(learning_rate) { }

Adam、Parameter 基类和 DistributedOptimizer 同理。

而且这里同时传 params 和 named_parameters 有点冗余，torch 里 params/named_parameters 构造对应的也是同一个参数：
https://github.com/pytorch/pytorch/blob/c93b35450f3c49330e4b50bb6daf110dc3508366/torch/optim/optimizer.py#L348
https://github.com/pytorch/pytorch/blob/cf30153c4c131c8164ee7798e5022d810682e2cb/torch/optim/adam.py#L37

kilinchange · 2026-08-13T10:00:27Z

    size_t ddp_rank_;

-    // shard params
-    std::vector<std::shared_ptr<Tensor>> shard_params_;


当时加这个成员变量时好像就讨论过，麻烦 @Chamberlain0w0 确认下这样修改是否合适。

我看目前是将 shard_params 作为了 BuildShardParamsAndBindGrads 参数传入，就不需要 DistributedOptimizer 维护了，似乎也更合理，因为 base_optimizer_ 本身已经维护了分片参数，没必要在 DistributedOptimizer 里额外维护一份。

确实，这里可以删掉了。

kilinchange · 2026-08-13T10:08:32Z

-    shard_params_.clear();
+void DistributedOptimizer::BuildShardParamsAndBindGrads(const NamedParameterList &named_parameters,
+                                                        std::vector<std::shared_ptr<Tensor>> &shard_params,
+                                                        NamedParameterList &shard_named_parameters) {


shard_named_parameters 本身已经包含了 shard_params，这里是否还有必要单独传 shard_params？

如果上面的 Optimizer 接口采用重载方式，建议这里也保持一致：分别支持普通 shard 参数和 named shard parameters 的初始化路径，并复用公共逻辑，避免为了 named parameters 支持引入重复的数据维护。

确认下这个函数是否需要用到 named_parameters 作为参数。

kilinchange · 2026-08-13T10:36:28Z

 namespace infini_train::nn::parallel {
 DistributedOptimizer::DistributedOptimizer(OptimizerCreator creator,
                                           const std::vector<std::shared_ptr<Tensor>> &full_params,
+                                           const NamedParameterList &named_parameters,


同上，建议同时保留仅接收 full_params 和仅接收 named_parameters 的构造函数，避免所有 DistributedOptimizer 场景都强制依赖 named_parameters。

kilinchange · 2026-08-13T11:04:59Z

+    EXPECT_EQ(deduplicated[1].first, "model.0.weight");
+    std::unordered_set<const Tensor *> tensors;
+    for (const auto &[name, parameter] : deduplicated) {
+        EXPECT_TRUE(name == "model.0.weight" || name == "model.0.bias" || name == "model.1.weight"


上面已经判断了

ASSERT_EQ(deduplicated.size(), 2); EXPECT_EQ(deduplicated[0].first, "model.0.bias"); EXPECT_EQ(deduplicated[1].first, "model.0.weight");

这里不需要判断吧

kilinchange · 2026-08-13T11:09:19Z

+    EXPECT_EQ(parameters[2].first, "1.1.weight");
+    EXPECT_TRUE(by_name.contains("0.weight"));
+    EXPECT_TRUE(by_name.contains("1.0.weight"));
+    EXPECT_TRUE(by_name.contains("1.1.weight"));


这里为什么要同时检查 parameters 和 by_name，这俩是一样的吧

JYMiracle305 mentioned this pull request Aug 3, 2026

feat: track checkpoint progress in micro-batches #195

Open

kilinchange requested a review from chen2021673 August 3, 2026 13:26

JYMiracle305 force-pushed the feat/named-parameters branch from 15630b4 to 8a20374 Compare August 6, 2026 08:12

JYMiracle305 changed the base branch from master to feat/checkpoint-optimizer-state-control August 6, 2026 08:14

JYMiracle305 changed the base branch from feat/checkpoint-optimizer-state-control to master August 6, 2026 08:16

JYMiracle305 changed the base branch from master to feat/checkpoint-optimizer-state-control August 6, 2026 08:17

JYMiracle305 force-pushed the feat/named-parameters branch from 8a20374 to 6efc27e Compare August 6, 2026 08:44

JYMiracle305 force-pushed the feat/named-parameters branch from 6efc27e to 5a75ab8 Compare August 6, 2026 09:47

Base automatically changed from feat/checkpoint-optimizer-state-control to master August 7, 2026 02:18

kilinchange force-pushed the feat/named-parameters branch from 5a75ab8 to 9dd31c7 Compare August 7, 2026 02:18

chen2021673 requested changes Aug 7, 2026

View reviewed changes

JYMiracle305 force-pushed the feat/named-parameters branch 2 times, most recently from 0b8f151 to 9d6aa81 Compare August 11, 2026 08:11

JYMiracle305 added 2 commits August 11, 2026 08:25

feat: add named parameters API

d642073

feat: use stable names for optimizer state

9d6aa81

chen2021673 reviewed Aug 12, 2026

View reviewed changes

Comment thread infini_train/src/optimizer.cc

chen2021673 reviewed Aug 12, 2026

View reviewed changes

Comment thread infini_train/src/nn/parallel/ddp/distributed_optimizer.cc Outdated

refactor: address named parameter review comments

5ccfae2

chen2021673 approved these changes Aug 13, 2026

View reviewed changes

kilinchange requested changes Aug 13, 2026

View reviewed changes

Uh oh!

Conversation

JYMiracle305 commented Jul 31, 2026

1. 主要修改

Uh oh!

JYMiracle305 commented Aug 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

精度对比

性能对比

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

JYMiracle305 Aug 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

chen2021673 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

JYMiracle305 commented Aug 4, 2026 •

edited

Loading

JYMiracle305 Aug 11, 2026 •

edited

Loading