DistilBERT preparation and export

The preparation flow still exports a fixed-shape PyTorch program, but the wrapper is different from GPT-2. Instead of keeping final-token logits, it extracts logits at the [MASK] token:

    

        
        
class _MaskLogits(torch.nn.Module):
    def __init__(self, model, mask_token_id):
        super().__init__()
        self.model = model
        self.mask_token_id = mask_token_id

    def forward(self, input_ids, attention_mask):
        # DistilBERT returns [batch, sequence, vocabulary] logits.
        logits = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
        ).logits
        mask_positions = (
            input_ids == self.mask_token_id
        ).to(torch.int64).argmax(dim=1)
        batch_indices = torch.arange(
            input_ids.shape[0],
            device=input_ids.device,
        )
        return logits[batch_indices, mask_positions, :]

    

self.model is the original Hugging Face DistilBERT model. mask_positions contains one token position for each batch item. Pairing it with batch_indices selects one [MASK] position from each item and returns [batch, vocabulary] logits.

The input is constrained to one mask token per batch item, so argmax(dim=1) returns a fixed [batch] position tensor.

The export example prompt must contain the tokenizer’s mask token. In this module, you will keep it to one mask token so the wrapper has a single position to extract:

    

        
        
tokenizer = AutoTokenizer.from_pretrained(args.model_id, use_fast=True)
prompt = args.example_prompt or f"The Raspberry Pi is a {tokenizer.mask_token} computer."
if tokenizer.mask_token not in prompt:
    raise ValueError(f"Example prompt must contain {tokenizer.mask_token}")

    

DistilBERT uses a shorter default sequence length in this module:

    

        
        
parser.add_argument("--seq-len", type=int, default=64)

    

Run the following command to perform model export on your development machine:

    

        
        
module-1-transformer-inference-with-xnnpack/scripts/aot/prepare_distilbert_fillmask.py \
  --out module-1-transformer-inference-with-xnnpack/artifacts/distilbert

    

This writes distilbert_fillmask_exported.pt2 and model_spec.json. The exported program still takes input_ids and attention_mask, but the output is [B, V] logits for the single mask position rather than next-token logits for a decoder.