Extend Cosmos3 CPU offloading to component and layerwise paths (#4695)

vllm_omni/diffusion/model_loader/diffusers_loader.py

@@ -636,6 +636,14 @@ class DiffusersPipelineLoader:
         # where the transformer lives at "pipe.transformer").
         discovered_modules = ModuleDiscovery.discover(model)
 
+        # Shard only the outermost DiTs. A pipeline may list a DiT and one of its
+        # submodules as separate DiTs (e.g. Cosmos3's transformer and the nested
+        # transformer.language_model) for offload's independent rings; for HSDP an
+        # inner DiT is already covered by its ancestor's _hsdp_shard_conditions, so
+        # sharding it again would double-wrap blocks and require the inner stack to
+        # declare its own conditions.
+        outer_dit_names, outer_dits = discovered_modules.outermost_dits()
+
         # Online FP8 quantization (Fp8OnlineLinearMethod) leaves layer weights
         # as non-contiguous transpose views (qweight.t()) so the Cutlass kernel
         # gets a column-major B. FSDP2 fully_shard rejects non-contiguous params.
@@ -647,14 +655,14 @@ class DiffusersPipelineLoader:
                 prepare_fp8_layers_for_fsdp,
             )
 
-            for trans in discovered_modules.dits:
+            for trans in outer_dits:
                 prepare_fp8_layers_for_fsdp(trans)
 
-        if not discovered_modules.dits:
+        if not outer_dits:
             raise ValueError("No DiT modules discovered for HSDP sharding")
 
-        # Apply HSDP sharding to all discovered DiT transformers
-        for name, trans in zip(discovered_modules.dit_names, discovered_modules.dits):
+        # Apply HSDP sharding to each outermost DiT transformer
+        for name, trans in zip(outer_dit_names, outer_dits):
             logger.debug("Applying HSDP to %s", name)
             apply_hsdp_to_model(trans, hsdp_config, target_device=target_device)