Upload batch 65: 500 files (0.01 GiB)

74da989 verified about 1 month ago

16 kB

import torch
import torch.distributed as dist
import math

import time


class ModulePlugin:
    def __init__(self, module, module_id, global_state=None):
        self.module = module
        self.module_id = module_id
        self.global_state = global_state
        self.enable = True
        self.implement_forward()

    @property
    def is_log_node(self):
        return (
            self.global_state.get("dist_controller").rank == 0
            and self.module_id[1] == 0
        )

    @property
    def t(self):
        return self.global_state.get("timestep")

    @property
    def p(self):
        return self.t / 1000

    def implement_forward(self):
        module = self.module
        if not hasattr(module, "old_forward"):
            module.old_forward = module.forward
        self.new_forward = self.get_new_forward()

        def forward(*args, **kwargs):
            self.update_config()  # update config
            return (
                self.new_forward(*args, **kwargs)
                if self.enable
                else self.old_forward(*args, **kwargs)
            )

        module.forward = forward

    def set_enable(self, enable=True):
        self.enable = enable

    def get_new_forward(self):
        raise NotImplementedError

    def update_config(self, config: dict = None):
        if config is None:
            config = self.global_state.get("plugin_configs", {}).get(
                self.module_id[0], {}
            )
        for key, value in config.items():
            setattr(self, key, value)

    def _calculate_padding(self, kernel_size, chunk_dim):
        """Calculate padding parameters for convolution layers.
        Args:
            kernel_size: kernel size tuple or int
            chunk_dim: dimension to chunk/pad along
        Returns:
            tuple: (padding_left, padding_right, padding_flag)
        """
        if isinstance(kernel_size, int):
            kernel_size = tuple([kernel_size] * (chunk_dim - 1))

        # For 3D conv, use width dimension (index 2)
        # For 2D conv, calculate based on chunk_dim
        if len(kernel_size) == 3:
            kernel_width = kernel_size[2]
        else:
            pad_kernel_dim = chunk_dim - 3
            kernel_width = kernel_size[pad_kernel_dim]

        d = kernel_width - 1
        padding_left = d // 2
        padding_right = d - padding_left
        padding_flag = padding_left if d > 0 else 0

        return padding_left, padding_right, padding_flag

    def _pad_context_bilateral_impl(
        self, h, padding_left, padding_right, slice_dim, cat_dim
    ):
        """Core implementation for bilateral padding with context from adjacent ranks.
        Args:
            h: input tensor to pad
            padding_left: number of elements to pad on left
            padding_right: number of elements to pad on right
            slice_dim: dimension to slice from (3 for 3D, 3 for 2D height/width)
            cat_dim: dimension to concatenate along
        Returns:
            padded tensor with context from adjacent ranks
        """
        if padding_left == 0:
            return h

        # Extract slices to share with neighbors
        if slice_dim == 3:
            share_to_left = h[:, :, :, :padding_left].contiguous()
            share_to_right = h[:, :, :, -padding_right:].contiguous()
        elif slice_dim == 4:
            share_to_left = h[:, :, :, :, :padding_left].contiguous()
            share_to_right = h[:, :, :, :, -padding_right:].contiguous()
        elif slice_dim == 2:
            share_to_left = h[:, :, :padding_left].contiguous()
            share_to_right = h[:, :, -padding_right:].contiguous()
        else:  # slice_dim == 1
            share_to_left = h[:, :, :, :padding_left].contiguous()
            share_to_right = h[:, :, :, -padding_right:].contiguous()

        rank = dist.get_rank()
        world_size = dist.get_world_size()
        adj_groups = self.global_state.get("dist_controller").adj_groups

        if rank % 2:
            # Odd rank: pad left first, then right
            if rank:
                padding_list = [torch.zeros_like(share_to_left) for _ in range(2)]
                dist.all_gather(padding_list, share_to_left, group=adj_groups[rank - 1])
                left_context = padding_list[0].to(h.device, non_blocking=True)
            else:
                left_context = torch.zeros_like(share_to_left).to(
                    h.device, non_blocking=True
                )

            if rank != world_size - 1:
                padding_list = [torch.zeros_like(share_to_right) for _ in range(2)]
                dist.all_gather(padding_list, share_to_right, group=adj_groups[rank])
                right_context = padding_list[1].to(h.device, non_blocking=True)
            else:
                right_context = torch.zeros_like(share_to_right).to(
                    h.device, non_blocking=True
                )
        else:
            # Even rank: pad right first, then left
            if rank != world_size - 1:
                padding_list = [torch.zeros_like(share_to_right) for _ in range(2)]
                dist.all_gather(padding_list, share_to_right, group=adj_groups[rank])
                right_context = padding_list[1].to(h.device, non_blocking=True)
            else:
                right_context = torch.zeros_like(share_to_right).to(
                    h.device, non_blocking=True
                )

            if rank:
                padding_list = [torch.zeros_like(share_to_left) for _ in range(2)]
                dist.all_gather(padding_list, share_to_left, group=adj_groups[rank - 1])
                left_context = padding_list[0].to(h.device, non_blocking=True)
            else:
                left_context = torch.zeros_like(share_to_left).to(
                    h.device, non_blocking=True
                )

        h_with_context = torch.cat([left_context, h, right_context], dim=cat_dim)
        return h_with_context

    def _pad_context_bilateral(self, h, padding_left, padding_right, chunk_dim):
        """Pad tensor with context from adjacent ranks using bilateral communication.
        Args:
            h: input tensor to pad
            padding_left: number of elements to pad on left
            padding_right: number of elements to pad on right
            chunk_dim: dimension to pad along
        Returns:
            padded tensor with context from adjacent ranks
        """
        slice_dim = 3 if chunk_dim == 3 else 4
        return self._pad_context_bilateral_impl(
            h, padding_left, padding_right, slice_dim, chunk_dim
        )

    def _pad_context_bilateral_2d(self, h, padding_left, padding_right, chunk_dim):
        """Pad 2D tensor with context from adjacent ranks using bilateral communication.
        Similar to _pad_context_bilateral but pads on chunk_dim-1 dimension for 2D convolutions.
        Args:
            h: input tensor to pad
            padding_left: number of elements to pad on left
            padding_right: number of elements to pad on right
            chunk_dim: base dimension reference (actual padding on chunk_dim-1)
        Returns:
            padded tensor with context from adjacent ranks
        """
        slice_dim = 2 if chunk_dim == 3 else 1
        return self._pad_context_bilateral_impl(
            h, padding_left, padding_right, slice_dim, chunk_dim - 1
        )


