[Bug Report] `convert_gpt_oss_weights` raises `TypeError: 'Tensor' object is not subscriptable` on MXFP4 experts

[hijohnnylin]

HookedTransformer.from_pretrained_no_processing("openai/gpt-oss-20b") fails during weight conversion.
On an MXFP4 gpt-oss checkpoint, experts.gate_up_proj is a triton_kernels.tensor.Tensor — the packed
4-bit payload plus its scales — not a torch.Tensor, so the converter's gate_up_proj[e, :, ::2] slice
raises. The wrapper class is also named Tensor, which makes the error look like a torch bug.

Reproduce

import torch
from transformer_lens import HookedTransformer

HookedTransformer.from_pretrained_no_processing("openai/gpt-oss-20b", device="cuda", dtype=torch.bfloat16)

Actual

  File ".../transformer_lens/loading_from_pretrained.py", line 1991, in get_pretrained_state_dict
    state_dict = convert_gpt_oss_weights(hf_model, cfg)
  File ".../transformer_lens/pretrained/weight_conversions/openai.py", line 92, in convert_gpt_oss_weights
    state_dict[f"blocks.{l}.mlp.experts.{e}.W_gate.weight"] = gate_up_proj[
                                                             ^^^^^^^^^^^^^
TypeError: 'Tensor' object is not subscriptable

The object being sliced, loaded through plain transformers:

>>> from transformers import AutoModelForCausalLM
>>> import torch
>>> m = AutoModelForCausalLM.from_pretrained("openai/gpt-oss-20b", dtype=torch.bfloat16, device_map="cuda")
>>> g = m.model.layers[0].mlp.experts.gate_up_proj
>>> type(g).__module__ + "." + type(g).__name__
'_gpt_oss_triton_kernels_cuda_8a730d9.tensor.Tensor'   # triton_kernels, fetched via the `kernels` hub
>>> isinstance(g, torch.Tensor)
False
>>> g[0, :, ::2]
TypeError: 'Tensor' object is not subscriptable
>>> [a for a in ("data", "storage", "shape", "dtype") if hasattr(g, a)]
['data', 'storage', 'shape', 'dtype']

Note this happens with dtype=torch.bfloat16: the MXFP4 weights stay packed in the triton-kernels
wrapper rather than being materialized as bf16 tensors. transformers itself loads and runs the
checkpoint fine, so the same script that produced the object above can be used as a reference forward.

Expected

Either a converted HookedTransformer, or a clear error saying MXFP4 gpt-oss checkpoints are not
supported by this path — 'Tensor' object is not subscriptable sends the reader looking at torch.

Suggested fix

In convert_gpt_oss_weights, unwrap or dequantize the expert weights before slicing. The wrapper
exposes .storage / .data, and transformers can also be asked for dequantized weights
(dequantize=True on the MXFP4 quantization config) so the converter always sees plain tensors. At
minimum, an isinstance(gate_up_proj, torch.Tensor) guard raising NotImplementedError and naming
MXFP4 would make the failure legible.

Environment

transformer_lens	3.6.0, and main at 69e98ab — same TypeError, same slice, no MXFP4 handling in that converter on either
transformers	5.14.1
torch	2.13.0
python	3.11.10
GPU	NVIDIA A100 80GB PCIe, driver 550.127.05
checkpoint	openai/gpt-oss-20b (GptOssForCausalLM, MXFP4 experts)