fix(replay): align BSHD data with Megatron token order

[Shikhargupta]

Summary

• convert BSHD replay tensors from rollout [B, S, ...] order to Megatron [S, B, ...] order before flattening
• apply sequence-parallel sharding on the sequence dimension, before flattening sequence and batch
• add CPU regression coverage for multi-sample and TP=2 layouts

Why this is needed

For B=2, S=4, a replay tensor containing:

sample 0: 0 1 2 3
sample 1: 4 5 6 7

was flattened directly as [0, 1, 2, 3, 4, 5, 6, 7]. Megatron's BSHD hidden states are flattened from [S, B, ...], so the corresponding route order is [0, 4, 1, 5, 2, 6, 3, 7]. With multiple samples, the old order can therefore replay a token's routing decision onto a different token.

For sequence parallelism with TP=2, slicing the already-flattened tensor also selected all of sample 0 on rank 0. The correct rank-local layouts are [0, 4, 1, 5] and [2, 6, 3, 7].

Validation

• pytest tests/fast/backends/training_utils/test_replay_data.py -q — 6 passed
• Miles pre-commit hooks on the changed files — passed
• git diff --check — passed

Related work

#2698 changes per-sample BSHD padding in the same path. This PR is independent: it fixes token ordering after padding/CP slicing and adds focused ordering tests.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.