Do not pin DDP device_ids for CPU reference models

[delock]

Problem

tests/unit/v1/zero/test_zero_user_backward.py builds torch DDP reference models (the known-good baseline that ZeRO results are compared against) at five sites:

model_ddp = DDP(model_ddp, device_ids=[rank], output_device=rank)

device_ids=[rank] assumes rank ↔ GPU index. torch's DDP contract only allows device_ids/output_device for single-device GPU modules; CPU modules live on one shared device and must omit them, so multi-rank CPU runs died inside the DDP constructor with:

ValueError: DistributedDataParallel device_ids and output_device arguments only work with single-device/multiple-device GPU modules or CPU modules, ...

Change

Route all five sites through one helper:

def wrap_ddp_reference(model, device, rank):
    # Only indexed devices take device_ids/output_device; CPU modules live on one shared device.
    if torch.device(device).type == 'cpu':
        return DDP(model)
    return DDP(model, device_ids=[rank], output_device=rank)

Design notes:

• the condition restates the exact precondition torch's own DDP constructor enforces, in torch's device vocabulary — it follows the model's actual device rather than the global accelerator configuration;
• a single helper means new reference-model sites cannot forget the branch (the first fix round in [DON'T MERGE] Run multi-rank CPU unit tests in CI via LOCAL_SIZE #8381 missed 4 of the 5 sites for exactly this reason);
• GPU behavior is unchanged.

Validation

Validated as part of the multi-rank CPU CI experiment in #8381: all DDP-constructor failures were eliminated (the file's few remaining failures there are unrelated — see the triage table in that PR), zero regressions vs the same-commit baseline.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 41e1dc720b

ℹ️ About Codex in GitHub

Codex has been enabled to automatically 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 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

Add the required sign-off trailer

This is a non-merge commit, but its message has no Signed-off-by trailer, so it does not satisfy the repository's commit requirements and may be rejected by the DCO/CI checks. Recreate the commit with git commit --signoff using the configured Git identity.

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

Useful? React with 👍 / 👎.