Support num_epochs in TrainingConfig

Is your feature request related to a problem? Please describe.

The current implementation relies on iteration-based training (train_iters). While precise for step-counting, it lacks a native mechanism to ensure training completes exactly at the end of a full data pass. Manually calculating the total steps required to cover a dataset is prone to errors and requires constant recalculation whenever the global_batch_size or dataset composition changes.

Describe the solution you'd like

I am requesting the addition of a --num_epochs parameter. This feature should: Automatically calculate the equivalent train_iters based on the total sample count of the loaded dataset and the global_batch_size. Support fractional values (e.g., 0.5 epochs).

Describe alternatives you've considered

The current workaround involves manual calculation using the following logic:

\text{train\_iters} = \frac{\text{total\_tokens\_to\_train}}{\text{batch\_size} \times \text{seq\_len}}

which is very inaccurate

---

Thanks!