[FEA]: Elastic Deterministic Resumption #222

Is this a new feature, an improvement, or a change to existing functionality?

New Feature

How would you describe the priority of this feature request

Critical (currently preventing usage)

Please provide a clear description of problem this feature solves

Long-running training jobs may be interrupted by preemption or hardware failures and need to resume with a different number of GPUs. For example, a run checkpointed on 16 GPUs may only have 8 GPUs available when it restarts.

We need to resume data loading from the checkpoint while preserving the remaining shuffled sample sequence, without unintentionally repeating or skipping samples. Changes to data-parallel world size or dataloader worker count should not change which samples training consumes next.

This would make recovery more flexible and improve reproducibility when GPU availability changes.

Feature Description

As a user training on a shared GPU cluster, I want to restore an Energon dataloader checkpoint with a different data-parallel world size or worker count, so that I can continue the same logical data stream using the resources available.

Ideally, this would support:

• Redistributing the remaining sample stream across the new ranks and workers.
• Preserving sample order and shuffle progression, including across epoch boundaries.
• Reconstructing the same ordered global batches when the global batch size stays fixed.
• Preserving sample-level randomness for preprocessing or augmentation where supported.
• Documenting the guarantees and limitations for dataset mixing, packing, and buffered or prefetched samples.

The requested determinism concerns the data stream and batch composition; bitwise-identical model training across different GPU configurations is outside the scope of this request.

Is this already supported through a recommended checkpointing workflow, or is it planned for a future release?

Describe your ideal solution

Support topology-independent dataloader checkpoints that can be restored with a different data-parallel world size or number of workers.

The checkpoint should capture logical progress through the global data stream, including shuffle state and any state needed for dataset mixing, packing, and preprocessing. On restore, Energon would redistribute the remaining work across the available ranks and workers while preserving:

• The remaining global sample sequence, without unintended duplication or skipping relative to the checkpoint.
• Shuffle progression across epoch boundaries.
• The same ordered global batches when the global batch size is unchanged.
• Deterministic sample-level preprocessing and augmentation where supported.

Prefetched but unconsumed samples should be handled consistently so that changing worker counts does not alter the resumed stream.

Ideally, this would work through the existing save/restore APIs without requiring users to manually transform checkpoint state. Unsupported configurations should produce a clear error explaining which guarantee cannot be preserved.

For example, a job checkpointed on 16 GPUs should be able to resume on 8 GPUs and consume the same subsequent global batches, provided the global batch size and data-processing configuration remain unchanged.

Describe any alternatives you have considered

No response

Additional context

No response

Code of Conduct

• I agree to follow the Code of Conduct
• I have searched the open feature requests and have found no duplicates for this feature request