def test_deepspeed_with_bfloat16_precision():
    """Test that the DeepSpeed strategy works with bfloat16 precision."""

    class Model(nn.Module):
        def __init__(self):
            super().__init__()
            self.layer = nn.Linear(32, 2)

        def forward(self, x):
            assert x.dtype == torch.bfloat16
            return self.layer(x)

    fabric = Fabric(accelerator="cuda", devices=2, strategy="deepspeed_stage_3", precision="bf16-mixed")
    assert isinstance(fabric._strategy.precision, DeepSpeedPrecision)
    assert fabric._strategy.precision.precision == "bf16-mixed"
    assert fabric._strategy.config["zero_optimization"]["stage"] == 3
    fabric.launch()

    with fabric.init_module():
        model = Model()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.1)
    model, optimizer = fabric.setup(model, optimizer)
    assert fabric._strategy.config["bf16"]["enabled"]
    assert model.layer.weight.dtype == torch.bfloat16

    batch = torch.rand(2, 32, device=fabric.device)
    assert batch.dtype == torch.float32
    loss = model(batch).sum()
    fabric.backward(loss)
    optimizer.step()
    optimizer.zero_grad()
Test that the DeepSpeed strategy works with bfloat16 precision.
tests/tests_fabric/strategies/test_deepspeed_integration.py