[Afloat16]

Problem

StableEmbedding casts its embedding output to the current global default dtype before calling LayerNorm. After a normal model.half(), model.bfloat16() or model.double(), the normalization parameters have a different dtype and forward fails with a mixed-dtype error. Changing the global default dtype after constructing the layer can also break an otherwise unchanged model.

Change

Cast to the normalization parameters' actual dtype, then retain the existing conversion back to the embedding weight dtype. A normalization layer created in full precision stays in full precision unless the caller explicitly converts it. Its normal Module call and hooks are preserved.

Reproducer

import torch
import bitsandbytes as bnb

model = torch.nn.Sequential(
    bnb.nn.StableEmbedding(32, 16),
    torch.nn.Linear(16, 4),
).bfloat16()
loss = model(torch.tensor([[1, 2, 3]])).float().square().mean()
loss.backward()

On unchanged main this fails in LayerNorm with RuntimeError: mixed dtype (CPU). The patch completes forward and backward.

Validation

Base: a616886cacf86701514f0cb1d6be760f967481c0; exact patch: df9928c5f5547da1a41148903fd62e09ab8e3c91.

• The same 15 new tests on unchanged source: 9 failed, 6 passed. Patched: 15 passed.
• New tests plus the complete existing tests/test_modules.py on CPU: 102 passed, 2 skipped.
• Native Embedding + LayerNorm + Linear reference models match outputs, every parameter gradient, parameters and momentum buffers through three SGD updates for float16, bfloat16, float32 and float64, including both to(dtype) and convenience conversion methods.
• Controls cover padding-row gradients, pretrained mixed-dtype embeddings with full-precision normalization, changes to the global default dtype, and a serialized model/optimizer checkpoint followed by another training update.
• Built the current CPU native library from source with CMake 4.4.4 and GCC; Python 3.12.14 / PyTorch 2.15.0+cpu.
• Repository-pinned Ruff 0.14.3 checks and formatting, plus git diff --check, pass.

BNB_TEST_DEVICE=cpu OMP_NUM_THREADS=2 python -m pytest \
  tests/test_stable_embedding_dtype.py tests/test_modules.py -q

GPU/distributed execution and the complete repository suite were not run. No model-quality or throughput improvement is claimed.

Checked all-state StableEmbedding/dtype/half/norm issues and PRs, the complete 132-entry module history, the latest PR catalog and current fork branches. The actual #87 patch introduces constructor dtype/device support and the global-default cast; it does not address subsequent module conversions or changes to the global default.