daniel/h3-blackwell-runtime
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Match H3 final AdaLN bias

Browse source
This commit is contained in:
Daniel Maddern 2026-08-13 12:40:44 +07:00
parent ae84a94159
commit a85fe116c1
2 changed files with 6 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
src/h3_blackwell_runtime/final.py
View file

@ -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),

6
tools/patch_comfy_h3_final_capture.py
View file

@ -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: