Alternative Weight-Loading Approach

converter/convert-tokenizer-hf.py

@@ -34,7 +34,6 @@ def __init__(self, dirPath, tokenizerConfig):
     def resolvePreTrainedTokenizerFast(self):
         utb = unicodeToBytes()
         tokenizer = PreTrainedTokenizerFast(tokenizer_file = os.path.join(self.dirPath, 'tokenizer.json'))
-        config = openJson(os.path.join(self.dirPath, 'config.json'))
         vocabLen = len(tokenizer.get_vocab())
         for i in range(vocabLen):
             tokenChars = list(tokenizer.convert_ids_to_tokens([i])[0])
@@ -47,18 +46,11 @@ def resolvePreTrainedTokenizerFast(self):
             self.tokens.append(bytes(tokenBytes))
             self.scores.append(-float(i))
 
-        # Pad tokenizer vocab to match model vocab_size if needed
-        targetVocabSize = config.get('vocab_size', vocabLen)
-        if targetVocabSize > vocabLen:
-            print(f'⚠️ Padding tokenizer vocab from {vocabLen} to {targetVocabSize}')
-            for i in range(vocabLen, targetVocabSize):
-                self.tokens.append(f'<|reserved_{i}|>'.encode('utf-8'))
-                self.scores.append(-float(i))
-
         self.bosId = tokenizer.bos_token_id
         if (tokenizer.eos_token_id):
             self.eosIds = [tokenizer.eos_token_id]
-        if (self.bosId is None or self.eosIds is None):
+        if (self.bosId is None or self.eosId is None):
+            config = openJson(os.path.join(self.dirPath, 'config.json'))
             if (self.bosId is None):
                 self.bosId = config['bos_token_id']
             if (self.eosIds is None):
@@ -91,8 +83,7 @@ def resolveLlamaTokenizer(self):
 
     def resolve(self):
         cls = self.tokenizerConfig['tokenizer_class']
-        if (cls == 'PreTrainedTokenizer' or
-            cls == 'PreTrainedTokenizerFast' or
+        if (cls == 'PreTrainedTokenizerFast' or
             cls == 'LlamaTokenizerFast' or
             cls == 'Qwen2Tokenizer'):
             return self.resolvePreTrainedTokenizerFast()

converter/requirements.txt

@@ -1,6 +1,5 @@
-# python>=3.9
+python>=3.9
 numpy==1.23.5
-torch==2.0.1+cpu --index-url https://download.pytorch.org/whl/cpu
+pytorch==2.0.1
 safetensors==0.4.2
-sentencepiece==0.1.99
-transformers==4.57.6
+sentencepiece==0.1.99

src/api-types.hpp

@@ -264,8 +264,6 @@ void to_json(json& j, const ModelList& models) {
         {"data", models.data}};
 }
 