class GroupNormPlugin(ModulePlugin):
    def __init__(self, module, module_id, global_state=None):
        super().__init__(module, module_id, global_state)

    def get_new_forward(self):
        module = self.module

        def new_forward(x):
            shape = x.shape
            N, C, G = shape[0], shape[1], module.num_groups
            assert C % G == 0

            x = x.reshape(N, G, -1)

            mean = x.mean(-1, keepdim=True).to(torch.float32)
            dist.all_reduce(mean)

            mean = mean / dist.get_world_size()

            var = ((x - mean.to(x.dtype)) ** 2).mean(-1, keepdim=True).to(torch.float32)

            dist.all_reduce(var)
            var = var / dist.get_world_size()

            x = (x - mean.to(x.dtype)) / (var.to(x.dtype) + module.eps).sqrt()
            x = x.view(shape)

            new_shape = [1 for _ in shape]
            new_shape[1] = -1

            return x * module.weight.view(new_shape) + module.bias.view(new_shape)

        return new_forward


class Conv3DSafeNewPligin(ModulePlugin):
    def __init__(self, module, module_id, global_state=None):
        super().__init__(module, module_id, global_state)

        self.kernel_size = getattr(module, "kernel_size", (1, 1, 1))

        if isinstance(self.kernel_size, int):
            self.kernel_size = (self.kernel_size, self.kernel_size, self.kernel_size)

        self.chunk_dim = global_state.get("chunk_dim")
        self.padding_left, self.padding_right, self.padding_flag = (
            self._calculate_padding(self.kernel_size, self.chunk_dim)
        )

    def pad_context(self, h):
        return self._pad_context_bilateral(
            h, self.padding_left, self.padding_right, self.chunk_dim
        )

    def get_new_forward(self):
        module = self.module

        def new_forward(hidden_states, cache_x=None, *args, **kwargs):
            if self.padding_flag == 0:
                # print(f"padding=0, return old_forward")
                return module.old_forward(hidden_states, cache_x, *args, **kwargs)

            hidden_states = self.pad_context(hidden_states)
            if cache_x is not None:
                cache_x = self.pad_context(cache_x)

            result = module.old_forward(hidden_states, cache_x, *args, **kwargs)
            if self.chunk_dim == 3:
                result = result[
                    :,
                    :,
                    :,
                    self.padding_left : (
                        -self.padding_right if self.padding_right > 0 else None
                    ),
                ]
            else:
                result = result[
                    :,
                    :,
                    :,
                    :,
                    self.padding_left : (
                        -self.padding_right if self.padding_right > 0 else None
                    ),
                ]

            return result

        return new_forward


