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

Fix issue with the way chroma handled gradient checkpointing.

Browse Source
This commit is contained in:
Jaret Burkett
2025-05-28 08:41:47 -06:00
parent 34f4c14cd6
commit ffaf2f154a
1 changed files with 7 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
extensions_built_in/diffusion_models/chroma/src/model.py
View File

@@ -96,6 +96,7 @@ class Chroma(nn.Module):
self.params = params self.params = params
self.in_channels = params.in_channels self.in_channels = params.in_channels
self.out_channels = self.in_channels self.out_channels = self.in_channels
self.gradient_checkpointing = False
if params.hidden_size % params.num_heads != 0: if params.hidden_size % params.num_heads != 0:
raise ValueError( raise ValueError(
f"Hidden size {params.hidden_size} must be divisible by num_heads {params.num_heads}" f"Hidden size {params.hidden_size} must be divisible by num_heads {params.num_heads}"
@@ -168,6 +169,9 @@ class Chroma(nn.Module):
# Get the device of the module (assumes all parameters are on the same device) # Get the device of the module (assumes all parameters are on the same device)
return next(self.parameters()).device return next(self.parameters()).device
def enable_gradient_checkpointing(self, enable: bool = True):
self.gradient_checkpointing = enable
def forward( def forward(
self, self,
img: Tensor, img: Tensor,
@@ -246,8 +250,7 @@ class Chroma(nn.Module):
txt_mod = mod_vectors_dict[f"double_blocks.{i}.txt_mod.lin"] txt_mod = mod_vectors_dict[f"double_blocks.{i}.txt_mod.lin"]
double_mod = [img_mod, txt_mod] double_mod = [img_mod, txt_mod]
# just in case in different GPU for simple pipeline parallel if torch.is_grad_enabled() and self.gradient_checkpointing:
if self.training:
img.requires_grad_(True) img.requires_grad_(True)
img, txt = ckpt.checkpoint( img, txt = ckpt.checkpoint(
block, img, txt, pe, double_mod, txt_img_mask block, img, txt, pe, double_mod, txt_img_mask
@@ -260,7 +263,7 @@ class Chroma(nn.Module):
img = torch.cat((txt, img), 1) img = torch.cat((txt, img), 1)
for i, block in enumerate(self.single_blocks): for i, block in enumerate(self.single_blocks):
single_mod = mod_vectors_dict[f"single_blocks.{i}.modulation.lin"] single_mod = mod_vectors_dict[f"single_blocks.{i}.modulation.lin"]
if self.training: if torch.is_grad_enabled() and self.gradient_checkpointing:
img.requires_grad_(True) img.requires_grad_(True)
img = ckpt.checkpoint(block, img, pe, single_mod, txt_img_mask) img = ckpt.checkpoint(block, img, pe, single_mod, txt_img_mask)
else: else: