Skip to content

Convert Pytest to Unittest for tests under test/dtypes/ #2216

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 4 commits into
base: main
Choose a base branch
from
Open

Convert Pytest to Unittest for tests under test/dtypes/ #2216

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
7c8978c
[migrate from pytest to unittest] test/dtypes/test_nf4.py
yuanjua May 16, 2025
ff16dea
[migrate from pytest to unittest] test/dtypes/test_uintx.py
yuanjua May 16, 2025
bf0d6e0
[migrate from pytest to unittest] test/dtypes/test_bitpacking.py
yuanjua May 16, 2025
f2579f2
run ruff fix
yuanjua May 16, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
run ruff fix
  • Loading branch information
@yuanjua
yuanjua committed May 16, 2025
commit f2579f208d2c4dd93fb66543c6632221c6b4070f
27 changes: 15 additions & 12 deletions test/dtypes/test_bitpacking.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,17 +4,17 @@
# This source code is licensed under the BSD 3-Clause license found in the
# LICENSE file in the root directory of this source tree.
import unittest
import torch
from torch.utils._triton import has_triton

from torchao.dtypes.uintx.bitpacking import pack, pack_cpu, unpack, unpack_cpu

import torch
from torch.testing._internal.common_utils import (
TestCase,
instantiate_parametrized_tests,
parametrize,
run_tests,
)
from torch.utils._triton import has_triton

from torchao.dtypes.uintx.bitpacking import pack, pack_cpu, unpack, unpack_cpu

bit_widths = (1, 2, 3, 4, 5, 6, 7)
dimensions = (0, -1, 1)
Expand All @@ -23,7 +23,7 @@
class TestBitpacking(TestCase):
def setUp(self):
torch._dynamo.reset()

def tearDown(self):
torch._dynamo.reset()

Expand All @@ -37,17 +37,17 @@ def test_CPU(self, bit_width, dim):
unpacked = unpack_cpu(packed, bit_width, dim=dim)
self.assertTrue(unpacked.allclose(test_tensor))


@parametrize("bit_width", bit_widths)
@parametrize("dim", dimensions)
@unittest.skipIf(not torch.cuda.is_available(), reason="CUDA not available")
def test_GPU(self, bit_width, dim):
test_tensor = torch.randint(0, 2**bit_width, (32, 32, 32), dtype=torch.uint8).cuda()
test_tensor = torch.randint(
0, 2**bit_width, (32, 32, 32), dtype=torch.uint8
).cuda()
packed = pack(test_tensor, bit_width, dim=dim)
unpacked = unpack(packed, bit_width, dim=dim)
self.assertTrue(unpacked.allclose(test_tensor))


@parametrize("bit_width", bit_widths)
@parametrize("dim", dimensions)
@unittest.skipIf(not torch.cuda.is_available(), reason="CUDA not available")
Expand All @@ -56,12 +56,13 @@ def test_compile(self, bit_width, dim):
torch._dynamo.config.specialize_int = True
torch.compile(pack, fullgraph=True)
torch.compile(unpack, fullgraph=True)
test_tensor = torch.randint(0, 2**bit_width, (32, 32, 32), dtype=torch.uint8).cuda()
test_tensor = torch.randint(
0, 2**bit_width, (32, 32, 32), dtype=torch.uint8
).cuda()
packed = pack(test_tensor, bit_width, dim=dim)
unpacked = unpack(packed, bit_width, dim=dim)
self.assertTrue(unpacked.allclose(test_tensor))


# these test cases are for the example pack walk through in the bitpacking.py file
@unittest.skipIf(not torch.cuda.is_available(), reason="CUDA not available")
def test_pack_example(self):
Expand All @@ -70,12 +71,13 @@ def test_pack_example(self):
).cuda()
shard_4, shard_2 = pack(test_tensor, 6)
print(shard_4, shard_2)
assert torch.tensor([0, 105, 151, 37], dtype=torch.uint8).cuda().allclose(shard_4)
assert (
torch.tensor([0, 105, 151, 37], dtype=torch.uint8).cuda().allclose(shard_4)
)
assert torch.tensor([39, 146], dtype=torch.uint8).cuda().allclose(shard_2)
unpacked = unpack([shard_4, shard_2], 6)
self.assertTrue(unpacked.allclose(test_tensor))


def test_pack_example_CPU(self):
test_tensor = torch.tensor(
[0x30, 0x29, 0x17, 0x5, 0x20, 0x16, 0x9, 0x22], dtype=torch.uint8
Expand All @@ -87,6 +89,7 @@ def test_pack_example_CPU(self):
unpacked = unpack([shard_4, shard_2], 6)
assert unpacked.allclose(test_tensor)


instantiate_parametrized_tests(TestBitpacking)

if __name__ == "__main__":
Expand Down
37 changes: 18 additions & 19 deletions test/dtypes/test_uintx.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,14 @@
# This source code is licensed under the BSD 3-Clause license found in the
# LICENSE file in the root directory of this source tree.
import unittest

import torch
from torch.testing._internal.common_quantization import QuantizationTestCase
from torch.testing._internal.common_utils import (
instantiate_parametrized_tests,
parametrize,
run_tests,
)

from torchao.dtypes.uintx.uintx_layout import to_uintx
from torchao.quantization.quant_api import quantize_, uintx_weight_only
Expand All @@ -19,13 +26,6 @@
TORCH_VERSION_AT_LEAST_2_3,
TORCH_VERSION_AT_LEAST_2_5,
)
from torch.testing._internal.common_utils import (
instantiate_parametrized_tests,
parametrize,
run_tests,
)
from torch.testing._internal.common_quantization import QuantizationTestCase


