github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Upcast adaln pruned layers to fp32 on h3 to prevent overflow.

Browse Source
This commit is contained in:
Jaret Burkett
2026-08-06 07:51:15 -06:00
parent 1e1418b22c
commit 139a38f5bd
1 changed files with 10 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
extensions_built_in/diffusion_models/minimax_h3/src/transformer.py
View File

@@ -192,8 +192,11 @@ class MiniMaxH3AdalnProj(nn.Module):
(M, time_embed_dim) -> ``expand`` tensors of (M * modalities, hidden); row (M, time_embed_dim) -> ``expand`` tensors of (M * modalities, hidden); row
layout ``[t0_mod0, t0_mod1, t0_mod2, t1_mod0, ...]``, addressed by layout ``[t0_mod0, t0_mod1, t0_mod2, t1_mod0, ...]``, addressed by
``timestep_index * MODALITY_NUM + tag``. The SiLU runs at temb's own ``timestep_index * MODALITY_NUM + tag``. The whole projection computes in
(float32) precision; only its result is cast to the projection dtype. float32 (pruned checkpoints store these linears fp16, whose 65504 ceiling
overflows here); the weights stay at their stored dtype and are only
upcast for the matmul, so whole-model casts can't undo this. Outputs are
float32; the use sites cast them to the block dtype.
""" """
def __init__( def __init__(
@@ -214,7 +217,11 @@ class MiniMaxH3AdalnProj(nn.Module):
def forward(self, temb: torch.Tensor): def forward(self, temb: torch.Tensor):
if self.apply_silu: if self.apply_silu:
temb = F.silu(temb) temb = F.silu(temb)
x = self.linear(temb.to(self.linear.weight.dtype)) x = F.linear(
temb.float(),
self.linear.weight.float(),
self.linear.bias.float() if self.linear.bias is not None else None,
)
x = x.view(x.shape[0] * self.modalities, self.expand * self.hidden) x = x.view(x.shape[0] * self.modalities, self.expand * self.hidden)
return x.chunk(self.expand, dim=-1) return x.chunk(self.expand, dim=-1)