Mamba在搭载ROCm的AMD GPU上的应用

于 2024-11-21 00:30:00 发布
阅读量962 收藏 11
点赞数 31
分类专栏： # ROCm 文章标签： 人工智能 学习
版权声明：本文为博主原创文章，遵循 CC 4.0 BY-SA 版权协议，转载请附上原文出处链接和本声明。
本文链接：https://blog.csdn.net/eidolon_foot/article/details/143864498
版权

Mamba on AMD GPUs with ROCm — ROCm Blogs

2024年6月28日 作者：Sean Song、Jassani Adeem、Moskvichev Arseny。

最近，Mamba引入了一种新颖的架构，不仅在模型效果上超越了Transformers，还实现了输入序列长度的线性扩展。

在这篇博客中，我们深入探讨了Mamba架构，并展示了如何在搭载ROCm平台的AMD GPU上使用Mamba。

Mamba 基础: 选择性SSM

Transformer架构是一个基础性模型，在包括语言、视觉和音频在内的各种模态中都取得了最新的性能。尽管取得了成功，由于注意力机制的推动，Transformer模型存在一些基本的局限性。首先，它们的计算复杂度随着序列长度呈二次方增长。其次，它们只能在有限的上下文窗口中进行建模。为了应对这些问题，已经提出了许多新颖的架构，例如线性注意力机制、门控卷积和循环模型以及结构化状态空间模型 (SSMs)。虽然这些替代方法通常在计算效率方面可以实现亚二次或线性时间，但在数据建模效果方面往往不如Transformer，这主要通过任务特定的指标（如准确性、精确度、召回率等）来评估。

序列处理问题可以看作是一个压缩任务。最佳的模型架构应有效地压缩序列的上下文并生成有效的结果。尽管具有注意力机制的Transformer架构在效果方面达到了最先进的水平，但在上下文压缩方面的效率却很低。这从自回归推理期间需要存储整个上下文（KV缓存）这一点上就可以看出。

