Optimize data fetching latency for full-hetero encoder

Author: Yangsx-1

loongforge/data/encoder_strided_sampler.py

@@ -0,0 +1,147 @@
+# Copyright 2026 The LoongForge Authors.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Strided data access for full_hetero_dp encoder.
+
+Provides two implementations:
+- EncoderStridedSampler: batch_sampler for indexable (map-style) datasets.
+- EncoderStridedIterator: iterator-level filter for streaming (Energon/WebDataset) dataloaders.
+
+Both yield only the microbatches assigned to a specific PP rank,
+maintaining data consistency with the decoder by using the same
+step-relative position filtering logic.
+"""
+
+import threading
+import queue
+
+from megatron.legacy.data.data_samplers import MegatronPretrainingRandomSampler
+
+_SENTINEL = object()
+
+
+class PrefetchIterator:
+    """Prefetches items from a source iterator in a background thread.
+
+    Allows the next training step's data to be loaded while the current
+    step's encoder/decoder computation is running.
+    """
+
+    def __init__(self, source_iter, prefetch_count):
+        self._queue = queue.Queue(maxsize=prefetch_count)
+        self._source = source_iter
+        self._thread = threading.Thread(target=self._worker, daemon=True)
+        self._thread.start()
+
+    def _worker(self):
+        try:
+            while True:
+                item = next(self._source)
+                self._queue.put(item)
+        except StopIteration:
+            self._queue.put(_SENTINEL)
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        item = self._queue.get()
+        if item is _SENTINEL:
+            raise StopIteration
+        return item
+
+
+class EncoderStridedSampler:
+    """Yields only microbatches assigned to this PP rank's encoder.
+
+    Internally iterates the same index sequence as the decoder's sampler,
+    but only yields batches at positions belonging to this PP rank.
+
+    Processes items in chunks of num_real_microbatch (one step's worth)
+    to ensure the position pattern is always step-relative, staying in
+    sync with the decoder's DataLoader across steps.
+    """
+
+    def __init__(self, dataset, total_samples, consumed_samples, micro_batch_size,
+                 data_parallel_rank, data_parallel_size, data_sharding,
+                 pp_rank, tp_size, model_size, num_real_microbatch):
+        self.dataset = dataset
+        self.total_samples = total_samples
+        self.consumed_samples = consumed_samples
+        self.micro_batch_size = micro_batch_size
+        self.data_parallel_rank = data_parallel_rank
+        self.data_parallel_size = data_parallel_size
+        self.data_sharding = data_sharding
+        self.pp_rank = pp_rank
+        self.tp_size = tp_size
+        self.model_size = model_size
+        self.num_real_microbatch = num_real_microbatch
+
+    def __len__(self):
+        return self.total_samples
+
+    def __iter__(self):
+        base_sampler = MegatronPretrainingRandomSampler(
+            self.dataset,
+            total_samples=self.total_samples,
+            consumed_samples=self.consumed_samples,
+            micro_batch_size=self.micro_batch_size,
+            data_parallel_rank=self.data_parallel_rank,
+            data_parallel_size=self.data_parallel_size,
+            data_sharding=self.data_sharding,
+        )
+        start = self.pp_rank * self.tp_size
+        end = start + self.tp_size
+        # Process in step-sized chunks to keep position pattern step-relative
+        step_buffer = []
+        for batch in base_sampler:
+            step_buffer.append(batch)
+            if len(step_buffer) == self.num_real_microbatch:
+                for i, b in enumerate(step_buffer):
+                    if start <= (i % self.model_size) < end:
+                        yield b
+                step_buffer = []
+        # Yield remaining partial step
+        for i, b in enumerate(step_buffer):
+            if start <= (i % self.model_size) < end:
+                yield b
+
+
+class EncoderStridedIterator:
+    """Iterator-level strided filter for streaming (Energon/WebDataset) dataloaders.
+
+    Wraps an EnergonDataloader, consumes all microbatches from it but only
+    yields those assigned to this PP rank. Logically equivalent to
+    EncoderStridedSampler but for iterable (non-indexable) datasets.
+
+    Buffers microbatches in step-sized chunks (num_real_microbatch) to
+    maintain correct position-based assignment, identical to the logic in
+    EncoderStridedSampler.
+    """
+
+    def __init__(self, energon_dataloader, pp_rank, tp_size, model_size, num_real_microbatch):
+        self._source = energon_dataloader
+        self._pp_rank = pp_rank
+        self._tp_size = tp_size
+        self._model_size = model_size
+        self._num_real_microbatch = num_real_microbatch
+        self._gen = self._filter()
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        return next(self._gen)
+
+    def _filter(self):
+        start = self._pp_rank * self._tp_size
+        end = start + self._tp_size
+        step_buffer = []
+        while True:
+            batch = next(self._source)
+            step_buffer.append(batch)
+            if len(step_buffer) == self._num_real_microbatch:
+                for i, b in enumerate(step_buffer):
+                    if start <= (i % self._model_size) < end:
+                        yield b
+                step_buffer = []

loongforge/data/multimodal/dataloader_provider.py

@@ -440,3 +440,30 @@ def cyclic_iter(iter):
     while True:
         for x in iter:
             yield x
+
+
+def build_full_hetero_encoder_energon_iterator(
+    task_encoder, collator, pp_rank, tp_size, model_size, num_real_microbatch
+):
+    """Build encoder iterator for full_hetero_dp with Energon dataloader.
+
+    Creates a separate EnergonDataloader (same dataset config as decoder) and
+    wraps it with EncoderStridedIterator to yield only microbatches assigned
+    to this PP rank.
+    """
+    from loongforge.data.encoder_strided_sampler import EncoderStridedIterator, PrefetchIterator
+    from loongforge.train.initialize import get_num_micro_batches_per_decoder_dp
+
+    encoder_dataset = get_train_dataset(task_encoder)
+    encoder_dataloader = get_train_loader(encoder_dataset, collator)
+
+    strided_iter = EncoderStridedIterator(
+        energon_dataloader=encoder_dataloader,
+        pp_rank=pp_rank,
+        tp_size=tp_size,
+        model_size=model_size,
+        num_real_microbatch=num_real_microbatch,
+    )
+    _, encoder_rounds = get_num_micro_batches_per_decoder_dp()
+    prefetch_count = tp_size * encoder_rounds
+    return PrefetchIterator(strided_iter, prefetch_count=prefetch_count)

loongforge/train/pretrain/pretrain_vlm.py

@@ -213,6 +213,7 @@ def _create_mock_batch(reference_batch):
 deepstack_grad_list = []
 
 _cpu_offload_manager = None
+_encoder_data_iterator = None
 
 
 def get_cpu_offload_manager():
@@ -227,6 +228,15 @@ def get_cpu_offload_manager():
         )
     return _cpu_offload_manager
 
+def get_encoder_data_iterator():
+    """Return the encoder data iterator for full_hetero_dp mode."""
+    return _encoder_data_iterator
+
+def set_encoder_data_iterator(iterator):
+    """Set the encoder data iterator for full_hetero_dp mode."""
+    global _encoder_data_iterator
+    _encoder_data_iterator = iterator
+
 def get_embedding_list():
     """Return the global embedding list."""
     return embedding_list
@@ -442,6 +452,26 @@ def train_valid_test_dataset_provider(train_val_test_num_samples, vp_stage=None)
         train_data_iterator, valid_data_iterator, test_data_iterator = (
             build_sft_cyclic_iterators(train_ds, None, None, collator)
         )
+
+        # Build encoder-specific iterator for full_hetero_dp
+        if getattr(args, 'enable_full_hetero_dp', False):
+            from loongforge.train.sft.utils import (
+                build_full_hetero_encoder_data_iterator,
+            )
+            from loongforge.train.initialize import (
+                get_model_size, get_num_real_micro_batches_per_decoder_dp,
+            )
+            pp_rank = mpu.get_pipeline_model_parallel_rank()
+            tp_size = mpu.get_tensor_model_parallel_world_size()
+            model_size = get_model_size()
+            num_real_microbatch = get_num_real_micro_batches_per_decoder_dp()
+            encoder_iter = build_full_hetero_encoder_data_iterator(
+                train_ds, args.consumed_train_samples, collator,
+                pp_rank=pp_rank, tp_size=tp_size, model_size=model_size,
+                num_real_microbatch=num_real_microbatch,
+            )
+            set_encoder_data_iterator(encoder_iter)
+
         return train_data_iterator, None, None
 
     else:
@@ -455,6 +485,26 @@ def train_valid_test_dataset_provider(train_val_test_num_samples, vp_stage=None)
 
         collator = build_sft_data_collator(VLMPretrainCollator)
         train_dataloader = get_train_loader(train_dataset, collator)
+
+        # Build encoder-specific iterator for full_hetero_dp (Energon path)
+        if getattr(args, 'enable_full_hetero_dp', False):
+            from loongforge.data.multimodal.dataloader_provider import (
+                build_full_hetero_encoder_energon_iterator,
+            )
+            from loongforge.train.initialize import (
+                get_model_size, get_num_real_micro_batches_per_decoder_dp,
+            )
+            pp_rank = mpu.get_pipeline_model_parallel_rank()
+            tp_size = mpu.get_tensor_model_parallel_world_size()
+            model_size = get_model_size()
+            num_real_microbatch = get_num_real_micro_batches_per_decoder_dp()
+            encoder_iter = build_full_hetero_encoder_energon_iterator(
+                task_encoder, collator,
+                pp_rank=pp_rank, tp_size=tp_size,
+                model_size=model_size, num_real_microbatch=num_real_microbatch,
+            )
+            set_encoder_data_iterator(encoder_iter)
+
         return train_dataloader, None, None
 
 

loongforge/train/sft/utils.py

@@ -233,6 +233,51 @@ def build_sft_cyclic_iterators(
     return train_iter, valid_iter, test_iter
 
 
+def build_encoder_data_iterator(
+    dataset: "Dataset",
+    consumed_samples: int,
+    data_collator: DataCollatorForSupervisedDataset,
+    pp_rank: int,
+    tp_size: int,
+    model_size: int,
+    num_real_microbatch: int,
+):
+    """Build a DataLoader iterator for the encoder in full_hetero_dp mode.
+
+    Uses EncoderStridedSampler to yield only microbatches assigned to this PP rank,
+    avoiding unnecessary disk IO for microbatches handled by other ranks.
+    """
+    from loongforge.data.encoder_strided_sampler import EncoderStridedSampler
+
+    args = get_args()
+    batch_sampler = EncoderStridedSampler(
+        dataset,
+        total_samples=len(dataset),
+        consumed_samples=consumed_samples,
+        micro_batch_size=args.micro_batch_size,
+        data_parallel_rank=mpu.get_data_parallel_rank(),
+        data_parallel_size=mpu.get_data_parallel_world_size(),
+        data_sharding=args.data_sharding,
+        pp_rank=pp_rank,
+        tp_size=tp_size,
+        model_size=model_size,
+        num_real_microbatch=num_real_microbatch,
+    )
+    dataloader = DataLoader(
+        dataset,
+        batch_sampler=batch_sampler,
+        collate_fn=data_collator,
+        num_workers=args.num_workers,
+        pin_memory=True,
+        persistent_workers=True if args.num_workers > 0 else False,
+    )
+    from loongforge.data.encoder_strided_sampler import PrefetchIterator
+    from loongforge.train.initialize import get_num_micro_batches_per_decoder_dp
+    _, encoder_rounds = get_num_micro_batches_per_decoder_dp()
+    prefetch_count = tp_size * encoder_rounds
+    return PrefetchIterator(iter(_cyclic_iter(dataloader)), prefetch_count=prefetch_count)
+
+
 ######## utils for get_batch ########
 def _get_position_ids(data: torch.Tensor):
     """create position ids"""

loongforge/train/training_utils.py

@@ -1459,51 +1459,69 @@ def train_step(
             get_batch, get_embedding_list,
             get_visual_pos_masks_list, get_deepstack_visual_embeds_list,
             get_deepstack_grad_list, _create_mock_batch,
+            get_encoder_data_iterator,
         )
 
         num_microbatch, encoder_rounds = get_num_micro_batches_per_decoder_dp()
         num_real_microbatch = get_num_real_micro_batches_per_decoder_dp()
         unwrapped_model = unwrap_model(model[0])
-        if isinstance(data_iterator, list):
-            first_iter, backup_iter = itertools.tee(data_iterator[0])
-            data_iterator = [RerunDataIterator(first_iter)] + data_iterator[1:]
+
+        encoder_iter = get_encoder_data_iterator()
+        if encoder_iter is not None:
+            if isinstance(data_iterator, list):
+                data_iterator = [RerunDataIterator(data_iterator[0])] + data_iterator[1:]
+            else:
+                data_iterator = RerunDataIterator(data_iterator)
         else:
-            data_iterator, backup_iter = itertools.tee(data_iterator)
-            data_iterator = RerunDataIterator(data_iterator)
+            if isinstance(data_iterator, list):
+                first_iter, backup_iter = itertools.tee(data_iterator[0])
+                data_iterator = [RerunDataIterator(first_iter)] + data_iterator[1:]
+            else:
+                data_iterator, backup_iter = itertools.tee(data_iterator)
+                data_iterator = RerunDataIterator(data_iterator)
 
         pp_layer = mpu.get_pipeline_model_parallel_rank()
         tp_size = mpu.get_tensor_model_parallel_world_size()
-
         model_size = num_microbatch // encoder_rounds
-        iter_count = 0
+
+        all_raw_batches = None
+        if encoder_iter is None:
+            all_raw_batches = [next(backup_iter) for _ in range(num_real_microbatch)]
+
         batch_list = []
+        last_real_batch = None
+        has_any_real = (pp_layer * tp_size < num_real_microbatch)
         embedding_list = get_embedding_list()
         visual_pos_masks_list = get_visual_pos_masks_list()
         deepstack_visual_embeds_list = get_deepstack_visual_embeds_list()
         for round in range(encoder_rounds):
             batch_list.clear()
             front = pp_layer * tp_size + round * model_size
-            end = (pp_layer + 1) * tp_size + round * model_size
             all_mock_in_range = (front >= num_real_microbatch)
-            # Skip batches before this PP rank's range, but only for real data.
-            # When entire range is mock, save the last skipped batch as reference.
-            skip_count = max(0, min(front, num_real_microbatch) - iter_count)
-            last_skipped_batch = None
-            for skip_i in range(skip_count):
-                if skip_i == skip_count - 1 and all_mock_in_range:
-                    last_skipped_batch = copy.deepcopy(get_batch(backup_iter))
-                else:
-                    next(backup_iter)
             for tp_idx in range(tp_size):
                 global_mb_idx = front + tp_idx
                 if global_mb_idx >= num_real_microbatch:
                     # This microbatch is beyond real data — use mock batch
-                    mock_ref = batch_list[-1] if batch_list else last_skipped_batch
+                    if not batch_list and last_real_batch is None:
+                        # No real batch seen yet — need a reference for mock shape
+                        if encoder_iter is not None and has_any_real:
+                            last_real_batch = get_batch(encoder_iter)
+                        elif all_raw_batches is not None:
+                            last_real_batch = copy.deepcopy(get_batch(iter([all_raw_batches[num_real_microbatch - 1]])))
+                        else:
+                            iter_arg = (
+                                data_iterator if not isinstance(data_iterator, list) else data_iterator[0]
+                            )
+                            last_real_batch = get_batch(iter_arg)
+                    mock_ref = batch_list[-1] if batch_list else last_real_batch
                     batch_list.append(_create_mock_batch(mock_ref))
                 else:
-                    local_batch = copy.deepcopy(get_batch(backup_iter))
-                    batch_list.append(local_batch)
-            iter_count = min(end, num_real_microbatch)
+                    if encoder_iter is not None:
+                        batch = get_batch(encoder_iter)
+                    else:
+                        batch = copy.deepcopy(get_batch(iter([all_raw_batches[global_mb_idx]])))
+                    last_real_batch = batch
+                    batch_list.append(batch)
 
             input_embeds_list = []
             for i in range(tp_size):