[Bug] torch.export canonicalizes DistilBERT input embeddings to the tied output projector

[thealgebraist]

System Info

• Python 3.13.5
• PyTorch 2.13.0+cu130
• Transformers 5.13.0
• Model: distilbert-base-uncased
• Task: masked-language modeling
• Backend: Linux host with CUDA-enabled PyTorch

Who can help?

The issue concerns torch.export parameter provenance for a Transformers
masked-language model. I am leaving maintainer tagging blank.

Information

This text was written by codex and the bug description based on a tool written by codex.
The problem arises when using an officially supported masked-language-model
task with a direct torch.export call.

Reproduction

import torch
from transformers import AutoModelForMaskedLM, AutoTokenizer

model = AutoModelForMaskedLM.from_pretrained(
    "distilbert-base-uncased").eval()
tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
inputs = tokenizer("the quick brown fox", return_tensors="pt")

exported = torch.export.export(
    model,
    (inputs.input_ids,),
    {"attention_mask": inputs.attention_mask},
)

for node in exported.graph_module.graph.nodes:
    if node.op == "call_function" and "embedding" in str(node.target):
        print("embedding operand:", node.args[0])

print(exported.graph_signature.inputs_to_parameters)

The native model uses distilbert.embeddings.word_embeddings for input token
embeddings and vocab_projector for masked-language-model output. The model
ties these weights, so eager and exported inference agree for the original
checkpoint (max error = 0.0).

However, the exported graph declares both
p_distilbert_embeddings_word_embeddings_weight and
p_vocab_projector_weight, while the input embedding operation consumes
p_vocab_projector_weight. The word-embedding placeholder is not used by that
operation.

As a controlled semantic check, I replaced only the native input embedding
with an independent zero-initialized embedding after export. Native logits
changed substantially: maximum absolute change 19.816930770874023, mean
absolute change 3.182511806488037. The exported graph still identifies the
projector weight as the input embedding operand.

Expected behavior

Please either preserve the original parameter identity in the exported graph
or record an explicit alias/equality constraint for this canonicalization. If
this behavior is intentional because tied weights are specialized, it would be
helpful for the export contract or documentation to state that the graph is
not valid after the tied parameters are independently rewritten.

Thanks for taking a look!

[SparshGarg999]

Proposed Plan to Preserve Parameter Provenance on Tied Embeddings in torch.export

When calling torch.export.export() on DistilBertForMaskedLM, weight tying between word_embeddings and vocab_projector causes torch.export graph tracing to canonicalize the input embedding operation to p_vocab_projector_weight.

Fix:
Add weight tying metadata preservation hints or export alias attributes in src/transformers/models/distilbert/modeling_distilbert.py so torch.export retains input embedding parameter identity.

[thealgebraist]

Thats a whole lot of deletions, is the code being moved elsewhere or ?