\text{Attention}(Q, K, V) = \text{Softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V

另一方面，循环神经网络 (RNN)由于其有限状态在效率方面表现很好，但它对上下文压缩的方法使其效果较差。

基于这一发现，Mamba模型提出了一种选择性状态空间模型 (SSM) 架构，该架构主要基于结构化状态空间序列模型 (S4)，这是一种广泛与RNN相关的状态空间模型 (SSM)。这种选择性SSM架构结合了RNN固有的效率和选择性机制，以实现Transformer的质量性能。在深入探讨Mamba之前，我们首先回顾一下SSM架构。

状态空间模型 (SSM)

SSM模型被训练来通过中间状态h(t)将输入序列x(t)映射到输出y(t)，中间状态应能正确捕捉序列的上下文并有效地生成输出。离散SSM模型可以通过以下方程定义：

\begin{aligned} \bar{h}_t &= \bar{A} \bar{h}_{t-1} + \bar{B} x_t \\ y_t &= C \bar{h}_t \end{aligned}

上述方程中，用横线表示离散化参数。要注意的是，从连续SSM方程到上述方程需要一个离散化步骤，这里未显示。为了全面了解，我们推荐阅读结构化状态空间序列模型 (S4)。

此SSM方程的一个重要特性是参数 $\bar{A}$、$\bar{B}$ 和 $C$ 随时间保持恒定。这使得模型的动态是恒定的，即线性时间不变 (LTI)。LTI架构在线性时间扩展上表现高效。然而，建模能力受限于LTI特性，因为LTI模型使用同一组参数处理每个输入标记，无法选择性地对不同数据/标记进行处理。这种非选择性LTI导致在长序列中聚合信息，引入噪声并降低建模效果。

通过比较Transformer和SSM，我们可以发现Transformer在效果上更优秀但效率较低，而SSM效率更高但效果较差。

名称

Transformer

SSM

效果

✓

✗

效率

✗

✓

选择性状态空间模型 (SSM)

无论是Transformer还是SSM都不是完美的。Mamba模型通过在SSM模型中融入一个选择机制，以平衡模型的有效性和效率。这种机制使得影响隐藏状态和输出的参数可以依赖于输入。

h_t = \bar{A}(x_t) h_{t-1} + \bar{B}(x_t) x_t

y_t = C(x_t) h_t

这种选择性机制克服了线性时间不变(LTI)模型在有效性上的不足，通过选择性地忽略或关注序列中的某个token。这类似于Transformer的注意力机制，通过为不同的token赋予不同的权重来决定应忽略或关注的程度。但与Transformer不同，后者不进行压缩且需要为每个token存储历史信息（K和V），选择性SSM的循环表示将整个历史数据压缩成一个状态，这可以有效地将上下文信息融入推理过程中。

png

图像来源:Mamba: Linear-Time Sequence Modeling with Selective State Spaces。

算法1和算法2依据SSM和选择性SSM在红色高亮处展示了差异。主要的差异是使参数 Δ_、_B 和 C 成为输入的函数。_Δ_ 控制着当前输入的关注程度或忽略程度。你会发现 A 被保持为常量而不依赖于输入，因为 A 是通过 Δ 离散化的，_Δ_ 中的选择性可以确保在\bar{A}中足够的选择性。通过使 B 和 C 变得动态和选择性，这允许对当前输入 x 如何贡献于中间状态 h 以及中间状态 h 对输出 y 的影响进行细粒度的控制。

选择性SSM的优化

为了在GPU上实现高效执行mamba，利用了两个层面的优化。

  • 选择性扫描算法:

    由于选择性SSM的动态特性（参数 Δ_、_B 和 C 是输入依赖的），无法像SSM（_A_、_B_、_C_ 是固定矩阵）那样用卷积表示来计算，请参见原论文中的方程式3a和3b。为了避免序列递归并实现并行化，对于给定的输入序列，选择性SSM的计算可以通过 并行扫描算法（Smith等，2023；Blelloch，1990）来优化，这类似于并行前缀和的概念。

  • 硬件感知优化:

    在现代GPU中，大多数操作（除了矩阵乘法）都受限于内存访问（FlashAttention Dao 等，2022）。为了克服这一瓶颈，选择性SSM的实现借鉴了FlashAttention中使用的类似策略（内核融合和重新计算），以针对GPU进行优化。

    • 内核融合: 而不是在GPU HBM中生成大小为 (B, L, D, N) 的 A¯ 和 B¯。它在SRAM中融合了步骤5和步骤6的操作，以避免将 A¯ 和 B¯ 储存和读取到/从 HBM，这样它们的大小比 X 和 Y (B, L, D) 大 N 倍。这有助于通过减少 _O_(N)（状态维度）的IO操作从而显著提高速度。 

    • 在训练反向传播过程中，有必要使用中间状态来更新权重。为了减少存储这些中间状态造成的内存占用，它在反向传递过程中在SRAM中重新计算这些状态。这有助于使选择性SSM层的内存占用与FlashAttention变压器层相同。

png


图片来源 Mamba: Linear-Time Sequence Modeling with Selective State Spaces

Mamba 架构

一个标准的 Mamba 块由门控多层感知器（MLP）和选择性状态空间模型（SSM）模块组成。一个 Mamba 模型重复这个块，并用标准归一化和残差连接交错进行。理解选择性 SSM 是理解 Mamba 架构的关键。

png

图片来源 Mamba: Linear-Time Sequence Modeling with Selective State Spaces

在 AMD GPU 上使用 ROCm 进行 Mamba 推理

Mamba repo 托管了 Mamba 模型的源代码。要在使用 ROCm 的 AMD GPU 上安装和运行 Mamba，您需要执行一个额外的步骤来使其正常工作。 csrc 文件夹中包含了 CUDA 源代码，该代码已经纳入了针对 Mamba 的硬件优化。 PyTorch 使用一个名为 Hipify_torch 的工具来将 CUDA 源代码转换为 HIP C++，从而使自定义内核适合在 ROCm 上运行。这个翻译过程在构建 CUDA 扩展时由 PyTorch 内部调用，确保在 ROCm 上使用自定义内核时能获得无缝体验。 有关更多详细信息，我们鼓励您阅读 HIPIFY 和Hipify_torch。幸运的是，我们已经为您完成了此步骤，因此您可以直接在使用 ROCm 的 AMD GPU 上安装和使用 Mamba 模型。

设置

本博客是使用以下设置创建的。有关设置的全面支持详情，请参考 ROCm 文档。

在本博客中，我们在配备 MI210 GPU 的 Linux 机器上利用  rocm/pytorch-nightly  docker 镜像，并使用 AMD GPU 驱动程序版本 6.7.0。

注意：_如果您的机器尚未安装 ROCm 或需要更新驱动程序，请按照 通过 AMDGPU 安装程序安装 ROCm 中显示的步骤进行操作。在本博客中，我们使用 amdgpu-install_6.1.60101-1_all.deb 包安装了包括 GPU 驱动程序在内的 ROCm。

入门指南

我们使用 rocm/pytorch-nightly Docker 镜像，并在容器中构建 Mamba。

docker pull rocm/pytorch-nightly:latest
docker run -it --name mamba --rm \
            --device=/dev/kfd --device=/dev/dri/ \
            --group-add=video --shm-size 8G \
            rocm/pytorch-nightly

在带有 ROCm 的 AMD GPU 上安装 Mamba。

git clone https://github.com/state-spaces/mamba.git
cd mamba
pip install .

Mamba 推理

你可以在 ROCm 博客仓库 中找到本文中使用的 mamba_transformer_evaluation.py 文件。

python mamba_transformer_evaluation.py --model-name "state-spaces/mamba-2.8b" --prompt "Please explain Turing test and provide your perspective on Turing test." --topp 0.9 --temperature 0.7 --repetition-penalty 1.2

输出:

    Loading model state-spaces/mamba-2.8b
    Number of parameters: 2768345600
    ['Please explain Turing test and provide your perspective on Turing test.\nI am not sure what is the difference between a human being, an AI or even another person? I have read about this in some books but still confused with it.. Please help me to understand more clearly... \n\nA:\n\nThe turing machine was invented by Alan M.Turing as part of his work into computability theory (which he called "theory A".) The idea behind that invention were twofold : first you can simulate any other computer using only one tape']
    Prompt length: 14, generation length: 100
    prompt processing + decoding time: 1.20s
    memory used: 5GB

Mamba 和 Transformer 的效率对比

Mamba 享有快速推理、序列长度的线性扩展和 Transformer 级别的性能。在本节中，我们将考察 Mamba 的推理能力，并与实现了 FlashAttention-2 的 Transformer 模型进行比较。关于困惑度等性能指标，请参阅原始论文的评估部分。为了与 Transformer 模型进行公平比较，我们选择  EleutherAI/gpt-neo-2.7B 作为基线。`gpt-neo-2.7B` 是一个 Transformer 模型，由 EleutherAI 仿制 GPT-3 架构设计。该模型支持 FlashAttention-2，并且其参数数量与 state-spaces/mamba-2.8b 相似。 

为Transformer安装Flash Attention

按照以下步骤在ROCm上安装FlashAttention。

git clone --recursive https://github.com/ROCm/flash-attention.git
cd flash-attention
MAX_JOBS=$((`nproc` - 1)) pip install -v .

安装`EleutherAI/gpt-neo-2.7B`模型所需的软件包。

!pip install accelerate git+https://github.com/huggingface/transformers matplotlib

测试FlashAttention是否在平台上安装正确。

首先运行没有FlashAttention的模型。

python mamba_transformer_evaluation.py --model-name "EleutherAI/gpt-neo-2.7B" --attn_implementation "eager" --prompt "Please explain Turing test and provide your perspective on Turing test." --topp 0.9 --temperature 0.7 --repetition-penalty 1.2

输出结果:

    Loading model EleutherAI/gpt-neo-2.7B
    EleutherAI/gpt-neo-2.7B model configuration attn_implementation:eager
    Number of parameters: 2651307520
    ['Please explain Turing test and provide your perspective on Turing test.\n\nA Turing test is a test that can be used to determine whether a computer can think. It is a test that can be used to determine whether a computer can think.\n\nA Turing test is a test that can be used to determine whether a computer can think. It is a test that can be used to determine whether a computer can think.\n\n−\n\nThe Turing test is a test that can be used to determine whether a computer can think. It is a test that']
    Prompt length: 12, generation length: 100
    prompt processing + decoding time: 2.93s
    memory used: 5GB

再运行启用FlashAttention的模型。

python mamba_transformer_evaluation.py --model-name "EleutherAI/gpt-neo-2.7B" --attn_implementation "flash_attention_2" --prompt "Please explain Turing test and provide your perspective on Turing test." --topp 0.9 --temperature 0.7 --repetition-penalty 1.2

输出结果:

    Loading model EleutherAI/gpt-neo-2.7B
    EleutherAI/gpt-neo-2.7B model configuration attn_implementation:flash_attention_2
    Number of parameters: 2651307520
    ['Please explain Turing test and provide your perspective on Turing test.\n\nA Turing test is a test that can be used to determine whether a computer can think. It is a test that can be used to determine whether a computer can think.\n\nA Turing test is a test that can be used to determine whether a computer can think. It is a test that can be used to determine whether a computer can think.\n\n−\n\nThe Turing test is a test that can be used to determine whether a computer can think. It is a test that']
    Prompt length: 12, generation length: 100
    prompt processing + decoding time: 2.46s
    memory used: 5GB

我们可以观察到处理时间从2.93秒下降到了2.46秒，这意味着FlashAttention已经成功安装。

接下来我们以相同的配置（批量大小、提示、提示长度（1k）、生成的令牌数（1k）等）运行Mamba-2.8B和GPT-Neo-2.7B模型。在每次迭代中，批量大小将会加倍。我们监控每次运行的推理时间和内存使用情况，从而对比Mamba模型和Transformer模型在AMD GPU上使用ROCm的效率。

import subprocess
sizes = [1, 2, 4, 8, 16, 32, 64, 128]
for batch_size in sizes:
    print(f"************ Evaluating mamba and transformed with Batch size: {batch_size} ***************")
    command_mamba = f"python mamba_transformer_evaluation.py --model-name 'state-spaces/mamba-2.8b' --batch {batch_size} --promptlen 1024 --genlen 1024 --topp 0.9 --temperature 0.7 --repetition-penalty 1.2"
    command_transformer = f"python mamba_transformer_evaluation.py --model-name 'EleutherAI/gpt-neo-2.7B' --attn_implementation 'flash_attention_2' --batch {batch_size} --promptlen 1024  --genlen 1024 --topp 0.9 --temperature 0.7 --repetition-penalty 1.2"
    subprocess.run(command_mamba, shell=True)
    subprocess.run(command_transformer, shell=True)

输出结果:

    ************ Evaluating mamba and transformed with Batch size: 1 ***************
    Loading model state-spaces/mamba-2.8b
    Number of parameters: 2768345600
    Prompt length: 1024, generation length: 1024
    prompt processing + decoding time: 12.12s
    memory used: 5GB

    Loading model EleutherAI/gpt-neo-2.7B
    EleutherAI/gpt-neo-2.7B model configuration attn_implementation:flash_attention_2
    Number of parameters: 2651307520
    Prompt length: 1024, generation length: 1024
    prompt processing + decoding time: 24.99s
    memory used: 7GB

    ************ Evaluating mamba and transformed with Batch size: 2 ***************
    Loading model state-spaces/mamba-2.8b
    Number of parameters: 2768345600
    Prompt length: 1024, generation length: 1024
    prompt processing + decoding time: 13.27s
    memory used: 6GB

    Loading model EleutherAI/gpt-neo-2.7B
    EleutherAI/gpt-neo-2.7B model configuration attn_implementation:flash_attention_2
    Number of parameters: 2651307520
    Prompt length: 1024, generation length: 1024
    prompt processing + decoding time: 24.87s
    memory used: 9GB
    ************ Evaluating mamba and transformed with Batch size: 4 ***************
    Loading model state-spaces/mamba-2.8b
    Number of parameters: 2768345600
    Prompt length: 1024, generation length: 1024
    prompt processing + decoding time: 13.78s
    memory used: 6GB

    Loading model EleutherAI/gpt-neo-2.7B
    EleutherAI/gpt-neo-2.7B model configuration attn_implementation:flash_attention_2
    Number of parameters: 2651307520
    Prompt length: 1024, generation length: 1024
    prompt processing + decoding time: 27.61s
    memory used: 13GB

    ************ Evaluating mamba and transformed with Batch size: 8 ***************
    Loading model state-spaces/mamba-2.8b
    Number of parameters: 2768345600
    Prompt length: 1024, generation length: 1024
    prompt processing + decoding time: 15.56s
    memory used: 7GB

    Loading model EleutherAI/gpt-neo-2.7B
    EleutherAI/gpt-neo-2.7B model configuration attn_implementation:flash_attention_2
    Number of parameters: 2651307520
    Prompt length: 1024, generation length: 1024
    prompt processing + decoding time: 38.54s
    memory used: 20GB

    ************ Evaluating mamba and transformed with Batch size: 16 ***************
    Loading model state-spaces/mamba-2.8b
    Number of parameters: 2768345600
    Prompt length: 1024, generation length: 1024
    prompt processing + decoding time: 18.57s
    memory used: 9GB

    Loading model EleutherAI/gpt-neo-2.7B
    EleutherAI/gpt-neo-2.7B model configuration attn_implementation:flash_attention_2
    Number of parameters: 2651307520
    Prompt length: 1024, generation length: 1024
    prompt processing + decoding time: 65.94s
    memory used: 35GB

    ************ Evaluating mamba and transformed with Batch size: 32 ***************
    Loading model state-spaces/mamba-2.8b
    Number of parameters: 2768345600
    Prompt length: 1024, generation length: 1024
    prompt processing + decoding time: 25.87s
    memory used: 12GB

    Traceback (most recent call last):
      File "mamba_transformer_evaluation.py", line 89, in <module>
        fn()
      File "mamba_transformer_evaluation.py", line 70, in <lambda>
        fn = lambda: model.generate(
      File "/opt/conda/envs/py_3.8/lib/python3.8/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
        return func(*args, **kwargs)
      ...
      File "/opt/conda/envs/py_3.8/lib/python3.8/site-packages/transformers/models/gpt_neo/modeling_gpt_neo.py", line 330, in forward
        value = torch.cat((past_value, value), dim=-2)
    torch.OutOfMemoryError: HIP out of memory. Tried to allocate 278.00 MiB. GPU 0 has a total capacity of 63.98 GiB of which 190.00 MiB is free. Of the allocated memory 59.51 GiB is allocated by PyTorch, and 3.95 GiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_HIP_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)


    Loading model EleutherAI/gpt-neo-2.7B
    EleutherAI/gpt-neo-2.7B model configuration attn_implementation:flash_attention_2
    Number of parameters: 2651307520

    ************ Evaluating mamba and transformed with Batch size: 64 ***************
    Loading model state-spaces/mamba-2.8b
    Number of parameters: 2768345600
    Prompt length: 1024, generation length: 1024
    prompt processing + decoding time: 42.82s
    memory used: 18GB

    Traceback (most recent call last):
      File "mamba_transformer_evaluation.py", line 82, in <module>
        out = fn()
      File "mamba_transformer_evaluation.py", line 70, in <lambda>
        fn = lambda: model.generate(
      ...
        key = torch.cat((past_key, key), dim=-2)
    torch.OutOfMemoryError: HIP out of memory. Tried to allocate 428.00 MiB. GPU 0 has a total capacity of 63.98 GiB of which 68.00 MiB is free. Of the allocated memory 54.31 GiB is allocated by PyTorch, and 9.26 GiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_HIP_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)


    Loading model EleutherAI/gpt-neo-2.7B
    EleutherAI/gpt-neo-2.7B model configuration attn_implementation:flash_attention_2
    Number of parameters: 2651307520

    ************ Evaluating mamba and transformed with Batch size: 128 ***************
    Loading model state-spaces/mamba-2.8b
    Number of parameters: 2768345600
    Prompt length: 1024, generation length: 1024
    prompt processing + decoding time: 75.77s
    memory used: 31GB

    Traceback (most recent call last):
      File "mamba_transformer_evaluation.py", line 82, in <module>
        out = fn()
      File "mamba_transformer_evaluation.py", line 70, in <lambda>
        fn = lambda: model.generate(
      ...
        key = torch.cat((past_key, key), dim=-2)
    torch.OutOfMemoryError: HIP out of memory. Tried to allocate 642.00 MiB. GPU 0 has a total capacity of 63.98 GiB of which 486.00 MiB is free. Of the allocated memory 61.38 GiB is allocated by PyTorch, and 1.80 GiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_HIP_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)


    Loading model EleutherAI/gpt-neo-2.7B
    EleutherAI/gpt-neo-2.7B model configuration attn_implementation:flash_attention_2
    Number of parameters: 2651307520

现在我们通过绘制批处理大小与推理时间和内存使用情况的图表来比较Mamba模型和Transformer模型的效率。

png

左图显示了GPT-Neo-2.7B的推理时间随着批处理大小的增加而急剧增加。相比之下，得益于选择性状态空间模型（SSM）的效率，Mamba模型的推理时间并没有如此大幅度地增加。GPT-Neo-2.7B在批处理大小为32时失败，无法在批处理大小≥32的情况下运行，因为它需要的内存超过了硬件（MI210有64GB内存）能提供的容量。

右图揭示了类似的趋势：由于注意力机制使用的KV缓存，GPT-Neo-2.7B模型的内存使用量显著高于Mamba模型，并且在批处理大小为32时因内存不足（OOM）停止运行。相比之下，Mamba模型可以继续在更大的批处理大小下运行。这使得Mamba能够在推理或训练时使用更大的批处理大小来提高吞吐量。

这两幅图表展示了Mamba模型的效率，与Mamba论文中提出的发现一致。

致谢

我们要感谢Tri Dao审查并合并对Mamba库支持ROCm的更改。

参考文献

Mamba: Linear-Time Sequence Modeling with Selective State Spaces
Efficiently Modeling Long Sequences with Structured State Spaces
Introduction to State Space Models (SSM)
GPT-Neo

评论 1
添加红包

请填写红包祝福语或标题

个

红包个数最小为10个

元

红包金额最低5元

当前余额3.43元 前往充值 >
需支付：10.00元
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

109702008

你的鼓励将是我创作的最大动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付：¥1
获取中
扫码支付

您的余额不足，请更换扫码支付或充值

打赏作者

实付元
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明：

1.余额是钱包充值的虚拟货币，按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载，可以购买VIP、付费专栏及课程。

余额充值