From eae6b3d1c18d126a75b6deee7b3afe892cedba4e Mon Sep 17 00:00:00 2001 From: Timothy Ngotiaoco Date: Thu, 17 Sep 2026 13:51:12 -0400 Subject: [PATCH 1/2] Reject sequence packing under pipeline parallelism PipelineStageModule.forward receives only hidden states, and pipeline_step never read batch["doc_ids"], so `pp > 1` with `data.pack_sequences = true` trained with cross-document attention: packed documents attended across each other while the labels still masked the boundary positions. The loss looked correct while attention leaked, which is why it went unnoticed. JobConfig.validate now raises on the combination. Packing without PP and PP without packing are both unchanged, and no shipped config sets pack_sequences, so no existing run is affected. --- CHANGELOG.md | 2 ++ kempnerforge/config/job.py | 10 ++++++++++ tests/unit/test_config.py | 12 ++++++++++++ 3 files changed, 24 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8c5e2cc..6526f65 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -122,6 +122,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Fixed +- **Sequence packing is now rejected under pipeline parallelism** instead of silently training on cross-document context. `PipelineStageModule.forward` receives only hidden states and `pipeline_step` never read `batch["doc_ids"]`, so with `pp > 1` and `data.pack_sequences = true` the packed documents attended across each other while the labels still masked the boundary positions -- the loss looked correct while attention leaked, which is why it went unnoticed. `JobConfig.validate` now raises. Packing without PP and PP without packing are both unchanged; no shipped config sets `pack_sequences`, so no existing run is affected. + - `kempnerforge/config/job.py` (+ `tests/unit/test_config.py`). - **Captions train an EOS stop token.** `_tokenize_and_mask` appends the tokenizer's EOS (when defined) so the last caption token learns to predict *stop*; previously captions (`add_special_tokens=False`) had no stop target and generation never learned to terminate. Unprompted, the first caption token stays unsupervised (standard LM) — use a prompt so its last token predicts it. - `kempnerforge/data/vlm_dataset.py` (+ `tests/unit/test_vlm_dataset.py`, `tests/unit/test_video_dataset.py`). - **VLM captioning labels are now next-token aligned.** `_tokenize_and_mask` labeled each position with its own input token (`labels[i] == input_ids[i]`); with the no-shift loss, each text logit was scored against the current, already-visible token, so the model learned to copy it and the captioning loss collapsed to `0`. Labels are now shifted (`labels[i] = input_ids[i+1]`, last real position `-100`), matching text pretraining; prompt-predicting positions stay masked. Fixes all VLM arches. diff --git a/kempnerforge/config/job.py b/kempnerforge/config/job.py index 90d129c..3b9863d 100644 --- a/kempnerforge/config/job.py +++ b/kempnerforge/config/job.py @@ -197,6 +197,16 @@ def validate(self, world_size: int = 1) -> None: "splitting. Use FSDP, TP, or EP instead." ) + if self.distributed.pp > 1 and self.data.pack_sequences: + raise ValueError( + "Sequence packing + Pipeline Parallelism is not supported. " + "PipelineStageModule.forward receives only hidden states, so doc_ids " + "never reaches the stages and packed documents would attend across " + "document boundaries while the labels still mask those positions -- " + "silently training on cross-document context. " + "Set data.pack_sequences=false, or train without pipeline parallelism." + ) + if self.distributed.ep > 1: if not self.model.is_moe: raise ValueError("ep > 1 requires an MoE model (num_experts > 0)") diff --git a/tests/unit/test_config.py b/tests/unit/test_config.py index f7378a7..67224e5 100644 --- a/tests/unit/test_config.py +++ b/tests/unit/test_config.py @@ -490,6 +490,18 @@ def test_validate_dense_with_pp_passes(self): ) config.validate(world_size=2) # Should not raise — dense + PP is fine + def test_validate_packing_with_pp_rejected(self): + config = JobConfig( + data=DataConfig(pack_sequences=True), + distributed=DistributedConfig(pp=2, dp_shard=1), + ) + with pytest.raises(ValueError, match="Sequence packing.*Pipeline Parallelism"): + config.validate(world_size=2) + + def test_validate_packing_without_pp_passes(self): + config = JobConfig(data=DataConfig(pack_sequences=True)) + config.validate(world_size=1) # Should not raise — packing is fine without PP + def test_validate_vlm_seq_len_too_short(self): config = JobConfig( model=ModelConfig(max_seq_len=1024), From 927d9499acb4ad5eaca5000685b72cc0536f2ec5 Mon Sep 17 00:00:00 2001 From: Timothy Ngotiaoco Date: Thu, 17 Sep 2026 14:09:47 -0400 Subject: [PATCH 2/2] Vectorize packed doc_id assignment _compute_packed_output carried a document counter in a Python for-loop over every token. It runs inside __getitem__, so it cost one interpreter iteration per token on every sample the loader yields -- seq_len iterations per sample, on the dataloader's critical path. doc_ids[i] is just the number of EOS tokens strictly before i, so a single np.cumsum over the shifted EOS mask replaces the loop. Behavior-preserving; TestComputePackedOutput pins the semantics and a new randomized test compares against the literal counter loop it replaces. --- kempnerforge/data/dataset.py | 11 +++++------ tests/unit/test_packing.py | 21 +++++++++++++++++++++ 2 files changed, 26 insertions(+), 6 deletions(-) diff --git a/kempnerforge/data/dataset.py b/kempnerforge/data/dataset.py index 066a73f..9956da0 100644 --- a/kempnerforge/data/dataset.py +++ b/kempnerforge/data/dataset.py @@ -41,13 +41,12 @@ def _compute_packed_output(tokens: np.ndarray, eos_token_id: int) -> dict[str, t cross-document boundaries), and ``doc_ids`` (seq_len, integer document assignment per input token for attention masking). """ - # Assign a document ID to each token: increment after every EOS + # Assign a document ID to each token: increment after every EOS, i.e. + # doc_ids[i] is the number of EOS tokens strictly before i. Vectorized + # because this runs per __getitem__, once per token -- a Python loop here + # costs seq_len interpreter iterations on every sample the loader yields. doc_ids = np.zeros(len(tokens), dtype=np.int64) - doc_id = 0 - for i in range(len(tokens)): - doc_ids[i] = doc_id - if tokens[i] == eos_token_id: - doc_id += 1 + doc_ids[1:] = np.cumsum(tokens[:-1] == eos_token_id) token_tensor = torch.from_numpy(tokens.copy()).long() doc_id_tensor = torch.from_numpy(doc_ids.copy()) diff --git a/tests/unit/test_packing.py b/tests/unit/test_packing.py index cd4ee57..6ac6ad8 100644 --- a/tests/unit/test_packing.py +++ b/tests/unit/test_packing.py @@ -89,6 +89,27 @@ def test_output_dtypes(self): assert result["labels"].dtype == torch.long assert result["doc_ids"].dtype == torch.long + def test_doc_ids_match_sequential_reference(self): + """Vectorized doc_id assignment matches a literal carry-the-counter loop.""" + + def reference(tokens: np.ndarray, eos_token_id: int) -> np.ndarray: + doc_ids = np.zeros(len(tokens), dtype=np.int64) + doc_id = 0 + for i in range(len(tokens)): + doc_ids[i] = doc_id + if tokens[i] == eos_token_id: + doc_id += 1 + return doc_ids + + rng = np.random.default_rng(0) + for _ in range(50): + # Small vocab so EOS is dense: boundaries, runs of consecutive EOS, + # and EOS at either end all show up across the draws. + tokens = rng.integers(0, 4, size=int(rng.integers(3, 64))).astype(np.int64) + result = _compute_packed_output(tokens, eos_token_id=0) + # _compute_packed_output returns the input half of the doc ids. + assert result["doc_ids"].tolist() == reference(tokens, 0)[:-1].tolist() + # --------------------------------------------------------------------------- # Attention mask for packed sequences