mindspore-ai / mindspore-ai/hyper-parallel

[Bug]: fully_shard comm_fusion keeps duplicate sharded parameter storage

Open
#231 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

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

  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

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.