Feature: add token-level gradient collection

Currently gradients are accumulated across sequences. We can accumulate targeted tokens using a labels mask but it would be nice to have the option to inspect individual gradients.

Low priority.