Reject sequence packing under pipeline parallelism

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.

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)")

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())

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),

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