class Conv2DSafeNewPligin(ModulePlugin):
    def __init__(self, module, module_id, global_state=None):
        super().__init__(module, module_id, global_state)

        self.kernel_size = getattr(module, "kernel_size", (1, 1))
        self.stride = getattr(module, "stride", (1, 1))

        if isinstance(self.kernel_size, int):
            self.kernel_size = (self.kernel_size, self.kernel_size)

        self.chunk_dim = global_state.get("chunk_dim")
        pad_kernel_dim = self.chunk_dim - 3
        kernel_height = self.kernel_size[pad_kernel_dim]
        d = kernel_height - 1
        self.padding_left = d // 2
        self.padding_right = d - self.padding_left
        self.padding = self.padding_left if d > 0 else 0

    def pad_context(self, h):
        return self._pad_context_bilateral_2d(
            h, self.padding_left, self.padding_right, self.chunk_dim
        )

    def get_new_forward(self):
        module = self.module

        def new_forward(hidden_states: torch.Tensor) -> torch.Tensor:
            if self.padding == 0:
                return module.old_forward(hidden_states)

            hidden_states = self.pad_context(hidden_states)

            if self.chunk_dim == 3:
                hidden_states = module.old_forward(hidden_states)[
                    :,
                    :,
                    self.padding_left : (
                        -self.padding_right if self.padding_right > 0 else None
                    ),
                ]
            else:
                hidden_states = module.old_forward(hidden_states)[
                    :,
                    :,
                    :,
                    self.padding_left : (
                        -self.padding_right if self.padding_right > 0 else None
                    ),
                ]
            return hidden_states

        return new_forward


class Conv2DSafeNewPliginStride2(ModulePlugin):
    def __init__(self, module, module_id, global_state=None):
        super().__init__(module, module_id, global_state)

        self.kernel_size = getattr(module, "kernel_size", (1, 1))
        self.stride = getattr(module, "stride", (1, 1))

        if isinstance(self.kernel_size, int):
            self.kernel_size = (self.kernel_size, self.kernel_size)

        self.chunk_dim = global_state.get("chunk_dim")
        assert self.kernel_size[0] == self.kernel_size[1]

        kernel_height = self.kernel_size[0]
        d = kernel_height - 1
        self.padding_left = d // 2
        self.padding_right = d - self.padding_left
        self.padding = self.padding_left if d > 0 else 0

    def pad_context(self, h):
        if self.padding == 0:
            return h

        rank = dist.get_rank()
        world_size = dist.get_world_size()

        if self.chunk_dim == 3:
            share_to_left = h[:, :, : self.padding_left].contiguous()
        else:
            share_to_left = h[:, :, :, : self.padding_left].contiguous()

        if rank < world_size - 1:
            right_context = torch.zeros_like(share_to_left)

            dist.recv(right_context, src=rank + 1)
        if rank > 0:
            dist.send(share_to_left, dst=rank - 1)
        # torch.cuda.synchronize()
        if rank < world_size - 1:
            h_with_context = torch.cat([h, right_context], dim=self.chunk_dim - 1)
        else:
            h_with_context = h
        return h_with_context

    def get_new_forward(self):
        module = self.module

        def new_forward(hidden_states: torch.Tensor) -> torch.Tensor:
            # assert self.global_state["chunk_dim"] == 3 # only support chunk height now
            if self.padding == 0:
                return module.old_forward(hidden_states)

            if self.chunk_dim == 3:
                hidden_states = hidden_states[:, :, :-1, :]
            else:
                hidden_states = hidden_states[:, :, :, :-1]

            hidden_states = self.pad_context(hidden_states)

            if self.chunk_dim == 3:
                hidden_states = torch.nn.functional.pad(hidden_states, (0, 0, 0, 1))
            else:
                hidden_states = torch.nn.functional.pad(hidden_states, (0, 1))
            hidden_states = module.old_forward(
                hidden_states
            )  # [:,:,self.padding_left:-self.padding_right if self.padding_right > 0 else None]
            return hidden_states

        return new_forward


