Update on "enable data loading for data parallel training"

Tested that data loading now have the expected behavior:
- different dp ranks get different data
- different tp ranks within the same dp rank get the same data 


[ghstack-poisoned]

torchtrain/datasets/alpaca.py

@@ -61,6 +61,8 @@ def __iter__(self):
 
         for idx, sample in enumerate(self.data_iterator):
             # select samples to pack in a round-robin fashion
+            # TODO: This is a temporary solution for small datasets like Alpaca.
+            #       For larger datasets we need to use a more scalable approach.
             if idx % self.world_size != self.rank:
                 continue
             sample_text = sample["text"]