From a85fe116c13ffd05d31bdd9ca6e304ff53df1db2 Mon Sep 17 00:00:00 2001 From: Daniel Maddern Date: Thu, 13 Aug 2026 12:40:44 +0700 Subject: [PATCH] Match H3 final AdaLN bias --- src/h3_blackwell_runtime/final.py | 2 +- tools/patch_comfy_h3_final_capture.py | 6 +++++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/src/h3_blackwell_runtime/final.py b/src/h3_blackwell_runtime/final.py index ed89e5e..c7f2b0d 100644 --- a/src/h3_blackwell_runtime/final.py +++ b/src/h3_blackwell_runtime/final.py @@ -43,7 +43,7 @@ class H3FinalLayer(nn.Module): 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.adaln_proj.linear.bias", 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.audio_out.weight", dtype=torch.bfloat16), diff --git a/tools/patch_comfy_h3_final_capture.py b/tools/patch_comfy_h3_final_capture.py index 21324ba..db7d259 100644 --- a/tools/patch_comfy_h3_final_capture.py +++ b/tools/patch_comfy_h3_final_capture.py @@ -37,7 +37,11 @@ new = ( " 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(), \"adaln_effective_weight\": capture_weight.detach().cpu(), \"adaln_effective_bias\": capture_bias.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" + " video_weight, video_bias, video_offload = comfy.ops.cast_bias_weight(self.video_out, hv, offloadable=True)\n" + " audio_weight, audio_bias, audio_offload = comfy.ops.cast_bias_weight(self.audio_out, ha, offloadable=True)\n" + " comfy.ops.uncast_bias_weight(self.video_out, video_weight, video_bias, video_offload)\n" + " comfy.ops.uncast_bias_weight(self.audio_out, audio_weight, audio_bias, audio_offload)\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(), \"adaln_effective_weight\": capture_weight.detach().cpu(), \"adaln_effective_bias\": capture_bias.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(), \"video_effective_weight\": video_weight.detach().cpu(), \"video_effective_bias\": video_bias.detach().cpu(), \"audio_weight\": self.audio_out.weight.detach().cpu(), \"audio_bias\": self.audio_out.bias.detach().cpu(), \"audio_effective_weight\": audio_weight.detach().cpu(), \"audio_effective_bias\": audio_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: