Class ContextParallelSplitHook

src/diffusers/hooks/context_parallel.py:129–217

class ContextParallelSplitHook(ModelHook):
    def __init__(self, metadata: ContextParallelModelPlan, parallel_config: ContextParallelConfig) -> None:
        super().__init__()
        self.metadata = metadata
        self.parallel_config = parallel_config
        self.module_forward_metadata = None

    def initialize_hook(self, module):
        cls = unwrap_module(module).__class__
        self.module_forward_metadata = ModuleForwardMetadata(_cls=cls)
        return module

    def pre_forward(self, module, *args, **kwargs):
        args_list = list(args)

        for name, cpm in self.metadata.items():
            if isinstance(cpm, ContextParallelInput) and cpm.split_output:
                continue

            # Maybe the parameter was passed as a keyword argument
            input_val, is_kwarg, index = self.module_forward_metadata._get_parameter_from_args_kwargs(
                name, args_list, kwargs
            )

            if input_val is None:
                continue

            # The input_val may be a tensor or list/tuple of tensors. In certain cases, user may specify to shard
            # the output instead of input for a particular layer by setting split_output=True
            if isinstance(input_val, torch.Tensor):
                input_val = self._prepare_cp_input(input_val, cpm)
            elif isinstance(input_val, (list, tuple)):
                if len(input_val) != len(cpm):
                    raise ValueError(
                        f"Expected input model plan to have {len(input_val)} elements, but got {len(cpm)}."
                    )
                sharded_input_val = []
                for i, x in enumerate(input_val):
                    if torch.is_tensor(x) and not cpm[i].split_output:
                        x = self._prepare_cp_input(x, cpm[i])
                    sharded_input_val.append(x)
                input_val = sharded_input_val
            else:
                raise ValueError(f"Unsupported input type: {type(input_val)}")

            if is_kwarg:
                kwargs[name] = input_val
            elif index is not None and index < len(args_list):
                args_list[index] = input_val
            else:
                raise ValueError(
                    f"An unexpected error occurred while processing the input '{name}'. Please open an "
                    f"issue at https://github.com/huggingface/diffusers/issues and provide a minimal reproducible "
                    f"example along with the full stack trace."
                )

        return tuple(args_list), kwargs