Fix BitNet packed-weight unpacking dtype (`F.linear` dtype mismatch)
Author: jiqing-feng
Pull request: yes

## Problem

Loading `microsoft/bitnet-b1.58-2B-4T` in `bfloat16` and running a forward pass crashes:

```
RuntimeError: expected mat1 and mat2 to have the same dtype, but got: BFloat16 != unsigned char
```

## Reproduction

```python
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "microsoft/bitnet-b1.58-2B-4T"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, dtype=torch.bfloat16).to("cuda")

messages = [
    {"role": "system", "content": "You are a helpful AI assistant."},
    {"role": "user", "content": "How are you?"},
]
prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
outputs = model.generate(**inputs, max_new_tokens=32)  # crashes before this fix
print(tokenizer.decode(outputs[0][inputs["input_ids"].shape[-1]:], skip_special_tokens=True))
```

## Cause

For `autobitlinear`/`offline` checkpoints, weights are unpacked in
`BitNetDeserialize.convert` via `unpack_weights(weight_uint8, dtype=target_dtype)`.
`target_dtype` was taken from the *packed* tensor (`uint8`), so the unpacked weights
stayed `uint8`. `AutoBitLinear.forward` then calls `F.linear` without re-casting,
giving the `bfloat16 != uint8` mismatch.

## Fix

Unpack into the module's compute dtype (`weight_scale.dtype`):

```python
if actual_out * VALUES_PER_ITEM == expected_out:
    needs_unpacking = True
    if hasattr(module, "weight_scale"):
        target_dtype = module.weight_scale.dtype
```

## Validation

With the fix, the model loads and generates coherent text:

```
I'm an AI, so I don't have feelings, but I'm here and ready to help you! How can I assist you today?
```


Comments:
This looks good, but can you explain why tests didn't catch it? Can we update/add a test to make sure we don't regress on this one?

> This looks good, but can you explain why tests didn't catch it? Can we update/add a test to make sure we don't regress on this one?

## Why the existing test didn't catch the bug

The old `test_unpack_packed_weights` could not reproduce the crash because of how
it built its input:

```python
# created without dtype, and packed weights cast to bfloat16
model.linear = AutoBitLinear(in_features=in_features, out_features=out_features, bias=False)
packed = pack_weights(original.clone().float()).to(torch.bfloat16)   # <-- cast to bfloat16
...
self.assertTrue(torch.equal(result["weight"], original))             # <-- only checks values
```

Two things masked the bug:

1. **The packed tensor was cast to `bfloat16`.** A real safetensors checkpoint
   keeps the packed BitNet weights as `uint8`. Because the input was already
   `bfloat16`, the buggy `target_dtype = weight.dtype` happened to pick the
   correct compute dtype, so the wrong code path produced a correct result.

2. **The assertion only checked shape and values, never the dtype.** It never
   asserted the unpacked dtype, nor that the result could be fed to `F.linear`
   against a `bfloat16` activation.

## The regression test

`test_unpack_packed_weights_uint8_checkpoint` mirrors the real load path:

- The `AutoBitLinear` module is built with `dtype=torch.bfloat16`, so its
  `weight_scale` (the compute dtype) is `bfloat16`.
- The packed weights stay `uint8`, exactly as stored in a checkpoint.
- It asserts the unpacked dtype equals the compute dtype (not `uint8`) and runs
  `F.linear` with a `bfloat16` activation to reproduce the original crash.

Verified: it fails on the pre-fix code (`torch.uint8 != torch.bfloat16`) and
passes with the fix.

**[For maintainers]** Suggested jobs to run (before merge)

run-slow: bitnet_integration

**CI Dashboard:** [View test results in Grafana](https://transformers-ci.lor-e.huggingface.cool/d/pytest-observability-by-pr/pytest-observability-branch?var-pr=46808)

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_46808). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.