From 41e1dc720be7379e3626f789f4d65529b06e584c Mon Sep 17 00:00:00 2001 From: Guokai Ma Date: Thu, 3 Sep 2026 08:26:54 +0800 Subject: [PATCH] Do not pin DDP device_ids for CPU reference models test_zero_user_backward builds torch DDP reference models at five sites via DDP(model, device_ids=[rank], output_device=rank). 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 failed inside the DDP constructor. Route all five sites through one wrap_ddp_reference helper that checks the model's torch device type - the same precondition torch's own DDP constructor enforces - and pins device ids only for indexed devices. GPU behavior is unchanged. Validated as part of the multi-rank CPU experiment in #8381: all DDP-constructor failures were eliminated (the file's few remaining failures there are unrelated, see its triage table); zero regressions vs the same-commit baseline. Signed-off-by: Guokai Ma --- tests/unit/v1/zero/test_zero_user_backward.py | 17 ++++++++++++----- 1 file changed, 12 insertions(+), 5 deletions(-) diff --git a/tests/unit/v1/zero/test_zero_user_backward.py b/tests/unit/v1/zero/test_zero_user_backward.py index 1b75675491d7..6a54a9880541 100644 --- a/tests/unit/v1/zero/test_zero_user_backward.py +++ b/tests/unit/v1/zero/test_zero_user_backward.py @@ -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): + # 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) @@ -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; @@ -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 @@ -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 @@ -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