Reference
Parity Across PyTorch Execution Modes
copy markdownThe trainer's claim that blockwise==resident equivalence is structural —
both paths execute identical layer math, only weight residency differs — is
tested empirically here. For each mode, scripts/parity_modes_sweep.py runs
the resident-vs-blockwise forward gate (torch.equal, atol=0) plus a
run-to-run reproducibility check and a 3-step training-loss sanity, each mode
in a fresh subprocess so global backend state cannot leak between modes.
Measured on NVIDIA A100 80GB PCIe, torch 2.10.0+cu128, transformers 5.5.0,
NF4, block_size=4. Raw records: validation/modes_sweep.jsonl.
Mode matrix
| Mode | qwen2 (0.5B) | qwen3 (0.6B) | gemma3 (1B) | gemma4 (31B) |
|---|---|---|---|---|
| baseline (deterministic=True, math SDPA) | 0.00e+00 | 0.00e+00 | 0.00e+00 | 0.00e+00 |
| default fast SDPA (deterministic=False) | 0.00e+00 | 0.00e+00 | 0.00e+00 | 0.00e+00 |
| forced math SDPA | 0.00e+00 | 0.00e+00 | 0.00e+00 | 0.00e+00 |
| TF32 enabled | 0.00e+00 | 0.00e+00 | 0.00e+00 | 0.00e+00 |
bf16 autocast (torch.autocast) | 0.00e+00 | 0.00e+00 | 0.00e+00 | 0.00e+00 |
torch.compile (per-layer) | 0.00e+00 | 0.00e+00 | 0.00e+00 | 0.00e+00 |
| forced flash SDPA | skipped* | skipped* | skipped* | skipped* |
| forced mem-efficient SDPA | skipped* | skipped* | skipped* | skipped* |
* sdpa_kernel([FLASH_ATTENTION]) / [EFFICIENT_ATTENTION] raise
RuntimeError: No available kernel on this stack because the models pass an
explicit attention mask those kernels do not accept — an upstream SDPA
backend constraint, not a Poros parity result. The unforced "default fast"
row covers what PyTorch actually dispatches in practice.
In every runnable mode, rerun_equal was also true: two blockwise runs with
the same seed are bitwise identical to each other.
Interpretation — what this does and does not claim
- Parity is mode-independent. Blockwise and resident execute the same kernels in the same mode, so whatever numerics a mode produces, both paths produce them identically. TF32, autocast, and compile change the numbers relative to the deterministic baseline (e.g. the TF32 loss differs from the fp32 loss) — but blockwise tracks resident exactly within each mode. That is the guarantee: Poros never adds divergence on top of whatever mode you chose.
- This is not a determinism claim for non-deterministic modes. With
deterministic=False, run-to-run reproducibility is a property of the kernels, not of Poros (it held in these runs, but is not guaranteed by PyTorch for all shapes/ops). - bf16 autocast here wraps the forward of an NF4-quantized frozen base; it
is distinct from
quantization="bf16", which remains experimental and carries no bitwise claim.
torch.compile note
Compile the transformer layers, not the streaming driver: the residency
engine swaps frozen-block storage between layer calls, which is data
movement, not graph math. Compiling each layer module (layer.compile())
works and preserves bitwise parity; compiling the whole blockwise forward
fails in Dynamo at the storage-swap (param.data = ...) inside
load_block_for_inference, by design.
Reproduce
# all modes, one model per invocation, JSONL out
python scripts/parity_modes_sweep.py --all --model Qwen/Qwen2.5-0.5B --out modes_sweep.jsonl
# torch-version gates (creates per-version venvs)
bash scripts/parity_torch_versions.sh torch_versions_parity.txt