Match H3 final output heads

This commit is contained in:
Daniel Maddern 2026-08-13 13:09:42 +07:00
parent a85fe116c1
commit 985a529fc9
2 changed files with 5 additions and 6 deletions

View file

@ -184,6 +184,7 @@ gate and is exact.
| Token refiner | Yes | Bit-exact from captured 5376-wide refiner input through both blocks and final RMSNorm; Qwen-to-refiner projection boundary is still not separately captured | | Token refiner | Yes | Bit-exact from captured 5376-wide refiner input through both blocks and final RMSNorm; Qwen-to-refiner projection boundary is still not separately captured |
| Prompt-only FL2VA packer | Yes | Bit-exact for the coherent captured text-only FL2VA DiT input; no keyframe/reference condition rows | | Prompt-only FL2VA packer | Yes | Bit-exact for the coherent captured text-only FL2VA DiT input; no keyframe/reference condition rows |
| H3 DiT backbone | Yes | Bit-exact through all 50 blocks from the coherent assembled FL2VA input; requires the standalone Comfy Kitchen fused Q/K RMSNorm + split-half RoPE operator | | H3 DiT backbone | Yes | Bit-exact through all 50 blocks from the coherent assembled FL2VA input; requires the standalone Comfy Kitchen fused Q/K RMSNorm + split-half RoPE operator |
| H3 final layer | Yes | Bit-exact final AdaLN, target-row modulation, and video/audio patch rows; Comfy materializes the AdaLN and output-head biases through BF16 |
| H3 DiT | Yes | Strict all-block numeric parity not achieved | | H3 DiT | Yes | Strict all-block numeric parity not achieved |
| Beta/RES sampler | Yes | Exact sigma/state/update parity not demonstrated; direct preview has wrong initial audio state | | Beta/RES sampler | Yes | Exact sigma/state/update parity not demonstrated; direct preview has wrong initial audio state |
| Video VAE decoder | Yes | No direct-versus-Comfy same-latent pixel comparison yet | | Video VAE decoder | Yes | No direct-versus-Comfy same-latent pixel comparison yet |
@ -195,10 +196,8 @@ gate and is exact.
Only these are outstanding. Do not recapture or revisit rows marked complete Only these are outstanding. Do not recapture or revisit rows marked complete
unless the checkpoint, Comfy version, prompt, or backend changes. unless the checkpoint, Comfy version, prompt, or backend changes.
1. **Final rows and video latent.** Final DiT hidden state, final AdaLN, and 1. **Final video latent layout.** Compare direct `unpatchify_video` output
final RMSNorm are exact. Repair the remaining final target-segment with Comfy before entering the sampler.
modulation/FP32-head boundary (`video_hidden` max `0.00598395`) before
`unpatchify_video` output with Comfy before entering the sampler.
2. **Sampler replay using Comfy state.** Feed `initial.pt` and captured 2. **Sampler replay using Comfy state.** Feed `initial.pt` and captured
`sigmas` directly to the sampler. Compare video and audio denoised/update `sigmas` directly to the sampler. Compare video and audio denoised/update
states for all 12 saved steps. Do not test seed equality until this passes. states for all 12 saved steps. Do not test seed equality until this passes.

View file

@ -45,9 +45,9 @@ class H3FinalLayer(nn.Module):
checkpoint.tensor("final_layer.adaln_proj.linear.weight", dtype=torch.bfloat16), checkpoint.tensor("final_layer.adaln_proj.linear.weight", dtype=torch.bfloat16),
checkpoint.tensor("final_layer.adaln_proj.linear.bias", dtype=torch.bfloat16), checkpoint.tensor("final_layer.adaln_proj.linear.bias", dtype=torch.bfloat16),
checkpoint.tensor("final_layer.video_out.weight", dtype=torch.bfloat16), checkpoint.tensor("final_layer.video_out.weight", dtype=torch.bfloat16),
checkpoint.tensor("final_layer.video_out.bias", dtype=torch.float32), checkpoint.tensor("final_layer.video_out.bias", dtype=torch.bfloat16),
checkpoint.tensor("final_layer.audio_out.weight", dtype=torch.bfloat16), checkpoint.tensor("final_layer.audio_out.weight", dtype=torch.bfloat16),
checkpoint.tensor("final_layer.audio_out.bias", dtype=torch.float32), checkpoint.tensor("final_layer.audio_out.bias", dtype=torch.bfloat16),
) )
def forward( def forward(