-static std::string normalizeMessageContent(const json &content);
-
 std::vector<ChatMessage> parseChatMessages(json &json){
     std::vector<ChatMessage> messages;
     messages.reserve(json.size());
@@ -274,7 +272,7 @@ std::vector<ChatMessage> parseChatMessages(json &json){
         ChatMessage msg;
         msg.role = item["role"].template get<std::string>();
         if (item.contains("content") && !item["content"].is_null())
-            msg.content = normalizeMessageContent(item["content"]);
+            msg.content = item["content"].template get<std::string>();
         if (item.contains("tool_call_id"))
             msg.tool_call_id = item["tool_call_id"].template get<std::string>();
         if (item.contains("tool_calls") && item["tool_calls"].is_array()) {
@@ -303,43 +301,6 @@ std::vector<ChatMessage> parseChatMessages(json &json){
     return messages;
 }
 
-static std::string normalizeMessageContent(const json &content) {
-    if (content.is_null())
-        return "";
-    if (content.is_string())
-        return content.template get<std::string>();
-    if (content.is_array()) {
-        std::string result;
-        for (const auto &part : content) {
-            std::string piece;
-            if (part.is_string()) {
-                piece = part.template get<std::string>();
-            } else if (part.is_object()) {
-                if (part.contains("type") && part["type"].is_string()) {
-                    const std::string type = part["type"].template get<std::string>();
-                    if ((type == "text" || type == "input_text") && part.contains("text") && part["text"].is_string()) {
-                        piece = part["text"].template get<std::string>();
-                    } else if (type == "text" && part.contains("content") && part["content"].is_string()) {
-                        piece = part["content"].template get<std::string>();
-                    }
-                } else if (part.contains("text") && part["text"].is_string()) {
-                    piece = part["text"].template get<std::string>();
-                }
-            }
-
-            if (piece.empty())
-                continue;
-            if (!result.empty() && result.back() != '\n')
-                result += ' ';
-            result += piece;
-        }
-        return result;
-    }
-    if (content.is_object())
-        return content.dump();
-    return content.dump();
-}
-
 InferenceParams parseInferenceParams(json &json, float defaultTemperature, float defaultTopp, unsigned long long defaultSeed) {
     InferenceParams params;
     params.temperature = defaultTemperature;
@@ -400,11 +361,7 @@ InferenceParams parseInferenceParams(json &json, float defaultTemperature, float
         }
     }
     if (json.contains("stop")) {
-        if (json["stop"].is_string()) {
-            params.stop = std::vector<std::string>{json["stop"].template get<std::string>()};
-        } else {
-            params.stop = json["stop"].template get<std::vector<std::string>>();
-        }
+        params.stop = json["stop"].template get<std::vector<std::string>>();
     } else {
         const std::string defaultStop = "<|eot_id|>";
         params.stop = std::vector<std::string>{defaultStop};

src/app.cpp

@@ -276,8 +276,14 @@ void runInferenceApp(AppCliArgs *args, void (*handler)(AppInferenceContext *cont
     std::vector<NnExecutorDevice> devices = resolveDevices(args, &net.netConfig, rootNodeConfig, &execution);
     NnExecutor executor(&net.netConfig, rootNodeConfig, &devices, &execution, synchronizer.get(), args->benchmark);
 
-    NnRootWeightLoader weightLoader(&executor, network, nNodes);
+    NnLocalWeightLoader weightLoader(&executor, 0, nNodes);
     loadLlmNetWeight(args->modelPath, &net, &weightLoader);
+    if (network != nullptr) {
+        printf("💿 Waiting for workers to load weights...\n");
+        for (NnUint socketIndex = 0; socketIndex < nNodes - 1; socketIndex++)
+            network->readAck(socketIndex);
+        printf("💿 All workers ready\n");
+    }
 
     RootLlmInference inference(&net, &execution, &executor, network);
 
@@ -322,8 +328,26 @@ void runWorkerApp(AppCliArgs *args) {
         NnNetworkNodeSynchronizer synchronizer(network, &execution, &netConfig, &nodeConfig);
         NnExecutor executor(&netConfig, &nodeConfig, &devices, &execution, &synchronizer, false);
 
-        NnWorkerWeightReader weightReader(&executor, network);
-        weightReader.read();
+        if (args->modelPath == nullptr)
+            throw std::runtime_error("--model is required for worker mode");
+
+        NnFloatType syncType = F_32;
+        for (NnUint i = 0; i < netConfig.nPipes; i++) {
+            if (std::strcmp(netConfig.pipes[i].name, "ZQ") == 0) {
+                syncType = netConfig.pipes[i].size.floatType;
+                break;
+            }
+        }
+
+        {
+            LlmHeader workerHeader = loadLlmHeader(args->modelPath, 0, syncType);
+            LlmNet workerNet = buildLlmNet(&workerHeader, netConfig.nNodes, netConfig.nBatches);
+            std::unique_ptr<LlmNet, void(*)(LlmNet *)> workerNetPtr(&workerNet, releaseLlmNet);
+            NnLocalWeightLoader weightLoader(&executor, nodeConfig.nodeIndex, netConfig.nNodes);
+            loadLlmNetWeight(args->modelPath, &workerNet, &weightLoader);
+        }
+
+        network->writeAck(ROOT_SOCKET_INDEX);
 
         WorkerLlmInference inference(&execution, network);
         bool isFirstAttempt = true;

src/llm.cpp

@@ -611,7 +611,7 @@ void releaseLlmNet(LlmNet *net) {
     delete[] net->nodeConfigs;
 }
 
-void loadLlmNetWeight(const char *path, LlmNet *net, NnRootWeightLoader *loader) {
+void loadLlmNetWeight(const char *path, LlmNet *net, NnLocalWeightLoader *loader) {
     MmapFile file;
     openMmapFile(&file, path, net->header->fileSize);
 #if DEBUG_USE_MMAP_FOR_WEIGHTS

src/llm.hpp

@@ -99,6 +99,6 @@ LlmHeader loadLlmHeader(const char* path, const unsigned int maxSeqLen, NnFloatT
 void printLlmHeader(LlmHeader *header);
 LlmNet buildLlmNet(LlmHeader *h, NnUint nNodes, NnUint nBatches);
 void releaseLlmNet(LlmNet *net);
-void loadLlmNetWeight(const char* path, LlmNet *net, NnRootWeightLoader *loader);
+void loadLlmNetWeight(const char* path, LlmNet *net, NnLocalWeightLoader *loader);
 
 #endif

src/nn/nn-network.cpp

@@ -935,3 +935,61 @@ void NnWorkerWeightReader::read() {
     }
     printf("💿 Weights loaded\n");
 }
+
+NnLocalWeightLoader::NnLocalWeightLoader(NnExecutor *executor, NnUint nodeIndex, NnUint nNodes) {
+    this->executor = executor;
+    this->nodeIndex = nodeIndex;
+    this->nNodes = nNodes;
+    this->tempSize = 0;
+}
+
+NnLocalWeightLoader::~NnLocalWeightLoader() {
+    if (tempSize > 0)
+        delete[] temp;
+}
+
+void NnLocalWeightLoader::finish() {}
+
+void NnLocalWeightLoader::allocate(NnSize size) {
+    if (tempSize < size) {
+        if (tempSize > 0)
+            delete[] temp;
+        tempSize = size;
+        temp = new NnByte[size];
+    }
+}
+
+NnSize NnLocalWeightLoader::loadRoot(const char *opName, NnUint opIndex, NnSize nBytes, NnByte *weight) {
+    if (nodeIndex == 0)
+        executor->loadWeight(opName, opIndex, 0u, nBytes, weight);
+    return nBytes;
+}
+
+NnSize NnLocalWeightLoader::loadAll(const char *opName, NnUint opIndex, NnSize nBytes, NnByte *weight) {
+    executor->loadWeight(opName, opIndex, 0u, nBytes, weight);
+    return nBytes;
+}
+
+NnSize NnLocalWeightLoader::loadRowMatmulSlices(const char *opName, const NnUint opIndex, const NnUint expertIndex, NnRowMatmulSlice *slice, NnByte *weight) {
+    const NnUint offset = expertIndex * slice->sliceSize.nBytes;
+    if (nNodes == 1u) {
+        executor->loadWeight(opName, opIndex, offset, slice->sliceSize.nBytes, weight);
+        return slice->size.nBytes;
+    }
+    allocate(slice->sliceSize.nBytes);
+    splitRowMatmulWeight(slice, nodeIndex, weight, temp);
+    executor->loadWeight(opName, opIndex, offset, slice->sliceSize.nBytes, temp);
+    return slice->size.nBytes;
+}
+
+NnSize NnLocalWeightLoader::loadColMatmulSlices(const char *opName, const NnUint opIndex, const NnUint expertIndex, NnColMatmulSlice *slice, NnByte *weight) {
+    const NnUint offset = expertIndex * slice->sliceSize.nBytes;
+    if (nNodes == 1u) {
+        executor->loadWeight(opName, opIndex, offset, slice->sliceSize.nBytes, weight);
+        return slice->size.nBytes;
+    }
+    allocate(slice->sliceSize.nBytes);
+    splitColMatmulWeight(slice, nodeIndex, weight, temp);
+    executor->loadWeight(opName, opIndex, offset, slice->sliceSize.nBytes, temp);
+    return slice->size.nBytes;
+}

src/nn/nn-network.hpp

@@ -133,4 +133,22 @@ class NnWorkerWeightReader {
     void allocate(NnUint size);
 };
 
+class NnLocalWeightLoader {
+private:
+    NnExecutor *executor;
+    NnUint nodeIndex;
+    NnUint nNodes;
+    NnByte *temp;
+    NnSize tempSize;
+    void allocate(NnSize size);
+public:
+    NnLocalWeightLoader(NnExecutor *executor, NnUint nodeIndex, NnUint nNodes);
+    ~NnLocalWeightLoader();
+    NnSize loadRoot(const char *opName, NnUint opIndex, NnSize nBytes, NnByte *weight);
+    NnSize loadAll(const char *opName, NnUint opIndex, NnSize nBytes, NnByte *weight);
+    NnSize loadRowMatmulSlices(const char *opName, const NnUint opIndex, const NnUint expertIndex, NnRowMatmulSlice *slice, NnByte *weight);
+    NnSize loadColMatmulSlices(const char *opName, const NnUint opIndex, const NnUint expertIndex, NnColMatmulSlice *slice, NnByte *weight);
+    void finish();
+};
+
 #endif