mindspore-ai / mindspore-ai/hyper-parallel

[Bug]: Tied weights lose Parameter identity after sharding and meta materialization

Open
#606 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Python
Stars
53
Forks
63
Avg merge
23h 45m
Merged PRs (30d)
63

Description

Checklist
  • 1. I have searched the existing issues (https://gitcode.com/mindspore/hyper-parallel/issues)
  • 2. I have read the relevant documentation.
  • 3. I have created a minimal reproduction case that clearly demonstrates the issue, including a complete code example and the error message with full traceback and error logs.
🐛 Describe the bug

问题描述

模型中的 tied weights 在参数分片或 meta device 物化后,可能只共享底层
storage,但不再是同一个 Parameter 对象。

例如 embedding 和 lm_head 完成分片后,原有逻辑只替换底层 tensor:

tensor_a = param_a.to_local() if isinstance(param_a, DTensor) else param_a.data

if isinstance(param_b, DTensor):
    param_b._local_tensor = tensor_a
else:
    param_b.data = tensor_a

此时虽然两端可能共享数据,但仍然存在两个独立的参数对象:

model.embed_tokens.weight is not model.lm_head.weight

这会影响 tied weights 的梯度聚合、优化器状态复用以及 HSDP 对共享参数
所有权的识别。

另外,meta 模型通过 to_empty 物化后,没有重新恢复 tied weights:

model.to_empty(device=device)
return model

to_empty 可能为同一绑定参数的不同 FQN 创建独立对象,导致 checkpoint
加载前的参数发现看到错误的参数关系。

期望行为

Expected behavior

参数分片及 meta materialization 完成后:

model.embed_tokens.weight is model.lm_head.weight

应继续成立。

对于 HSDP 模型,需要在 checkpoint target discovery 之前刷新 HSDP
参数状态并恢复共享参数;同时应校验绑定参数的 shape、dtype、DeviceMesh和 placements 是否一致,避免错误布局被静默绑定。

Additional context
Environment info

atlas a2/, torch 2.9.0

Thanks for contributing 🎉!

schema_version: 1
source: gitcode
gitcode_repo: mindspore/hyper-parallel
gitcode_issue: 388
source_url: https://gitcode.com/mindspore/hyper-parallel/issues/388

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Trace the parameter sharding path shown in the issue and the meta-model to_empty materialization path first. Reproduce the embedding/lm_head case, then inspect checkpoint target discovery and HSDP parameter-state handling. Done means tied names reference the same Parameter after sharding and materialization, with incompatible shape, dtype, DeviceMesh, or placements rejected.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
distributed-systems, machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
38/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.