Function test_embedding_error

tests/test_modules.py:377–413

(device, embedding_class, input_shape, embedding_dim, quant_storage)

    ids=lambda x: x.__name__ if inspect.isclass(x) else str(x),
)
def test_embedding_error(device, embedding_class, input_shape, embedding_dim, quant_storage):
    if device == "hpu":
        if embedding_class is bnb.nn.EmbeddingFP4:
            pytest.skip("FP4 is not supported on HPU")
        elif embedding_class is bnb.nn.EmbeddingNF4 and not is_supported_on_hpu("nf4", torch.float32, quant_storage):
            pytest.skip("This configuration is not supported on HPU")

    is_8bit = embedding_class is bnb.nn.Embedding8bit

    num_embeddings = 128

    src_weight = torch.rand((num_embeddings, embedding_dim), dtype=torch.float32)

    emb_base = nn.Embedding(
        num_embeddings=num_embeddings,
        embedding_dim=embedding_dim,
        _freeze=True,
        _weight=src_weight,
    )
    if is_8bit:
        e = embedding_class(num_embeddings=num_embeddings, embedding_dim=embedding_dim)
    else:
        e = embedding_class(num_embeddings=num_embeddings, embedding_dim=embedding_dim, quant_storage=quant_storage)

    e.load_state_dict(emb_base.state_dict())

    emb_base.to(device)
    e.to(device)

    input_tokens = torch.randint(low=0, high=num_embeddings, size=input_shape, device=device)

    torch.testing.assert_close(
        actual=e(input_tokens),
        expected=emb_base(input_tokens),
        atol=0.05 if is_8bit else 0.20,
        rtol=0.0,
    )


@pytest.mark.parametrize("device", get_available_devices())