[diffusion] fix: fix ci on amd (#18716)

Co-authored-by: Mick <mickjagger19@icloud.com>

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

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

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)