def test_deepspeed_multigpu_stage_3():
    """Test to ensure ZeRO Stage 3 works with a parallel model."""
    fabric = Fabric(
        strategy=DeepSpeedStrategy(stage=3),
        accelerator="cuda",
        devices=2,
        precision="16-mixed",
    )
    fabric.launch()

    def _make_block():
        return nn.Sequential(nn.Linear(32, 32, bias=False), nn.ReLU())

    with fabric.init_module():
        model = nn.Sequential(*(_make_block() for _ in range(5)), nn.Linear(32, 3))

    optimizer = torch.optim.Adam(model.parameters(), lr=0.1)
    model, optimizer = fabric.setup(model, optimizer)

    x = torch.rand(2, 32, device=fabric.device)
    y = torch.ones(x.size(0), device=x.device, dtype=torch.long)
    x = model(x)
    x = x.float()  # Ensure output is in float32 for softmax operation
    logits = F.softmax(x, dim=1)
    loss = F.cross_entropy(logits, y)
    fabric.backward(loss)
    optimizer.step()
    optimizer.zero_grad()
Test to ensure ZeRO Stage 3 works with a parallel model.
tests/tests_fabric/strategies/test_deepspeed_integration.py