大模型部署框架FastLLM实现细节解析

2023-7-27 21:36| 发布者: 开心| 查看: 579| 评论: 2

摘要: 0x0. 前言接着大模型部署框架 FastLLM 简要解析这篇文章首先梳理了一下FastLLM的调用链和关键的数据结构，然后解析了 FastLLM 的一些实现细节和CPU/GPU后端实现采用的 ...

    0x0. 前言
  接着大模型部署框架 FastLLM 简要解析这篇文章首先梳理了一下FastLLM的调用链和关键的数据结构，然后解析了 FastLLM 的一些实现细节和CPU/GPU后端实现采用的优化技巧。
  0x1. 调用链和数据结构解析
  以chatglm-6b的支持为例，函数入口在 https://github.com/ztxz16/fastllm/blob/master/src/models/chatglm.cpp#L626 ，这里的 input 就是输入的 context（string类型）。然后 https://github.com/ztxz16/fastllm/blob/master/src/models/chatglm.cpp#L633 这行代码对 input 进行 tokenizer encode并构造好inputIds，再构造好attentionMask之后就可以给Forward函数推理，拿到推理结果之后再使用tokenizer进行decode得到输出。
  在这里，inputIds和attentionMask都是Data数据类型，类比于PyTorch的Tensor，来对输入数据以及device，shape等信息进行统一管理。下面的代码展示了Data数据结构的定义，源码在：https://github.com/ztxz16/fastllm/blob/master/include/fastllm.h#L201-L286

