Add cuda graph support for MLP memcpy

d2/runtime/dispatch_fn.py

@@ -1,3 +1,4 @@
+from typing import Iterable
 import torch
 
 from d2.runtime.attn_kernels.dispatch import (
@@ -104,6 +105,102 @@ def backward(ctx, signal_grad):
             return (grad_attn_out,) + (None,) * 6
 
 
+# No stream because this should always run on the compute stream.
+# FIXME: currently duplicating above for easy debugging
+class pre_all2all_layout_transfer_for_cuda_graph_fwd(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
+        send_seqlens: tuple[torch.Tensor, torch.Tensor],
+        kv_replica_mask: torch.Tensor,
+        send_memcpy_metadata: Iterable[torch.Tensor],
+        # metadata: AlltoAllMetadata, bwd_metadata: AlltoAllMetadata,
+        dispatcher_id: int, is_qkv: bool,
+    ):
+        # a signal tensor output to maintain the autograd graph dependency.
+        signal = torch.empty((1,), dtype=q.dtype, device=q.device)
+        save_tensors = []
+
+        if is_qkv:
+            q = q.contiguous()
+            k = k.contiguous()
+            v = v.contiguous()
+            q_seq_lens, k_seq_lens = send_seqlens
+            pre_a2a_qkv(
+                q, k, v, kv_replica_mask, q_seq_lens, k_seq_lens,
+                *send_memcpy_metadata,
+                is_fwd=True, instance_id=dispatcher_id,
+            )
+        else:
+            q = q.contiguous()
+            assert k is None and v is None
+            pre_a2a_attn_out(
+                q, metadata.seq_lens[0].send_seqlens,
+                *metadata.send_memcpy_metadata, instance_id=dispatcher_id
+            )
+        return q, k, v
+
+    @staticmethod
+    def backward(ctx, q_grad, k_grad, v_grad):
+        return (q_grad, k_grad, v_grad) + (None,) * 5
+
+
+class pre_all2all_layout_transfer_for_cuda_graph_bwd(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
+        recv_seqlens: tuple[torch.Tensor, torch.Tensor],
+        kv_replica_mask: torch.Tensor,
+        recv_memcpy_metadata: Iterable[torch.Tensor],
+        tensor_shapes: Iterable[tuple[int, ...]],
+        dispatcher_id: int, is_qkv: bool,
+    ):
+        signal = torch.empty((1,), dtype=q.dtype, device=q.device)
+        ctx.bwd_recv_shapes = tuple(tensor_shapes)
+        ctx.dispatcher_id = dispatcher_id
+        ctx.is_qkv = is_qkv
+        ctx.save_for_backward(
+            kv_replica_mask,
+            *recv_seqlens,
+            *recv_memcpy_metadata,
+        )
+        return signal
+
+    @staticmethod
+    def backward(ctx, signal_grad):
+        switch_buffer = ctx.dispatcher_id is None
+        if ctx.is_qkv:
+            grad_q_shape = ctx.bwd_recv_shapes[0]
+            grad_k_shape = ctx.bwd_recv_shapes[1]
+            grad_q = torch.empty(grad_q_shape, dtype=signal_grad.dtype, device=signal_grad.device)
+            grad_k = torch.zeros(grad_k_shape, dtype=signal_grad.dtype, device=signal_grad.device)  # must be zero not empty
+            grad_v = torch.zeros_like(grad_k)
+            post_a2a_qkv(
+                grad_q, grad_k, grad_v,
+                # qkv dispatch, q_seq_tokens, k_seq_tokens, v_seq_tokens, q_offset, k_offset, v_offset
+                *ctx.saved_tensors,
+                is_fwd=False,
+                switch_buffer=switch_buffer,
+                instance_id=ctx.dispatcher_id,
+            )
+            grad_k = grad_k.sum(dim=0)
+            grad_v = grad_v.sum(dim=0)
+            return (grad_q, grad_k, grad_v) + (None,) * 6
+        else:
+            grad_attn_out_shape = ctx.bwd_recv_shapes[0]
+            grad_attn_out = torch.empty(
+                grad_attn_out_shape, dtype=signal_grad.dtype, device=signal_grad.device
+            )
+            post_a2a_attn_out(
+                grad_attn_out,
+                *ctx.saved_tensors,
+                switch_buffer=switch_buffer,
+                instance_id=ctx.dispatcher_id,
+            )
+            # the dummy position k,v should have a None gradient.
+            return (grad_attn_out,) + (None,) * 8
+
+
 # we have to add the arg stream here because backward cannot be assigned
 # a stream outside this function.
 class all_to_all(torch.autograd.Function):

d2/runtime/megatron/base_transformer_layer.py

@@ -1,8 +1,9 @@
-from typing import Any, Optional
+from typing import Any, Iterable, Optional
 import time
 import torch
 from torch import Tensor
 import os
+from d2.runtime.dispatch_fn import pre_all2all_layout_transfer_for_cuda_graph_fwd
 from megatron.core import tensor_parallel
 from megatron.core.models.common.embeddings.rope_utils import (
     apply_rotary_pos_emb,
@@ -38,6 +39,235 @@ def __init__(
         from megatron.core.extensions.transformer_engine import TEDotProductAttention
         self.self_attention: SelfAttention
         assert isinstance(self.self_attention.core_attention, TEDotProductAttention)
+        self.pre_attn_cuda_graph = [None] * 2
+        self.post_attn_cuda_graph = None
+
+    def init_pre_attn_cuda_graph(
+        self,
+        prev_layer: "Optional[TransformerLayer]",
+        seq_len: int,
+        max_num_seq: int,
+        max_cp_degree: int,
+        device: torch.device,
+        dtype: torch.dtype
+    ):
+        for dispatcher_id in range(2):
+            self.dispatcher_id = dispatcher_id
+            static_send_memcpy_metadata = (
+                [torch.zeros((max_num_seq,), dtype=int, device=device), torch.zeros((max_num_seq,), dtype=int, device=device)],
+                torch.zeros((max_num_seq, max_cp_degree), device=device, dtype=torch.int8),
+                (
+                    torch.full((max_num_seq,), 0, device=device),
+                    torch.full((max_cp_degree, max_num_seq), -1, device=device),
+                    torch.full((max_cp_degree, max_num_seq), -1, device=device),
+                ),
+            )
+            static_send_memcpy_metadata[0][0][0] = static_send_memcpy_metadata[0][1][0] = seq_len
+            self.qkv_bwd_tensor_shapes = [(seq_len, self.self_attention.query_projection_size), (max_cp_degree, seq_len, self.self_attention.kv_projection_size)]
+            if prev_layer is None:
+                static_input = torch.zeros((seq_len, 1, self.config.hidden_size), device=device, dtype=dtype, requires_grad=True)
+                self.pre_attn_cuda_graph[dispatcher_id] = torch.cuda.make_graphed_callables(self.__pre_attn_cuda_graph, (static_input, *static_send_memcpy_metadata))
+            else:
+                static_core_attn_out = torch.zeros((seq_len, 1, prev_layer.self_attention.query_projection_size), device=device, dtype=dtype, requires_grad=True)
+                static_residual = torch.zeros((seq_len, 1, self.config.hidden_size), device=device, dtype=dtype, requires_grad=True)
+                def post_then_pre_core_attn_cuda_graph(core_attn_out: Tensor, residual: Tensor, *send_memcpy_metadata):
+                    hidden_states = prev_layer._post_attn_cuda_graph(core_attn_out, residual)
+                    return self.__pre_attn_cuda_graph(hidden_states, *send_memcpy_metadata)
+                self.pre_attn_cuda_graph[dispatcher_id] = torch.cuda.make_graphed_callables(post_then_pre_core_attn_cuda_graph, (static_core_attn_out, static_residual, *static_send_memcpy_metadata))
+
+    def init_post_attn_cuda_graph(self, seq_len: int, device: torch.device, dtype: torch.dtype):
+        static_core_attn_out = torch.zeros((seq_len, 1, self.self_attention.query_projection_size), device=device, dtype=dtype, requires_grad=True)
+        static_residual = torch.zeros((seq_len, 1, self.config.hidden_size), device=device, dtype=dtype, requires_grad=True)
+        self.post_attn_cuda_graph = torch.cuda.make_graphed_callables(self._post_attn_cuda_graph, (static_core_attn_out, static_residual))
+
+    def _forward_pre_attn_cuda_graph(
+        self,
+        args: Iterable[torch.Tensor],
+        rotary_pos_emb: Optional[Tensor] = None,
+        rotary_pos_cos: Optional[Tensor] = None,
+        rotary_pos_sin: Optional[Tensor] = None,
+        packed_seq_params: Optional[PingPangSingleStepPackedSeqParams] = None,
+        sequence_len_offset: Optional[Tensor] = None
+    ):
+        # torch.distributed.breakpoint()
+        # TODO: do we need self.dispatcher_id = dispatcher_id?
+        query, key, value, residual = self.pre_attn_cuda_graph[packed_seq_params.dispatcher_id](
+            *args,
+            [seq_len.send_seqlens for seq_len in packed_seq_params.qkv_fwd_metadata.seq_lens],
+            packed_seq_params.qkv_fwd_metadata.kv_replica_mask,
+            packed_seq_params.qkv_fwd_metadata.send_memcpy_metadata,
+        )
+
+        # log_memory_usage(f"(L{self.layer_number}) _forward_pre_core_attn:(init, before input layernorm)")
+
+        # assert rotary_pos_cos is None and rotary_pos_sin is None
+
+        # # For self attention we just duplicate the rotary_pos_emb if it isn't already
+        # if rotary_pos_emb is not None and not isinstance(rotary_pos_emb, tuple):
+        #     rotary_pos_emb = (rotary_pos_emb,) * 2
+
+        # #### Some code in core_attention. This is because we don't want the pos embedding
+        # # being handled in the attention layout (the pos id will be hard to handle)
+        # inference_context = None
+
+        # log_memory_usage(f"(L{self.layer_number}) _forward_pre_core_attn:(after  qkv, before adjust_key_value_for_inference)")
+
+        # query, key, value, rotary_pos_emb, attn_mask_type = self.self_attention._adjust_key_value_for_inference(
+        #     inference_context,
+        #     query,
+        #     key,
+        #     value,
+        #     rotary_pos_emb,
+        #     rotary_pos_cos,
+        #     rotary_pos_sin,
+        #     sequence_len_offset,
+        # )
+        # if packed_seq_params is not None:
+        #     query = query.squeeze(1)
+        #     key = key.squeeze(1)
+        #     value = value.squeeze(1)
+
+        # log_memory_usage(f"(L{self.layer_number}) _forward_pre_core_attn:(after adjust_key_value_for_inference, before rope)")
+
+        # # ================================================
+        # # relative positional embedding (rotary embedding)
+        # # ================================================
+        # if rotary_pos_emb is not None and not self.config.flash_decode:
+        #     q_pos_emb, k_pos_emb = rotary_pos_emb
+
+        #     if packed_seq_params is not None:
+        #         if packed_seq_params.cu_seqlens_q_padded is not None:
+        #             cu_seqlens_q = packed_seq_params.cu_seqlens_q_padded
+        #         else:
+        #             cu_seqlens_q = packed_seq_params.cu_seqlens_q
+        #         if packed_seq_params.cu_seqlens_kv_padded is not None:
+        #             cu_seqlens_kv = packed_seq_params.cu_seqlens_kv_padded
+        #         else:
+        #             cu_seqlens_kv = packed_seq_params.cu_seqlens_kv
+        #     else:
+        #         cu_seqlens_q = cu_seqlens_kv = None
+
+        #     if q_pos_emb is not None:
+        #         # TODO VIJAY: simplify
+        #         query = apply_rotary_pos_emb(
+        #             query, q_pos_emb, config=self.config, cu_seqlens=cu_seqlens_q
+        #         )
+        #     if k_pos_emb is not None:
+        #         key = apply_rotary_pos_emb(
+        #             key, k_pos_emb, config=self.config, cu_seqlens=cu_seqlens_kv
+        #         )
+
+        #     # TODO, can apply positional embedding to value_layer so it has
+        #     # absolute positional embedding.
+        #     # otherwise, only relative positional embedding takes effect
+        #     # value_layer = apply_rotary_pos_emb(value_layer, k_pos_emb)
+        # log_memory_usage(f"(L{self.layer_number}) _forward_pre_core_attn:(after rope, before return)")
+        attn_mask_type = self.self_attention.attn_mask_type
+        return query, key, value, residual, attn_mask_type
+
+    def _forward_post_attn_cuda_graph(
+        self, core_attn_out: Tensor, residual: Tensor,
+        context: Optional[Tensor] = None, context_mask: Optional[Tensor] = None,
+    ):
+        assert context is None and context_mask is None, "not supported in cudagraph"
+        return self.post_attn_cuda_graph(core_attn_out, residual)
+
+
+    def __pre_attn_cuda_graph(
+        self, hidden_states: Tensor,
+        send_seqlens: tuple[Tensor, Tensor],
+        kv_replica_mask: Tensor,
+        send_memcpy_metadata: Iterable[Tensor],
+    ):
+        if self.recompute_input_layernorm:
+            self.input_layernorm_checkpoint = tensor_parallel.CheckpointWithoutOutput()
+            input_layernorm_output = self.input_layernorm_checkpoint.checkpoint(
+                self.input_layernorm, hidden_states
+            )
+        else:
+            input_layernorm_output = self.input_layernorm(hidden_states)
+
+        log_memory_usage(f"(L{self.layer_number}) _forward_pre_core_attn:(after input layernorm, before qkv)")
+
+        # q, k, v
+        query, key, value = self.self_attention.get_query_key_value_tensors(input_layernorm_output, None)
+        query = query.reshape(query.shape[0], self.self_attention.query_projection_size // self.config.tensor_model_parallel_size)
+        key = key.reshape(query.shape[0], self.self_attention.kv_projection_size // self.config.tensor_model_parallel_size)
+        value = value.reshape_as(key)
+
+        query, key, value = pre_all2all_layout_transfer_for_cuda_graph_fwd.apply(
+            query, key, value,
+            send_seqlens, kv_replica_mask, send_memcpy_metadata, self.dispatcher_id, True,
+        )
+
+        return query, key, value, hidden_states
+
+    def _post_attn_cuda_graph(self, core_attn_out: Tensor, residual: Tensor):
+        attention_output_with_bias = self.self_attention.linear_proj(core_attn_out)
+
+        log_memory_usage(f"(L{self.layer_number}) _forward_post_core_attn:(before layernorm)")
+        if self.recompute_input_layernorm:
+            # discard the output of the input layernorm and register the recompute
+            # as a gradient hook of attention_output_with_bias[0]
+            self.input_layernorm_checkpoint.discard_output_and_register_recompute(
+                attention_output_with_bias[0]
+            )
+
+        log_memory_usage(f"(L{self.layer_number}) _forward_post_core_attn:(after layernorm)")
+
+        # TODO: could we move `bias_dropout_add_exec_handler` itself
+        # inside the module provided in the `bias_dropout_add_spec` module?
+        with self.bias_dropout_add_exec_handler():
+            hidden_states = self.self_attn_bda(self.training, self.config.bias_dropout_fusion)(
+                attention_output_with_bias, residual, self.hidden_dropout
+            )
+
+        # Residual connection.
+        residual = hidden_states
+
+        # # Optional Layer norm after self-attention
+        # pre_cross_attn_layernorm_output = self.pre_cross_attn_layernorm(hidden_states)
+
+        # # Cross attention.
+        # log_memory_usage(f"(L{self.layer_number}) _forward_post_core_attn:(cross attention)")
+        # attention_output_with_bias = self.cross_attention(
+        #     pre_cross_attn_layernorm_output,
+        #     attention_mask=context_mask,
+        #     key_value_states=context,
+        #     inference_context=inference_context,
+        # )
+        # log_memory_usage(f"(L{self.layer_number}) _forward_post_core_attn:(after cross attention)")
+
+        # if isinstance(attention_output_with_bias, dict) and "context" in attention_output_with_bias:
+        #     context = attention_output_with_bias["context"]
+
+        attention_output_with_bias = self.pre_cross_attn_layernorm(hidden_states)
+
+        # TODO: could we move `bias_dropout_add_exec_handler` itself
+        # inside the module provided in the `bias_dropout_add_spec` module?
+        with self.bias_dropout_add_exec_handler():
+            hidden_states = self.cross_attn_bda(self.training, self.config.bias_dropout_fusion)(
+                attention_output_with_bias, residual, self.hidden_dropout
+            )
+        log_memory_usage(f"(L{self.layer_number}) _forward_post_core_attn:(after cross attn bda)")
+
+        # Residual connection.
+        residual = hidden_states
+
+        # Optional Layer norm post the cross-attention.
+        if self.recompute_pre_mlp_layernorm:
+            self.pre_mlp_norm_checkpoint = tensor_parallel.CheckpointWithoutOutput()
+            pre_mlp_layernorm_output = self.pre_mlp_norm_checkpoint.checkpoint(
+                self.pre_mlp_layernorm, hidden_states
+            )
+        else:
+            pre_mlp_layernorm_output = self.pre_mlp_layernorm(hidden_states)
+
+        log_memory_usage(f"(L{self.layer_number}) _forward_post_core_attn:(after pre mlp layernorm)")
+
+        mlp_output = self._forward_mlp(pre_mlp_layernorm_output, residual)
+        log_memory_usage(f"(L{self.layer_number}) _forward_post_core_attn:(after mlp)")
+        return mlp_output
 
     def _forward_pre_core_attn(
         self,

d2/runtime/megatron/forward_backward_func.py

@@ -90,7 +90,7 @@ def send_forward_backward_recv_forward_backward(output_tensors, input_tensor_gra
 forward_backward_pipelining_without_interleaving_first_run = True
 
 
-import wlbllm.registry
+# import wlbllm.registry
 def wlb_swap_next_forward_metadata():
     # Call this function before entering forward step.
     swap_metadata_fn = wlbllm.registry.get("swap_metadata_fn")