class WanAttentionPlugin(ModulePlugin):
    def __init__(self, module, module_id, global_state=None):
        self.rank = dist.get_rank()
        self.world_size = dist.get_world_size()
        self.gather_dim = global_state.get("chunk_dim")

        super().__init__(module, module_id, global_state)

    def get_new_forward(self):
        module = self.module
        rank = self.rank
        world_size = self.world_size

        def new_forward(hidden_states: torch.Tensor) -> torch.Tensor:
            gathered_tensors = [
                torch.zeros_like(hidden_states) for _ in range(world_size)
            ]
            dist.all_gather(gathered_tensors, hidden_states)

            combined_tensor = torch.cat(gathered_tensors, dim=self.gather_dim)
            forward_output = module.old_forward(combined_tensor)

            # chunk_sizes = [t.size(self.gather_dim) for t in gathered_tensors]
            # start_idx = sum(chunk_sizes[:rank])
            # end_idx = start_idx + chunk_sizes[rank]
            # local_output = forward_output[:, :, :, start_idx:end_idx].contiguous()
            local_output = forward_output.chunk(chunks=world_size, dim=self.gather_dim)[
                rank
            ]

            return local_output

        return new_forward

	import torch
	import torch.distributed as dist
	import math

	import time


	class ModulePlugin:
	def __init__(self, module, module_id, global_state=None):
	self.module = module
	self.module_id = module_id
	self.global_state = global_state
	self.enable = True
	self.implement_forward()

	@property
	def is_log_node(self):
	return (
	self.global_state.get("dist_controller").rank == 0
	and self.module_id[1] == 0
	)

	@property
	def t(self):
	return self.global_state.get("timestep")

	@property
	def p(self):
	return self.t / 1000

	def implement_forward(self):
	module = self.module
	if not hasattr(module, "old_forward"):
	module.old_forward = module.forward
	self.new_forward = self.get_new_forward()

	def forward(args, *kwargs):
	self.update_config() # update config
	return (
	self.new_forward(args, *kwargs)
	if self.enable
	else self.old_forward(args, *kwargs)
	)

	module.forward = forward

	def set_enable(self, enable=True):
	self.enable = enable

	def get_new_forward(self):
	raise NotImplementedError

	def update_config(self, config: dict = None):
	if config is None:
	config = self.global_state.get("plugin_configs", {}).get(
	self.module_id[0], {}
	)
	for key, value in config.items():
	setattr(self, key, value)

	def _calculate_padding(self, kernel_size, chunk_dim):
	"""Calculate padding parameters for convolution layers.

	Args:
	kernel_size: kernel size tuple or int
	chunk_dim: dimension to chunk/pad along

	Returns:
	tuple: (padding_left, padding_right, padding_flag)
	"""
	if isinstance(kernel_size, int):
	kernel_size = tuple([kernel_size] * (chunk_dim - 1))

	# For 3D conv, use width dimension (index 2)
	# For 2D conv, calculate based on chunk_dim
	if len(kernel_size) == 3:
	kernel_width = kernel_size[2]
	else:
	pad_kernel_dim = chunk_dim - 3
	kernel_width = kernel_size[pad_kernel_dim]

	d = kernel_width - 1
	padding_left = d // 2
	padding_right = d - padding_left
	padding_flag = padding_left if d > 0 else 0

	return padding_left, padding_right, padding_flag

	def _pad_context_bilateral_impl(
	self, h, padding_left, padding_right, slice_dim, cat_dim
	):
	"""Core implementation for bilateral padding with context from adjacent ranks.

	Args:
	h: input tensor to pad
	padding_left: number of elements to pad on left
	padding_right: number of elements to pad on right
	slice_dim: dimension to slice from (3 for 3D, 3 for 2D height/width)
	cat_dim: dimension to concatenate along

	Returns:
	padded tensor with context from adjacent ranks
	"""
	if padding_left == 0:
	return h

	# Extract slices to share with neighbors
	if slice_dim == 3:
	share_to_left = h[:, :, :, :padding_left].contiguous()
	share_to_right = h[:, :, :, -padding_right:].contiguous()
	elif slice_dim == 4:
	share_to_left = h[:, :, :, :, :padding_left].contiguous()
	share_to_right = h[:, :, :, :, -padding_right:].contiguous()
	elif slice_dim == 2:
	share_to_left = h[:, :, :padding_left].contiguous()
	share_to_right = h[:, :, -padding_right:].contiguous()
	else: # slice_dim == 1
	share_to_left = h[:, :, :, :padding_left].contiguous()
	share_to_right = h[:, :, :, -padding_right:].contiguous()

	rank = dist.get_rank()
	world_size = dist.get_world_size()
	adj_groups = self.global_state.get("dist_controller").adj_groups

	if rank % 2:
	# Odd rank: pad left first, then right
	if rank:
	padding_list = [torch.zeros_like(share_to_left) for _ in range(2)]
	dist.all_gather(padding_list, share_to_left, group=adj_groups[rank - 1])
	left_context = padding_list[0].to(h.device, non_blocking=True)
	else:
	left_context = torch.zeros_like(share_to_left).to(
	h.device, non_blocking=True
	)

	if rank != world_size - 1:
	padding_list = [torch.zeros_like(share_to_right) for _ in range(2)]
	dist.all_gather(padding_list, share_to_right, group=adj_groups[rank])
	right_context = padding_list[1].to(h.device, non_blocking=True)
	else:
	right_context = torch.zeros_like(share_to_right).to(
	h.device, non_blocking=True
	)
	else:
	# Even rank: pad right first, then left
	if rank != world_size - 1:
	padding_list = [torch.zeros_like(share_to_right) for _ in range(2)]
	dist.all_gather(padding_list, share_to_right, group=adj_groups[rank])
	right_context = padding_list[1].to(h.device, non_blocking=True)
	else:
	right_context = torch.zeros_like(share_to_right).to(
	h.device, non_blocking=True
	)

	if rank:
	padding_list = [torch.zeros_like(share_to_left) for _ in range(2)]
	dist.all_gather(padding_list, share_to_left, group=adj_groups[rank - 1])
	left_context = padding_list[0].to(h.device, non_blocking=True)
	else:
	left_context = torch.zeros_like(share_to_left).to(
	h.device, non_blocking=True
	)

	h_with_context = torch.cat([left_context, h, right_context], dim=cat_dim)
	return h_with_context

	def _pad_context_bilateral(self, h, padding_left, padding_right, chunk_dim):
	"""Pad tensor with context from adjacent ranks using bilateral communication.

	Args:
	h: input tensor to pad
	padding_left: number of elements to pad on left
	padding_right: number of elements to pad on right
	chunk_dim: dimension to pad along

	Returns:
	padded tensor with context from adjacent ranks
	"""
	slice_dim = 3 if chunk_dim == 3 else 4
	return self._pad_context_bilateral_impl(
	h, padding_left, padding_right, slice_dim, chunk_dim
	)

	def _pad_context_bilateral_2d(self, h, padding_left, padding_right, chunk_dim):
	"""Pad 2D tensor with context from adjacent ranks using bilateral communication.

	Similar to _pad_context_bilateral but pads on chunk_dim-1 dimension for 2D convolutions.

	Args:
	h: input tensor to pad
	padding_left: number of elements to pad on left
	padding_right: number of elements to pad on right
	chunk_dim: base dimension reference (actual padding on chunk_dim-1)

	Returns:
	padded tensor with context from adjacent ranks
	"""
	slice_dim = 2 if chunk_dim == 3 else 1
	return self._pad_context_bilateral_impl(
	h, padding_left, padding_right, slice_dim, chunk_dim - 1
	)


	class GroupNormPlugin(ModulePlugin):
	def __init__(self, module, module_id, global_state=None):
	super().__init__(module, module_id, global_state)

	def get_new_forward(self):
	module = self.module

	def new_forward(x):
	shape = x.shape
	N, C, G = shape[0], shape[1], module.num_groups
	assert C % G == 0

	x = x.reshape(N, G, -1)

	mean = x.mean(-1, keepdim=True).to(torch.float32)
	dist.all_reduce(mean)

	mean = mean / dist.get_world_size()

	var = ((x - mean.to(x.dtype)) ** 2).mean(-1, keepdim=True).to(torch.float32)

	dist.all_reduce(var)
	var = var / dist.get_world_size()

	x = (x - mean.to(x.dtype)) / (var.to(x.dtype) + module.eps).sqrt()
	x = x.view(shape)

	new_shape = [1 for _ in shape]
	new_shape[1] = -1

	return x * module.weight.view(new_shape) + module.bias.view(new_shape)

	return new_forward


	class Conv3DSafeNewPligin(ModulePlugin):
	def __init__(self, module, module_id, global_state=None):
	super().__init__(module, module_id, global_state)

	self.kernel_size = getattr(module, "kernel_size", (1, 1, 1))

	if isinstance(self.kernel_size, int):
	self.kernel_size = (self.kernel_size, self.kernel_size, self.kernel_size)

	self.chunk_dim = global_state.get("chunk_dim")
	self.padding_left, self.padding_right, self.padding_flag = (
	self._calculate_padding(self.kernel_size, self.chunk_dim)
	)

	def pad_context(self, h):
	return self._pad_context_bilateral(
	h, self.padding_left, self.padding_right, self.chunk_dim
	)

	def get_new_forward(self):
	module = self.module

	def new_forward(hidden_states, cache_x=None, args, *kwargs):
	if self.padding_flag == 0:
	# print(f"padding=0, return old_forward")
	return module.old_forward(hidden_states, cache_x, args, *kwargs)

	hidden_states = self.pad_context(hidden_states)
	if cache_x is not None:
	cache_x = self.pad_context(cache_x)

	result = module.old_forward(hidden_states, cache_x, args, *kwargs)
	if self.chunk_dim == 3:
	result = result[
	:,
	:,
	:,
	self.padding_left : (
	-self.padding_right if self.padding_right > 0 else None
	),
	]
	else:
	result = result[
	:,
	:,
	:,
	:,
	self.padding_left : (
	-self.padding_right if self.padding_right > 0 else None
	),
	]

	return result

	return new_forward


	class Conv2DSafeNewPligin(ModulePlugin):
	def __init__(self, module, module_id, global_state=None):
	super().__init__(module, module_id, global_state)

	self.kernel_size = getattr(module, "kernel_size", (1, 1))
	self.stride = getattr(module, "stride", (1, 1))

	if isinstance(self.kernel_size, int):
	self.kernel_size = (self.kernel_size, self.kernel_size)

	self.chunk_dim = global_state.get("chunk_dim")
	pad_kernel_dim = self.chunk_dim - 3
	kernel_height = self.kernel_size[pad_kernel_dim]
	d = kernel_height - 1
	self.padding_left = d // 2
	self.padding_right = d - self.padding_left
	self.padding = self.padding_left if d > 0 else 0

	def pad_context(self, h):
	return self._pad_context_bilateral_2d(
	h, self.padding_left, self.padding_right, self.chunk_dim
	)

	def get_new_forward(self):
	module = self.module

	def new_forward(hidden_states: torch.Tensor) -> torch.Tensor:
	if self.padding == 0:
	return module.old_forward(hidden_states)

	hidden_states = self.pad_context(hidden_states)

	if self.chunk_dim == 3:
	hidden_states = module.old_forward(hidden_states)[
	:,
	:,
	self.padding_left : (
	-self.padding_right if self.padding_right > 0 else None
	),
	]
	else:
	hidden_states = module.old_forward(hidden_states)[
	:,
	:,
	:,
	self.padding_left : (
	-self.padding_right if self.padding_right > 0 else None
	),
	]
	return hidden_states

	return new_forward


	class Conv2DSafeNewPliginStride2(ModulePlugin):
	def __init__(self, module, module_id, global_state=None):
	super().__init__(module, module_id, global_state)

	self.kernel_size = getattr(module, "kernel_size", (1, 1))
	self.stride = getattr(module, "stride", (1, 1))

	if isinstance(self.kernel_size, int):
	self.kernel_size = (self.kernel_size, self.kernel_size)

	self.chunk_dim = global_state.get("chunk_dim")
	assert self.kernel_size[0] == self.kernel_size[1]

	kernel_height = self.kernel_size[0]
	d = kernel_height - 1
	self.padding_left = d // 2
	self.padding_right = d - self.padding_left
	self.padding = self.padding_left if d > 0 else 0

	def pad_context(self, h):
	if self.padding == 0:
	return h

	rank = dist.get_rank()
	world_size = dist.get_world_size()

	if self.chunk_dim == 3:
	share_to_left = h[:, :, : self.padding_left].contiguous()
	else:
	share_to_left = h[:, :, :, : self.padding_left].contiguous()

	if rank < world_size - 1:
	right_context = torch.zeros_like(share_to_left)

	dist.recv(right_context, src=rank + 1)
	if rank > 0:
	dist.send(share_to_left, dst=rank - 1)
	# torch.cuda.synchronize()
	if rank < world_size - 1:
	h_with_context = torch.cat([h, right_context], dim=self.chunk_dim - 1)
	else:
	h_with_context = h
	return h_with_context

	def get_new_forward(self):
	module = self.module

	def new_forward(hidden_states: torch.Tensor) -> torch.Tensor:
	# assert self.global_state["chunk_dim"] == 3 # only support chunk height now
	if self.padding == 0:
	return module.old_forward(hidden_states)

	if self.chunk_dim == 3:
	hidden_states = hidden_states[:, :, :-1, :]
	else:
	hidden_states = hidden_states[:, :, :, :-1]

	hidden_states = self.pad_context(hidden_states)

	if self.chunk_dim == 3:
	hidden_states = torch.nn.functional.pad(hidden_states, (0, 0, 0, 1))
	else:
	hidden_states = torch.nn.functional.pad(hidden_states, (0, 1))
	hidden_states = module.old_forward(
	hidden_states
	) # [:,:,self.padding_left:-self.padding_right if self.padding_right > 0 else None]
	return hidden_states

	return new_forward


	class WanAttentionPlugin(ModulePlugin):
	def __init__(self, module, module_id, global_state=None):
	self.rank = dist.get_rank()
	self.world_size = dist.get_world_size()
	self.gather_dim = global_state.get("chunk_dim")

	super().__init__(module, module_id, global_state)

	def get_new_forward(self):
	module = self.module
	rank = self.rank
	world_size = self.world_size

	def new_forward(hidden_states: torch.Tensor) -> torch.Tensor:
	gathered_tensors = [
	torch.zeros_like(hidden_states) for _ in range(world_size)
	]
	dist.all_gather(gathered_tensors, hidden_states)

	combined_tensor = torch.cat(gathered_tensors, dim=self.gather_dim)
	forward_output = module.old_forward(combined_tensor)

	# chunk_sizes = [t.size(self.gather_dim) for t in gathered_tensors]
	# start_idx = sum(chunk_sizes[:rank])
	# end_idx = start_idx + chunk_sizes[rank]
	# local_output = forward_output[:, :, :, start_idx:end_idx].contiguous()
	local_output = forward_output.chunk(chunks=world_size, dim=self.gather_dim)[
	rank
	]

	return local_output

	return new_forward