Support CP > 1 with inter-document masking in pretraining
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 6h
- Merged PRs (30d)
- 271
Description
## Problem
PR #5298 added THD-based inter-document masking through `--dataloader-inter-document-masking`, but it is not currently supported with `--context-parallel-size > 1`. PR #5531 proposes enforcing `CP == 1`.
A user may require inter-document masking for long-context pretraining, where `CP == 1` is unlikely to be sufficient.
## Requested behavior
Support `--dataloader-inter-document-masking` with `CP > 1`, including correct per-document THD partitioning/padding and coverage for CP sizes 2 and 4.
## Related
- #5298
- #5531
- #1878
Contributor guide
Research direction
Start by reading PR #5298 for the existing THD-based inter-document masking and PR #5531 for the CP==1 restriction. Trace how context parallelism partitions and pads per-document THD data, then add coverage for CP sizes 2 and 4. Done means --dataloader-inter-document-masking works correctly with --context-parallel-size greater than 1.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- distributed-systems, machine-learning
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 45/100