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.
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.