mindspore-ai / mindspore-ai/hyper-parallel
[Bug]: fully_shard comm_fusion keeps duplicate sharded parameter storage
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
When fully_shard enables comm_fusion, HSDPParamGroup creates a contiguous flat parameter buffer and rebases each hsdp_param shard onto slices of that buffer. Before the fix, the old per-parameter sharded storages were still active after the rebase. Memory snapshot showed both the old hsdp_param.sharded_param shard storages and the new flat buffer storage alive at the same time.
For the reproduction modeled after the zero3 comm_fusion precision case, each rank kept two old 4 MiB shard blocks plus one new 8 MiB flat buffer block after fully_shard, causing redundant device memory usage. The model parameter and optimizer parameter identities were consistent, so the issue is stale shard storage lifetime rather than optimizer using a different parameter object.
Expected behavior
After comm_fusion rebases sharded parameters into the flat parameter buffer, only the flat buffer storage should remain active for the local shards. Access through module parameters and optimizer updates should continue to target the same managed parameter object, and training precision should remain unchanged.
Additional context
The fix should also work when parameters are initialized lazily through defer-init/meta-init flows. The validation should compare non-fusion and fusion training results and verify that sharded local tensors alias the flat buffer after lazy initialization.
Environment info
HyperParallel PyTorch fully_shard/HSDP path on Ascend NPU multi-card distributed test environment.
schema_version: 1
source: gitcode
gitcode_repo: mindspore/hyper-parallel
gitcode_issue: 271
source_url: https://gitcode.com/mindspore/hyper-parallel/issues/271
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
Start in the fully_shard/HSDP comm_fusion path and reproduce the memory snapshot showing both old shard storage and the flat buffer. Compare fusion and non-fusion training, then exercise lazy defer-init/meta-init initialization. Done means only flat-buffer storage remains, local shards alias it, optimizer and module access remain consistent, and training precision is unchanged.
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
- 50/100