BitEmbedding: the bit-bias embedding as a plain module

[AmitMY]

BitEmbedding: the bit-bias embedding as a plain module

PatchedBitEmbeddings puts the bit-bias (a byte table plus a zero-initialized projection of each byte's 8 bits) behind a tieable façade weight, refreshed by hooks, to patch HF models' embeddings.

WeLT on Megatron-Bridge owns its byte embeddings (no HF tying) and runs under torch.compile, so it uses the same math as a plain module. BitEmbedding moves that module here, so WeLT can import it instead of keeping its own copy (sign/WeLT ByteEmbedding):

• BitEmbedding(weight): the table, initialized from weight, plus bit_proj.
• folded_weight(): the equivalent plain table, e.g. to export to inference engines.

Tests

A new test checks that the embedding starts as its table, that both parameters get gradients, and that the bits are folded into the table correctly. All 14 byte-embedding tests pass, and ruff is clean.

🤖 Generated with Claude Code

https://claude.ai/code/session_015Xm62qoeaTunFSU5GYCQ1f

---

Note

Low Risk
Additive API and tests only; existing patched HF embedding path is untouched.

Overview
Adds BitEmbedding, a hook-free variant of the byte bit-bias embedding used by PatchedBitEmbeddings, so stacks like Megatron that own embeddings and use torch.compile can use the same math without the tieable façade and refresh hooks.

The module keeps a 256×D table plus a zero-initialized D×8 bit_proj, folds bits into the table via folded_weight(), and looks up with F.embedding. PatchedBitEmbeddings and HF patching are unchanged.

A new unit test checks identity with the initial table, gradients on both parameters, and correct folding when bit_proj is non-zero.

^{Reviewed by Cursor Bugbot for commit e5377cd. Bugbot is set up for automated code reviews on this repo. Configure here.}

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit e5377cd. Configure here.}

[cursor]

Mixed devices in BitEmbedding init

Medium Severity

BitEmbedding clones weight onto that tensor's device but allocates bit_proj and the bits buffer on the default device. A non-CPU table therefore leaves the module split across devices, and folded_weight only casts dtype, so the first forward raises.

Additional Locations (1)

• utf8_tokenizer/byte_embeddings.py#L162-L163

 

^{Reviewed by Cursor Bugbot for commit e5377cd. Configure here.}