Compare commits
No commits in common. "807b54ca0e09aa3734fb9c4956cfda1ed4c0cda8" and "eed3a3d951edd9b84b467a4b7ca2731767458cbd" have entirely different histories.
807b54ca0e
...
eed3a3d951
26 changed files with 43 additions and 904 deletions
227
PARITY.md
227
PARITY.md
|
|
@ -1,227 +0,0 @@
|
|||
# H3 Parity Evidence Ledger
|
||||
|
||||
ComfyUI is the oracle. This ledger records completed work and measured facts so
|
||||
we do not repeat investigations. A source review is never recorded as a passed
|
||||
numeric gate.
|
||||
|
||||
## Reference Runs And Immutable Artifacts
|
||||
|
||||
| Asset | Location on Spark | Contents | Status |
|
||||
| --- | --- | --- | --- |
|
||||
| Qwen trace | `/tmp/fl2va-qwen-trace` | `qwen_input_embeds.pt`, 50 decoder outputs in `qwen_layers/00.pt` through `49.pt`, and `qwen_layer50.pt` | Complete, immutable Comfy capture |
|
||||
| Qwen layer-0 trace | `/tmp/fl2va-qwen0-capture` | `qwen0_norm1.pt`, attention, post-attention, norm2, MLP, output | Complete, immutable Comfy capture |
|
||||
| H3 full trace | `/tmp/fl2va-full-capture` | assembled DiT input/output and `blocks/00.pt` through `49.pt` | Complete, immutable Comfy capture |
|
||||
| Working-copy capture mirror | `/home/daniel/aeon-spark-test/h3/h3-blackwell-runtime/artifacts/capture` | Qwen layer-0 and layer-50 assets plus H3 input/output | Present; do not treat as immutable if the working checkout is modified |
|
||||
| Sampler reference | `/home/daniel/aeon-spark-test/h3/h3-blackwell-runtime/artifacts/fl2va-sampler-reference` | `initial.pt`, exact `sigmas`, `step_00.pt` through `step_11.pt` | Complete, immutable Comfy capture |
|
||||
|
||||
All captures above use the established prompt-only FL2VA reference:
|
||||
|
||||
```text
|
||||
prompt: A brass-and-paper dragon flies above a rain-washed old city at blue hour.
|
||||
size: 320x192, length: 22 frames
|
||||
seed: 440204
|
||||
sampler: beta / res_multistep / 12 steps
|
||||
Comfy attention baseline: --use-sage-attention (SageAttention2)
|
||||
```
|
||||
|
||||
This reference is **not** the 72-frame cat preview. Do not overwrite it to
|
||||
diagnose preview output.
|
||||
|
||||
## Completed Findings
|
||||
|
||||
| Boundary | Evidence | Result | Conclusion |
|
||||
| --- | --- | --- | --- |
|
||||
| Direct runtime build | GB10 CUDA 13 image, Torch `2.9.1+cu130`, committed SageAttention3 wheel | Repeated direct runs complete | Environment is operational; PyTorch's `sm_121` warning is stale metadata in this environment |
|
||||
| H3 attention baseline | Comfy `--use-sage-attention` versus direct backend investigation | SageAttention2 (`sage2`) was established as the H3 correctness baseline | Use `sage2` for all H3 parity gates; Sage3 remains experimental |
|
||||
| Qwen tokenizer | Direct raw-prompt tokenizer versus MiniMax H3 source contract | Raw text, no chat template, `add_special_tokens=False` | Correct for nonempty text-only prompts; empty prompt differs: direct substitutes a space while Comfy injects token `151643` |
|
||||
| Qwen embedding loading, before fix | Direct versus Comfy input embedding capture | Example direct values were approximately `11` where Comfy was approximately `0.017`; downstream refined-text mean absolute delta `1.79263` | Root cause was omitted per-row `model.embed_tokens.weight_scale` |
|
||||
| Qwen embedding loading, after fix | Direct rows multiplied by `weight_scale` in FP32 | Input embeddings matched the Comfy capture exactly | Passed for the captured nonempty text-only reference |
|
||||
| Qwen layer-50, before embedding fix | Direct versus Comfy layer-50 capture | Mean absolute delta `23.50596`, max `38977.09375` | Failed; invalidated direct text conditioning |
|
||||
| Qwen layer-50, after embedding fix but before projection fix | Direct versus Comfy layer-50 capture | Mean absolute delta improved to about `5.63` | Embedding issue fixed; remaining error was material |
|
||||
| Qwen all-layer replay, after FP32 fix | `compare_qwen_layer_trace.py` replayed `/tmp/fl2va-qwen-trace` without new Comfy inference | Layers `00` through `49` and `qwen_layer50.pt` all reported mean/max `0` / `0` | **Passed: prompt-only Qwen conditioning is bit-exact for the immutable 17-token FL2VA reference.** |
|
||||
| Qwen layer-0 normalization | Direct `F.rms_norm` versus `qwen0_norm1.pt` | Mean absolute delta about `8.96e-06` | Passed to expected BF16-level tolerance |
|
||||
| Qwen layer-0 attention, before projection fix | Direct versus `qwen0_attention.pt` | Mean absolute delta about `0.0357`; layer output about `0.1496` | First material Qwen divergence was at projection/attention boundary, not tokenization, embeddings, RMSNorm, or RoPE |
|
||||
| Qwen attention backend | Comfy source inspection | Qwen chooses Comfy small-input attention, i.e. SDPA path, not Sage2 | Direct Qwen was changed to BF16 causal-mask SDPA with GQA |
|
||||
| Qwen SDPA dispatch audit and fix | Actual Spark Comfy block audit found direct raw SDPA omitted Comfy's four-dimensional mask wrapper, backend priority, and masked-GQA fallback. Direct ported that policy without importing Comfy. | Replaying immutable layer-0 capture after the change: attention mean/max `0.000244372` / `0.0161562`; post-attention `0.00026845` / `0.0213833`; layer output `0.000808998` / `0.0526199`. Previous attention mean was about `0.0357`; previous layer-output mean about `0.1496`. | Active code mismatch identified and substantially repaired. The remaining layer output error is now downstream of projection/elementwise BF16 boundaries, not an unreviewed SDPA dispatcher difference. |
|
||||
| Qwen NVFP4 metadata | Qwen layer-0 checkpoint sidecars inspected | Every Q/K/V/O and MLP projection has `full_precision_matrix_mult: true`; `o_proj` and `down_proj` also have `pre_quant_scale` | Direct was incorrectly quantizing activations for all Qwen projections |
|
||||
| Qwen NVFP4 fix | Commit `0cad1db` | Direct honors `full_precision_matrix_mult` and loads/applies `pre_quant_scale`; Spark smoke passed finite expected-shape Q and O projections | Structural fix passed; post-fix numeric QKV/MLP/layer-50 comparisons are still required |
|
||||
| Qwen layer-0 projections, after NVFP4 fix | Fresh Comfy projection capture at `/tmp/fl2va-qwen0-projections`; direct compared from captured input embedding and post-attention state | Q mean/max `6.6547e-05` / `0.00195485`; K `9.11704e-05` / `0.00241077`; V `6.0298e-05` / `0.00102112`; gate `0.000365695` / `0.0082469`; up `0.000335494` / `0.00792789`; activated `4.76882e-05` / `0.031621`; down `0.000557594` / `0.0749016` | The full-precision metadata fix removed the previous large projection failure. Projection parity is close but not strict; `down_proj` is the largest remaining layer-0 projection boundary. Preserve these numbers and do not recapture this case. |
|
||||
| Qwen projection code audit, post-fix | Statement-level comparison against actual Spark Comfy commit `43cb4fffc89bba20ab7bd61467a36d0339338dab`, Torch `2.9.1+cu130`, and `/h3-models/text_encoders/qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors` | All 350 Qwen projection sidecars set `full_precision_matrix_mult: true`; exactly 100 (`o_proj` and `down_proj` in each layer) have BF16 `pre_quant_scale`; no Qwen projection has bias. Packed weight is contiguous uint8, block scale is native E4M3, tensor scale is scalar FP32. | Direct and Comfy agree on the active full-precision branch, NVFP4 scale association, logical shapes, no-bias behavior, pre-scale order, BF16 compute dtype, and use the same installed Comfy Kitchen NVFP4 dequantizer. Do not change projection math based on speculation. |
|
||||
| H3 DiT block trace | `/tmp/fl2va-full-capture` | Block 0 was close, around mean absolute delta `0.00018`; small independent differences accumulated across blocks | Full 50-block strict parity remains failed/unresolved |
|
||||
| Sampler initial state | Comfy `initial.pt` inspected against direct preview | Comfy carries joint AV state; direct preview initializes video noise only and leaves audio zero | Known direct mismatch; identical integer seed is not sampler parity |
|
||||
| Preview output, before Qwen fixes | Dragon/locomotive prompt previews | Flower-like output | Confirmed incorrect conditioning had visible semantic impact |
|
||||
| Preview output, after Qwen metadata fix | 72-frame cat, 12-step direct preview | Cat semantic structure appears, with persistent 32-pixel-like square artifacts | Conditioning improved; denoiser/sampler/VAE boundary remains non-parity |
|
||||
| Video encoder | Direct ffmpeg raw RGB to H.264/yuv420p, no audio | Different from Comfy `CreateVideo`/`SaveVideo` | Not a valid latent-quality gate; cannot explain model-space corruption by itself |
|
||||
|
||||
## Qwen Projection Code Audit: Remaining Non-Active Differences
|
||||
|
||||
These are real implementation differences, but the audit established they are
|
||||
not active numerical branches for the captured Qwen checkpoint. They must not
|
||||
be presented as an explanation for the current `down_proj` delta without new
|
||||
evidence.
|
||||
|
||||
| Direct location | Difference from Comfy | Active for current capture? |
|
||||
| --- | --- | --- |
|
||||
| `nvfp4.py:47-51` | Calls `.contiguous()` while Comfy preserves loaded storage/view identity. In particular Comfy treats E4M3 block scales as a dtype view, not a numeric cast. | No: checkpoint tensors are already contiguous and direct also uses a dtype view. |
|
||||
| `nvfp4.py:75-76` | Always flattens and contiguates activation rows; Comfy full-precision `F.linear` retains N-D input/stride. | No known value effect: elements and BF16 dtype are unchanged. |
|
||||
| `nvfp4.py:82` | Uses `packed_weight.dequantize().to(flat_x)` rather than an explicitly named Comfy compute-dtype dequantization call. | No: input and output compute dtype are BF16, and both call the same installed Kitchen layout/dequantizer. |
|
||||
| `nvfp4.py:46,81` | Does not implement Comfy global full-precision overrides, disabled-format controls, forced casts, weight functions, or LoRA/offload dispatch. | No: every Qwen sidecar itself forces full precision, no projection has bias, and no patches/LoRAs are attached. |
|
||||
|
||||
**Audit conclusion:** projection source code contains cleanup/generalization gaps,
|
||||
but no identified active branch mismatch explains the current captured values.
|
||||
The next code audit must move outward to the exact Qwen block boundary:
|
||||
Comfy input dtype, full-precision weight construction, SDPA backend selection,
|
||||
and residual/BF16 rounding boundaries. Do not run another full inference merely
|
||||
to rediscover this projection result.
|
||||
|
||||
## Qwen Layer-0 Arithmetic Audit After `4050afe`
|
||||
|
||||
The follow-up statement audit covered every remaining active layer-0 arithmetic
|
||||
boundary outside the already-fixed SDPA wrapper: full-precision NVFP4 dispatch,
|
||||
weight/scales, `pre_quant_scale`, BF16 GEMM, activation layout, Q/K RMSNorm,
|
||||
O projection/residual, post-attention norm, SiLU/gate/up product, down
|
||||
projection, and final residual.
|
||||
|
||||
**Result: no further source-proven active mismatch was found.** For the actual
|
||||
checkpoint, Comfy and direct both use BF16 input, BF16 NVFP4 dequantization,
|
||||
BF16 `F.linear`, no projection biases, identical pre-scale ordering, identical
|
||||
MLP evaluation order, and ordinary BF16 residual additions. Explicit flattening
|
||||
and `.contiguous()` in direct are not a demonstrated numerical difference for
|
||||
the captured contiguous batch-one input.
|
||||
|
||||
Do not modify `nvfp4.py` or Qwen MLP/residual code to chase the remaining
|
||||
layer-0 delta. The next justified work is a bounded component experiment that
|
||||
identifies the actual selected SDPA kernel/GQA fallback and verifies capture
|
||||
consistency; it is not another full denoising or video run.
|
||||
|
||||
### Bounded SDPA/Layer-0 Result
|
||||
|
||||
The bounded replay was completed with the fresh projection-capture tensors.
|
||||
For the reference's 17-token layer-0 attention, both paths have native cuDNN
|
||||
GQA available (`flash=False`, `cudnn=True`, `efficient=False`). Therefore the
|
||||
direct SDPA wrapper takes the same native-GQA branch Comfy takes.
|
||||
|
||||
| Boundary | Mean absolute error | Maximum absolute error |
|
||||
| --- | ---: | ---: |
|
||||
| O-projected attention | `0.000244372` | `0.0161562` |
|
||||
| First residual | `0.00026845` | `0.0213833` |
|
||||
| Post-attention RMSNorm | `3.87714e-05` | `0.000426032` |
|
||||
| MLP output | `0.000674194` | `0.0500984` |
|
||||
| Layer-0 output | `0.000808998` | `0.0526199` |
|
||||
|
||||
This confirms the residual error is already present in the captured projection
|
||||
outputs and is amplified by the MLP. The NVFP4 and MLP code audits found no
|
||||
remaining active source-level branch difference. Further edits to flattening,
|
||||
contiguity, dequantization dtype, `pre_quant_scale`, SiLU ordering, or residual
|
||||
adds would be speculative and are prohibited until a new code-path difference
|
||||
is demonstrated.
|
||||
|
||||
## Loaded-Comfy Projection Module Probe
|
||||
|
||||
The source audit was followed by a direct replay against tensors consumed and
|
||||
produced by the actual loaded Comfy layer-0 modules. This proves the remaining
|
||||
difference is present **before attention** and is not inferred from source.
|
||||
|
||||
| Item | Evidence |
|
||||
| --- | --- |
|
||||
| Capture location | `/tmp/fl2va-qwen0-loaded-modules` on Spark |
|
||||
| QKV capture SHA-256 | `8998722e73ba24c40c09a86133ccdb22b708998a3e8b8fb4c4293227345f19db` |
|
||||
| MLP capture SHA-256 | `e52ed70f4aa76f1ba26302b9955631d8f490fb6dc269d96a0ce3e1ea981e7ac2` |
|
||||
| Loaded Comfy class | `comfy.ops.mixed_precision_ops.<locals>.MixedPrecisionOps.Linear` |
|
||||
| Loaded Comfy weight | `comfy_kitchen.tensor.base.QuantizedTensor`, logical dtype BF16, `TensorCoreNVFP4Layout`, full precision enabled |
|
||||
| Q/K/V dimensions | `(8192,5120)`, `(1024,5120)`, `(1024,5120)`; no pre-scale |
|
||||
| Gate/up/down dimensions | `(25600,5120)`, `(25600,5120)`, `(5120,25600)`; only down has pre-scale |
|
||||
|
||||
Direct replay on the **identical loaded-Comfy module inputs**:
|
||||
|
||||
| Projection | Mean absolute error | Maximum absolute error |
|
||||
| --- | ---: | ---: |
|
||||
| Q | `6.11124e-05` | `0.0019514` |
|
||||
| K | `8.47002e-05` | `0.00167805` |
|
||||
| V | `5.4459e-05` | `0.00098893` |
|
||||
| Gate | `0.000308973` | `0.0082469` |
|
||||
| Up | `0.00028138` | `0.00792789` |
|
||||
| Down, using Comfy's captured activated input | `0.000424157` | `0.0500984` |
|
||||
|
||||
**Conclusion:** same checkpoint tensors, same BF16 inputs, and the same named
|
||||
Kitchen layout do not currently produce identical projection outputs through
|
||||
direct `Nvfp4Linear` and Comfy's loaded `MixedPrecisionOps.Linear`. The next
|
||||
investigation must capture/compare the effective `QuantizedTensor` parameters
|
||||
and direct packed-weight wrapper at dispatch time. Do not revisit attention,
|
||||
sampler, or MLP ordering until that exact module-dispatch difference is found.
|
||||
|
||||
### Effective Weight Root Cause And Resolution
|
||||
|
||||
The storage-safe dispatch probe found the active difference. Comfy's text
|
||||
encoder calls Qwen with `dtype=torch.float32`; its full-precision NVFP4 branch
|
||||
casts each projection to an effective FP32 matrix before `F.linear`. Direct had
|
||||
constructed the Qwen encoder with BF16, causing its full-precision NVFP4 branch
|
||||
to dequantize the same packed data into BF16.
|
||||
|
||||
The probe proved that packed qdata, tensor scale, and block scale are exact for
|
||||
all Q/K/V/gate/up/down projections. With direct Qwen switched to FP32:
|
||||
|
||||
| Projection check | Result |
|
||||
| --- | --- |
|
||||
| Effective Q/K/V/gate/up/down FP32 weights | Bit-exact to Comfy |
|
||||
| Q/K/V/gate/up linear outputs | Bit-exact to Comfy |
|
||||
| Layer-0 O-projected attention | Mean/max `0` / `0` |
|
||||
| Layer-0 post-attention residual | Mean/max `0` / `0` |
|
||||
| Layer-0 post-attention RMSNorm | Mean/max `0` / `0` |
|
||||
| Layer-0 MLP output | Mean/max `0` / `0` |
|
||||
| Layer-0 final output | Mean/max `0` / `0` |
|
||||
|
||||
The special `down` dispatch sub-probe that supplied Comfy's already-activated
|
||||
input directly to `F.linear` is intentionally invalid because it bypasses the
|
||||
module's required `pre_quant_scale`; the full layer replay above is the valid
|
||||
gate and is exact.
|
||||
|
||||
## Current Runtime Scope
|
||||
|
||||
| Component | Implemented | Known limitation |
|
||||
| --- | --- | --- |
|
||||
| Text-only Qwen | Yes | No vision encoder, MRoPE, image/video expansion, reference labels, or modality tags |
|
||||
| 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 |
|
||||
| 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 | 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 |
|
||||
| Video VAE decoder | Yes | No direct-versus-Comfy same-latent pixel comparison yet |
|
||||
| Audio VAE/decode/mux | No | Preview intentionally produces video only |
|
||||
| Full Ref2VA | No | References, vision conditioning, VAE encode, audio, and muxing are unimplemented |
|
||||
|
||||
## Remaining Gates, In Dependency Order
|
||||
|
||||
Only these are outstanding. Do not recapture or revisit rows marked complete
|
||||
unless the checkpoint, Comfy version, prompt, or backend changes.
|
||||
|
||||
1. **Final rows and video latent.** Final DiT hidden state, final AdaLN, and
|
||||
final RMSNorm are exact. Repair the remaining final target-segment
|
||||
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
|
||||
`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.
|
||||
3. **VAE pixel gate.** Decode the identical captured final video latent in
|
||||
direct and Comfy and compare pixels before ffmpeg. This decides whether any
|
||||
residual grid comes from decoder behavior.
|
||||
4. **End-to-end FL2VA preview.** Generate the reference dragon only after
|
||||
gates 1-8 pass. Compare its raw final latent first, then video.
|
||||
5. **Feature/performance work.** Only then add audio, Ref2VA/reference paths,
|
||||
Sage3, CUDA graphs, and multi-GPU execution.
|
||||
|
||||
## Existing Tools And Their Intended Gate
|
||||
|
||||
| Gate | Tools |
|
||||
| --- | --- |
|
||||
| Qwen captures | `patch_comfy_qwen_output_capture.py`, `patch_comfy_qwen_layer_trace.py`, `patch_comfy_qwen_layer0_sublayers.py`, `patch_comfy_qwen_layer0_projections.py` |
|
||||
| Qwen projection gate | `compare_qwen0_projections.py` |
|
||||
| Qwen local diagnosis | `trace_qwen0_attention.py` needs repair before reuse; it currently passes an obsolete encoder argument |
|
||||
| H3 block capture | `patch_comfy_h3_capture.py`, `patch_comfy_h3_block_capture.py`, `patch_comfy_h3_block0_sublayers.py` |
|
||||
| H3 local diagnosis | `compare_block0_qkv.py`, `compare_attention_backends.py`, `compare_block0_mlp_projections.py`, `localize_block_sublayers.py`, `localize_block_mismatch.py` |
|
||||
| Sampler | `patch_comfy_sampler_capture.py`, `compare_fl2va_steps.py` |
|
||||
| VAE | `smoke_h3_vae_decoder.py`; a same-latent direct-versus-Comfy pixel comparator still needs to be added |
|
||||
|
||||
The `patch_comfy_*` tools modify `/opt/ComfyUI` and are only safe against the
|
||||
restorable Comfy reference image. Capture provenance must be appended here at
|
||||
the time a new capture is made.
|
||||
|
|
@ -56,24 +56,6 @@ def apply_split_half_rope(x: torch.Tensor, rotation: torch.Tensor) -> torch.Tens
|
|||
return torch.cat((pair[..., 0], pair[..., 1], x[..., rotated_width:]), dim=-1)
|
||||
|
||||
|
||||
def rms_rope_split_half_(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
rotation: torch.Tensor,
|
||||
q_weight: torch.Tensor,
|
||||
k_weight: torch.Tensor,
|
||||
eps: float,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Run Comfy Kitchen's standalone fused H3 Q/K normalization and RoPE."""
|
||||
import comfy_kitchen # Registers the independent CUDA extension operators.
|
||||
|
||||
del comfy_kitchen
|
||||
torch.ops.comfy_kitchen.rms_rope_split_half_(
|
||||
q, k, rotation, q_weight, k_weight, eps, rotation.shape[-3] * 2
|
||||
)
|
||||
return q, k
|
||||
|
||||
|
||||
class H3SageAttention(nn.Module):
|
||||
"""One MiniMax H3 attention module, independent of ComfyUI and Raylight."""
|
||||
|
||||
|
|
@ -119,13 +101,12 @@ class H3SageAttention(nn.Module):
|
|||
sequence = x.shape[0]
|
||||
inner = self.heads * self.head_dim
|
||||
q, k, v = self.qkv_proj(x).split(inner, dim=-1)
|
||||
q = q.view(1, sequence, self.heads, self.head_dim)
|
||||
k = k.view(1, sequence, self.heads, self.head_dim)
|
||||
q = rms_norm(q.view(1, sequence, self.heads, self.head_dim), self.q_norm_weight, self.eps)
|
||||
k = rms_norm(k.view(1, sequence, self.heads, self.head_dim), self.k_norm_weight, self.eps)
|
||||
v = v.view(1, sequence, self.heads, self.head_dim)
|
||||
|
||||
q, k = rms_rope_split_half_(q, k, rope_rotation, self.q_norm_weight, self.k_norm_weight, self.eps)
|
||||
q = q.transpose(1, 2).contiguous()
|
||||
k = k.transpose(1, 2).contiguous()
|
||||
q = apply_split_half_rope(q, rope_rotation).transpose(1, 2).contiguous()
|
||||
k = apply_split_half_rope(k, rope_rotation).transpose(1, 2).contiguous()
|
||||
v = v.transpose(1, 2).contiguous()
|
||||
|
||||
out = run_attention(q, k, v, backend=self.backend, is_causal=False)
|
||||
|
|
|
|||
|
|
@ -42,11 +42,11 @@ class H3FinalLayer(nn.Module):
|
|||
return cls(
|
||||
checkpoint.tensor("adaln_t_table", dtype=torch.float32),
|
||||
checkpoint.tensor("final_layer.norm.weight", dtype=output_dtype),
|
||||
checkpoint.tensor("final_layer.adaln_proj.linear.weight", dtype=torch.bfloat16),
|
||||
checkpoint.tensor("final_layer.adaln_proj.linear.bias", dtype=torch.float16),
|
||||
checkpoint.tensor("final_layer.video_out.weight", dtype=torch.bfloat16),
|
||||
checkpoint.tensor("final_layer.adaln_proj.linear.weight", dtype=torch.float32),
|
||||
checkpoint.tensor("final_layer.adaln_proj.linear.bias", dtype=torch.float32),
|
||||
checkpoint.tensor("final_layer.video_out.weight", dtype=torch.float32),
|
||||
checkpoint.tensor("final_layer.video_out.bias", dtype=torch.float32),
|
||||
checkpoint.tensor("final_layer.audio_out.weight", dtype=torch.bfloat16),
|
||||
checkpoint.tensor("final_layer.audio_out.weight", dtype=torch.float32),
|
||||
checkpoint.tensor("final_layer.audio_out.bias", dtype=torch.float32),
|
||||
)
|
||||
|
||||
|
|
@ -64,9 +64,10 @@ class H3FinalLayer(nn.Module):
|
|||
|
||||
video_start, video_stop, video_row = video_segment
|
||||
audio_start, audio_stop, audio_row = audio_segment
|
||||
video_hidden = (rms_norm(hidden[video_start:video_stop], self.norm_weight, self.eps) * (1.0 + scale[video_row]) + shift[video_row]).to(torch.float32)
|
||||
audio_hidden = (rms_norm(hidden[audio_start:audio_stop], self.norm_weight, self.eps) * (1.0 + scale[audio_row]) + shift[audio_row]).to(torch.float32)
|
||||
normalized = rms_norm(hidden, self.norm_weight, self.eps)
|
||||
video_hidden = normalized[video_start:video_stop] * (1 + scale[video_row].to(hidden.dtype)) + shift[video_row].to(hidden.dtype)
|
||||
audio_hidden = normalized[audio_start:audio_stop] * (1 + scale[audio_row].to(hidden.dtype)) + shift[audio_row].to(hidden.dtype)
|
||||
return (
|
||||
functional.linear(video_hidden, self.video_weight, self.video_bias),
|
||||
functional.linear(audio_hidden, self.audio_weight, self.audio_bias),
|
||||
functional.linear(video_hidden.float(), self.video_weight, self.video_bias),
|
||||
functional.linear(audio_hidden.float(), self.audio_weight, self.audio_bias),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -67,8 +67,8 @@ class Nvfp4Linear(nn.Module):
|
|||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
if x.shape[-1] != self.in_features:
|
||||
raise ValueError(f"Expected feature width {self.in_features}, received {x.shape[-1]}.")
|
||||
if x.dtype not in (torch.float16, torch.bfloat16, torch.float32):
|
||||
raise ValueError("NVFP4 linear accepts FP16, BF16, or FP32 activations.")
|
||||
if x.dtype not in (torch.float16, torch.bfloat16):
|
||||
raise ValueError("NVFP4 linear accepts FP16 or BF16 activations.")
|
||||
|
||||
from comfy_kitchen.tensor import QuantizedTensor
|
||||
|
||||
|
|
@ -82,8 +82,6 @@ class Nvfp4Linear(nn.Module):
|
|||
weight = packed_weight.dequantize().to(flat_x)
|
||||
output = functional.linear(flat_x, weight, bias)
|
||||
return output.reshape(*original_shape, self.out_features)
|
||||
if x.dtype == torch.float32:
|
||||
raise ValueError("Quantized NVFP4 activation GEMM requires FP16 or BF16 activations.")
|
||||
packed_x = QuantizedTensor.from_float(flat_x, "TensorCoreNVFP4Layout")
|
||||
output = functional.linear(packed_x, packed_weight, bias)
|
||||
return output[:flat_x.shape[0], :self.out_features].reshape(*original_shape, self.out_features)
|
||||
|
|
|
|||
|
|
@ -20,7 +20,7 @@ def patchify_video(latent: torch.Tensor) -> torch.Tensor:
|
|||
|
||||
|
||||
def pack_audio(latent: torch.Tensor) -> torch.Tensor:
|
||||
return latent[0].permute(1, 2, 0).reshape(-1, latent.shape[1]).transpose(0, 1).contiguous().transpose(0, 1)
|
||||
return latent[0].permute(1, 2, 0).reshape(-1, latent.shape[1])
|
||||
|
||||
|
||||
def unpatchify_video(rows: torch.Tensor, frames: int, latent_height: int, latent_width: int) -> torch.Tensor:
|
||||
|
|
@ -58,10 +58,10 @@ class H3PromptPacker:
|
|||
"""Build `[text | audio | video]` tokens for prompt-only H3 T2V."""
|
||||
|
||||
def __init__(self, checkpoint):
|
||||
self.video_weight = checkpoint.tensor("video_patch_proj.weight", dtype=torch.bfloat16).to(torch.float32)
|
||||
self.video_bias = checkpoint.tensor("video_patch_proj.bias", dtype=torch.bfloat16).to(torch.float32)
|
||||
self.audio_weight = checkpoint.tensor("audio_patch_proj.weight", dtype=torch.bfloat16).to(torch.float32)
|
||||
self.audio_bias = checkpoint.tensor("audio_patch_proj.bias", dtype=torch.bfloat16).to(torch.float32)
|
||||
self.video_weight = checkpoint.tensor("video_patch_proj.weight", dtype=torch.float32)
|
||||
self.video_bias = checkpoint.tensor("video_patch_proj.bias", dtype=torch.float32)
|
||||
self.audio_weight = checkpoint.tensor("audio_patch_proj.weight", dtype=torch.float32)
|
||||
self.audio_bias = checkpoint.tensor("audio_patch_proj.bias", dtype=torch.float32)
|
||||
self.text_weight = checkpoint.tensor("condition_proj.weight", dtype=torch.bfloat16)
|
||||
self.text_bias = checkpoint.tensor("condition_proj.bias", dtype=torch.bfloat16)
|
||||
|
||||
|
|
@ -84,7 +84,7 @@ class H3PromptPacker:
|
|||
row = {value: index for index, value in enumerate(unique_times)}
|
||||
video_row, audio_row = row[video_time] * 3, row[audio_time] * 3
|
||||
times = torch.tensor(unique_times, device=hidden.device, dtype=torch.float32)
|
||||
positions = torch.cat((torch.stack((torch.arange(text_length, dtype=torch.float64), torch.zeros(text_length), torch.zeros(text_length)), dim=-1), _audio_positions(audio.shape[-1], float(text_length), video.shape[-1], video.shape[-2]), _video_positions(video.shape[2], video.shape[-2], video.shape[-1], float(text_length))))
|
||||
positions = torch.cat((torch.stack((torch.arange(text_length, dtype=torch.float64), torch.zeros(text_length), torch.zeros(text_length)), dim=-1), _audio_positions(audio.shape[-1], float(text_length), video.shape[-1], video.shape[-2]), _video_positions(video.shape[2], video.shape[-2], video.shape[-1], float(text_length) + audio.shape[-1])))
|
||||
block_video_segment = (text_length + audio_length, hidden.shape[0], video_row)
|
||||
block_audio_segment = (text_length, text_length + audio_length, audio_row + 2)
|
||||
final_video_segment = (text_length + audio_length, hidden.shape[0], row[video_time])
|
||||
|
|
|
|||
|
|
@ -66,29 +66,6 @@ def _rope(query: torch.Tensor, key: torch.Tensor, theta: float) -> tuple[torch.T
|
|||
return query_output.to(query.dtype), key_output.to(key.dtype)
|
||||
|
||||
|
||||
def _qwen_attention(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
|
||||
"""Match Comfy's small-input SDPA wrapper for masked Qwen GQA."""
|
||||
kwargs = {"attn_mask": mask.unsqueeze(0).unsqueeze(0), "dropout_p": 0.0, "is_causal": False, "enable_gqa": True}
|
||||
if query.numel() >= 1024 * 128:
|
||||
from torch.nn.attention import SDPBackend, sdpa_kernel
|
||||
|
||||
priority = [SDPBackend.FLASH_ATTENTION, SDPBackend.CUDNN_ATTENTION, SDPBackend.EFFICIENT_ATTENTION, SDPBackend.MATH]
|
||||
params = torch.backends.cuda.SDPAParams(query, key, value, kwargs["attn_mask"], 0.0, False, True)
|
||||
supports_native_gqa = (
|
||||
torch.backends.cuda.can_use_flash_attention(params)
|
||||
or torch.backends.cuda.can_use_cudnn_attention(params)
|
||||
or torch.backends.cuda.can_use_efficient_attention(params)
|
||||
)
|
||||
if not supports_native_gqa:
|
||||
repeats = query.shape[-3] // key.shape[-3]
|
||||
key = key.repeat_interleave(repeats, dim=-3)
|
||||
value = value.repeat_interleave(repeats, dim=-3)
|
||||
kwargs["enable_gqa"] = False
|
||||
with sdpa_kernel(priority, set_priority=True):
|
||||
return F.scaled_dot_product_attention(query, key, value, **kwargs)
|
||||
return F.scaled_dot_product_attention(query, key, value, **kwargs)
|
||||
|
||||
|
||||
class _Qwen3VLBlock(nn.Module):
|
||||
def __init__(self, checkpoint: H3Checkpoint, prefix: str, config: Qwen3VL32BTextConfig, dtype: torch.dtype):
|
||||
super().__init__()
|
||||
|
|
@ -122,7 +99,14 @@ class _Qwen3VLBlock(nn.Module):
|
|||
dtype=query.dtype,
|
||||
device=query.device,
|
||||
).triu_(1)
|
||||
attention = _qwen_attention(query, key, value, causal_mask)
|
||||
attention = F.scaled_dot_product_attention(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
attn_mask=causal_mask,
|
||||
is_causal=False,
|
||||
enable_gqa=True,
|
||||
)
|
||||
hidden_states = residual + self.o_proj(attention.transpose(1, 2).reshape(batch, sequence, -1))
|
||||
residual = hidden_states
|
||||
x = self.post_attention_layernorm(hidden_states)
|
||||
|
|
@ -133,7 +117,7 @@ class Qwen3VL32BTextEncoder(nn.Module):
|
|||
"""Mounted-checkpoint Qwen3-VL prompt conditioner returning layer-50 states."""
|
||||
config = Qwen3VL32BTextConfig()
|
||||
|
||||
def __init__(self, checkpoint_path: str | Path, *, device: str | torch.device = "cuda", dtype: torch.dtype = torch.float32):
|
||||
def __init__(self, checkpoint_path: str | Path, *, device: str | torch.device = "cuda", dtype: torch.dtype = torch.bfloat16):
|
||||
super().__init__()
|
||||
self.checkpoint_path = str(checkpoint_path)
|
||||
self.device_name = str(device)
|
||||
|
|
@ -195,7 +179,7 @@ class Qwen3VL32BTextEncoder(nn.Module):
|
|||
|
||||
class Qwen3VLPromptConditioner:
|
||||
"""Tokenize raw H3 prompt text and produce Qwen layer-50 conditioning."""
|
||||
def __init__(self, checkpoint_path: str | Path, tokenizer_dir: str | Path | None = None, *, device: str | torch.device = "cuda", dtype: torch.dtype = torch.float32):
|
||||
def __init__(self, checkpoint_path: str | Path, tokenizer_dir: str | Path | None = None, *, device: str | torch.device = "cuda", dtype: torch.dtype = torch.bfloat16):
|
||||
from .conditioning import H3PromptTokenizer
|
||||
|
||||
tokenizer_dir = tokenizer_dir or Path(__file__).with_name("qwen25_tokenizer")
|
||||
|
|
|
|||
|
|
@ -1,7 +1,5 @@
|
|||
"""Minimal direct prompt-only H3 video sampler for visual smoke previews."""
|
||||
|
||||
import time
|
||||
|
||||
import torch
|
||||
|
||||
from .packing import H3PromptPacker, unpatchify_video
|
||||
|
|
@ -54,12 +52,8 @@ def sample_video_res_multistep(model, packer: H3PromptPacker, text: torch.Tensor
|
|||
audio_carried = audio * 4.0
|
||||
video_history = audio_history = None
|
||||
video_history_sigma = audio_history_sigma = None
|
||||
total_steps = len(sigmas) - 1
|
||||
started = time.perf_counter()
|
||||
for index, sigma in enumerate(sigmas[:-1], start=1):
|
||||
step_started = time.perf_counter()
|
||||
previous_index = index - 1
|
||||
sigma_down = sigmas[index]
|
||||
for index, sigma in enumerate(sigmas[:-1]):
|
||||
sigma_down = sigmas[index + 1]
|
||||
sigma_audio = _audio_sigma(sigma)
|
||||
native_audio = audio_carried * (sigma_audio / sigma)
|
||||
hidden, times, segments, positions, video_segment, audio_segment = packer(text, video, native_audio, float(sigma))
|
||||
|
|
@ -70,18 +64,11 @@ def sample_video_res_multistep(model, packer: H3PromptPacker, text: torch.Tensor
|
|||
velocity_audio = (1.0 - 4.0) * (audio_carried * carry) + (1.0 + 3.0 * sigma_audio) * velocity_audio_native
|
||||
video_denoised = video - sigma * velocity_video
|
||||
audio_denoised = audio_carried - sigma * velocity_audio
|
||||
previous_sigma = sigmas[previous_index - 1] if previous_index else None
|
||||
previous_sigma = sigmas[index - 1] if index else None
|
||||
video = res_multistep_update(video, video_denoised, sigma, sigma_down, video_history, video_history_sigma, previous_sigma)
|
||||
audio_carried = res_multistep_update(audio_carried, audio_denoised, sigma, sigma_down, audio_history, audio_history_sigma, previous_sigma)
|
||||
video_history, audio_history = video_denoised, audio_denoised
|
||||
video_history_sigma = audio_history_sigma = sigma_down
|
||||
elapsed = time.perf_counter() - started
|
||||
eta = elapsed / index * (total_steps - index)
|
||||
print(
|
||||
f"sampling step {index}/{total_steps}: "
|
||||
f"{time.perf_counter() - step_started:.1f}s, elapsed {elapsed:.1f}s, eta {eta:.1f}s",
|
||||
flush=True,
|
||||
)
|
||||
return video
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -1,49 +0,0 @@
|
|||
"""Compare direct prompt-only FL2VA packing with an immutable Comfy DiT input."""
|
||||
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as functional
|
||||
|
||||
from h3_blackwell_runtime.checkpoint import H3Checkpoint
|
||||
from h3_blackwell_runtime.packing import H3PromptPacker, pack_audio, patchify_video
|
||||
from h3_blackwell_runtime.sampler import _audio_sigma
|
||||
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--capture", type=Path, required=True)
|
||||
parser.add_argument("--refiner-trace", type=Path, required=True)
|
||||
parser.add_argument("--checkpoint", default="/models/minimax_h3_fl2va_pruned_nvfp4.safetensors")
|
||||
args = parser.parse_args()
|
||||
|
||||
captured = torch.load(args.capture, map_location="cuda", weights_only=False)
|
||||
text = torch.load(args.refiner_trace / "refiner_output.pt", map_location="cuda", weights_only=False)
|
||||
if text.ndim == 2:
|
||||
text = text.unsqueeze(0)
|
||||
checkpoint = H3Checkpoint(args.checkpoint)
|
||||
packer = H3PromptPacker(checkpoint)
|
||||
|
||||
video_shape = (1, 24, 7, 12, 20)
|
||||
audio_shape = (1, 32, 2, 37)
|
||||
video = captured["video_x"].to("cuda")
|
||||
audio = captured["audio_x"].to("cuda")
|
||||
sigma = 1 - captured["timesteps"].min()
|
||||
|
||||
with torch.inference_mode():
|
||||
hidden, times, segments, positions, _, _ = packer(text, video, audio, float(sigma))
|
||||
text_rows = text[0].to(torch.bfloat16)
|
||||
audio_rows = functional.linear(pack_audio(audio).float(), packer.audio_weight, packer.audio_bias).to(torch.bfloat16)
|
||||
video_rows = functional.linear(patchify_video(video).float(), packer.video_weight, packer.video_bias).to(torch.bfloat16)
|
||||
|
||||
for name, actual, expected in (
|
||||
("text", text_rows, captured["hidden"][:17]),
|
||||
("audio", audio_rows, captured["hidden"][17:91]),
|
||||
("video", video_rows, captured["hidden"][91:]),
|
||||
("hidden", hidden, captured["hidden"]),
|
||||
("times", times, captured["timesteps"]),
|
||||
("positions", positions, captured["position_ids"]),
|
||||
):
|
||||
delta = (actual.float() - expected.to(actual.device).float()).abs()
|
||||
print(f"{name} shape={tuple(actual.shape)} max_abs={delta.max().item():.6g} mean_abs={delta.mean().item():.6g}")
|
||||
print(f"segments direct={segments} comfy={captured['segments']}")
|
||||
|
|
@ -1,30 +0,0 @@
|
|||
"""Compare direct patch GEMMs with saved outputs from loaded Comfy modules."""
|
||||
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as functional
|
||||
|
||||
from h3_blackwell_runtime.checkpoint import H3Checkpoint
|
||||
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--capture", type=Path, required=True)
|
||||
parser.add_argument("--checkpoint", default="/models/minimax_h3_fl2va_pruned_nvfp4.safetensors")
|
||||
args = parser.parse_args()
|
||||
|
||||
captured = torch.load(args.capture, map_location="cuda", weights_only=False)
|
||||
checkpoint = H3Checkpoint(args.checkpoint)
|
||||
for name, data in captured.items():
|
||||
rows = data["rows"].to("cuda")
|
||||
weight = checkpoint.tensor(f"{name}_patch_proj.weight", dtype=torch.bfloat16).to(torch.float32)
|
||||
bias = checkpoint.tensor(f"{name}_patch_proj.bias", dtype=torch.float32)
|
||||
with torch.inference_mode():
|
||||
direct = functional.linear(rows, weight, bias)
|
||||
row_delta = (rows.float() - data["rows"].to(rows.device).float()).abs()
|
||||
output_delta = (direct.float() - data["output"].to(direct.device).float()).abs()
|
||||
forward_delta = (data["forward_cast"].float() - data["output"].float()).abs()
|
||||
print(f"{name}.rows stride={tuple(rows.stride())} comfy_stride={data['rows_stride']} max_abs={row_delta.max().item():.6g}")
|
||||
print(f"{name}.direct max_abs={output_delta.max().item():.6g} mean_abs={output_delta.mean().item():.6g}")
|
||||
print(f"{name}.comfy_forward_cast max_abs={forward_delta.max().item():.6g} mean_abs={forward_delta.mean().item():.6g}")
|
||||
|
|
@ -1,37 +0,0 @@
|
|||
"""Compare direct FP32 patch GEMMs under Comfy's captured matmul policy."""
|
||||
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as functional
|
||||
|
||||
from h3_blackwell_runtime.checkpoint import H3Checkpoint
|
||||
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--capture", type=Path, required=True)
|
||||
parser.add_argument("--checkpoint", default="/models/minimax_h3_fl2va_pruned_nvfp4.safetensors")
|
||||
args = parser.parse_args()
|
||||
|
||||
captured = torch.load(args.capture, map_location="cuda", weights_only=False)
|
||||
checkpoint = H3Checkpoint(args.checkpoint)
|
||||
previous_tf32 = torch.backends.cuda.matmul.allow_tf32
|
||||
previous_precision = torch.get_float32_matmul_precision()
|
||||
torch.backends.cuda.matmul.allow_tf32 = captured["matmul"]["allow_tf32"]
|
||||
torch.set_float32_matmul_precision(captured["matmul"]["precision"])
|
||||
try:
|
||||
for name in ("video", "audio"):
|
||||
rows = captured[f"{name}_rows"].to("cuda")
|
||||
weight = checkpoint.tensor(f"{name}_patch_proj.weight", dtype=torch.bfloat16).to(torch.float32)
|
||||
bias = checkpoint.tensor(f"{name}_patch_proj.bias", dtype=torch.bfloat16).to(torch.float32)
|
||||
output = functional.linear(rows, weight, bias)
|
||||
for label, actual, expected in (("weight", weight, captured[f"{name}_weight"]), ("bias", bias, captured[f"{name}_bias"]), ("fp32", output, captured[f"{name}_embed_fp32"])):
|
||||
delta = (actual.float() - expected.to(actual.device).float()).abs()
|
||||
print(f"{name}.{label} max_abs={delta.max().item():.6g} mean_abs={delta.mean().item():.6g}")
|
||||
bf16_delta = (output.to(torch.bfloat16).float() - captured[f"{name}_embed_fp32"].to(torch.bfloat16).float()).abs()
|
||||
print(f"{name}.bf16 max_abs={bf16_delta.max().item():.6g} mean_abs={bf16_delta.mean().item():.6g}")
|
||||
print(f"matmul={captured['matmul']}")
|
||||
finally:
|
||||
torch.backends.cuda.matmul.allow_tf32 = previous_tf32
|
||||
torch.set_float32_matmul_precision(previous_precision)
|
||||
|
|
@ -1,30 +0,0 @@
|
|||
"""Replay direct layer-0 projections from actual loaded-Comfy module inputs."""
|
||||
|
||||
import argparse
|
||||
|
||||
import torch
|
||||
|
||||
from h3_blackwell_runtime.qwen3vl_text import Qwen3VL32BTextEncoder
|
||||
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--capture-dir", required=True)
|
||||
parser.add_argument("--checkpoint", default="/text-encoders/qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors")
|
||||
args = parser.parse_args()
|
||||
|
||||
encoder = Qwen3VL32BTextEncoder(args.checkpoint).eval()
|
||||
layer = encoder.layers[0]
|
||||
qkv = torch.load(f"{args.capture_dir}/qwen0_loaded_qkv.pt", map_location="cuda", weights_only=False)
|
||||
mlp = torch.load(f"{args.capture_dir}/qwen0_loaded_mlp.pt", map_location="cuda", weights_only=False)
|
||||
|
||||
with torch.inference_mode():
|
||||
qkv_actual = {"q": layer.q_proj(qkv["input"].to(encoder.dtype)), "k": layer.k_proj(qkv["input"].to(encoder.dtype)), "v": layer.v_proj(qkv["input"].to(encoder.dtype))}
|
||||
gate = layer.gate_proj(mlp["input"].to(encoder.dtype))
|
||||
up = layer.up_proj(mlp["input"].to(encoder.dtype))
|
||||
mlp_actual = {"gate": gate, "up": up, "down": layer.down_proj(mlp["output"]["activated"].to(encoder.dtype))}
|
||||
|
||||
for group, actual, captured in (("qkv", qkv_actual, qkv), ("mlp", mlp_actual, mlp)):
|
||||
print(f"{group}.loaded_metadata={captured['metadata']}")
|
||||
for name, value in actual.items():
|
||||
delta = (value.float() - captured["output"][name].float()).abs()
|
||||
print(f"{group}.{name} max_abs={delta.max().item():.6g} mean_abs={delta.mean().item():.6g}")
|
||||
|
|
@ -1,35 +0,0 @@
|
|||
"""Compare direct Qwen layer-0 projections with a Comfy capture."""
|
||||
|
||||
import argparse
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as functional
|
||||
|
||||
from h3_blackwell_runtime.qwen3vl_text import Qwen3VL32BTextEncoder
|
||||
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--capture-dir", required=True)
|
||||
parser.add_argument("--checkpoint", default="/text-encoders/qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors")
|
||||
args = parser.parse_args()
|
||||
|
||||
encoder = Qwen3VL32BTextEncoder(args.checkpoint).eval()
|
||||
layer = encoder.layers[0]
|
||||
embeds = torch.load(f"{args.capture_dir}/qwen_input_embeds.pt", map_location="cuda", weights_only=False)
|
||||
qkv_expected = torch.load(f"{args.capture_dir}/qwen0_qkv.pt", map_location="cuda", weights_only=False)
|
||||
mlp_expected = torch.load(f"{args.capture_dir}/qwen0_mlp_projections.pt", map_location="cuda", weights_only=False)
|
||||
|
||||
with torch.inference_mode():
|
||||
norm1 = layer.input_layernorm(embeds.to(encoder.dtype))
|
||||
qkv_actual = {"q": layer.q_proj(norm1), "k": layer.k_proj(norm1), "v": layer.v_proj(norm1)}
|
||||
post_attention = torch.load(f"{args.capture_dir}/qwen0_post_attention.pt", map_location="cuda", weights_only=False)
|
||||
norm2 = layer.post_attention_layernorm(post_attention.to(encoder.dtype))
|
||||
gate = layer.gate_proj(norm2)
|
||||
up = layer.up_proj(norm2)
|
||||
activated = functional.silu(gate) * up
|
||||
mlp_actual = {"gate": gate, "up": up, "activated": activated, "down": layer.down_proj(activated)}
|
||||
|
||||
for group, actual, expected in (("qkv", qkv_actual, qkv_expected), ("mlp", mlp_actual, mlp_expected)):
|
||||
for name, value in actual.items():
|
||||
delta = (value.float() - expected[name].float()).abs()
|
||||
print(f"{group}.{name} shape={tuple(value.shape)} max_abs={delta.max().item():.6g} mean_abs={delta.mean().item():.6g}")
|
||||
|
|
@ -1,46 +0,0 @@
|
|||
"""Compare direct NVFP4 wrapper and effective Comfy BF16 weights."""
|
||||
|
||||
import argparse
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as functional
|
||||
|
||||
from h3_blackwell_runtime.qwen3vl_text import Qwen3VL32BTextEncoder
|
||||
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--capture-dir", required=True)
|
||||
parser.add_argument("--checkpoint", default="/text-encoders/qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors")
|
||||
parser.add_argument("--dtype", choices=("bfloat16", "float32"), default="bfloat16")
|
||||
args = parser.parse_args()
|
||||
|
||||
encoder = Qwen3VL32BTextEncoder(args.checkpoint, dtype=getattr(torch, args.dtype)).eval()
|
||||
layer = encoder.layers[0]
|
||||
qkv = torch.load(f"{args.capture_dir}/qwen0_quantized_qkv.pt", map_location="cuda", weights_only=False)
|
||||
mlp = torch.load(f"{args.capture_dir}/qwen0_quantized_mlp.pt", map_location="cuda", weights_only=False)
|
||||
|
||||
for group, capture, modules, inputs in (
|
||||
("qkv", qkv, {"q": layer.q_proj, "k": layer.k_proj, "v": layer.v_proj}, {"q": qkv["input"], "k": qkv["input"], "v": qkv["input"]}),
|
||||
("mlp", mlp, {"gate": layer.gate_proj, "up": layer.up_proj, "down": layer.down_proj}, {"gate": mlp["input"], "up": mlp["input"], "down": mlp["activated"]}),
|
||||
):
|
||||
for name, module in modules.items():
|
||||
expected = capture["modules"][name]
|
||||
packed = module._packed_weight()
|
||||
params = packed._params
|
||||
qdata, scale, block_scale = packed.layout_cls.get_plain_tensors(packed)
|
||||
reference_qdata = expected["qdata"].as_subclass(torch.Tensor).cpu()
|
||||
print(f"{group}.{name}.qdata exact={torch.equal(qdata.cpu(), reference_qdata)} shape={tuple(qdata.shape)}")
|
||||
for field, actual in (("scale", scale), ("block_scale", block_scale)):
|
||||
reference = expected[field].to(actual.device)
|
||||
print(f"{group}.{name}.{field} exact={torch.equal(actual, reference)} shape={tuple(actual.shape)} dtype={actual.dtype}")
|
||||
value = inputs[name].to(encoder.dtype)
|
||||
with torch.inference_mode():
|
||||
dequantized = packed.dequantize()
|
||||
direct = module(value)
|
||||
effective = expected["effective_weight"].to("cuda")
|
||||
comfy_effective = functional.linear(value.to(effective.dtype), effective, None)
|
||||
print(f"{group}.{name}.effective dtype={expected['effective_dtype']} shape={expected['effective_shape']} direct_dtype={encoder.dtype}")
|
||||
for variant, actual in (("weight", dequantized), ("linear", direct), ("effective_linear", comfy_effective)):
|
||||
reference = expected["effective_weight"].to(actual.device) if variant == "weight" else expected["output"].to(actual.device)
|
||||
delta = (actual.float() - reference.float()).abs()
|
||||
print(f"{group}.{name}.{variant}_vs_comfy max_abs={delta.max().item():.6g} mean_abs={delta.mean().item():.6g}")
|
||||
|
|
@ -1,27 +0,0 @@
|
|||
"""Replay direct Qwen against an immutable per-layer Comfy trace."""
|
||||
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
|
||||
from h3_blackwell_runtime.qwen3vl_text import Qwen3VL32BTextEncoder
|
||||
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--trace-dir", type=Path, required=True)
|
||||
parser.add_argument("--checkpoint", default="/text-encoders/qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors")
|
||||
args = parser.parse_args()
|
||||
|
||||
encoder = Qwen3VL32BTextEncoder(args.checkpoint).eval()
|
||||
hidden = torch.load(args.trace_dir / "qwen_input_embeds.pt", map_location="cuda", weights_only=False).to(encoder.dtype)
|
||||
|
||||
with torch.inference_mode():
|
||||
for index, layer in enumerate(encoder.layers):
|
||||
hidden = layer(hidden)
|
||||
expected = torch.load(args.trace_dir / "qwen_layers" / f"{index:02d}.pt", map_location="cuda", weights_only=False)
|
||||
delta = (hidden.float() - expected.float()).abs()
|
||||
print(f"layer={index:02d} max_abs={delta.max().item():.6g} mean_abs={delta.mean().item():.6g}")
|
||||
expected = torch.load(args.trace_dir / "qwen_layer50.pt", map_location="cuda", weights_only=False)
|
||||
delta = (hidden.float() - expected.float()).abs()
|
||||
print(f"layer50 max_abs={delta.max().item():.6g} mean_abs={delta.mean().item():.6g}")
|
||||
|
|
@ -1,29 +0,0 @@
|
|||
"""Replay the direct token refiner against a Comfy boundary trace."""
|
||||
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
from h3_blackwell_runtime.checkpoint import H3Checkpoint
|
||||
from h3_blackwell_runtime.token_refiner import H3TokenRefiner
|
||||
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--trace-dir", type=Path, required=True)
|
||||
parser.add_argument("--checkpoint", default="/models/minimax_h3_fl2va_pruned_nvfp4.safetensors")
|
||||
args = parser.parse_args()
|
||||
|
||||
refiner = H3TokenRefiner(H3Checkpoint(args.checkpoint), attention_backend="sage2").eval()
|
||||
hidden = torch.load(args.trace_dir / "refiner_input.pt", map_location="cuda", weights_only=False)
|
||||
|
||||
with torch.inference_mode():
|
||||
for index, block in enumerate(refiner.blocks):
|
||||
hidden = block(hidden)
|
||||
expected = torch.load(args.trace_dir / f"refiner_block{index}.pt", map_location="cuda", weights_only=False)
|
||||
delta = (hidden.float() - expected.float()).abs()
|
||||
print(f"block={index} max_abs={delta.max().item():.6g} mean_abs={delta.mean().item():.6g}")
|
||||
output = refiner.final_norm
|
||||
output = torch.nn.functional.rms_norm(hidden, output.shape, weight=output.to(hidden), eps=1e-5).unsqueeze(0)
|
||||
expected = torch.load(args.trace_dir / "refiner_output.pt", map_location="cuda", weights_only=False)
|
||||
delta = (output.float() - expected.float()).abs()
|
||||
print(f"output max_abs={delta.max().item():.6g} mean_abs={delta.mean().item():.6g}")
|
||||
|
|
@ -24,14 +24,14 @@ replace_once(
|
|||
" if capture_dir and not H3_CAPTURE_ACTIVE:\n"
|
||||
" H3_CAPTURE_ACTIVE = True\n"
|
||||
" os.makedirs(capture_dir, exist_ok=True)\n"
|
||||
" torch.save({\"hidden\": h.detach().cpu(), \"timesteps\": t_vals.detach().cpu(), \"t_emb\": t_emb.detach().cpu(), \"position_ids\": layout.position_ids, \"segments\": mod_segments, \"video_x\": video_x.detach().cpu(), \"audio_x\": audio_x.detach().cpu()}, os.path.join(capture_dir, \"input.pt\"))\n\n"
|
||||
" torch.save({\"hidden\": h.detach().cpu(), \"timesteps\": t_vals.detach().cpu(), \"position_ids\": layout.position_ids, \"segments\": mod_segments}, os.path.join(capture_dir, \"input.pt\"))\n\n"
|
||||
" # blocks\n patches_replace = transformer_options.get(\"patches_replace\", {})\n",
|
||||
)
|
||||
replace_once(
|
||||
model,
|
||||
" return [-video_out.to(video_x.dtype), -audio_out.to(audio_x.dtype)]\n",
|
||||
" if capture_dir and H3_CAPTURE_ACTIVE:\n"
|
||||
" torch.save({\"video\": video_out.detach().cpu(), \"audio\": audio_out.detach().cpu(), \"video_segment\": video_seg, \"audio_segment\": audio_seg, \"adaln_t_table\": self.adaln_t_table.detach().cpu()}, os.path.join(capture_dir, \"output.pt\"))\n"
|
||||
" torch.save({\"video\": video_out.detach().cpu(), \"audio\": audio_out.detach().cpu(), \"video_segment\": video_seg, \"audio_segment\": audio_seg}, os.path.join(capture_dir, \"output.pt\"))\n"
|
||||
" H3_CAPTURE_ACTIVE = False\n"
|
||||
" return [-video_out.to(video_x.dtype), -audio_out.to(audio_x.dtype)]\n",
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1,42 +0,0 @@
|
|||
"""Capture H3 final-layer intermediates from one reference inference."""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
model = Path("/opt/ComfyUI/comfy/ldm/minimax/model.py")
|
||||
source = model.read_text(encoding="utf-8")
|
||||
old = (
|
||||
" shift, scale = self.adaln_proj(t_emb)\n"
|
||||
" va, vb, vrow = video_seg\n"
|
||||
" aa, ab, arow = audio_seg\n"
|
||||
" hv = (self.norm(x[va:vb]) * (1.0 + scale[vrow]) + shift[vrow]).to(torch.float32)\n"
|
||||
" ha = (self.norm(x[aa:ab]) * (1.0 + scale[arow]) + shift[arow]).to(torch.float32)\n"
|
||||
" return self.video_out(hv), self.audio_out(ha)\n"
|
||||
)
|
||||
new = (
|
||||
" capture_input = t_emb.clone() if H3_CAPTURE_ACTIVE else None\n"
|
||||
" adaln_output = self.adaln_proj(t_emb)\n"
|
||||
" capture_replay = self.adaln_proj.linear(capture_input.clone()) if H3_CAPTURE_ACTIVE else None\n"
|
||||
" capture_forward = self.adaln_proj.linear._forward(capture_input.clone(), self.adaln_proj.linear.weight, self.adaln_proj.linear.bias) if H3_CAPTURE_ACTIVE else None\n"
|
||||
" shift, scale = adaln_output\n"
|
||||
" va, vb, vrow = video_seg\n"
|
||||
" aa, ab, arow = audio_seg\n"
|
||||
" norm_v = self.norm(x[va:vb])\n"
|
||||
" norm_a = self.norm(x[aa:ab])\n"
|
||||
" hv = (norm_v * (1.0 + scale[vrow]) + shift[vrow]).to(torch.float32)\n"
|
||||
" ha = (norm_a * (1.0 + scale[arow]) + shift[arow]).to(torch.float32)\n"
|
||||
" video = self.video_out(hv)\n"
|
||||
" audio = self.audio_out(ha)\n"
|
||||
" capture_dir = os.getenv(\"H3_CAPTURE_DIR\") if H3_CAPTURE_ACTIVE else None\n"
|
||||
" if capture_dir:\n"
|
||||
" torch.cuda.synchronize(x.device)\n"
|
||||
" linear = self.adaln_proj.linear\n"
|
||||
" torch.save({\"hidden\": x.detach().cpu(), \"adaln_input\": capture_input.detach().cpu(), \"adaln_output\": torch.cat((shift, scale), dim=-1).detach().cpu(), \"adaln_replay\": capture_replay.detach().cpu(), \"adaln_forward\": capture_forward.detach().cpu(), \"t_emb\": t_emb.detach().cpu(), \"norm_v\": norm_v.detach().cpu(), \"norm_a\": norm_a.detach().cpu(), \"shift\": shift.detach().cpu(), \"scale\": scale.detach().cpu(), \"adaln_weight\": linear.weight.detach().cpu(), \"adaln_bias\": linear.bias.detach().cpu(), \"adaln_metadata\": {\"class\": f\"{type(linear).__module__}.{type(linear).__qualname__}\", \"weight_class\": f\"{type(linear.weight).__module__}.{type(linear.weight).__qualname__}\", \"quant_format\": getattr(linear, \"quant_format\", None), \"layout_type\": getattr(linear, \"layout_type\", None), \"full_precision_mm\": getattr(linear, \"_full_precision_mm\", None), \"input_scale\": getattr(linear, \"input_scale\", None)}, \"video_hidden\": hv.detach().cpu(), \"audio_hidden\": ha.detach().cpu(), \"video_weight\": self.video_out.weight.detach().cpu(), \"video_bias\": self.video_out.bias.detach().cpu(), \"audio_weight\": self.audio_out.weight.detach().cpu(), \"audio_bias\": self.audio_out.bias.detach().cpu(), \"video\": video.detach().cpu(), \"audio\": audio.detach().cpu(), \"matmul\": {\"allow_tf32\": torch.backends.cuda.matmul.allow_tf32, \"precision\": torch.get_float32_matmul_precision()}}, os.path.join(capture_dir, \"final.pt\"))\n"
|
||||
" return video, audio\n"
|
||||
)
|
||||
if source.count(old) == 1:
|
||||
source = source.replace(old, new)
|
||||
elif new not in source:
|
||||
raise RuntimeError("Unable to locate H3 FinalLayer.forward.")
|
||||
model.write_text(source, encoding="utf-8")
|
||||
print("Applied H3 final-layer capture patch.")
|
||||
|
|
@ -1,33 +0,0 @@
|
|||
"""Probe loaded Comfy H3 patch projections on their captured FP32 rows."""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
path = Path("/opt/ComfyUI/comfy/ldm/minimax/model.py")
|
||||
source = path.read_text(encoding="utf-8")
|
||||
if "import os\n" not in source:
|
||||
source = source.replace("import math\n", "import math\nimport os\n", 1)
|
||||
old = (
|
||||
" video_embed = self.video_patch_proj(all_video_rows).to(dtype)\n"
|
||||
" audio_embed = self.audio_patch_proj(all_audio_rows).to(dtype)\n"
|
||||
)
|
||||
new = (
|
||||
" video_embed_fp32 = self.video_patch_proj(all_video_rows)\n"
|
||||
" audio_embed_fp32 = self.audio_patch_proj(all_audio_rows)\n"
|
||||
" capture_dir = os.getenv(\"H3_CAPTURE_DIR\")\n"
|
||||
" if capture_dir and not os.path.exists(os.path.join(capture_dir, \"h3_loaded_patch_projection_probe.pt\")):\n"
|
||||
" def probe(layer, rows, output):\n"
|
||||
" weight, bias, state = comfy.ops.cast_bias_weight(layer, rows, offloadable=True)\n"
|
||||
" try:\n"
|
||||
" cast_output = torch.nn.functional.linear(rows, weight, bias)\n"
|
||||
" return {\"rows\": rows.detach().cpu(), \"output\": output.detach().cpu(), \"forward_cast\": layer.forward_comfy_cast_weights(rows).detach().cpu(), \"effective_weight\": weight.detach().cpu(), \"effective_bias\": None if bias is None else bias.detach().cpu(), \"effective_weight_dtype\": str(weight.dtype), \"effective_bias_dtype\": None if bias is None else str(bias.dtype), \"cast_output\": cast_output.detach().cpu(), \"rows_dtype\": str(rows.dtype), \"rows_stride\": tuple(rows.stride()), \"weight_dtype\": str(layer.weight.dtype), \"weight_stride\": tuple(layer.weight.stride()), \"bias_dtype\": None if layer.bias is None else str(layer.bias.dtype), \"bias_stride\": None if layer.bias is None else tuple(layer.bias.stride()), \"autocast\": torch.is_autocast_enabled(), \"force_cast\": getattr(layer, \"comfy_force_cast_weights\", False), \"weight_functions\": len(layer.weight_function), \"bias_functions\": len(layer.bias_function)}\n"
|
||||
" finally:\n"
|
||||
" comfy.ops.uncast_bias_weight(layer, weight, bias, state)\n"
|
||||
" torch.save({\"video\": probe(self.video_patch_proj, all_video_rows, video_embed_fp32), \"audio\": probe(self.audio_patch_proj, all_audio_rows, audio_embed_fp32)}, os.path.join(capture_dir, \"h3_loaded_patch_projection_probe.pt\"))\n"
|
||||
" video_embed = video_embed_fp32.to(dtype)\n"
|
||||
" audio_embed = audio_embed_fp32.to(dtype)\n"
|
||||
)
|
||||
if source.count(old) != 1:
|
||||
raise RuntimeError("Unable to locate H3 patch-projection calls.")
|
||||
path.write_text(source.replace(old, new), encoding="utf-8")
|
||||
print("Applied loaded H3 patch-projection probe.")
|
||||
|
|
@ -1,24 +0,0 @@
|
|||
"""Capture Comfy H3 FP32 patch-projection inputs, outputs, and matmul policy."""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
path = Path("/opt/ComfyUI/comfy/ldm/minimax/model.py")
|
||||
source = path.read_text(encoding="utf-8")
|
||||
old = (
|
||||
" video_embed = self.video_patch_proj(all_video_rows).to(dtype)\n"
|
||||
" audio_embed = self.audio_patch_proj(all_audio_rows).to(dtype)\n"
|
||||
)
|
||||
new = (
|
||||
" video_embed_fp32 = self.video_patch_proj(all_video_rows)\n"
|
||||
" audio_embed_fp32 = self.audio_patch_proj(all_audio_rows)\n"
|
||||
" capture_dir = os.getenv(\"H3_CAPTURE_DIR\")\n"
|
||||
" if capture_dir and not os.path.exists(os.path.join(capture_dir, \"h3_patch_projections.pt\")):\n"
|
||||
" torch.save({\"video_rows\": all_video_rows.detach().cpu(), \"audio_rows\": all_audio_rows.detach().cpu(), \"video_embed_fp32\": video_embed_fp32.detach().cpu(), \"audio_embed_fp32\": audio_embed_fp32.detach().cpu(), \"video_weight\": self.video_patch_proj.weight.detach().cpu(), \"video_bias\": self.video_patch_proj.bias.detach().cpu(), \"audio_weight\": self.audio_patch_proj.weight.detach().cpu(), \"audio_bias\": self.audio_patch_proj.bias.detach().cpu(), \"matmul\": {\"allow_tf32\": torch.backends.cuda.matmul.allow_tf32, \"precision\": torch.get_float32_matmul_precision()}}, os.path.join(capture_dir, \"h3_patch_projections.pt\"))\n"
|
||||
" video_embed = video_embed_fp32.to(dtype)\n"
|
||||
" audio_embed = audio_embed_fp32.to(dtype)\n"
|
||||
)
|
||||
if source.count(old) != 1:
|
||||
raise RuntimeError("Unable to locate H3 patch-projection calls.")
|
||||
path.write_text(source.replace(old, new), encoding="utf-8")
|
||||
print("Applied H3 patch-projection probe.")
|
||||
|
|
@ -1,34 +0,0 @@
|
|||
"""Capture Comfy H3 token-refiner boundaries once for direct parity replay."""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
path = Path("/opt/ComfyUI/comfy/ldm/minimax/model.py")
|
||||
source = path.read_text(encoding="utf-8")
|
||||
if "import os\n" not in source:
|
||||
source = source.replace("import math\n", "import math\nimport os\n", 1)
|
||||
|
||||
old = (
|
||||
" def forward(self, x, transformer_options={}):\n"
|
||||
" for block in self.blocks:\n"
|
||||
" x = block(x, transformer_options=transformer_options)\n"
|
||||
" return self.final_norm(x)\n"
|
||||
)
|
||||
new = (
|
||||
" def forward(self, x, transformer_options={}):\n"
|
||||
" capture_dir = os.getenv(\"H3_CAPTURE_DIR\")\n"
|
||||
" if capture_dir:\n"
|
||||
" torch.save(x.detach().cpu(), os.path.join(capture_dir, \"refiner_input.pt\"))\n"
|
||||
" for index, block in enumerate(self.blocks):\n"
|
||||
" x = block(x, transformer_options=transformer_options)\n"
|
||||
" if capture_dir:\n"
|
||||
" torch.save(x.detach().cpu(), os.path.join(capture_dir, f\"refiner_block{index}.pt\"))\n"
|
||||
" output = self.final_norm(x)\n"
|
||||
" if capture_dir:\n"
|
||||
" torch.save(output.detach().cpu(), os.path.join(capture_dir, \"refiner_output.pt\"))\n"
|
||||
" return output\n"
|
||||
)
|
||||
if source.count(old) != 1:
|
||||
raise RuntimeError("Unable to locate TokenRefiner.forward.")
|
||||
path.write_text(source.replace(old, new), encoding="utf-8")
|
||||
print("Applied H3 token-refiner trace patch.")
|
||||
|
|
@ -1,29 +0,0 @@
|
|||
"""Capture H3 curve embeddings after every transformer block."""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
model = Path("/opt/ComfyUI/comfy/ldm/minimax/model.py")
|
||||
source = model.read_text(encoding="utf-8")
|
||||
old = (
|
||||
" else:\n"
|
||||
" h = block(h, t_emb, mod_segments, rope_freqs, transformer_options=transformer_options)\n"
|
||||
" if prefetch_queue is not None:\n"
|
||||
)
|
||||
new = (
|
||||
" else:\n"
|
||||
" if capture_dir and H3_CAPTURE_ACTIVE and i == 0:\n"
|
||||
" torch.save({\"t_emb\": t_emb.clone().detach().cpu(), \"timesteps\": t_vals.clone().detach().cpu()}, os.path.join(capture_dir, \"t_emb_before_block0.pt\"))\n"
|
||||
" h = block(h, t_emb, mod_segments, rope_freqs, transformer_options=transformer_options)\n"
|
||||
" if capture_dir and H3_CAPTURE_ACTIVE:\n"
|
||||
" temb_dir = os.path.join(capture_dir, \"t_emb\")\n"
|
||||
" os.makedirs(temb_dir, exist_ok=True)\n"
|
||||
" torch.save({\"value\": t_emb.detach().cpu(), \"class\": f\"{type(block.adaln_proj.linear).__module__}.{type(block.adaln_proj.linear).__qualname__}\", \"data_ptr\": t_emb.data_ptr(), \"stride\": t_emb.stride()}, os.path.join(temb_dir, f\"{i:02d}.pt\"))\n"
|
||||
" if prefetch_queue is not None:\n"
|
||||
)
|
||||
if source.count(old) == 1:
|
||||
source = source.replace(old, new)
|
||||
elif new not in source:
|
||||
raise RuntimeError("Unable to locate the H3 transformer block loop.")
|
||||
model.write_text(source, encoding="utf-8")
|
||||
print("Applied H3 t_emb trace patch.")
|
||||
|
|
@ -1,32 +0,0 @@
|
|||
"""Capture Comfy Qwen layer-0 projection boundaries for direct parity checks."""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
path = Path("/opt/ComfyUI/comfy/text_encoders/llama.py")
|
||||
source = path.read_text(encoding="utf-8")
|
||||
old = " xq = self.q_proj(hidden_states)\n xk = self.k_proj(hidden_states)\n xv = self.v_proj(hidden_states)\n"
|
||||
new = old + (
|
||||
" capture_dir = os.getenv(\"H3_CAPTURE_DIR\") if getattr(self, \"_h3_trace_index\", -1) == 0 else None\n"
|
||||
" if capture_dir:\n"
|
||||
" torch.save({\"q\": xq.detach().cpu(), \"k\": xk.detach().cpu(), \"v\": xv.detach().cpu()}, os.path.join(capture_dir, \"qwen0_qkv.pt\"))\n"
|
||||
)
|
||||
if source.count(old) != 1:
|
||||
raise RuntimeError("Unable to locate Qwen QKV projections.")
|
||||
source = source.replace(old, new)
|
||||
|
||||
old = " return self.down_proj(self.activation(self.gate_proj(x)) * self.up_proj(x))\n"
|
||||
new = (
|
||||
" gate = self.gate_proj(x)\n"
|
||||
" up = self.up_proj(x)\n"
|
||||
" activated = self.activation(gate) * up\n"
|
||||
" output = self.down_proj(activated)\n"
|
||||
" capture_dir = os.getenv(\"H3_CAPTURE_DIR\") if getattr(self, \"_h3_trace_index\", -1) == 0 else None\n"
|
||||
" if capture_dir:\n"
|
||||
" torch.save({\"gate\": gate.detach().cpu(), \"up\": up.detach().cpu(), \"activated\": activated.detach().cpu(), \"down\": output.detach().cpu()}, os.path.join(capture_dir, \"qwen0_mlp_projections.pt\"))\n"
|
||||
" return output\n"
|
||||
)
|
||||
if source.count(old) != 1:
|
||||
raise RuntimeError("Unable to locate Qwen MLP projections.")
|
||||
path.write_text(source.replace(old, new), encoding="utf-8")
|
||||
print("Applied Qwen layer-0 projection capture patch.")
|
||||
|
|
@ -57,12 +57,7 @@ if source.count(old) != 1:
|
|||
source = source.replace(old, new)
|
||||
|
||||
old = " for i, layer in enumerate(self.layers):\n"
|
||||
new = (
|
||||
" for i, layer in enumerate(self.layers):\n"
|
||||
" layer._h3_trace_index = i\n"
|
||||
" layer.self_attn._h3_trace_index = i\n"
|
||||
" layer.mlp._h3_trace_index = i\n"
|
||||
)
|
||||
new = " for i, layer in enumerate(self.layers):\n layer._h3_trace_index = i\n"
|
||||
if source.count(old) != 1:
|
||||
raise RuntimeError("Unable to locate Qwen decoder layer loop.")
|
||||
path.write_text(source.replace(old, new), encoding="utf-8")
|
||||
|
|
|
|||
|
|
@ -1,36 +0,0 @@
|
|||
"""Capture actual loaded Comfy Qwen layer-0 projection module behavior."""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
path = Path("/opt/ComfyUI/comfy/text_encoders/llama.py")
|
||||
source = path.read_text(encoding="utf-8")
|
||||
old = " xq = self.q_proj(hidden_states)\n xk = self.k_proj(hidden_states)\n xv = self.v_proj(hidden_states)\n"
|
||||
new = old + (
|
||||
" capture_dir = os.getenv(\"H3_CAPTURE_DIR\") if getattr(self, \"_h3_trace_index\", -1) == 0 else None\n"
|
||||
" if capture_dir:\n"
|
||||
" modules = {\"q\": self.q_proj, \"k\": self.k_proj, \"v\": self.v_proj}\n"
|
||||
" metadata = {name: {\"class\": f\"{type(module).__module__}.{type(module).__qualname__}\", \"weight_type\": f\"{type(module.weight).__module__}.{type(module.weight).__qualname__}\", \"weight_dtype\": str(module.weight.dtype), \"weight_shape\": tuple(module.weight.shape), \"layout_type\": getattr(module, \"layout_type\", None), \"full_precision\": getattr(module, \"_full_precision_mm\", None), \"pre_quant_scale\": getattr(module, \"pre_quant_scale\", None) is not None} for name, module in modules.items()}\n"
|
||||
" torch.save({\"input\": hidden_states.detach().cpu(), \"output\": {\"q\": xq.detach().cpu(), \"k\": xk.detach().cpu(), \"v\": xv.detach().cpu()}, \"metadata\": metadata}, os.path.join(capture_dir, \"qwen0_loaded_qkv.pt\"))\n"
|
||||
)
|
||||
if source.count(old) != 1:
|
||||
raise RuntimeError("Unable to locate Qwen QKV projections.")
|
||||
source = source.replace(old, new)
|
||||
|
||||
old = " return self.down_proj(self.activation(self.gate_proj(x)) * self.up_proj(x))\n"
|
||||
new = (
|
||||
" gate = self.gate_proj(x)\n"
|
||||
" up = self.up_proj(x)\n"
|
||||
" activated = self.activation(gate) * up\n"
|
||||
" output = self.down_proj(activated)\n"
|
||||
" capture_dir = os.getenv(\"H3_CAPTURE_DIR\") if getattr(self, \"_h3_trace_index\", -1) == 0 else None\n"
|
||||
" if capture_dir:\n"
|
||||
" modules = {\"gate\": self.gate_proj, \"up\": self.up_proj, \"down\": self.down_proj}\n"
|
||||
" metadata = {name: {\"class\": f\"{type(module).__module__}.{type(module).__qualname__}\", \"weight_type\": f\"{type(module.weight).__module__}.{type(module.weight).__qualname__}\", \"weight_dtype\": str(module.weight.dtype), \"weight_shape\": tuple(module.weight.shape), \"layout_type\": getattr(module, \"layout_type\", None), \"full_precision\": getattr(module, \"_full_precision_mm\", None), \"pre_quant_scale\": getattr(module, \"pre_quant_scale\", None) is not None} for name, module in modules.items()}\n"
|
||||
" torch.save({\"input\": x.detach().cpu(), \"output\": {\"gate\": gate.detach().cpu(), \"up\": up.detach().cpu(), \"activated\": activated.detach().cpu(), \"down\": output.detach().cpu()}, \"metadata\": metadata}, os.path.join(capture_dir, \"qwen0_loaded_mlp.pt\"))\n"
|
||||
" return output\n"
|
||||
)
|
||||
if source.count(old) != 1:
|
||||
raise RuntimeError("Unable to locate Qwen MLP projections.")
|
||||
path.write_text(source.replace(old, new), encoding="utf-8")
|
||||
print("Applied loaded Qwen layer-0 projection probe.")
|
||||
|
|
@ -1,58 +0,0 @@
|
|||
"""Capture loaded Comfy Qwen QuantizedTensor fields and dispatch variants."""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
path = Path("/opt/ComfyUI/comfy/text_encoders/llama.py")
|
||||
source = path.read_text(encoding="utf-8")
|
||||
helper = (
|
||||
"def _h3_owned_cpu(tensor):\n"
|
||||
" return tensor.detach().to(device=\"cpu\", copy=True).contiguous().clone()\n\n"
|
||||
)
|
||||
if "def _h3_owned_cpu(tensor):\n" not in source:
|
||||
source = helper + source
|
||||
old = " xq = self.q_proj(hidden_states)\n xk = self.k_proj(hidden_states)\n xv = self.v_proj(hidden_states)\n"
|
||||
new = old + (
|
||||
" capture_dir = os.getenv(\"H3_CAPTURE_DIR\") if getattr(self, \"_h3_trace_index\", -1) == 0 else None\n"
|
||||
" if capture_dir:\n"
|
||||
" modules = {\"q\": self.q_proj, \"k\": self.k_proj, \"v\": self.v_proj}\n"
|
||||
" def probe(module):\n"
|
||||
" stored = module.weight\n"
|
||||
" qdata, scale, block_scale = stored.layout_cls.get_plain_tensors(stored)\n"
|
||||
" weight, bias, offload_stream = comfy.ops.cast_bias_weight(module, hidden_states, offloadable=True)\n"
|
||||
" try:\n"
|
||||
" output = torch.nn.functional.linear(hidden_states, weight, bias)\n"
|
||||
" return {\"qdata\": _h3_owned_cpu(qdata), \"scale\": _h3_owned_cpu(scale), \"block_scale\": _h3_owned_cpu(block_scale), \"layout_class\": stored._layout_cls, \"orig_dtype\": stored._params.orig_dtype, \"orig_shape\": tuple(stored._params.orig_shape), \"effective_weight\": _h3_owned_cpu(weight), \"effective_dtype\": str(weight.dtype), \"effective_shape\": tuple(weight.shape), \"output\": _h3_owned_cpu(output)}\n"
|
||||
" finally:\n"
|
||||
" comfy.ops.uncast_bias_weight(module, weight, bias, offload_stream)\n"
|
||||
" torch.save({\"input\": hidden_states.detach().cpu(), \"modules\": {name: probe(module) for name, module in modules.items()}}, os.path.join(capture_dir, \"qwen0_quantized_qkv.pt\"))\n"
|
||||
)
|
||||
if source.count(old) != 1:
|
||||
raise RuntimeError("Unable to locate Qwen QKV projections.")
|
||||
source = source.replace(old, new)
|
||||
|
||||
old = " return self.down_proj(self.activation(self.gate_proj(x)) * self.up_proj(x))\n"
|
||||
new = (
|
||||
" gate = self.gate_proj(x)\n"
|
||||
" up = self.up_proj(x)\n"
|
||||
" activated = self.activation(gate) * up\n"
|
||||
" output = self.down_proj(activated)\n"
|
||||
" capture_dir = os.getenv(\"H3_CAPTURE_DIR\") if getattr(self, \"_h3_trace_index\", -1) == 0 else None\n"
|
||||
" if capture_dir:\n"
|
||||
" modules = {\"gate\": self.gate_proj, \"up\": self.up_proj, \"down\": self.down_proj}\n"
|
||||
" def probe(module, value):\n"
|
||||
" stored = module.weight\n"
|
||||
" qdata, scale, block_scale = stored.layout_cls.get_plain_tensors(stored)\n"
|
||||
" weight, bias, offload_stream = comfy.ops.cast_bias_weight(module, value, offloadable=True)\n"
|
||||
" try:\n"
|
||||
" output = torch.nn.functional.linear(value, weight, bias)\n"
|
||||
" return {\"qdata\": _h3_owned_cpu(qdata), \"scale\": _h3_owned_cpu(scale), \"block_scale\": _h3_owned_cpu(block_scale), \"layout_class\": stored._layout_cls, \"orig_dtype\": stored._params.orig_dtype, \"orig_shape\": tuple(stored._params.orig_shape), \"effective_weight\": _h3_owned_cpu(weight), \"effective_dtype\": str(weight.dtype), \"effective_shape\": tuple(weight.shape), \"output\": _h3_owned_cpu(output)}\n"
|
||||
" finally:\n"
|
||||
" comfy.ops.uncast_bias_weight(module, weight, bias, offload_stream)\n"
|
||||
" torch.save({\"input\": x.detach().cpu(), \"activated\": activated.detach().cpu(), \"modules\": {\"gate\": probe(self.gate_proj, x), \"up\": probe(self.up_proj, x), \"down\": probe(self.down_proj, activated)}}, os.path.join(capture_dir, \"qwen0_quantized_mlp.pt\"))\n"
|
||||
" return output\n"
|
||||
)
|
||||
if source.count(old) != 1:
|
||||
raise RuntimeError("Unable to locate Qwen MLP projections.")
|
||||
path.write_text(source.replace(old, new), encoding="utf-8")
|
||||
print("Applied Qwen QuantizedTensor dispatch probe.")
|
||||
|
|
@ -4,7 +4,7 @@ import argparse
|
|||
|
||||
import torch
|
||||
|
||||
from h3_blackwell_runtime.attention import rms_norm, rms_rope_split_half_
|
||||
from h3_blackwell_runtime.attention import apply_split_half_rope, rms_norm
|
||||
from h3_blackwell_runtime.block import gate_segments, modulate_segments
|
||||
from h3_blackwell_runtime.checkpoint import H3Checkpoint
|
||||
from h3_blackwell_runtime.denoiser import H3PackedDenoiser
|
||||
|
|
@ -30,17 +30,8 @@ with torch.inference_mode():
|
|||
shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = adaln(inputs["timesteps"])
|
||||
norm1 = modulate_segments(rms_norm(inputs["hidden"], block.norm1_weight, block.norm_eps), shift_msa, scale_msa, inputs["segments"])
|
||||
q, k, v = block.attention.qkv_proj(norm1).split(7168, dim=-1)
|
||||
raw_q, raw_k, raw_v = q.clone(), k.clone(), v.clone()
|
||||
q_prepared, k_prepared = rms_rope_split_half_(
|
||||
q.view(1, -1, 56, 128),
|
||||
k.view(1, -1, 56, 128),
|
||||
rotation,
|
||||
block.attention.q_norm_weight,
|
||||
block.attention.k_norm_weight,
|
||||
1e-5,
|
||||
)
|
||||
q_prepared = q_prepared.transpose(1, 2).contiguous()
|
||||
k_prepared = k_prepared.transpose(1, 2).contiguous()
|
||||
q_prepared = apply_split_half_rope(rms_norm(q.view(1, -1, 56, 128), block.attention.q_norm_weight, 1e-5), rotation).transpose(1, 2).contiguous()
|
||||
k_prepared = apply_split_half_rope(rms_norm(k.view(1, -1, 56, 128), block.attention.k_norm_weight, 1e-5), rotation).transpose(1, 2).contiguous()
|
||||
v_prepared = v.view(1, -1, 56, 128).transpose(1, 2).contiguous()
|
||||
from sageattention import sageattn
|
||||
attention = block.attention.out_proj(sageattn(q_prepared, k_prepared, v_prepared, is_causal=False, tensor_layout="HND", smooth_k=False).transpose(1, 2).reshape(norm1.shape[0], -1))
|
||||
|
|
@ -51,9 +42,9 @@ with torch.inference_mode():
|
|||
|
||||
for name, actual, expected in (
|
||||
("norm1", norm1, capture["norm1"]),
|
||||
("raw_q", raw_q, capture["qkv_raw"]["q"]),
|
||||
("raw_k", raw_k, capture["qkv_raw"]["k"]),
|
||||
("raw_v", raw_v, capture["qkv_raw"]["v"]),
|
||||
("raw_q", q, capture["qkv_raw"]["q"]),
|
||||
("raw_k", k, capture["qkv_raw"]["k"]),
|
||||
("raw_v", v, capture["qkv_raw"]["v"]),
|
||||
("q", q_prepared, capture["qkv_prepared"]["q"]),
|
||||
("k", k_prepared, capture["qkv_prepared"]["k"]),
|
||||
("v", v_prepared, capture["qkv_prepared"]["v"]),
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue