In case anyone in the future is reading this, the above does work pretty well, but only for training. For evaluation, the (Distributed)LengthGroupedSampler is not used. Furthermore, even if it was, by nature of throwing some data away in the collator, we skip some of the evaluation set (which is a no-no for comparing results between methods).
I provided an example implementation of a batch sampler that groups based on a provided feature in the following comment:
Option for `(Distributed)LengthGroupedSampler` to treat groups as a hard constraint · Issue #12995 · huggingface/transformers · GitHub 5
It also requires a change to Transformers itself to support a batch sampler, which is in a PR linked to that issue.