feat: support llama 3.1. (#106)

b4rtaz · Jul 25, 2024 · 4b8a0ca · 4b8a0ca
1 parent 8c57298
commit 4b8a0ca
Show file tree

Hide file tree

Showing 15 changed files with 249 additions and 55 deletions.
diff --git a/.gitignore b/.gitignore
@@ -4,6 +4,7 @@
 *.0
 *.dSYM
 *.data
+*.temp
 __pycache__
 
 *-test

diff --git a/converter/convert-hf.py b/converter/convert-hf.py
@@ -143,6 +143,14 @@ def parseHiddenAct(act: str):
         raise Exception(f'Unsupported hidden act: {act}')
     return hiddenAct
 
+def parseRopeType(rt: str):
+    ropeType = {
+        'llama3': 2, # LLAMA3_1
+    }.get(rt)
+    if (ropeType is None):
+        raise Exception(f'Unsupported rope type: {ropeType}')
+    return ropeType
+
 def loadConfig(folderPath: str, weightsFloatType: int):
     allFiles = os.listdir(folderPath)
     allFiles.sort()
@@ -178,6 +186,14 @@ def loadConfig(folderPath: str, weightsFloatType: int):
     ropeTheta = config.get('rope_theta')
     if (ropeTheta is not None):
         result['rope_theta'] = int(ropeTheta)
+
+    ropeScaling = config.get('rope_scaling')
+    if (ropeScaling is not None):
+        result['rope_scaling_factor'] = int(ropeScaling['factor'])
+        result['rope_scaling_low_freq_factor'] = int(ropeScaling['low_freq_factor'])
+        result['rope_scaling_high_freq_factory'] = int(ropeScaling['high_freq_factor'])
+        result['rope_scaling_orig_max_seq_len'] = int(ropeScaling['original_max_position_embeddings'])
+        result['rope_type'] = parseRopeType(ropeScaling['rope_type'])
     return result
 
 def printUsage():

diff --git a/converter/writer.py b/converter/writer.py
@@ -125,7 +125,12 @@ def writeHeader(file, params):
         'max_seq_len': 10,
         'hidden_act': 11,
         'rope_theta': 12,
-        'weights_float_type': 13
+        'weights_float_type': 13,
+        'rope_scaling_factor': 14,
+        'rope_scaling_low_freq_factor': 15,
+        'rope_scaling_high_freq_factory': 16,
+        'rope_scaling_orig_max_seq_len': 17,
+        'rope_type': 18,
     }
     header = struct.pack('i', 0xA00ABCD)
 

diff --git a/launch.py b/launch.py
@@ -18,7 +18,12 @@
         'https://huggingface.co/b4rtaz/Llama-3-8B-Q40-Instruct-Distributed-Llama/resolve/main/dllama_model_lama3_instruct_q40.m?download=true',
         'https://huggingface.co/b4rtaz/Llama-3-8B-Q40-Instruct-Distributed-Llama/resolve/main/dllama_tokenizer_llama3.t?download=true',
         'q40', 'q80', 'chat'
-    ]
+    ],
+    'llama3_1_8b_instruct_q40': [
+        'https://huggingface.co/b4rtaz/Llama-3_1-8B-Q40-Instruct-Distributed-Llama/resolve/main/dllama_model_llama3.1_instruct_q40.m?download=true',
+        'https://huggingface.co/b4rtaz/Llama-3_1-8B-Q40-Instruct-Distributed-Llama/resolve/main/dllama_tokenizer_llama_3_1.t?download=true',
+        'q40', 'q80', 'chat'
+    ],
 }
 
 def downloadFile(url: str, path: str):

