feat: group training rows by length to improve DP and batching efficiency

[Dogacel]

In multi-GPU or batch size > 1 training, training samples are padded to match the longest sample + multiple of 256 (to prevent extra compiles). However this slows down the overall training if samples are coming in irregular batches.

I.e. assume we get sequences [10, 4500], [30, 5000], in a 2-GPU setting. Since we have to pad to the longest sequence, we are paying 4500t + 5000t training time. However if you re-arrange them by sorting [10, 30], [4500, 5000]; training cost is almost halved, 30t + 5000t.

Tests

[TODO]

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: dfb3588f4b

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

Backfill seq_len for old offline manifests

When opening an offline dataset materialized by the previous schema-1 writer, manifest rows only contain split, data_id, and file, so indexing row["seq_len"] raises KeyError before replay/training can start. Since OFFLINE_SCHEMA_VERSION is unchanged and no migration/fallback derives the length from the saved input_ids, this breaks existing offline datasets even though they still advertise the supported version; please tolerate missing seq_len or bump/migrate the schema.

Useful? React with 👍 / 👎.