Torch

check_fp8_support(device_id=0)

Check if FP8 is supported on the current GPU.

FP8 requires compute capability 8.9+ (Ada Lovelace/Hopper architecture or newer).

Source code in bionemo/testing/torch.py

def check_fp8_support(device_id: int = 0) -> tuple[bool, str, str]:
    """Check if FP8 is supported on the current GPU.

    FP8 requires compute capability 8.9+ (Ada Lovelace/Hopper architecture or newer).
    """
    if not torch.cuda.is_available():
        return False, "0.0", "CUDA not available"
    device_props = torch.cuda.get_device_properties(device_id)
    compute_capability = f"{device_props.major}.{device_props.minor}"
    device_name = device_props.name
    # FP8 is supported on compute capability 8.9+ (Ada Lovelace/Hopper architecture)
    is_supported = (device_props.major > 8) or (device_props.major == 8 and device_props.minor >= 9)
    return is_supported, compute_capability, f"Device: {device_name}, Compute Capability: {compute_capability}"

recursive_assert_approx_equal(x, y, atol=0.0001, rtol=0.0001)

Assert that all tensors in a nested structure are approximately equal.

Source code in bionemo/testing/torch.py

def recursive_assert_approx_equal(x, y, atol=1e-4, rtol=1e-4):
    """Assert that all tensors in a nested structure are approximately equal."""
    if isinstance(x, torch.Tensor):
        torch.testing.assert_close(x, y, atol=atol, rtol=rtol)
    elif isinstance(x, np.ndarray):
        np.testing.assert_allclose(x, y, atol=atol, rtol=rtol)
    elif isinstance(x, (list, tuple)):
        assert len(x) == len(y), f"Length mismatch: {len(x)} vs {len(y)}"
        for x_item, y_item in zip(x, y):
            recursive_assert_approx_equal(x_item, y_item, atol=atol, rtol=rtol)
    elif isinstance(x, dict):
        assert x.keys() == y.keys()
        for key in x:
            recursive_assert_approx_equal(x[key], y[key], atol=atol, rtol=rtol)
    else:
        assert x == y

recursive_detach(x)

Detach all tensors in a nested structure.

Source code in bionemo/testing/torch.py

def recursive_detach(x):
    """Detach all tensors in a nested structure."""
    if isinstance(x, torch.Tensor):
        return x.detach().cpu()
    elif isinstance(x, (list, tuple)):
        return type(x)(recursive_detach(item) for item in x)
    elif isinstance(x, dict):
        return {key: recursive_detach(value) for key, value in x.items()}
    else:
        return x