def test_fsdp_precision_forward_context():
    """Test to ensure that the context manager correctly is set to bfloat16."""
    from torch.distributed.fsdp.sharded_grad_scaler import ShardedGradScaler

    precision = FSDPPrecision(precision="16-mixed")
    assert isinstance(precision.scaler, ShardedGradScaler)
    assert torch.get_default_dtype() == torch.float32
    with precision.forward_context():
        assert torch.get_autocast_gpu_dtype() == torch.float16
    assert isinstance(precision.forward_context(), torch.autocast)
    assert precision.forward_context().fast_dtype == torch.float16

    precision = FSDPPrecision(precision="16-true")
    assert precision.scaler is None
    assert torch.get_default_dtype() == torch.float32
    with precision.forward_context():
        assert torch.get_default_dtype() == torch.float16
    assert isinstance(precision.forward_context(), _DtypeContextManager)
    assert precision.forward_context()._new_dtype == torch.float16

    precision = FSDPPrecision(precision="bf16-mixed")
    assert precision.scaler is None
    with precision.forward_context():
        assert torch.get_autocast_gpu_dtype() == torch.bfloat16
    assert isinstance(precision.forward_context(), torch.autocast)
    assert precision.forward_context().fast_dtype == torch.bfloat16

    precision = FSDPPrecision(precision="bf16-true")
    assert precision.scaler is None
    with precision.forward_context():  # forward context is not using autocast ctx manager
        assert torch.get_default_dtype() == torch.bfloat16
    assert isinstance(precision.forward_context(), _DtypeContextManager)
    assert precision.forward_context()._new_dtype == torch.bfloat16
Test to ensure that the context manager correctly is set to bfloat16.
tests/tests_fabric/plugins/precision/test_fsdp.py