Class ContextParallelGatherHook

src/diffusers/hooks/context_parallel.py:220–251

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

    def post_forward(self, module, output):
        is_tensor = isinstance(output, torch.Tensor)

        if is_tensor:
            output = [output]
        elif not (isinstance(output, (list, tuple)) and all(isinstance(x, torch.Tensor) for x in output)):
            raise ValueError(f"Expected output to be a tensor or a list/tuple of tensors, but got {type(output)}.")

        output = list(output)

        if len(output) != len(self.metadata):
            raise ValueError(f"Expected output to have {len(self.metadata)} elements, but got {len(output)}.")

        for i, cpm in enumerate(self.metadata):
            if cpm is None:
                continue
            if self.parallel_config.ulysses_anything or self.parallel_config.ring_anything:
                output[i] = PartitionAnythingSharder.unshard_anything(
                    output[i], cpm.gather_dim, self.parallel_config._flattened_mesh
                )
            else:
                output[i] = EquipartitionSharder.unshard(
                    output[i], cpm.gather_dim, self.parallel_config._flattened_mesh
                )

        return output[0] if is_tensor else tuple(output)


class AllGatherFunction(torch.autograd.Function):