Fix data_load_start position (#1481)

# Fix incorrect data loading time measurement

This PR fixes the timing of data_loading_times measurement in
batch_generator. Previously, the timer started after calling
next(data_iterator), which excluded the actual data fetching time from
the measurement.
Now, the timer starts before the next() call to correctly capture the
full DataLoader latency.

Author: speed1313

torchtitan/train.py

@@ -373,13 +373,13 @@ def batch_generator(
         data_iterator = iter(data_iterable)
 
         while True:
+            data_load_start = time.perf_counter()
             try:
                 batch = next(data_iterator)
             except StopIteration as ex:
                 # If data runs out during gradient accumulation, that
                 # entire step will not be executed.
                 raise DataloaderStopIteration() from ex
-            data_load_start = time.perf_counter()
             input_dict, labels = batch
             ntokens_batch = labels.numel()
             self.ntokens_seen += ntokens_batch