Module neuralops.tests.test_compression_enhancements

Tests for compression engine enhancements (GPTQ, AWQ, SmoothQuant, Wanda, LoRA/QLoRA/DoRA).

Functions

def test_awq_quantizer()

Expand source code

def test_awq_quantizer():
    layer = nn.Linear(16, 16)
    awq = AWQQuantizer(layer, bits=4)
    x = torch.randn(8, 16)
    awq.observe(x)
    res = awq.quantize()
    assert res.shape == (16, 16)

def test_gptq_quantizer()

Expand source code

def test_gptq_quantizer():
    layer = nn.Linear(16, 16)
    q = GPTQQuantizer(layer, bits=4)
    x = torch.randn(8, 16)
    q.add_batch(x)
    res = q.quantize()
    assert res.shape == (16, 16)

def test_lora_merge_and_extract()

Expand source code

def test_lora_merge_and_extract():
    lin = nn.Linear(16, 16)
    orig_w = lin.weight.data.clone()
    A = torch.randn(4, 16)
    B = torch.randn(16, 4)
    merge_lora_into_linear(lin, A, B, scaling=1.0)
    assert not torch.equal(lin.weight.data, orig_w)

    diff = B @ A
    A_ext, B_ext = extract_lora_from_diff(diff, rank=4)
    assert torch.allclose(diff, B_ext @ A_ext, atol=1e-4)

def test_qlora_and_dora()

Expand source code

def test_qlora_and_dora():
    x = torch.randn(2, 16)
    qlora = QLoRALinear(16, 16, rank=4)
    out_q = qlora(x)
    assert out_q.shape == (2, 16)

    dora = DoRALinear(16, 16, rank=4)
    out_d = dora(x)
    assert out_d.shape == (2, 16)

def test_smoothquant()

Expand source code

def test_smoothquant():
    layer = nn.Linear(16, 16)
    act_scales = torch.ones(16) * 2.0
    scales = smoothquant_layer(layer, act_scales)
    assert scales.shape == (16,)

def test_wanda_pruning()

Expand source code

def test_wanda_pruning():
    layer = nn.Linear(16, 16)
    act_norm = torch.ones(16)
    mask = wanda_prune_layer(layer, act_norm, sparsity=0.5)
    assert mask.shape == (16, 16)
    assert torch.sum(layer.weight.data == 0) > 0