fp8 medium DiT: calibrated tier — same bf16-fused speed, better worst-step fidelity (calibration by @ryanontheinside, #47) - #86
Merged
Conversation
Captures the FP8 calibration .npz for the medium DiT by driving the model's own pingpong generate() and recording the six DiT engine inputs at every sampling step. Native by construction: loads via the repo's own load_diffusion_cond, prompts from interface/reprompt.py, the real pingpong sigma schedule (no re-implementation of the inference path). From #47 (feat/dit-fp8); included here to support the calibrated fp8 tier.
Grafts @ryanontheinside's calibrated fp8 scales (#47) onto the shipped bakedmin fp8 structure (baked fp32 RoPE constant + bf16 fused MHA + weakly-typed), so the `dit_fp8` tier keeps its bf16-fused speed while gaining much better worst-step fidelity. The build recipe is unchanged — calibration lives entirely in the ONNX scale values. - transplant_scales.py: match every quantized Linear by weight-initializer name and swap only the fp8 scale VALUES (per-tensor activation amax + per-channel weight scales) from #47's calibrated ONNX onto our RoPE-baked graph; the 5.8 GB fp32 weights are untouched (TRT re-quantizes at build). - build_from_onnx.py: update the sa3-m-fp8 recipe notes — dit_fp8 is now the calibrated engine. - build/README.md + tensorRT/README.md: document the calibrated tier, the bake -> transplant producer flow, and that a from-scratch recalibration uses #47's full pipeline (make_calib.py -> build_dit_fp8.py). Measured H200 (L=4096, same run): speed unchanged (31.2 vs uncalibrated 30.8 ms/fwd, within noise); worst-step velocity-cos vs fp32 on adversarial seeds 0.52/0.57/0.64 -> 0.92/0.94/0.92, with steps 1-7 tracking the fully-calibrated #47 reference within ~0.001. Calibration approach + make_calib.py tooling by @ryanontheinside (#47). Co-authored-by: RyanOnTheInside <7623207+ryanontheinside@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Upgrades the
fp8medium-DiT tier to a calibrated engine — same bf16-fused speed, meaningfully better worst-case per-step fidelity.What changed. The shipped
dit_fp8(#84) used uncalibrated fp8 scales. This grafts the calibrated activation + per-channel weight scales from @ryanontheinside's fp8 work (#47) onto the existing bakedmin structure (baked fp32 RoPE + bf16 fused MHA + weakly-typed). Real-conditioning calibration capture is hismake_calib.py.Result (H200, L=4096, same-run).
Still medium-only, ≤4096, default stays
fp16mixed. Speed/quality knob: @ryanontheinside's fp16-attention variant (#47) buys the last ~0.03 of step-0 fidelity for +2 ms/fwd; this tier keeps bf16 attention for the speed.Calibration approach +
make_calib.pytooling by @ryanontheinside (#47).🤖 Generated with Claude Code