Skip to content

[data] switch pretrain from broadcast to replicated load#491

Merged
ananthsub merged 1 commit intoNVIDIA-NeMo:mainfrom
ananthsub:tp-replicated-load
Aug 26, 2025
Merged

[data] switch pretrain from broadcast to replicated load#491
ananthsub merged 1 commit intoNVIDIA-NeMo:mainfrom
ananthsub:tp-replicated-load

Conversation

@ananthsub
Copy link
Contributor

@ananthsub ananthsub commented Aug 26, 2025 •
edited
Loading

Don't load the batch in one TP rank and broadcast. instead, initialize the dataloader on all ranks and unconditionally call get_batch_from_iterator

Broadcast uses less file I/O and PCIe BW but hurts GPU memcpy and NVL BW

@ananthsub
switch pretrain to replicated load from broadcast
fe02f57 
Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
@copy-pr-bot
Copy link

copy-pr-bot bot commented Aug 26, 2025

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@ananthsub
Copy link
Contributor Author

ananthsub commented Aug 26, 2025

/ok to test fe02f57

sanandaraj5597
sanandaraj5597 approved these changes Aug 26, 2025
View reviewed changes
Copy link
Contributor

@sanandaraj5597 sanandaraj5597 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM.

@ananthsub ananthsub marked this pull request as ready for review August 26, 2025 21:01
@ananthsub ananthsub enabled auto-merge (squash) August 26, 2025 21:13
@ananthsub ananthsub merged commit 7612573 into NVIDIA-NeMo:main Aug 26, 2025
33 checks passed
ko3n1g pushed a commit that referenced this pull request Aug 26, 2025
@ananthsub
switch pretrain to replicated load from broadcast (#491)
74cc8d1 
Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
ananthsub added a commit that referenced this pull request Aug 26, 2025
@ko3n1g @ananthsub
switch pretrain to replicated load from broadcast (#491) (#492)
672a90c 
Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
Co-authored-by: Ananth Subramaniam <ansubramania@nvidia.com>
@ananthsub ananthsub mentioned this pull request Sep 10, 2025
Open
@ananthsub ananthsub deleted the tp-replicated-load branch February 17, 2026 07:16
@izikgo izikgo mentioned this pull request Feb 17, 2026
Open
5 tasks
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@sanandaraj5597 sanandaraj5597 sanandaraj5597 approved these changes

@erhoo82 erhoo82 Awaiting requested review from erhoo82

@malay-nagda malay-nagda Awaiting requested review from malay-nagda

Assignees

No one assigned

Labels

r0.1.0

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@ananthsub @sanandaraj5597