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

Fix layer offloaded adaln projection layer move

Browse Source
This commit is contained in:
Jaret Burkett
2026-08-06 08:10:20 -06:00
parent 139a38f5bd
commit 7309db4d74
1 changed files with 7 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -217,11 +217,13 @@ 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 = F.linear( # follow temb's device: the raw params may be CPU-resident under layer
temb.float(), # offloading / low_vram, and F.linear bypasses the paging hooks
self.linear.weight.float(), weight = self.linear.weight.to(device=temb.device, dtype=torch.float32)
self.linear.bias.float() if self.linear.bias is not None else None, bias = self.linear.bias
) if bias is not None:
bias = bias.to(device=temb.device, dtype=torch.float32)
x = F.linear(temb.float(), weight, bias)
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)