mirror of
https://github.com/elder-plinius/OBLITERATUS.git
synced 2026-08-17 16:37:30 +02:00
37 lines
1.3 KiB
Python
37 lines
1.3 KiB
Python
"""Real CUDA and bitsandbytes placement/operation probes."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from obliteratus import device
|
|
|
|
|
|
pytestmark = pytest.mark.gpu
|
|
|
|
|
|
def test_cuda_discovery_dtype_placement_and_operation():
|
|
if not torch.cuda.is_available():
|
|
pytest.skip("requires a CUDA runner; set ENABLE_CUDA_GATE and attach the cuda label")
|
|
assert device.is_cuda()
|
|
assert device.get_device("auto") == "cuda"
|
|
assert device.default_dtype("cuda") is torch.float16
|
|
tensor = torch.arange(16, device="cuda", dtype=torch.float32).reshape(4, 4)
|
|
result = tensor @ tensor.T
|
|
assert result.device.type == "cuda"
|
|
assert torch.isfinite(result).all()
|
|
|
|
|
|
def test_bitsandbytes_quantization_operation():
|
|
if not torch.cuda.is_available():
|
|
pytest.skip("requires a CUDA runner; set ENABLE_CUDA_GATE and attach the cuda label")
|
|
bnb = pytest.importorskip("bitsandbytes", reason="install locked bitsandbytes>=0.46.1")
|
|
assert device.supports_bitsandbytes("cuda")
|
|
source = torch.randn(16, 16, device="cuda", dtype=torch.float16)
|
|
quantized, state = bnb.functional.quantize_4bit(source, quant_type="nf4")
|
|
restored = bnb.functional.dequantize_4bit(quantized, state)
|
|
assert restored.device.type == "cuda"
|
|
assert restored.shape == source.shape
|
|
assert torch.isfinite(restored).all()
|