Repository: LAION-AI/Open-Assistant
File: model/model_training/models/patching.py (Python)

Code before:

    out = flash_attn(qkv, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen)
    # out: [bs, seq_len, num_attention_heads, attn_head_size]

    if attention_mask is not None:
        seqs = [out[start:end] for start, end in zip(cu_seqlens[:-1], cu_seqlens[1:])]
        out = pad_sequence(seqs, batch_first=True)
        # restore original sequence length
        out = F.pad(out, (0, 0) * (out.dim() - 2) + (0, max_len - out.size(1)), value=0.0)
    out = out.transpose(1, 2).to(out_dtype)
    return out, None


def add_dropout(module: nn.Module, patched_fwd: Callable, p_dropout: float = 0.1):
    dropout = nn.Dropout(p=p_dropout)
    module.old_forward = module.forward
    module.forward = partial(patched_fwd, dropout, module)


def add_flash_attn(module: nn.Module, causal: bool = True):
    """
    Replaces the standard attention implementation with Flash Attention [1].
    Limitations:
      - Only works for fp16 or bf16 inputs
      - Requires inputs to be on CUDA
      - `output_attentions=True` does work after patching, attention weights will be None
      - Non-contiguous attention masks are not supported (e.g. [1, 1, 0, 1, 1, 0, 0] will just become [1, 1, 1, 1, 1, 0, 0]).

    [1] https://github.com/HazyResearch/flash-attention
    """

    if not hasattr(module, "_attn"):
        warnings.warn("Provided module doesn't have a _attn() function to be patched.")
    flash_attn = FlashSelfAttention(causal=causal)
    module._attn = partial(_patched_gpt_neox_attn, module, flash_attn)


def patch_model(model: GPTNeoXModel, resid_pdrop: Optional[float] = 0.1, flash_attention: bool = True):
    """
    Helper function for patching HF language models.
    Currently supports: GPTNeoX-based models

    Limitations:
      - Flash attention requires CUDA and fp16/bf16 training. It also requires contiguous attention masks.
      - Residual dropout does not support multi-GPU training without DeepDpeed.
    """
    global FlashSelfAttention
    if flash_attention:
        try:
            from flash_attn.modules.mha import FlashSelfAttention  # pyright: reportMissingImports=false
        except ModuleNotFoundError:

Review comment (suggestion):
typo from my side: 
```suggestion
      - `output_attentions=True` does not work after patching, attention weights will be None
```

Code after:

    out = flash_attn(qkv, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen)
    # out: [bs, seq_len, num_attention_heads, attn_head_size]

    if attention_mask is not None:
        seqs = [out[start:end] for start, end in zip(cu_seqlens[:-1], cu_seqlens[1:])]
        out = pad_sequence(seqs, batch_first=True)
        # restore original sequence length
        out = F.pad(out, (0, 0) * (out.dim() - 2) + (0, max_len - out.size(1)), value=0.0)
    out = out.transpose(1, 2).to(out_dtype)
    return out, None


def add_dropout(module: nn.Module, patched_fwd: Callable, p_dropout: float = 0.1):
    dropout = nn.Dropout(p=p_dropout)
    module.old_forward = module.forward
    module.forward = partial(patched_fwd, dropout, module)


def add_flash_attn(module: nn.Module, causal: bool = True):
    """
    Replaces the standard attention implementation with Flash Attention [1].
    Limitations:
      - Only works for fp16 or bf16 inputs
      - Requires inputs to be on CUDA
      - `output_attentions=True` does not work after patching, attention weights will be None
      - Non-contiguous attention masks are not supported (e.g. [1, 1, 0, 1, 1, 0, 0] will just become [1, 1, 1, 1, 1, 0, 0]).

    [1] https://github.com/HazyResearch/flash-attention
    """

    if not hasattr(module, "_attn"):
        warnings.warn("Provided module doesn't have a _attn() function to be patched.")
    flash_attn = FlashSelfAttention(causal=causal)
    module._attn = partial(_patched_gpt_neox_attn, module, flash_attn)


def patch_model(model: GPTNeoXModel, resid_pdrop: Optional[float] = 0.1, flash_attention: bool = True):
    """
    Helper function for patching HF language models.
    Currently supports: GPTNeoX-based models

    Limitations:
      - Flash attention requires CUDA and fp16/bf16 training. It also requires contiguous attention masks.
      - Residual dropout does not support multi-GPU training without DeepDpeed.
    """
    global FlashSelfAttention
    if flash_attention:
        try:
            from flash_attn.modules.mha import FlashSelfAttention  # pyright: reportMissingImports=false
        except ModuleNotFoundError: