Forward free_data through partition() instead of hardcoding True

[vineethsaivs]

Problem

#6694 added a free_data flag so the ZeRO-3 leaf-module fast-sharding path could partition a parameter without calling free_param() on it, and then release the buffer itself after the whole submodule was done (param.data = empty_buffer), avoiding a per-parameter synchronisation.

The flag was added to all three signatures, but the two intermediate hops pass the literal True instead of the value they were given:

def partition(param_list=None, hierarchy=0, has_been_updated=False, free_data=True):
    ...
    self._partition(param_list, has_been_updated=has_been_updated, free_data=True)   # not free_data

def _partition(self, param_list, force=False, has_been_updated=False, free_data=True):
    ...
    self._partition_param(param, has_been_updated=has_been_updated, free_data=True)  # not free_data

So free_data=False never reaches the only place that reads it:

# _partition_param
if free_data:
    free_param(param)

Consequences:

• The if free_data: guard is unreachable as False; free_param() always runs.
• PartitionedParameterCoordinator.release_sub_module() computes free_data = not z3_leaf_module(submodule) or not self.fast_sharding_for_leaf_module and threads it through __release_param() -> param.partition(free_data=free_data), but the buffer is freed anyway. The fast path enabled by stage3_module_granularity_threshold > 0 is a no-op, and the param.data = empty_buffer that follows just overwrites an already-freed tensor.

Reproducible directly on the parameter API (single process, CPU):

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

l.weight.all_gather()
print(l.weight.numel())            # 18
l.weight.partition(free_data=False)
print(l.weight.numel())            # 0  <- expected 18

Fix

Pass free_data down at both hops. Two lines; no signature or default changes, so the default path (free_data=True) is byte-for-byte the same as before.

Test

Added TestPartitionWithoutFreeingData to tests/unit/runtime/zero/test_zero_context.py (world_size=1), next to the existing TestZeroGatheredParametersFree. It gathers a parameter, calls partition(free_data=False) and asserts param.data survives, then gathers again and calls plain partition() to assert the default still frees.

# new test against the unpatched source
1 failed, 2 passed        (TestZeroGatheredParametersFree and TestMiCSGatheredParametersFree are the controls)

# after the fix
3 passed

Whole-file and feature-path controls, clean tree vs patched:

tests/unit/runtime/zero/            (whole directory)
  clean:   1 failed, 57 passed, 250 skipped
  patched: 1 failed, 58 passed, 250 skipped     (net difference is exactly the new test)

tests/unit/v1/zero/test_zero_activation_checkpoint_lifecycle.py   (parametrises module_granularity_threshold)
  clean:   9 passed, 2 xfailed
  patched: 9 passed, 2 xfailed

The one failure is the same on both sides, test_zero_nesting_init.py::TestNestedParallelInit::test_nested_parallel_init, and is unrelated to this change.

yapf --style .style.yapf -d is clean on both files. flake8 --config .flake8 reports the same two pre-existing F824 global reuse_buffers warnings before and after; nothing new.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 1b22deacfa

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

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 👍 / 👎.

[vineethsaivs]

This one is a false positive: the commit does carry the trailer.

$ gh api repos/deepspeedai/DeepSpeed/pulls/8305/commits --jq '.[].commit.message' | tail -1
Signed-off-by: Vineeth Sai <vineethsai4444@gmail.com>

The DCO check on this PR passes, which is the same thing measured independently. No change needed.