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..419e36f 100644 --- a/kempnerforge/config/job.py +++ b/kempnerforge/config/job.py @@ -197,6 +197,19 @@ def validate(self, world_size: int = 1) -> None: "splitting. Use FSDP, TP, or EP instead." ) + # Why this is an error rather than a warning: PipelineStageModule.forward + # receives only hidden states, so doc_ids never reaches the stages and + # packed documents attend across each other. The labels still carry -100 + # at the boundaries, so the loss looks correct while attention leaks -- + # there is no signal in the training curve that would reveal it. + # Measured on 2 GPUs: pp=2 output matches unpacked causal attention + # exactly, rather than the packed reference. + if self.distributed.pp > 1 and self.data.pack_sequences: + raise ValueError( + "Sequence packing + Pipeline Parallelism is not supported. " + "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/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_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), 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