Reject sequence packing under pipeline parallelism

[timothyngo]

Summary

Two independent, self-contained commits on the sequence-packing path. Neither depends on FlexAttention; this is the first of three PRs (see Follow-ups).

• Reject pp > 1 + data.pack_sequences in JobConfig.validate. PipelineStageModule.forward receives only hidden states (pipeline_parallel.py:192, layer call at :216) and pipeline_step never read batch["doc_ids"] (loop.py:194-226), so the combination trained with cross-document attention. Packing's label half still worked — labels carry -100 at boundaries from the dataset — so the loss looked correct while attention leaked, which is why it went unnoticed.
• Vectorize _compute_packed_output's doc-id assignment. It carried a counter in a Python for loop over every token, inside __getitem__ — seq_len interpreter iterations per sample, on the dataloader's critical path. doc_ids[i] is just the number of EOS tokens strictly before i, so one np.cumsum over the shifted EOS mask replaces it.

Kept as separate commits so the refactor is not bundled with the bug fix.

Behavior change

The validation error is the only user-visible change. Packing without PP, and PP without packing, are both unchanged. No config in configs/train/ sets pack_sequences, so no existing run is affected — the shipped PP configs (13b_32gpu_tp4_pp2.toml, 29b_32gpu_tp4_pp2.toml, 70b_32gpu_tp4_pp4.toml) all still load, which test_shipped_config_loads covers.

The vectorization is behavior-preserving: TestComputePackedOutput pins the semantics, and a new randomized test compares against the literal counter loop it replaces. I verified that test has teeth by mutating the implementation off-by-one — it fails, along with three of the existing cases.

Testing

• uv run ruff check kempnerforge/ tests/ passes
• uv run ruff format --check passes
• uv run pytest tests/unit/ passes — 1808 passed, 13 skipped
• pre-commit (ruff, ruff-format) passes on all changed files
• uv run pyright kempnerforge/ — 1 error, pre-existing and unrelated: video_io.py:94 import "av" could not be resolved. av lives in the optional video dependency group, which uv sync does not install by default. Zero errors in the files this PR touches.

Not run: integration/distributed/e2e suites — this PR changes config validation and a numpy expression, touching no GPU or distributed code path.

One environment note for anyone reproducing on the cluster: pytest --timeout=60 from CONTRIBUTING.md gives spurious failures in tests/unit/test_packing.py on a cold NFS cache, because import transformers walks its whole models/ tree and can exceed 60s on its own. Those tests pass with a larger timeout; nothing to do with this change.

Follow-ups

• PR 2 — FlexAttention backend for packed sequences (model.attention_backend = "flex"), replacing the dense (B, 1, S, S) mask at attention.py:214-227 that currently drops SDPA off FlashAttention whenever packing is on.
• PR 3 — GPU validation matrix and throughput benchmarks for that path.
• Actually supporting packing under PP is separate work: the stage module's forward signature plus per-microbatch alignment, since the schedule splits along dim 0 and does not pass a chunk id. This PR's rejection is the assertion that change would invert.

Closes #204

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

Files with missing lines	Coverage Δ
kempnerforge/config/job.py	89.07% <100.00%> (+0.18%)	⬆️
kempnerforge/data/dataset.py	92.18% <100.00%> (-0.13%)	⬇️

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.