Token X-ray

The gradient weight each token got in one real training step, read from each library's own loop.

pip install token-xray-ml Start with the equation

Every token gets a vote

Most language-model trainers optimise the average loss over the tokens in a batch. With N target tokens, each one's gradient is scaled by exactly the same amount:

On a radiograph of one training step, where each cell is a token and its brightness is the weight it received, that is an even, featureless grey. Featureless is healthy. Rows are sequences; their length is the number of tokens they contribute.

Splitting the batch should change nothing

Large batches rarely fit in memory, so trainers split them into micro-batches and add up the gradients. Done right, the radiograph cannot tell the difference.

  1. The transformers Trainer, four accumulation steps

    Four micro-batches through the real Trainer: every token within 2×10−7 of its intended weight.

  2. Average each micro-batch first, then average the averages

    Small micro-batches now count as much as big ones. The real Trainer of transformers 4.45.2.

  3. Split across two processes, average on each

    With average_tokens_across_devices off, the rank with fewer tokens gets louder. The default is flat.

  4. Keep only the first micro-batch

    The other three micro-batches go dark. Emulates a DeepSpeed ZeRO offload bug (Limozin et al., 2026) that acts after backward, out of sight.

Some objectives tilt the vote on purpose

Reinforcement learning on language models scores whole answers. Six answers to one prompt, three right and three wrong, and three ways to turn their rewards into token weights. Brightness is how hard each token is pushed; the label says which way.

Small settings rewrite the objective

Each is a setting that quietly changes which tokens count, and each shows up on the film. All three are real runs of the library as configured.

Documents that read each other

Packing documents into one row saves padding, as long as they cannot see each other. Each grid gives one document a loss (row) and shows the gradient every other document received (column), which should be zero. Whether a leak exists does not depend on the model; its size does, and these sizes are for this untrained probe.

What preference and RL trainers apply

The same X-ray runs through TRL's trainers. These objectives are signed: a bright cell is a token the step pushes hard, up for an answer the loss favours and down for the others. The row label says which way. A dark cell received no gradient at all.

The completions are scripted tracer sequences with fixed rewards, so the trainer's own code does the slicing, masking, scoring and reduction, and every number can be checked against its equation. The sizes (the DPO margin, GRPO's ratios and clipping) depend on the model's state: these are for this untrained probe.

X-ray your own training step

Write a function that batches the probe's sequences exactly as your trainer does and runs one backward pass. Token X-ray swaps in a tiny model whose every token is a unique tracer and reads back the weight each one received. 

from token_xray import xray

def my_step(model, sequences):
    batch = my_collator(sequences)       # your code
    loss = my_loss(model, batch)         # your code
    loss.backward()

report = xray(my_step, lengths=[5, 9, 3, 12, 7, 4, 10, 6])
report.assert_token_mean()               # fails loudly if any token is off

Every check is mutation-tested: scripts/mutate.py plants 46 faults one at a time, 38 in the instrument and 8 realistic bugs in transformers and TRL. All are caught, each of the 54 tests catches at least one, and it runs on every push to main.