diff --git a/src/app.cpp b/src/app.cpp
@@ -41,24 +41,27 @@ AppArgs AppArgs::parse(int argc, char** argv, bool hasMode) {
     args.steps = 0;
     args.seed = (unsigned long long)time(NULL);
     args.chatTemplateType = TEMPLATE_UNKNOWN;
+    args.useDiscForKvCache = false;
 
     int i = 1;
     if (hasMode && argc > 1) {
         args.mode = argv[1];
         i++;
     }
     for (; i + 1 < argc; i += 2) {
-        if (strcmp(argv[i], "--model") == 0) {
-            args.modelPath = argv[i + 1];
-        } else if (strcmp(argv[i], "--tokenizer") == 0) {
-            args.tokenizerPath = argv[i + 1];
-        } else if (strcmp(argv[i], "--prompt") == 0) {
-            args.prompt = argv[i + 1];
-        } else if (strcmp(argv[i], "--weights-float-type") == 0) {
-            args.weightsFloatType = parseFloatType(argv[i + 1]);
-        } else if (strcmp(argv[i], "--buffer-float-type") == 0) {
-            args.bufferFloatType = parseFloatType(argv[i + 1]);
-        } else if (strcmp(argv[i], "--workers") == 0) {
+        char* name = argv[i];
+        char* value = argv[i + 1];
+        if (strcmp(name, "--model") == 0) {
+            args.modelPath = value;
+        } else if (strcmp(name, "--tokenizer") == 0) {
+            args.tokenizerPath = value;
+        } else if (strcmp(name, "--prompt") == 0) {
+            args.prompt = value;
+        } else if (strcmp(name, "--weights-float-type") == 0) {
+            args.weightsFloatType = parseFloatType(value);
+        } else if (strcmp(name, "--buffer-float-type") == 0) {
+            args.bufferFloatType = parseFloatType(value);
+        } else if (strcmp(name, "--workers") == 0) {
             int j = i + 1;
             for (; j < argc && argv[j][0] != '-'; j++);
             int count = j - i - 1;
@@ -82,22 +85,24 @@ AppArgs AppArgs::parse(int argc, char** argv, bool hasMode) {
             }
 
             i += count - 1;
-        } else if (strcmp(argv[i], "--port") == 0) {
-            args.port = atoi(argv[i + 1]);
-        } else if (strcmp(argv[i], "--nthreads") == 0) {
-            args.nThreads = atoi(argv[i + 1]);
-        } else if (strcmp(argv[i], "--steps") == 0) {
-            args.steps = atoi(argv[i + 1]);
-        } else if (strcmp(argv[i], "--temperature") == 0) {
-            args.temperature = atof(argv[i + 1]);
-        } else if (strcmp(argv[i], "--topp") == 0) {
-            args.topp = atof(argv[i + 1]);
-        } else if (strcmp(argv[i], "--seed") == 0) {
-            args.seed = atoll(argv[i + 1]);
-        } else if (strcmp(argv[i], "--chat-template") == 0) {
-            args.chatTemplateType = parseChatTemplateType(argv[i + 1]);
+        } else if (strcmp(name, "--port") == 0) {
+            args.port = atoi(value);
+        } else if (strcmp(name, "--nthreads") == 0) {
+            args.nThreads = atoi(value);
+        } else if (strcmp(name, "--steps") == 0) {
+            args.steps = atoi(value);
+        } else if (strcmp(name, "--temperature") == 0) {
+            args.temperature = atof(value);
+        } else if (strcmp(name, "--topp") == 0) {
+            args.topp = atof(value);
+        } else if (strcmp(name, "--seed") == 0) {
+            args.seed = atoll(value);
+        } else if (strcmp(name, "--chat-template") == 0) {
+            args.chatTemplateType = parseChatTemplateType(value);
+        } else if (strcmp(name, "--kv-cache-storage") == 0) {
+            args.useDiscForKvCache = strcmp(value, "disc") == 0;
         } else {
-            printf("Unknown option %s\n", argv[i]);
+            printf("Unknown option %s\n", name);
             exit(EXIT_FAILURE);
         }
     }
@@ -131,7 +136,10 @@ void App::run(AppArgs* args, void (*program)(Inference* inference, SocketPool* s
         args->steps = spec.seqLen;
     }
 
-    Transformer transformer = Transformer::loadRootFromFile(args->modelPath, &spec, socketPool);
+    TransformerConfig config;
+    config.useDiscForKvCache = args->useDiscForKvCache;
+
+    Transformer transformer = Transformer::loadRootFromFile(args->modelPath, &spec, &config, socketPool);
     socketPool->setTurbo(true);
 
     Inference inference = Inference(&arch, args->nThreads, &transformer, socketPool);

diff --git a/src/app.hpp b/src/app.hpp
@@ -16,7 +16,8 @@
 class AppArgs {
 public:
     char* mode;
-    int nThreads; 
+    int nThreads;
+    bool useDiscForKvCache;
 
     // inference
     char* modelPath;

diff --git a/src/apps/dllama/dllama.cpp b/src/apps/dllama/dllama.cpp
@@ -160,7 +160,7 @@ class Chat {
             int nInputTokens;
             tokenizer->encode((char*)inputPrompt.c_str(), inputTokens, &nInputTokens, true, false);
 
-            pos_t userPromptEndPos = (pos_t)std::min(spec->seqLen, pos + nInputTokens - 1);
+            pos_t userPromptEndPos = (pos_t)std::min(spec->seqLen, (int)pos + nInputTokens - 1);
             for (pos_t i = 0; pos < userPromptEndPos; pos++, i++) {
                 inference->infer(inputTokens[i], pos);
                 token = inputTokens[i + 1];
@@ -207,10 +207,13 @@ void worker(AppArgs* args) {
         throw std::runtime_error("Invalid port number");
     }
 
+    TransformerConfig config;
+    config.useDiscForKvCache = args->useDiscForKvCache;
+
     SocketServer server(args->port);
     Socket socket = server.accept();
     TransformerSpec spec;
-    Transformer transformer = Transformer::loadSlice(&spec, &socket);
+    Transformer transformer = Transformer::loadSlice(&spec, &config, &socket);
     TransformerArch arch = TransformerArchFactory::create(&spec);
 
     Worker worker = Worker(&arch, args->nThreads, &transformer, &socket);

diff --git a/src/commands.cpp b/src/commands.cpp
@@ -1,5 +1,8 @@
 #include <cassert>
 #include <cstring>
+#ifdef _WIN32
+    #define _USE_MATH_DEFINES
+#endif
 #include <cmath>
 #include "utils.hpp"
 #include "funcs.hpp"
@@ -167,6 +170,54 @@ void LlamaRopeCommand::forward(bool isQ, float* qOrK, pos_t pos, unsigned int nT
     }
 }
 
+Llama3_1RopeCommand::Llama3_1RopeCommand(RopeSlice *slice, float ropeScalingFactor, float ropeScalingLowFreqFactor, float ropeScalingHighFreqFactory, int ropeScalingOrigMaxSeqLen) {
+    this->slice = slice;
+    this->ropeScalingFactor = ropeScalingFactor;
+    this->ropeScalingLowFreqFactor = ropeScalingLowFreqFactor;
+    this->ropeScalingHighFreqFactory = ropeScalingHighFreqFactory;
+    this->ropeScalingOrigMaxSeqLen = ropeScalingOrigMaxSeqLen;
+    printf("🕒 ropeScalingFactor: %f\n", ropeScalingFactor);
+    printf("🕒 ropeScalingLowFreqFactor: %f\n", ropeScalingLowFreqFactor);
+    printf("🕒 ropeScalingHighFreqFactory: %f\n", ropeScalingHighFreqFactory);
+    printf("🕒 ropeScalingOrigMaxSeqLen: %d\n", ropeScalingOrigMaxSeqLen);
+}
+
+float Llama3_1RopeCommand::scale(float freq) {
+    float waveLen = 2.0f * M_PI * freq;
+    float lowFreqWavelen = ropeScalingOrigMaxSeqLen / ropeScalingLowFreqFactor;
+    float highFreqWavelen = ropeScalingOrigMaxSeqLen / ropeScalingHighFreqFactory;
+    if (waveLen < highFreqWavelen) {
+        return freq;
+    } else if (waveLen > lowFreqWavelen) {
+        return freq / ropeScalingFactor;
+    } else {
+        float smooth = (ropeScalingOrigMaxSeqLen / waveLen - ropeScalingLowFreqFactor) / (ropeScalingHighFreqFactory - ropeScalingLowFreqFactor);
+        return (1 - smooth) * freq / ropeScalingFactor + smooth * freq;
+    }
+}
+
+void Llama3_1RopeCommand::forward(bool isQ, float* qOrK, pos_t pos, unsigned int nThreads, unsigned int threadIndex) {
+    const unsigned int dim0Half = (isQ ? slice->qDim0 : slice->kvDim0) / 2;
+    const unsigned int shift = isQ ? slice->qShift : 0;
+    SPLIT_RANGE_TO_THREADS(s, e, 0, dim0Half, nThreads, threadIndex);
+    const unsigned int iStart = s * 2;
+    const unsigned int iEnd = e * 2;
+
+    for (unsigned int i = iStart; i < iEnd; i += 2) {
+        const unsigned int headDim = i % slice->headSize;
+        const float freq = 1.0f / powf(slice->ropeTheta, headDim / (float)slice->headSize);
+        const float val = pos * freq;
+        const float fcr = cosf(val);
+        const float fci = sinf(val);
+
+        float v0 = qOrK[i];
+        float v1 = qOrK[i + 1];
+
+        qOrK[i]     = scale(v0 * fcr - v1 * fci);
+        qOrK[i + 1] = scale(v0 * fci + v1 * fcr);
+    }
+}
+
 FalconRopeCommand::FalconRopeCommand(RopeSlice *slice) {
     this->slice = slice;
 }

diff --git a/src/commands.hpp b/src/commands.hpp
@@ -9,7 +9,7 @@
 // *Slice - calculates sizes, offsets, slice sizes etc. It is not responsible for memory allocation. It may help in the loading of data.
 // *Command - allocates memory for weights, performs calculations.
 
-typedef unsigned short pos_t;
+typedef unsigned int pos_t;
 typedef uint8_t slice_index_t;
 
 class MatmulSlice {
@@ -106,6 +106,19 @@ class LlamaRopeCommand : public RopeCommand {
     void forward(bool isQ, float* qOrK, pos_t pos, unsigned int nThreads, unsigned int threadIndex);
 };
 
+class Llama3_1RopeCommand : public RopeCommand {
+private:
+    RopeSlice* slice;
+    float ropeScalingFactor;
+    float ropeScalingLowFreqFactor;
+    float ropeScalingHighFreqFactory;
+    int ropeScalingOrigMaxSeqLen;
+public:
+    Llama3_1RopeCommand(RopeSlice *slice, float ropeScalingFactor, float ropeScalingLowFreqFactor, float ropeScalingHighFreqFactory, int ropeScalingOrigMaxSeqLen);
+    void forward(bool isQ, float* qOrK, pos_t pos, unsigned int nThreads, unsigned int threadIndex);
+    float scale(float freq);
+};
+
 class FalconRopeCommand : public RopeCommand {
 private:
     RopeSlice* slice;

diff --git a/src/grok1-tasks-test.cpp b/src/grok1-tasks-test.cpp
@@ -30,6 +30,7 @@ int main() {
     TransformerSpec spec;
     spec.headerSize = sizeof(TransformerFileOldHeader) + sizeof(int);
     spec.archType = GROK1;
+    spec.ropeType = ROPE_FALCON;
     spec.dim = 6144;
     spec.nLayers = 1;
     spec.nHeads = 48;
@@ -47,6 +48,9 @@ int main() {
     spec.hiddenAct = GELU;
     spec.ropeTheta = 10000.0f;
 
+    TransformerConfig config;
+    config.useDiscForKvCache = false;
+
     size_t beforeBlockBytes = spec.dim * spec.vocabSize * sizeof(float);
     size_t blockBytes = 956596224;
     size_t afterBlockBytes  = (spec.dim + spec.dim * spec.vocabSize) * sizeof(float);
@@ -60,7 +64,7 @@ int main() {
     for (int f = 0; f < nFloats; f++) block[f] = randomF32(&state) / 100.0;
 
     SocketPool socketPool(0, NULL);
-    Transformer transformer = Transformer::loadRoot(weights, &spec, &socketPool);
+    Transformer transformer = Transformer::loadRoot(weights, &spec, &config, &socketPool);
     transformer.pos = 0;
 
     float* x = transformer.x;

diff --git a/src/llama2-tasks-test.cpp b/src/llama2-tasks-test.cpp
@@ -528,6 +528,7 @@ int main() {
     TransformerSpec spec;
     spec.headerSize = sizeof(TransformerFileOldHeader) + sizeof(int);
     spec.archType = LLAMA;
+    spec.ropeType = ROPE_LLAMA;
     spec.dim = 4096;
     spec.nLayers = 1;
     spec.headSize = 128;
@@ -545,6 +546,9 @@ int main() {
     spec.hiddenAct = SILU;
     spec.ropeTheta = 10000.0f;
 
+    TransformerConfig config;
+    config.useDiscForKvCache = false;
+
     size_t beforeBlockBytes = /* embedding */ 524288000;
     size_t blockBytes       = 809533440;
     size_t afterBlockBytes  = /* norm */ 16384 + /* embedding */ 524288000;
@@ -562,7 +566,7 @@ int main() {
     for (int i = 0; i < mm; i++) mmData[i] = randomF32(&state) / 120.0;
 
     SocketPool socketPool(0, NULL);
-    Transformer transformer = Transformer::loadRoot((char*)data, &spec, &socketPool);
+    Transformer transformer = Transformer::loadRoot((char*)data, &spec, &config, &socketPool);
     transformer.pos = 0;
 
     float* x = transformer.x;