From 42530450deaaef630f95a89bbeaac07a9118177c Mon Sep 17 00:00:00 2001 From: Afloat Date: Fri, 13 Nov 2026 03:34:37 +0800 Subject: [PATCH] Forward StableEmbedding factory device and dtype by keyword --- bitsandbytes/nn/modules.py | 4 +- tests/test_stable_embedding_constructor.py | 110 +++++++++++++++++++++ 2 files changed, 112 insertions(+), 2 deletions(-) create mode 100644 tests/test_stable_embedding_constructor.py diff --git a/bitsandbytes/nn/modules.py b/bitsandbytes/nn/modules.py index ebc0b0943..f6ba918af 100644 --- a/bitsandbytes/nn/modules.py +++ b/bitsandbytes/nn/modules.py @@ -92,8 +92,8 @@ def __init__( scale_grad_by_freq, sparse, _weight, - device, - dtype, + device=device, + dtype=dtype, ) self.norm = torch.nn.LayerNorm(embedding_dim, device=device) GlobalOptimManager.get_instance().register_module_override(self, "weight", {"optim_bits": 32}) diff --git a/tests/test_stable_embedding_constructor.py b/tests/test_stable_embedding_constructor.py new file mode 100644 index 000000000..2300c841f --- /dev/null +++ b/tests/test_stable_embedding_constructor.py @@ -0,0 +1,110 @@ +import copy + +import pytest +import torch + +from bitsandbytes.nn import StableEmbedding + + +@pytest.mark.parametrize("device", [None, "cpu", torch.device("cpu")]) +@pytest.mark.parametrize("dtype", [None, torch.float16, torch.bfloat16, torch.float32, torch.float64]) +def test_stable_embedding_factory_kwargs_train(device, dtype): + torch.manual_seed(41) + embedding = StableEmbedding(19, 8, padding_idx=0, device=device, dtype=dtype) + expected_dtype = dtype or torch.get_default_dtype() + assert embedding.weight.device.type == "cpu" + assert embedding.weight.dtype == expected_dtype + assert embedding.weight.requires_grad + assert embedding.norm.weight.device == embedding.weight.device + assert embedding.norm.weight.dtype == torch.get_default_dtype() + + reference = torch.nn.Embedding(19, 8, padding_idx=0, device=device, dtype=dtype) + reference.load_state_dict({"weight": embedding.weight.detach().clone()}) + norm = copy.deepcopy(embedding.norm) + head = torch.nn.Linear(8, 3, dtype=expected_dtype) + reference_head = copy.deepcopy(head) + parameters = list(embedding.parameters()) + list(head.parameters()) + reference_parameters = list(reference.parameters()) + list(norm.parameters()) + list(reference_head.parameters()) + optimizer = torch.optim.SGD(parameters, lr=0.04, momentum=0.9) + reference_optimizer = torch.optim.SGD(reference_parameters, lr=0.04, momentum=0.9) + initial_weight = embedding.weight.detach().clone() + tokens = torch.tensor([[0, 1, 2, 1], [3, 4, 0, 3]]) + targets = torch.tensor([[0, 1, 2, 0], [1, 2, 0, 1]]) + + for _ in range(3): + optimizer.zero_grad() + reference_optimizer.zero_grad() + actual = head(embedding(tokens)) + expected = reference_head(norm(reference(tokens).to(norm.weight.dtype)).to(expected_dtype)) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + torch.nn.functional.cross_entropy(actual.float().reshape(-1, 3), targets.reshape(-1)).backward() + torch.nn.functional.cross_entropy(expected.float().reshape(-1, 3), targets.reshape(-1)).backward() + for parameter, reference_parameter in zip(parameters, reference_parameters): + assert parameter.grad is not None + torch.testing.assert_close(parameter.grad, reference_parameter.grad, rtol=0, atol=0) + assert torch.count_nonzero(embedding.weight.grad[0]) == 0 + optimizer.step() + reference_optimizer.step() + for parameter, reference_parameter in zip(parameters, reference_parameters): + torch.testing.assert_close(parameter, reference_parameter, rtol=0, atol=0) + torch.testing.assert_close( + optimizer.state[parameter]["momentum_buffer"], + reference_optimizer.state[reference_parameter]["momentum_buffer"], + rtol=0, + atol=0, + ) + assert not torch.equal(embedding.weight[1:], initial_weight[1:]) + assert torch.count_nonzero(embedding.weight[0]) == 0 + + +@pytest.mark.parametrize("dtype", [None, torch.float32, torch.bfloat16]) +def test_stable_embedding_factory_kwargs_meta(dtype): + embedding = StableEmbedding(19, 8, device="meta", dtype=dtype) + assert embedding.weight.device.type == "meta" + assert embedding.norm.weight.device.type == "meta" + assert embedding.weight.dtype == (dtype or torch.get_default_dtype()) + assert embedding.weight.requires_grad + + +@pytest.mark.parametrize("device", ["cpu", torch.device("cpu")]) +def test_stable_embedding_pretrained_weight_remains_trainable(device): + weight = torch.arange(19 * 8, dtype=torch.float64).reshape(19, 8) / 100 + embedding = StableEmbedding(19, 8, _weight=weight.clone(), device=device, dtype=torch.float32) + reference = torch.nn.Embedding(19, 8, _weight=weight.clone(), device=device, dtype=torch.float32) + # Factory dtype does not recast an explicitly supplied pretrained weight. + assert embedding.weight.dtype == reference.weight.dtype == torch.float64 + assert embedding.weight.requires_grad == reference.weight.requires_grad + embedding(torch.tensor([1, 2, 3])).square().sum().backward() + assert embedding.weight.grad is not None + assert torch.count_nonzero(embedding.weight.grad) > 0 + + +@pytest.mark.parametrize( + "options", + [ + {"max_norm": 0.5, "norm_type": 1.0}, + {"scale_grad_by_freq": True}, + {"sparse": True}, + ], +) +def test_stable_embedding_constructor_preserves_embedding_options(options): + torch.manual_seed(17) + embedding = StableEmbedding(19, 8, padding_idx=-1, device="cpu", dtype=torch.float32, **options) + reference = torch.nn.Embedding(19, 8, padding_idx=-1, device="cpu", dtype=torch.float32, **options) + reference.load_state_dict({"weight": embedding.weight.detach().clone()}) + norm = copy.deepcopy(embedding.norm) + tokens = torch.tensor([1, 1, 2, 18]) + actual = embedding(tokens) + expected = norm(reference(tokens)) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + coefficients = torch.arange(actual.numel(), dtype=actual.dtype).reshape_as(actual) + (actual * coefficients).sum().backward() + (expected * coefficients).sum().backward() + torch.testing.assert_close(embedding.weight, reference.weight, rtol=0, atol=0) + actual_gradient = embedding.weight.grad + expected_gradient = reference.weight.grad + if options.get("sparse"): + assert actual_gradient.is_sparse + actual_gradient = actual_gradient.to_dense() + expected_gradient = expected_gradient.to_dense() + torch.testing.assert_close(actual_gradient, expected_gradient, rtol=0, atol=0)