From 5e84bf0d0b333ed4ddb8be681a7c45a735c95a3a Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Thu, 28 May 2026 12:48:24 -0600 Subject: [PATCH] Fixed issue with hidream-01 that could cause a weird nan state. Took forever to track down as it was 1 in 10 starts. --- .../src/hidream_o1/qwen3_vl_transformers.py | 29 ------------------- 1 file changed, 29 deletions(-) diff --git a/extensions_built_in/diffusion_models/hidream/src/hidream_o1/qwen3_vl_transformers.py b/extensions_built_in/diffusion_models/hidream/src/hidream_o1/qwen3_vl_transformers.py index 17884fc..87d8c57 100644 --- a/extensions_built_in/diffusion_models/hidream/src/hidream_o1/qwen3_vl_transformers.py +++ b/extensions_built_in/diffusion_models/hidream/src/hidream_o1/qwen3_vl_transformers.py @@ -1740,35 +1740,6 @@ class Qwen3VLModel(Qwen3VLPreTrainedModel): input_ids, inputs_embeds=inputs_embeds, image_features=image_embeds ) inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) - elif torch.is_grad_enabled(): - # t2i task: no pixel_values, but we must run the vision encoder with a - # tiny dummy input so that EVERY rank has non-None (zero) gradients for - # vision-encoder parameters. This keeps the FSDP reduce-scatter and the - # replicate-group all-reduce symmetric across t2i and ref-task ranks, - # preventing collective hangs at backward / clip_grad_norm_. - # The dummy output is zeroed out before being added to inputs_embeds, so - # the forward result is numerically identical to the no-pixel_values path. - pe = self.visual.patch_embed # PatchEmbed - t_sz = pe.temporal_patch_size # e.g. 2 - m_sz = self.visual.spatial_merge_size # e.g. 2 - n_patches = t_sz * m_sz * m_sz - patch_dim = pe.in_channels * t_sz * pe.patch_size * pe.patch_size - fake_pv = torch.zeros( - n_patches, - patch_dim, - device=inputs_embeds.device, - dtype=pe.proj.weight.dtype, - ) - fake_grid = torch.tensor( - [[t_sz, m_sz, m_sz]], dtype=torch.long, device=inputs_embeds.device - ) - fake_embs, _ = self.get_image_features(fake_pv, fake_grid) - fake_embs = torch.cat(fake_embs, dim=0).to(inputs_embeds.dtype) - # Multiply by a zero tensor (same dtype/device) so the gradient path - # through the vision encoder is live but the numerical contribution is 0. - inputs_embeds = inputs_embeds + fake_embs.sum() * inputs_embeds.new_zeros( - [] - ) if pixel_values_videos is not None: video_embeds, _ = self.get_video_features(