Skip to content

[train, data] perf: optimize data fetching latency for full-hetero encoder - #76

Merged
nullnonenilNULL merged 1 commit into
baidu-baige:masterfrom
Yangsx-1:data-pull
Jun 1, 2026
Merged

[train, data] perf: optimize data fetching latency for full-hetero encoder#76
nullnonenilNULL merged 1 commit into
baidu-baige:masterfrom
Yangsx-1:data-pull

Conversation

@Yangsx-1

Copy link
Copy Markdown
Collaborator

Summary

Optimizes data loading for the encoder in full-hetero DP mode by eliminating redundant data fetching. Previously, the encoder iterated through all microbatches sequentially (including those assigned to other PP ranks), causing unnecessary IO and latency. Now each PP rank only fetches its own assigned microbatches.

Changes

File Change
loongforge/data/encoder_strided_sampler.py New module: EncoderStridedSampler (map-style), EncoderStridedIterator (streaming), PrefetchIterator (background prefetch)
loongforge/data/multimodal/dataloader_provider.py Add build_encoder_energon_iterator for Energon/WebDataset path
loongforge/train/sft/utils.py Add build_encoder_data_iterator for SFT map-style datasets
loongforge/train/pretrain/pretrain_vlm.py Wire up encoder iterator construction; add get/set_encoder_data_iterator
loongforge/train/training_utils.py Refactor train_step to consume from dedicated encoder iterator, removing itertools.tee + sequential skip pattern

@Yangsx-1 Yangsx-1 changed the title Optimize data fetching latency for full-hetero encoder [train, data] perf: optimize data fetching latency for full-hetero encoder May 29, 2026
@Yangsx-1
Yangsx-1 force-pushed the data-pull branch 3 times, most recently from 03aa349 to 3464bc1 Compare May 29, 2026 07:14
@nullnonenilNULL
nullnonenilNULL merged commit 5414969 into baidu-baige:master Jun 1, 2026
5 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants