fix(megatron): correct rollout weight reload, resume offsets and sampler seed

[taking-lying-flat]

PR type

• Bug Fix
• New Feature
• Document Updates
• More Models or Datasets Support

PR information

Fixes #10098 and #10099 with changes to four production files.

Megatron colocated weight reload passes an integer device ordinal into vLLM's post-load processing, which expects a torch.device. The existing exception handler hides this failure. Normalize the device and enable strict error propagation for the Megatron caller. Older vLLM releases without the helper retain the compatibility skip; errors inside an existing helper propagate in strict mode. Other callers retain non-strict behavior.

The on-policy checkpoint counts generated training samples, while the dataloader indexes unrepeated prompts. Convert the offset at dataloader construction for GRPO/GKD, without modifying the persisted counter, and reject offsets not divisible by num_generations. For example, 32 consumed rollout rows with four generations correspond to eight prompt rows. GKD currently forces num_generations=1, so its conversion is a no-op.

Pass data_seed into the random sampler and use seed + epoch in both sharding branches, including length grouping. The optional sampler argument defaults to zero for existing direct callers.

Compatibility note: honoring data_seed changes the shuffle order used by old checkpoints, including non-on-policy callers of this sampler. Seed zero reproduces the old sampler permutation only if dataset ordering is identical; changing data_seed can also affect upstream dataset preparation. This patch does not restore rollout buffers or guarantee identical generation/RNG state after a resume.

Experiment results

Validation scripts were kept local; the PR contains production code only.

• 56 local regression cases passed: real PyTorch samplers and DataLoaders, two simulated DP ranks, both sharding branches, shuffled/unshuffled resume, cross-epoch resume, incomplete dataset tails, length grouping, seed behavior, validation/streaming preservation, device conversion, legacy helper absence, and strict/non-strict error propagation. The harness loads the actual production functions/classes in isolation; it does not launch a distributed Megatron trainer.
• Real vLLM/GPU integration passed on an NVIDIA RTX A1000 Laptop GPU: a small probe module's CPU parameter moves to CUDA for post-load processing and returns to CPU with the expected value; an intentional processing error propagates. PyTorch 2.13.0+cu130; installed vLLM development build 0.1.dev564+ga3561ef8e.d20260823.
• The initial 33-case suite against the original code produced 22 failures, 8 passes, and 3 GPU skips in the sandbox.
• pre-commit run --files swift/rlhf_trainers/utils.py swift/megatron/trainers/rollout_mixin.py swift/megatron/trainers/base.py swift/megatron/trainers/batch_sampler.py: passed.
• git diff --check: passed.

No full distributed Megatron training run or NPU hardware validation was performed.