classData{public:boollockInCPU=false;//如果lock在CPU上，那么不允许移动到其余设备WeightTypeweightType=WeightType::NONE;//权重类型，NONE代表非权重（或未知权重）DataTypedataType=DataType::FLOAT32;//数据类型intunitSize,unitSizeDiv=1;//单个元素的字节数=unitSIze/unitSizeDivstd::vectordims;//数据形状std::vectorstrides;//跨度uint64_texpansionSize=0;//扩容后的尺寸uint64_texpansionBytes=0;//扩容后的字节数std::vectorexpansionDims;//预扩容的形状uint8_t*cpuData=nullptr;//数据指针void*cudaData=nullptr;std::vectorextraCudaData;void*deviceData=nullptr;std::vectorextraDeviceData;DataDevicedataDevice=DataDevice::CPU;//这两个参数用于量化，对FLOAT数据不适用intperChannelAxis=-1;//沿哪个轴分通道量化，-1代表没有分通道std::vectorperChannelsConfigs;//perChannelsConfigs[i]代表第i个通道的min,max;如果没有分通道，perChannelsConfigs[0]代表全局min,maxstd::vectorscales,mins;std::vectorzeros;std::vectorweightSum;//作为权重时，有时候需要存一些和加速计算std::stringfileName;longlongfilePos;std::shared_ptrm_file;Data(){};Data(DataTypetype);Data(DataTypetype,conststd::vector&dims);//构造函数//构造函数，创建好之后从data复制数据//data中是原始数据，如果type不是float那么需要量化Data(DataTypetype,conststd::vector&dims,conststd::vector&data);~Data();//析构函数Data(constData&ori);//深拷贝voidCopyFrom(constData&ori);//复制uint64_tGetBytes()const;//获取总字节数voidAllocate();//分配内存voidAllocate(floatv);//分配内存并初始化voidExpansion(conststd::vector&dims);//预扩容到相应尺寸voidMallocSpace(uint64_tsize);//在设备上分配voidFreeSpace();//回收设备上的内存voidUpdateUnitSize();//更新unitSizevoidResize(conststd::vector&dims);//更改尺寸voidReshape(conststd::vector&dims);//更改尺寸,但不修改数据uint64_tCount(inti)const;//dims[i]*strides[i]voidPrintShape()const;//输出形状voidPrint()const;//输出voidCalcWeightSum();//计算WeightSumvoidToDevice(DataDevicedevice);//移动到指定devicevoidToDevice(void*device);voidset_file(std::shared_ptrfile){m_file=file;}};

  在Forward函数里面，以Data为核心载体，运行chatglm-6b模型的流程，具体包含如下的一些算子：https://github.com/ztxz16/fastllm/blob/master/include/fastllm.h#L346-L408 。以Permute为例我们浏览下它的实现：

voidPermute(constData&input,conststd::vector&axis,Data&output){DataaxisData=Data(DataType::INT32PARAM,{(int)axis.size()});axisData.Allocate();for(inti=0;iRun("Permute",{{"input",(Data*)&input},{"axis",&axisData},{"output",(Data*)&output}},{},{});}

  这里的curExecutor负责根据FastLLM编译开启的后端选项把算子Dispatch到不同的device进行执行，{"input", (Data*)&input}, {"axis", &axisData}, {"output", (Data*)&output}} 这行代码表示的是一个DataDict对象，也就是一个值为data的字典，原始定义为typedef std::map  DataDict;。接着我们看一下curExecutor的定义和实现：

namespacefastllm{classExecutor{private:std::vectordevices;std::mapprofiler;public:Executor();//创建默认的Executor~Executor();//析构voidClearDevices();//清空devicesvoidAddDevice(BaseDevice*device);//增加一个device//运行一个opvoidRun(conststd::string&opType,constfastllm::DataDict&datas,constfastllm::FloatDict&floatParams,constfastllm::IntDict&intParams);voidClearProfiler();voidPrintProfiler();};}

  从Executor类的定义我们可以判断它负责了在设定的devices上根据opType和输入数据等执行Op的前向计算，也就是Run这个接口。由于Executor类是FastLLM的调度核心实现，所以我们来详细解析一下它的实现。

namespacefastllm{Executor::Executor(){this->devices.clear();#ifdefUSE_CUDA//将一个指向CudaDevice类对象的指针插入到devices向量的末尾。//这里通过new运算符创建了一个CudaDevice对象，并将返回的指针进行类型转换为BaseDevice*类型。this->devices.push_back((BaseDevice*)newCudaDevice());#endifthis->devices.push_back((BaseDevice*)newCpuDevice());}Executor::~Executor(){//释放devices向量中的每个指针元素所占用的内存。for(inti=0;idevices指的是当前对象的devices成员，即指向BaseDevice类对象的指针向量。this->devices.clear();}//该函数用于向devices向量中添加一个指向BaseDevice类对象的指针。voidExecutor::AddDevice(fastllm::BaseDevice*device){this->devices.push_back(device);}voidExecutor::Run(conststd::string&opType,constfastllm::DataDict&datas,constfastllm::FloatDict&floatParams,constfastllm::IntDict&intParams){//创建一个st变量，用于记录函数开始执行的时间。autost=std::now();//创建一个布尔变量lockInCPU，用于记录是否将数据锁定在CPU上。boollockInCPU=false;//在第一个for循环中，遍历数据字典datas，查找是否有"___batch"后缀的参数，//并根据情况设置lockInCPU的值。it.first是数据字典中的键（key），it.second//是对应的值（value）。如果存在"___batch"后缀的参数，则将lockInCPU设置为//对应数据的lockInCPU属性（布尔值），否则设置为当前数据的lockInCPU属性。for(auto&it:datas){if(intParams.find(it.first+"___batch")!=intParams.end()){intbatch=intParams.find(it.first+"___batch")->second;for(inti=0;ilockInCPU;}}else{lockInCPU|=it.second->lockInCPU;}}//第二个for循环遍历devices向量中的所有设备指针device。//在循环中，首先检查lockInCPU是否为真，并且当前设备的类型不是"cpu"，//如果是，则跳过当前设备（continue）。这个检查是为了保证数据锁定在CPU上时，只执行CPU设备上的操作。for(autodevice:devices){if(lockInCPU&&device->deviceType!="cpu"){continue;}//然后，通过调用device->CanRun(opType,datas,floatParams,intParams)//检查当前设备是否可以运行指定的操作opType。如果可以运行，则进行以下操作：if(device->CanRun(opType,datas,floatParams,intParams)){//第三个for循环遍历数据字典datas，如果存在"___batch"后缀的参数，//则将对应数据转移到当前设备上；否则，将当前数据转移到当前设备上。for(auto&it:datas){if(intParams.find(it.first+"___batch")!=intParams.end()){intbatch=intParams.find(it.first+"___batch")->second;for(inti=0;iToDevice((void*)device);}}else{it.second->ToDevice((void*)device);}}//调用device->Reshape(opType,datas,floatParams,intParams)//进行形状推导，device上的形状推导调用了opType对应的op的形状推导，//并且被各个不同的op重写。device->Reshape(opType,datas,floatParams,intParams);//对opType对应的这个算子进行推理。device->Run(opType,datas,floatParams,intParams);break;}}//最后，计算操作运行时间，并将其加入profiler成员变量，用于性能分析。floatspend=GetSpan(st,std::now());profiler[opType]+=spend;}//清除profile的信息voidExecutor::ClearProfiler(){profiler.clear();}//打印profile信息，也即输出每个层的运行时间和模型的总运行时间voidExecutor::PrintProfiler(){floatsum=0.0;for(auto&it:profiler){printf("%sspend%f",it.first.c_str(),it.second);sum+=it.second;}printf("totalspend%f",sum);}}

  自此，前向计算就顺利完成了，再把推理结果给 tokenizer 解码就结束了，整体的调度执行流程是很简单明了的。
  0x2. tokenizer 解析
  接着，我们来解析一下tokenizer的实现。先看一下tokenizer的定义（https://github.com/ztxz16/fastllm/blob/master/include/fastllm.h#L287-L310）：

structTokenizer{structTrieNode{inttokenId;std::mapnext;TrieNode();};TrieNode*root;std::unordered_maptokenToStringDict;Tokenizer();~Tokenizer();voidClear();//清空分词器voidInsert(conststd::string&s,inttokenId);//插入一个tokenDataEncode(conststd::string&s);//编码std::stringDecode(constData&data);//解码std::stringDecodeTokens(conststd::vector&tokens);//解码};

  我们从实现来看tokenizer的细节：

//这是Tokenizer类的嵌套结构TrieNode的构造函数的实现。//在构造函数中，将tokenId成员变量的值初始化为-999999。//这个值在构造函数中被硬编码，它是作为一个特殊标记来使用的。Tokenizer::TrieNode(){this->tokenId=-999999;}//Tokenizer类的构造函数的实现。//在构造函数中，通过new运算符创建一个新的TrieNode对象，//并将其指针赋值给root成员变量。这样，构造函数创建了一个空的字典树，//并将其根节点指针存储在root中。Tokenizer::Tokenizer(){root=newTrieNode();}//Tokenizer类的析构函数的实现。//在析构函数中，首先调用Clear()函数，用于释放动态分配的资源和清空数据。//然后，调用delete运算符释放通过new运算符创建的root对象的内存，从而释放整个字典树的内存。Tokenizer::~Tokenizer(){Clear();deleteroot;}//这是Tokenizer类的成员函数Clear()的定义，用于清空分词器并释放动态分配的资源。voidTokenizer::Clear(){//创建一个指向TrieNode的指针向量q，用于辅助遍历字典树。std::vectorq;//将字典树的根节点root加入q向量，作为遍历的起始点。q.push_back(root);//开始遍历q向量中的节点，这是一个广度优先搜索（BFS）的过程。for(inti=0;inext){//将当前节点now的子节点加入q向量中，以便继续遍历子节点的子节点。q.push_back(it.second);}}//当遍历完成后，q向量中包含了字典树中的所有节点。//创建一个新的TrieNode对象，并将其指针赋值给root成员变量，表示创建了一个空的字典树。root=newTrieNode();//清空tokenToStringDict映射表，以确保所有token的映射被清空。tokenToStringDict.clear();}//这是Tokenizer类的成员函数Insert的定义，用于向分词器中插入一个token。voidTokenizer::Insert(conststd::string&s,inttokenId){//创建一个指向TrieNode的指针now，并将其初始化为指向字典树的根节点root。TrieNode*now=this->root;//开始遍历输入的字符串s中的每个字符。for(inti=0;inext中添加新的子节点，该子节点的键为当前字符s[i]的编码值，//值为指向新创建的TrieNode对象的指针。这表示在字典树中添加了一个新的字符节点。if(now->next.find(s[i])==now->next.end()){now->next[s[i]]=newTrieNode();}//将now移动到下一个字符s[i]对应的节点，以便继续处理下一个字符。now=now->next[s[i]];}//遍历完成后，now将指向字典树中最后一个字符的节点。//设置当前节点的tokenId成员变量，表示当前节点代表一个token，//并使用传入的tokenId值来标识该token。now->tokenId=tokenId;//将传入的tokenId和对应的字符串s添加到tokenToStringDict//映射表中，用于后续的解码过程。tokenToStringDict[tokenId]=s;}//这是Tokenizer类的成员函数Encode的定义，用于对输入的字符串s进行编码。DataTokenizer::Encode(conststd::string&s){//创建一个浮点数向量v，用于存储编码结果。该向量将存储找到的token对应的tokenId值。std::vectorv;//开始遍历输入的字符串s中的每个字符。for(inti=0;iroot;//从当前字符s[i]开始继续遍历字符串s。for(intj=i;jnext.find(s[j])!=now->next.end()){//将now移动到下一个字符s[j]对应的节点。now=now->next[s[j]];//检查当前节点now是否代表一个token，即它的tokenId是否有效。if(now->tokenId!=-999999){//如果当前节点代表一个token，将tokenId和当前位置j存储到//tokenId和pos变量中，以便记录找到的token的信息。tokenId=now->tokenId;pos=j;}}else{//如果当前字符不再是token的一部分，退出内层循环，继续外层循环。break;}}//如果pos大于等于当前位置i，表示找到了一个token。//这里pos存储了找到的token的结束位置，i移动到pos处，以便继续遍历下一个字符。if(pos>=i){i=pos;v.push_back(tokenId);//printf("%d",tokenId);}}//printf("");//遍历完成后，v向量中存储了输入字符串中所有找到的token对应的tokenId值。//创建一个Data对象并返回，表示编码的结果。这里Data是一个数据结构，//用于存储数据及其相关信息。编码结果是一个一维浮点数数组，//表示输入字符串中所有找到的token对应的tokenId值。returnData(DataType::FLOAT32,{1,(int)v.size()},v);}//这是Tokenizer类的成员函数DecodeTokens的定义，//用于对输入的token数组进行解码，将token转换回原始的字符串。std::stringTokenizer::DecodeTokens(conststd::vector&tokens){//创建一个空字符串ret，用于存储解码结果。std::stringret="";//开始遍历输入的token数组tokens。for(inti=0;isecond);Data&output=*(datas.find("output")->second);Data&weight=*(datas.find("weight")->second);;output.Allocate();//这行代码为output分配内存。//这行代码从weight的维度中提取词汇大小（vocabSize）和嵌入大小（embSize）。intvocabSize=weight.dims[0],embSize=weight.dims[1];//这行代码计算input的长度。uint64_tinputLen=input.Count(0);//这行代码获取input的数据，并将其转换为浮点数的指针。float*inputData=(float*)input.cpuData;//接下来的代码根据内存模式和权重的数据类型的不同，分别处理了四种情况。//这四种情况可以归纳为两个大类：内存模式和权重的数据类型。//内存模式：如果GetLowMemMode()返回true，则表示处于低内存模式。//在这种模式下，权重数据不会一次性全部加载到内存中，而是每次只加载需要的部分。//否则，权重数据会全部加载到内存中。if(GetLowMemMode()){FILE*fi=fopen(weight.fileName.c_str(),"rb");//权重的数据类型：如果权重的数据类型为FLOAT32，则使用浮点数进行计算。//如果权重的数据类型为BFLOAT16，则使用16位浮点数进行计算。if(weight.dataType==DataType::FLOAT32){float*outputData=(float*)output.cpuData;for(inti=0;i

路过

雷人

握手

鲜花

鸡蛋

收藏分享邀请

发表评论

		自动登录	找回密码
密码			请使用中文注册

大模型部署框架FastLLM实现细节解析

最新评论

相关分类