# torch.uintx dtypes are introduced in 2.3
if TORCH_VERSION_AT_LEAST_2_3:
Expand Down Expand Up @@ -75,7 +75,8 @@ def tearDown(self):
@parametrize("group_size", group_sizes)
@unittest.skipIf(not torch.cuda.is_available(), reason="CUDA not available")
@unittest.skipIf(
not TORCH_VERSION_AT_LEAST_2_5, reason="only works with fix in the nightly build"
not TORCH_VERSION_AT_LEAST_2_5,
reason="only works with fix in the nightly build",
)
def test_uintx_quant_on_cpu_then_move_to_cuda(self, dtype, group_size):
scale = 512
Expand All @@ -88,18 +89,18 @@ def test_uintx_quant_on_cpu_then_move_to_cuda(self, dtype, group_size):
output_on_cuda = fp16_mod_on_cuda(test_input_on_cuda)
self.assertTrue(
torch.allclose(output_on_cpu, output_on_cuda.cpu(), atol=1.0e-3),
"The output of the model on CPU and CUDA should be close"
"The output of the model on CPU and CUDA should be close",
)


@parametrize("dtype", dtypes)
@parametrize("group_size", group_sizes)
@parametrize("device", devices)
@unittest.skipIf(not torch.cuda.is_available(), reason="CUDA not available")
@unittest.skipIf(
not TORCH_VERSION_AT_LEAST_2_5, reason="only works with fix in the nightly build"
not TORCH_VERSION_AT_LEAST_2_5,
reason="only works with fix in the nightly build",
)
def test_uintx_weight_only_model_quant(self,dtype, group_size, device):
def test_uintx_weight_only_model_quant(self, dtype, group_size, device):
scale = 512
fp16 = Linear16(scale, device)
quantize_(fp16, uintx_weight_only(dtype, group_size=group_size))
Expand All @@ -108,15 +109,15 @@ def test_uintx_weight_only_model_quant(self,dtype, group_size, device):
output = uintx.forward(test_input)
self.assertTrue(output is not None, "model quantization failed")


@parametrize("dtype", dtypes)
@parametrize("group_size", group_sizes)
@parametrize("device", devices)
@unittest.skipIf(not torch.cuda.is_available(), reason="CUDA not available")
@unittest.skipIf(
not TORCH_VERSION_AT_LEAST_2_5, reason="only works with fix in the nightly build"
not TORCH_VERSION_AT_LEAST_2_5,
reason="only works with fix in the nightly build",
)
def test_uintx_weight_only_quant(self,dtype, group_size, device):
def test_uintx_weight_only_quant(self, dtype, group_size, device):
input_float = torch.randn((1, 256), dtype=torch.float16, device=device)
mapping_type = MappingType.SYMMETRIC
eps = torch.finfo(torch.float32).eps
Expand Down Expand Up @@ -153,21 +154,19 @@ def test_uintx_weight_only_quant(self,dtype, group_size, device):
)
self.assertTrue(deqaunt is not None, "deqauntization failed")


@parametrize("dtype", dtypes)
@unittest.skipIf(not torch.cuda.is_available(), reason="Need CUDA available")
@unittest.skipIf(
not TORCH_VERSION_AT_LEAST_2_3, reason="sub byte dtype requires torch 2.3+"
)
def test_uintx_target_dtype(self,dtype):
def test_uintx_target_dtype(self, dtype):
from torchao.quantization.quant_api import uintx_weight_only

linear = torch.nn.Linear(128, 256, dtype=torch.bfloat16, device="cuda")
# make sure it runs
quantize_(linear, uintx_weight_only(dtype))
linear(torch.randn(1, 128, dtype=torch.bfloat16, device="cuda"))


@parametrize("dtype", dtypes)
@unittest.skipIf(not torch.cuda.is_available(), reason="Need CUDA available")
@unittest.skipIf(
Expand All @@ -183,7 +182,6 @@ def test_uintx_target_dtype_compile(self, dtype):
linear = torch.compile(linear)
linear(torch.randn(1, 128, dtype=torch.bfloat16, device="cuda"))


@parametrize("dtype", dtypes)
@unittest.skipIf(not torch.cuda.is_available(), reason="Need CUDA available")
@unittest.skipIf(
Expand Down Expand Up @@ -214,6 +212,7 @@ def test_uintx_model_size(self, dtype):
quantized_size = get_model_size_in_bytes(linear)
self.assertTrue(bf16_size * _dtype_to_ratio[dtype] == quantized_size)


instantiate_parametrized_tests(TestUintx)

if __name__ == "__main__":
Expand Down
Loading