已开启
[Bug]: fully_shard comm_fusion keeps duplicate sharded parameter storage #271
MengXY107创建于 7月1日
7月1日 添加了label:bug
7月1日 关联了pull request:fix: release rebased FSDP flat buffer shard storage
7月6日 关联了pull request:fix: rebase DTensor storage for comm fusion
Checklist
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.