mindspore-ai / mindspore-ai/hyper-parallel
[Bug]: Tied weights lose Parameter identity after sharding and meta materialization
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
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- 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