def test_deepspeed_with_bfloat16_precision(tmp_path):
    """Test that deepspeed works with bfloat16 precision."""
    model = BoringModel()
    trainer = Trainer(
        default_root_dir=tmp_path,
        strategy="deepspeed_stage_3",
        accelerator="gpu",
        devices=2,
        fast_dev_run=True,
        precision="bf16-mixed",
        num_sanity_val_steps=0,
        enable_progress_bar=False,
        enable_model_summary=False,
    )

    trainer.fit(model)
    assert isinstance(trainer.strategy.precision_plugin, DeepSpeedPrecision)
    assert trainer.strategy.precision_plugin.precision == "bf16-mixed"
    assert trainer.strategy.config["zero_optimization"]["stage"] == 3
    assert trainer.strategy.config["bf16"]["enabled"]
    assert model.layer.weight.dtype == torch.bfloat16
Test that deepspeed works with bfloat16 precision.
tests/tests_pytorch/strategies/test_deepspeed.py