From 6555b2a71d9f79eecb66c0baacbaf1da99b8d161 Mon Sep 17 00:00:00 2001 From: Bingxu Chen Date: Fri, 13 Feb 2026 11:51:24 +0800 Subject: [PATCH] [diffusion] fix: fix ci on amd (#18716) Co-authored-by: Mick --- .../component_loaders/text_encoder_loader.py | 29 +++++++++++++++++-- .../runtime/loader/fsdp_load.py | 13 +++++++-- .../pipelines_core/stages/denoising_dmd.py | 1 + 3 files changed, 39 insertions(+), 4 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py index 3aeb6093c..d14845059 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py @@ -279,9 +279,34 @@ class TextEncoderLoader(ComponentLoader): # if loaded_weights is not None: weights_not_loaded = weights_to_load - loaded_weights if weights_not_loaded: + # NOTE: + # If we silently continue with uninitialized weights, the text encoder can + # produce NaNs/garbage embeddings that later fail stage verification in a + # hard-to-debug way (e.g., `prompt_embeds` fails the NaN check). + # + # We allow a small set of known-optional parameters to be missing, but + # default to strict behavior for the rest. + allowed_missing_patterns = ( + getattr(model, "_allowed_missing_weights_patterns", []) or [] + ) + unexpected_missing = { + n + for n in weights_not_loaded + if not any(pat in n for pat in allowed_missing_patterns) + } + if unexpected_missing: + raise ValueError( + "Following text encoder weights were not initialized from checkpoint: " + f"{sorted(unexpected_missing)}. " + "This usually indicates a checkpoint/model-arch mismatch or a broken " + "weight-name mapping. If these are truly optional, set " + "`model._allowed_missing_weights_patterns` to whitelist patterns." + ) logger.warning( - "Following model weights were not initialized from " - f"checkpoint: {weights_not_loaded}" + "Following (allowed) text encoder weights were not initialized from " + "checkpoint: %s (allowed patterns: %s)", + sorted(weights_not_loaded), + allowed_missing_patterns, ) return model diff --git a/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py b/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py index 06f80f281..2c723ab27 100644 --- a/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py +++ b/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py @@ -272,7 +272,16 @@ def load_model_from_full_model_state_dict( else: sharded_tensor = full_tensor - if cpu_offload: + # Important: `cpu_offload` is intended for FSDP-managed parameter movement. + # If a parameter is not sharded into a DTensor (i.e., no `device_mesh`), FSDP + # will NOT manage it. Offloading it here would leave CPU parameters that + # later participate in GPU kernels (e.g., conv/embedding), causing device/dtype + # mismatches like "Input type (CUDABFloat16Type) and weight type (CPUBFloat16Type)". + # + # Therefore: + # - For non-FSDP models, keep the historical behavior (allow CPU offload). + # - For FSDP models, do NOT offload non-sharded parameters here. + if cpu_offload and not is_fsdp_model: sharded_tensor = sharded_tensor.cpu() else: full_tensor = full_tensor.to(device=device, dtype=param_dtype) @@ -309,7 +318,7 @@ def load_model_from_full_model_state_dict( sharded_tensor = torch.zeros_like( meta_sharded_param, device=device, dtype=param_dtype ) - if cpu_offload: + if cpu_offload and not is_fsdp_model: sharded_tensor = sharded_tensor.cpu() else: # Initialize with zeros and distribute diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_dmd.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_dmd.py index 057ba2d03..a47477388 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_dmd.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_dmd.py @@ -119,6 +119,7 @@ class DmdDenoisingStage(DenoisingStage): ) else: current_model = self.transformer + self._manage_device_placement(current_model, None, server_args) # Expand latents for I2V noise_latents = latents.clone() latent_model_input = latents.to(target_dtype)