fix: add transformer. prefix in GPT-2 and Falcon/Bloom preprocess_weights (#113)

## Problem

HF safetensors for GPT-2 and Bloom store layer weights **without** the
outer `transformer.` prefix (e.g. `h.N.*`, `wte.*`, `ln_f.*`,
`word_embeddings.*`). Our ONNX initialisers include it because the outer
module attribute is `self.transformer`. As a result,
`preprocess_weights()` produced keys that never matched any ONNX
initialiser — every weight was left unloaded, and ORT rejected the model
with `Node input 'transformer.wte.weight' is not a graph input,
initializer, or output of a previous node.`

## Root Cause

- **GPT-2 / OpenAI-GPT / GPT-SW3**: safetensors keys are `h.N.attn.*`,
`wte.*`, `wpe.*`, `ln_f.*` (no prefix). GPT-Neo and GPT-BigCode already
include `transformer.`.
- **Bloom**: safetensors keys are `h.N.*`, `word_embeddings.*`, `ln_f.*`
(no prefix).

## Fix

- **`gpt2.py`**: in step 1 of the weight-normalisation loop, add
`transformer.` for plain GPT-2/GPT-SW3 keys (biogpt.\*, model.\*,
output_projection.\* and lm_head.\* are handled separately).
- **`falcon.py`**: after the Falcon/Bloom QKV-split and rename loop,
re-key entries that lack the `transformer.` or `lm_head.` prefix.

## Testing

Golden tests confirmed passing after the fix:
```
tests/e2e_golden_test.py::TestL4CheckpointVerified::test_prefill_argmax_matches_golden[text-generation/gpt2] PASSED
tests/e2e_golden_test.py::TestL4CheckpointVerified::test_prefill_argmax_matches_golden[text-generation/bloom-560m] PASSED
```

Also removes now-stale xfail/skip marks from `integration_test.py`:
- gpt2 xfail `tie_word_embeddings graph reference issue in ORT`
- bloom-560m skip `word_embeddings_layernorm not implemented`

All 2277 unit tests pass. Lintrunner clean.

Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: justinchuby

src/mobius/models/falcon.py

@@ -577,6 +577,15 @@ def preprocess_weights(
             new_key = new_key.replace(".mlp.dense_4h_to_h.", ".mlp.down_proj.")
             new_state_dict[new_key] = value
 
+        # HF Falcon / Bloom safetensors omit the "transformer." prefix
+        # (e.g. "h.N.*", "word_embeddings.*", "ln_f.*").  ONNX initialiser
+        # names include it because the outer attribute is ``self.transformer``.
+        # ``lm_head.*`` is a top-level attribute and stays as-is.
+        new_state_dict = {
+            (k if k.startswith(("transformer.", "lm_head.")) else "transformer." + k): v
+            for k, v in new_state_dict.items()
+        }
+
         # Handle weight tying
         if self.config.tie_word_embeddings:
             embed_key = "transformer.word_embeddings.weight"

src/mobius/models/gpt2.py

@@ -123,6 +123,15 @@ def preprocess_weights(
                 name = "transformer.h." + name[len("biogpt.layers.") :]
             elif name.startswith("model.layers."):
                 name = "transformer.h." + name[len("model.layers.") :]
+            # GPT-2 / OpenAI-GPT / GPT-SW3 HF safetensors omit the
+            # "transformer." prefix (e.g. "h.N.*", "wte.*", "ln_f.*").
+            # GPT-Neo and GPT-BigCode already include it.
+            # biogpt.* / model.* / output_projection.* / lm_head.* are
+            # handled separately below and must not be prefixed here.
+            elif not name.startswith(
+                ("transformer.", "biogpt.", "model.", "output_projection.", "lm_head.")
+            ):
+                name = "transformer." + name
 
             # ── 2. Top-level embedding / norm renames ────────────────────────
             # OpenAI-GPT

tests/integration_test.py

@@ -97,9 +97,6 @@ def _model_accessible(model_id: str) -> bool:
         "openai-community/gpt2",
         False,
         id="gpt2",
-        marks=pytest.mark.xfail(
-            reason="tie_word_embeddings graph reference issue in ORT", strict=False
-        ),
     ),
     # OPT (learned positional embeddings)
     pytest.param(
@@ -113,10 +110,6 @@ def _model_accessible(model_id: str) -> bool:
         "bigscience/bloom-560m",
         False,
         id="bloom-560m",
-        marks=pytest.mark.skip(
-            reason="Bloom word_embeddings_layernorm not implemented "
-            "in FalconCausalLMModel — weights silently dropped"
-        ),
     ),
     # Falcon (ALiBi attention, multi-query)
     pytest.param(