Skip to content

Do not pin DDP device_ids for CPU reference models #8399

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

Merged
delock merged 2 commits into from
Sep 3, 2026
Merged

Do not pin DDP device_ids for CPU reference models #8399

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
41e1dc7
Do not pin DDP device_ids for CPU reference models
delock Sep 3, 2026
92f3b5f
Merge branch 'master' into pr-c-ddp-cpu-reference
delock Sep 3, 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
17 changes: 12 additions & 5 deletions tests/unit/v1/zero/test_zero_user_backward.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -111,11 +111,18 @@ def create_ddp_model(model_class, device, rank, dtype, seed=42, lr=1e-3, **model
torch.manual_seed(seed)
model = model_class(**model_kwargs)
model = model.to(device=device, dtype=dtype)
model = DDP(model, device_ids=[rank], output_device=rank)
model = wrap_ddp_reference(model, device, rank)
optimizer = torch.optim.Adam(model.parameters(), lr=lr)
return model, optimizer


def wrap_ddp_reference(model, device, rank):

@chatgpt-codex-connector chatgpt-codex-connector Bot Sep 3, 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.

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

# 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)


def create_deepspeed_engine(model_class, zero_stage, seed=42, gradient_accumulation_steps=1, **model_kwargs):
torch.manual_seed(seed)
model = model_class(**model_kwargs)
Expand Down Expand Up @@ -280,7 +287,7 @@ def run_frozen_checkpoint_comparison(model_cls,
torch.manual_seed(42)
model_ddp = model_cls(hidden_dim=hidden_dim, use_reentrant=use_reentrant, **model_kwargs)
model_ddp = model_ddp.to(device=device, dtype=dtype)
model_ddp = DDP(model_ddp, device_ids=[rank], output_device=rank)
model_ddp = wrap_ddp_reference(model_ddp, device, rank)
optimizer_ddp = torch.optim.Adam([p for p in model_ddp.parameters() if p.requires_grad], lr=1e-3)

# DeepSpeed engine with ZeRO partitioning. Only trainable params go to the optimizer;
Expand Down Expand Up @@ -1370,7 +1377,7 @@ def test_checkpointed_non_scalar_backward(self, zero_stage, use_reentrant):
torch.manual_seed(42)
model_ddp = CheckpointedModel(hidden_dim=hidden_dim, use_reentrant=use_reentrant)
model_ddp = model_ddp.to(device=device, dtype=dtype)
model_ddp = DDP(model_ddp, device_ids=[rank], output_device=rank)
model_ddp = wrap_ddp_reference(model_ddp, device, rank)
optimizer_ddp = torch.optim.Adam(model_ddp.parameters(), lr=1e-3)

# Create DeepSpeed model with ZeRO-3
Expand Down Expand Up @@ -1444,7 +1451,7 @@ def test_checkpointed_scalar_backward(self, zero_stage, use_reentrant):
torch.manual_seed(42)
model_ddp = CheckpointedModel(hidden_dim=hidden_dim, use_reentrant=use_reentrant)
model_ddp = model_ddp.to(device=device, dtype=dtype)
model_ddp = DDP(model_ddp, device_ids=[rank], output_device=rank)
model_ddp = wrap_ddp_reference(model_ddp, device, rank)
optimizer_ddp = torch.optim.Adam(model_ddp.parameters(), lr=1e-3)

# Create DeepSpeed model with ZeRO-3
Expand Down Expand Up @@ -1516,7 +1523,7 @@ def test_checkpointed_multiple_backward(self, zero_stage, use_reentrant):
torch.manual_seed(42)
model_ddp = CheckpointedModel(hidden_dim=hidden_dim, use_reentrant=use_reentrant)
model_ddp = model_ddp.to(device=device, dtype=dtype)
model_ddp = DDP(model_ddp, device_ids=[rank], output_device=rank)
model_ddp = wrap_ddp_reference(model_ddp, device, rank)
optimizer_ddp = torch.optim.Adam(model_ddp.parameters(), lr=1e-3)

# Create DeepSpeed model WITH checkpointing, using PyTorch Adam
Expand Down
Loading