fix: MTP CP-aware roll, segment-aware roll, and synthetic packing support - #4623
fix: MTP CP-aware roll, segment-aware roll, and synthetic packing support#4623chiaotung97 wants to merge 1 commit into
Conversation
- Add CP-aware left shift (_shift_left_one_cp_aware) using ppermute - Add segment-aware roll (roll_and_mask_by_segment) to prevent cross-document target leakage in MTP - Generate packed segment IDs for synthetic data when packing=True - Remove guard against synthetic+packing+CP in train_utils.py - Add 17 unit tests covering segment roll, CP shift, and CP+segment
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
|
Hi reviewers, just an FYI that this PR is from a forked repo. Due to secret key limitations, the Gemini review couldn't be triggered. |
| local_rolled = local_rolled.at[:, -1:, ...].set(0) | ||
|
|
||
| try: | ||
| cp_size = jax.lax.psum(1, axis_name=axis_name) |
There was a problem hiding this comment.
This is a very interesting way to collect cp_size. I suggest using
jax.lax.axis_size(axis_name)
OR
mesh.shape.get(axis_name, 1)
NuojCheng
left a comment
NuojCheng
left a comment
There was a problem hiding this comment.
LGTM and thank you for the contribution! Some minor comments:
- the cpu test failure can be solved by rebasing to main
- I don't think we need design doc included in this PR, but it is great for PR review. Can we move the design doc to the PR description?
| Returns: | ||
| Array shaped like x, left-shifted by 1 across CP boundaries. | ||
| """ | ||
| local_rolled = jnp.roll(x, -1, axis=1) |
There was a problem hiding this comment.
Would this still work with load balancing on? If not perhaps we should handle it differently or reject LB?
| mtp_xent, _ = max_utils.cross_entropy_with_logits( | ||
| mtp_logits, jax.nn.one_hot(rolled_target_ids, cfg.vocab_size), 0.0 | ||
| ) | ||
| mtp_xent_masked = mtp_xent * rolled_target_mask |
There was a problem hiding this comment.
Is rolled_target_mask a boolean mask?
There was a problem hiding this comment.
target_mask enters MultiTokenPredictionBlock.__call__ as jnp.bool_ (standard training pipeline behavior: True = valid position, False = padded/ignore).
We will cast it to jnp.int32 before the rolling loop because roll_and_mask_by_segment internally zeros boundary positions via jnp.where(mask, 0, rolled) — when rolled is booland 0 is int32, JAX auto-promotes to int32 with a scatter-input FutureWarning (will become a hard error in a future JAX release).
The effective value is 0/1 either way, and downstream mtp_xent * mask works identically for int32 and bool. Verified on TPU v6e. A follow-up commit will add an explicit bool → int32 cast at the entry of __call__.
There was a problem hiding this comment.
Thank you for your response.
target_mask enters MultiTokenPredictionBlock.__call__ as jnp.bool_ (standard training pipeline behavior: True = valid position, False = padded/ignore).
I think this is not true if you're doing packing right? e.g. targets_segmentation is an integer array
| rolled_target_mask = roll_and_mask(rolled_target_mask) | ||
| rolled_position_id = roll_and_mask(rolled_position_id) | ||
| rolled_input_ids = roll_and_mask_by_segment(rolled_input_ids, rolled_segment_ids) | ||
| rolled_target_ids = roll_and_mask_by_segment(rolled_target_ids, rolled_segment_ids) |
There was a problem hiding this comment.
Is rolled_segment_ids aligned with rolled_target_ids here? target_ids is already shifted by 1 so after this roll the actual target is 2 tokens ahead? segment IDs passed to the helper seem to track only 1 token
OK, I will remove the design doc later, thanks for these comments. |
4 participants
Summary
Fixes 3 issues when combining Multi-Token Prediction (MTP) with
All-Gather Context Parallelism (AG-CP) and Packing.
Checklist
_shift_left_one_cp_aware) with ppermute backward ringroll_and_mask_by_segment) to prevent cross-document leakage_make_packed_segment_ids)Problem
jnp.rolldoes not cross CP rank boundariesjnp.rollonly operates within the local shardroll_and_maskis segment-unawaresegment_idsis always all-ones even thoughbase.ymldefaults topacking: truetrain_utils.pyalso blocks the combination outrightSolution
1. CP-Aware Left Shift (
_shift_left_one_cp_aware)Uses
jax.lax.ppermutein a backward ring: rank r sends its first token torank r−1, receives rank r+1's first token, and places it in its last
position. Degrades to
jnp.rollwhen no CP axis is in scope — zero overhead.2. Segment-Aware Roll (
roll_and_mask_by_segment)Shifts both
xandsegment_idsvia_shift_left_one_cp_aware, then maskspositions where
seg_current != seg_next(document boundary) orseg_current == 0(padding). Falls back toroll_and_maskwhensegment_ids=None. All rolling variables inMultiTokenPredictionBlock.__call__now use this function.3. Synthetic Data with Packed Segment IDs
_make_packed_segment_idssplits each row into 2..N randomly-sized segmentswith sequential integer IDs starting from 1. Removed the
train_utils.pyguard that rejected
synthetic + packing + CP.Files Changed
layers/multi_token_prediction.py_shift_left_one_cp_aware,roll_and_mask_by_segment, wiringinput_pipeline/synthetic_data_processing.py_make_packed_segment_idsutils/train_utils.pytests/unit/multi_token_prediction_test.pydocs/design/ag_cp_mtp_packing_fix.mdUnit Tests
30 tests, all passing (TPU v6e-4, 4 devices):
TestRollAndMaskTestRollAndMaskBySegmentTestMakePackedSegmentIdsTestShiftLeftOneCpAwareTestRollAndMaskBySegmentWithCpBackward Compatibility
_shift_left_one_cp_aware: degrades tojnp.rollwhen CP=1 or no"context"axisroll_and_mask_by_segment: degrades toroll_and_maskwhensegment_ids=Noneroll_and_mask(shift=-1): equivalent to original path when CP is offsynthetic_data_processing: segment IDs remainjnp.oneswhenpacking=FalseVerification