skip to content

Reference

Parity Across PyTorch Execution Modes

copy markdown

The 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

Modeqwen2 (0.5B)qwen3 (0.6B)gemma3 (1B)gemma4 (31B)
baseline (deterministic=True, math SDPA)0.00e+000.00e+000.00e+000.00e+00
default fast SDPA (deterministic=False)0.00e+000.00e+000.00e+000.00e+00
forced math SDPA0.00e+000.00e+000.00e+000.00e+00
TF32 enabled0.00e+000.00e+000.00e+000.00e+00
bf16 autocast (torch.autocast)0.00e+000.00e+000.00e+000.00e+00
torch.compile (per-layer)0.00e+000.00e+000.00e+000.00e+00
forced flash SDPAskipped*skipped*skipped*skipped*
forced mem-efficient SDPAskipped*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