Quintus/docs/pipeline_hardening.md

# Pipeline Hardening

![Pipeline Hardening Flow](../assets/pipeline_hardening_flow.png)

This page summarizes the correctness and reliability lessons that shaped the Quintus codebase. Most of these are silent-failure classes: the pipeline can appear to run while producing invalid or misleading artifacts.

## Silent Serialization Bugs

Teacher token IDs must be stored in a dtype that can represent the tokenizer vocabulary.

An early offline-KD path stored top-k token IDs too narrowly. Qwen token IDs exceed signed 16-bit range, so IDs could wrap negative and later be clamped into valid-looking but wrong positions. Training could continue, but the KL support was corrupted.

Hardening rule:

- Store token IDs as `int32` or wider.
- Validate IDs on load.
- Reject negative IDs.
- Reject IDs outside the student vocabulary.
- Treat dtype as part of the shard contract.

## Row-Order Preservation

Teacher-logit extraction often sorts samples by length for throughput. Training usually expects logits to match the original tokenized row order.

If sorted extraction writes shards in sorted order without restoring original indices, the student receives teacher logits for the wrong sample. This is a model-poisoning bug, not a performance issue.

Hardening rule:

- Batch by sorted length if useful.
- Preserve `original_idx`.
- Write final shards in original dataset order.
- Verify teacher-logit length against the tokenized row length at training time.

## Dataset Schema And Decoding

Public instruction datasets do not share a single row schema. Some rows arrive as `messages`; others use Alpaca-style `instruction`, `input`, and `output` fields. Some content fields contain nested dict/list payloads that need structured coercion before templating.

Dataset streaming can also fail late when a compression codec or file decoder is missing. That failure should remain visible instead of being replaced by a generic "zero samples" result.

Hardening rule:

- Detect Alpaca-style instruction/output rows before chat-message conversion.
- Coerce nested dict/list content through structured serialization, then normalize to text.
- Normalize common role aliases before applying a chat template.
- Preserve the first real dataset exception when streaming fails.
- Validate dataset decoding and schema mapping before large model downloads.

## Zero-Data And Data-Erasure Guards

Data preparation should fail when no usable rows are produced. It should also distinguish "download only" from "tokenize and overwrite output".

Hardening rule:

- Abort if filtering retains zero samples.
- Abort if tokenization writes zero rows.
- Do not open tokenized output in write mode for asset-only setup.
- Use explicit flags for model-only or data-only phases.

## Missing Shards Must Fail

Replacing missing teacher-logit shards with zero tensors makes the training loop look healthy while removing the KD signal.

Hardening rule:

- Missing shard means hard failure.
- Stale shard directories are cleaned before extraction.
- `_provenance.json` is required for KD.
- Shard count, sample count, max sequence length, temperature, top-k, and schema version are checked before training.

## Provenance Contracts

Path equality is weak provenance because paths change across machines. Data identity should come from content and model contracts.

Useful provenance fields:

- schema version
- dataset fingerprint or SHA-256
- sample count
- shard count
- max sequence length
- top-k or full-vocab mode
- temperature
- teacher model ID and revision
- student model ID and revision
- tokenizer sizes
- tokenizer fingerprints
- shard dtypes

Tokenizer fingerprints can drift across library versions. Vocab size and schema compatibility should remain hard gates; fingerprint drift can be a warning when stronger invariants still match.

## Assistant-Only Loss Masks

Supervising prompt and chat-template tokens can teach formatting before substance. It can also make chat-mode behavior fragile.

Hardening rule:

- Tokenized rows must include `loss_mask`.
- Loss mask must be binary.
- Rows with zero assistant targets are rejected.
- User prompts, system prompts, separators, and padding are not targets.
- Assistant response tokens are the supervised region.

Prefix-stable mask derivation is useful when tokenizer-provided assistant masks are unavailable.

## Gradient Accumulation Semantics

DeepSpeed and non-DeepSpeed paths need different step-accounting logic.

DeepSpeed accumulation is global across the full run, not local to each epoch. Epoch-end remainder branches should not create phantom optimizer steps.

Non-DeepSpeed accumulation needs an explicit final flush when a leftover accumulation window exists. That flush must rescale gradients so the update represents the mean over the remainder, not a shrunken `remainder / grad_accum` update.

Hardening rule:

- Advance `global_step` only after a real optimizer update.
- Align scheduler steps with real updates.
- Log flush steps.
- Include flush steps in training-loss CSVs.
- Prefer validation split sizes that align with effective batch size.

## Checkpoint Semantics

`init_from_checkpoint` and `resume_from_checkpoint` are different operations.

- Initialization starts a new phase from an existing model.
- Resume continues an interrupted phase from training state.

Mixing the two can skip training, restart from the wrong model, or reuse stale state.

Hardening rule:

- Forbid simultaneous init and resume.
- Save trainer state and scheduler state.
- Search both `step_*` and `epoch_*` checkpoints for resume.
- Store batch offset for mid-epoch resume.
- Keep final model-loading checkpoints portable.

## Compiler Portability

Compiled PyTorch modules can save weights with `_orig_mod.` prefixes if not unwrapped. Standard Transformers and vLLM loaders do not expect those keys.

Hardening rule:

- Keep `torch.compile` opt-in.
- Treat dynamic-shape recompile overhead as a throughput risk, not just a startup cost.
- Unwrap compiled modules before saving.
- Strip `_orig_mod.` only as a repair path, not as the normal release path.
- Verify saved checkpoints load through standard APIs.

## Artifact Hygiene

Stale outputs are a real ML correctness problem. Old result JSONs, old plots, or old sample logs can make a failed run look successful.

Hardening rule:

- Clean evaluation output directories before a new run.
- Clean stale plots before rendering.
- Select result files by clear recency rules.
- Fail if expected task outputs are incomplete.
- Fail if a requested checkpoint is missing; do not fall back to older local weights.
- Include runtime versions in result summaries.

## Environment Contracts

Notebook and cloud images often contain mixed binary packages. Import success for `torch` alone does not prove the stack is healthy.

Hardening rule:

- Treat `torch`, `torchvision`, and `torchaudio` as one binary compatibility family.
- Use staged dependency manifests instead of ad hoc installs.
- Keep vLLM dependencies separate from HF-only evaluation dependencies.
- Prefer clear preflight errors over late framework crashes.
- Print exception chains, not only the outer error.

## Remote Code And Revisions

Model loading should be reproducible and explicit.

Hardening rule:

- Pin teacher, student, and tokenizer revisions when possible.
- Default remote-code trust to false.
- Provide an explicit override for models that need custom code.
- Explain remote-code failures clearly.

## Safe Logging

Training logs should be rich enough for issue diagnosis without dumping config internals.

Hardening rule:

- Avoid logging authentication values or full config payloads.
- Disable traceback local-variable dumps in rich tracebacks.
- Strip ANSI sequences from file logs while keeping colored notebook output if desired.
- Use UTF-8 file logs and replacement-safe console output for generated model text.
- Log checkpoint save/upload intent, output size, duration, and destination path without sensitive values.

## Public Release Rule

A project can be release-ready without every possible production safeguard. The line is crossed when:

- known silent corruption paths are removed,
- remaining tradeoffs are documented,
- artifacts are reproducible enough to audit,
- public docs focus on decisions, methods, and release artifacts,
- evaluation claims are tied to clear methodology.

For Quintus, the release surface should describe the engineering decisions and results.