diff --git a/tests/diffusion/test_compile.py b/tests/diffusion/test_compile.py new file mode 100644 index 000000000..522ef343f --- /dev/null +++ b/tests/diffusion/test_compile.py @@ -0,0 +1,35 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch.nn as nn + +from vllm_omni.diffusion.compile import regionally_compile + + +class _WrappedBlock(nn.Module): + def __init__(self) -> None: + super().__init__() + self.compile_called = False + + def compile(self, *args, **kwargs): + self.compile_called = True + return self + + +class _ModelWithWrappedRepeatedBlocks(nn.Module): + _repeated_blocks = ["OriginalBlock"] + _layerwise_offload_blocks_attrs = ["transformer_blocks"] + + def __init__(self) -> None: + super().__init__() + self.transformer_blocks = nn.ModuleList([_WrappedBlock(), _WrappedBlock()]) + self.other_blocks = nn.ModuleList([_WrappedBlock()]) + + +def test_regionally_compile_matches_wrapped_blocks_by_declared_container_attr(): + model = _ModelWithWrappedRepeatedBlocks() + + regionally_compile(model) + + assert all(block.compile_called for block in model.transformer_blocks) + assert not model.other_blocks[0].compile_called diff --git a/vllm_omni/diffusion/compile.py b/vllm_omni/diffusion/compile.py index 66a4675f3..0e4a3ff50 100644 --- a/vllm_omni/diffusion/compile.py +++ b/vllm_omni/diffusion/compile.py @@ -8,6 +8,25 @@ from vllm.logger import init_logger logger = init_logger(__name__) +def _matches_repeated_block( + name: str, + module: nn.Module, + repeated_blocks: list[str], + repeated_block_attrs: list[str], +) -> bool: + class_name = module.__class__.__name__ + if class_name in repeated_blocks: + return True + + for attr in ("_fsdp_wrapped_module", "module", "_orig_mod"): + wrapped = getattr(module, attr, None) + if wrapped is not None and wrapped.__class__.__name__ in repeated_blocks: + return True + + parts = name.split(".") + return len(parts) >= 2 and parts[-2] in repeated_block_attrs and parts[-1].isdigit() + + def regionally_compile(model: nn.Module, *compile_args: Any, **compile_kwargs: Any) -> nn.Module: """ Apply regional compilation to a PyTorch model. @@ -27,15 +46,25 @@ def regionally_compile(model: nn.Module, *compile_args: Any, **compile_kwargs: A logger.warning("Regional compilation skipped because the model does not define `_repeated_blocks`.") return model + repeated_block_attrs = getattr(model, "_layerwise_offload_blocks_attrs", []) + # Check if we have modules with the specified class names has_compiled_region = False - for submod in model.modules(): - if submod.__class__.__name__ in repeated_blocks: + compiled_region_count = 0 + for name, submod in model.named_modules(): + if _matches_repeated_block(name, submod, repeated_blocks, repeated_block_attrs): # Compile this submodule submod.compile(*compile_args, **compile_kwargs) has_compiled_region = True + compiled_region_count += 1 if not has_compiled_region: logger.warning(f"Regional compilation skipped because {repeated_blocks} classes are not found in the model.") + else: + logger.info( + "Regional compilation applied to %d module(s) for repeated blocks %s.", + compiled_region_count, + repeated_blocks, + ) return model