Skip to content

Forward free_data through partition() instead of hardcoding True #8305

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
vineethsaivs wants to merge 2 commits into
base: master
Choose a base branch
from
+24 −2
Open

Forward free_data through partition() instead of hardcoding True #8305

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
1b22dea
Forward free_data through partition() instead of hardcoding True
vineethsaivs Aug 24, 2026
61c2c59
Merge branch 'master' into fix/zero3-forward-free-data
sfc-gh-truwase Aug 26, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions deepspeed/runtime/zero/partition_parameters.py
100755 → 100644
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1579,7 +1579,7 @@ def partition(param_list=None, hierarchy=0, has_been_updated=False, free_data=Tr
force=False)
if param_list is None:
param_list = [cls]
self._partition(param_list, has_been_updated=has_been_updated, free_data=True)
self._partition(param_list, has_been_updated=has_been_updated, free_data=free_data)

@chatgpt-codex-connector chatgpt-codex-connector Bot Aug 24, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Add the missing sign-off trailer

This is a non-merge commit, but its commit message has no Signed-off-by: trailer. The repository requires every non-merge commit to be signed off, so compliance checks can reject this until the commit is recreated with git commit --signoff using the configured name and email.

AGENTS.md reference: AGENTS.md:L8-L8

Useful? React with 👍 / 👎.


def reduce_gradients_at_owner(param_list=None, hierarchy=0):
cls = param
Expand Down Expand Up @@ -1739,7 +1739,7 @@ def _partition(self, param_list, force=False, has_been_updated=False, free_data=
print_rank_0(f"Before Partitioning Param {param.ds_id}", force=False)
if self.zero_param_process_group is not None:
self._partition_param_sec(param, has_been_updated=has_been_updated)
self._partition_param(param, has_been_updated=has_been_updated, free_data=True)
self._partition_param(param, has_been_updated=has_been_updated, free_data=free_data)

param.ds_status = ZeroParamStatus.NOT_AVAILABLE
# if param.ds_tensor is not None:
Expand Down
22 changes: 22 additions & 0 deletions tests/unit/runtime/zero/test_zero_context.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -120,6 +120,28 @@ def __init__(self, hidden_dim):
assert model.l1.weight.numel() == 0, "outside of GatheredParameters the param should go back to be 0-sized"


class TestPartitionWithoutFreeingData(DistributedTest):
world_size = 1

def test(self):
with deepspeed.zero.Init():
l = torch.nn.Linear(6, 3, bias=False)

full_numel = l.in_features * l.out_features
l.weight.all_gather()
assert l.weight.numel() == full_numel

# The leaf-module fast-sharding path releases the buffer itself once the whole
# submodule is done, so partition() has to leave param.data alone when asked to.
l.weight.partition(free_data=False)
assert l.weight.ds_status == ZeroParamStatus.NOT_AVAILABLE
assert l.weight.numel() == full_numel, "partition(free_data=False) should not free param.data"

l.weight.all_gather()
l.weight.partition()
assert l.weight.numel() == 0, "partition() should free param.data by default"


class TestMiCSGatheredParametersFree(DistributedTest):
world_size = 1

Expand Down
Loading