在ara_soc裸机 (bare metal, RTL模型) 上运行llama2大型语言模型

硬件:ara是一个支持RISC-V 向量扩展1.0 (RVV 1.0)的向量协处理器。

软件:llama2.c是一个C语言实现的llama2模型。

一、动机

验证硬件的正确性

后续会在ara的基础上进行修改,调整硬件,而软件可以作为一个很好的测试向量用来验证SoC硬件的正确性。

大语言模型是目标应用之一,运行目标应用能更好地进行验证。

性能分析与软硬件协同优化

仿真结果得到的性能参数可以分析瓶颈,指导后续硬件的设计。

而直接在RTL模型上运行,优点是得到的结果比较准确,而缺点则是运行较慢。相比之下,更高抽象的仿真可以运行地更快,在探索设计空间和架构搜索时更常用。

二、如何在裸机上运行

裸机上运行需要

  1. 初始化程序(通常是一个汇编)

  2. 提供必要的库函数

  3. 划分好地址空间

  4. 将初始化汇编、应用程序、数据(可选)、库函数正确链接加载到地址空间(如RAM)

初始化程序

  • 初始化程序会初始化寄存器

  • 设置栈

  • 设计控制寄存器(如对于这个应用开启浮点、向量扩展)

  • 设置中断和异常向量处理函数

  • 跳转到main函数

必要的库函数

常见的有

  • 打印输出

  • 内存管理

  • 文件操作

这里打印输出是一个mock uart (当然不可综合就是了) RTL设备,在C程序中对特定地址进行写数据,地址请求会通过axi bar进行路由,对访问uart的地址请求经过axi、apb转到这个设备。
然后在其中判断写地址,然后将写数据写入控制寄存器或通过verilog函数$write打印到控制台。

内存管理我实现了简单的malloc、calloc、free等函数。空闲内存块以链表形式组织,在空闲内存块前存放空闲信息、大小、指针等数据。
这里遇到的坑有:地址对齐问题、连续malloc导致性能急剧下降。前者可以根据硬件特性将分配的地址对齐。

#define MG_ALIGN(size, align) (((size) + (align) - 1) & ~((align) - 1))

后者可以采用从最后分配位置继续分配,而避免每次从头遍历找空闲内存块。

文件操作看起来也可以实现一个mock file设备来进行文件读写,不过现在没时间,就通过稍微修改源代码的方式指定文件内容地址的方式来实现文件的读取。

补充说明地址对齐问题

初始化程序中会设置异常处理handler:

trap_vector:
    // Jump to the mtvec_handler, if it exists
    la t5, mtvec_handler
    beqz t5, 1f
    jr t5
1:  csrr a0, mcause
    j _fail

    .align 2
_eoc:
    la t0, eoc_address_reg
    sd a0, 0(t0)
    jal x0, _eoc

    .align 2
_fail:
    la t0, eoc_address_reg
    sd a0, 0(t0)
    jal x0, _eoc

当发生异常的时候会将csr寄存器mcause (Machine cause register )的值写入a0,作为返回值。然后后面会将a0返回值写入eoc_address_reg控制寄存器,在testbench中会判断这个值

  /*********
   *  EOC  *
   *********/

  always @(posedge clk_i) begin
    if (exit_o[0]) begin
      if (exit_o >> 1) begin
        $warning("Core Test ", $sformatf("*** FAILED *** (tohost = %0d)", (exit_o >> 1)));
      end else begin
        // Print vector HW runtime
        $display("[hw-cycles]: %d", int'(dut.runtime_buf_q));
        $info("Core Test ", $sformatf("*** SUCCESS *** (tohost = %0d)", (exit_o >> 1)));
      end

      $finish(exit_o >> 1);
    end
  end

比如这里报的错是返回值6。

查看RISCV mcause值的含义可知是Store/AMO address misaligned Exception。

在硬件中:在lsu会判断地址的低n位,如果不全为0,则产生未对齐信号。这个信号会经过mmu、load unit、store unit在commit阶段产生一个exception,然后会写入csr寄存器文件。

地址空间

这里我采用的地址空间如下

地址空间

链接与加载

我使用的Verilator RTL仿真器,它可以加载elf文件到指定的存储器,这里将链接得到的elf文件加载到0x80000000,也就是SoC的boot address。

CPP程序中会解析elf,将数据放入buffer。CPP程序会在编译硬件时一起进行编译。

数据部分有多种方式放到地址空间,由于这次采用的修改源码指定文件内容地址的方式,因此可以将权重文件输出到汇编文件(简单python脚本),指定到.data区,和程序一起链接,在C程序中可以作为全局变量访问文件内容。

三、结果

使用180K的权重和对应tokenizer文件,使用Verilator仿真RTL运行llama2,指定prompt,生成模式:

bowen@bowen-virtual-machine-ic:~/Projects/ara/hardware$ app=llama2 make simv
Makefile:83: "Specified QuestaSim version (questa-2021.2) not found in PATH /home/bowen/Projects/riscv/bin:/home/bowen/.vscode-server/cli/servers/Stable-dc96b837cf6bb4af9cd736aa3af08cf8279f7685/server/bin/remote-cli:/home/bowen/.local/bin:/home/bowen/Projects/riscv/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/usr/games:/usr/local/games:/snap/bin:/home/bowen/Software/Modelsim/modeltech/linux_x86_64:/home/bowen/Software/Modelsim/modeltech/linux_x86_64"
build/verilator/Vara_tb_verilator  -l ram,/home/bowen/Projects/ara/apps/bin/llama2,elf
Set `ram TOP.ara_tb_verilator.dut.i_ara_soc.i_dram 10 0x80000000 0x1000000 write with offset: 0x0 write with size: 0xc5f88
Registered memory regions:
        'ram' (80bits) at location: 'TOP.ara_tb_verilator.dut.i_ara_soc.i_dram' (LMA range [0x80000000, 0x80ffffff])
Simulation of Ara
=================


Simulation running, end by pressing CTRL-c.
---------------------
start
--------------------
 ---------------------
 build_transformer
---------------------
 do memory_init !!!
---------------------
 build_tokenizer
---------------------
>>>>>>>>>>>>>>>>>>>>>>>>>>>
 tokenizer: 800BF418
 0A 
<<<<<<<<<<<<<<<<<<<<<<<<<<<
 ---------------------
 build_sampler
---------------------
---------------------
generate
---------------------

Once upon a time, there was a little girl

Received stop request, shutting down simulation.

Simulation statistics
=====================
Executed cycles:  27449425
Wallclock time:   7993 s
Simulation speed: 3434.18 cycles/s (3.43418 kHz)

四、展望

  1. 实现一个mock file的RTL模块,以更方便实现文件读写,尽量少得改动软件源代码
  2. 是否可以直接使用标准库的函数,进行系统调用时代理到host端(类似spike pk的仿真方式?)
posted @ 2024-06-01 01:54  一鸣惊人_001  阅读(223)  评论(1)    收藏  举报