Stash gen of LLAMA mlx

src/inference.rs

@@ -0,0 +1,19 @@
+use crate::node_service::{InferenceState, Tensor};
+use crate::Shard;
+
+#[derive(Debug)]
+pub struct InferenceEngine {
+    state_cache: HashMap<String, _>,
+}
+
+impl InferenceEngine {
+    pub(crate) fn infer_tensor(
+        &self,
+        request_id: String,
+        shard: Shard,
+        tensor: Option<Tensor>,
+        inference_state: Option<InferenceState>,
+    ) -> Tensor {
+
+    }
+}

src/llama_module.rs

@@ -0,0 +1,226 @@
+use mlx_rs::macros::ModuleParameters;
+use mlx_rs::module::Module;
+use mlx_rs::module::ModuleParameters;
+use mlx_rs::module::Param;
+use mlx_rs::nested::NestedHashMap;
+use mlx_rs::nn;
+use mlx_rs::Array;
+use std::rc::Rc;
+
+// Define Shard struct to mirror Python dataclass
+#[derive(Debug, Clone)]
+pub struct Shard {
+    pub name: String,
+    pub start_layer: usize,
+    pub end_layer: usize,
+}
+
+impl Shard {
+    pub fn new(name: String, start_layer: usize, end_layer: usize) -> Self {
+        Shard {
+            name,
+            start_layer,
+            end_layer,
+        }
+    }
+
+    pub fn is_first_layer(&self) -> bool {
+        self.start_layer == 0
+    }
+
+    pub fn is_last_layer(&self) -> bool {
+        // Assuming end_layer is inclusive and represents the last layer index in the shard
+        // and num_hidden_layers is the total number of layers.
+        // We would need num_hidden_layers to accurately determine the last layer.
+        // For now, let's assume if end_layer is very large, it's the last layer in shard.
+        self.end_layer > 9999 // A large number as a placeholder, adjust as needed
+    }
+}
+
+// Define ModelArgs struct to mirror Python dataclass ModelArgs
+#[derive(Debug, Clone)]
+pub struct ModelArgsRs {
+    pub vocab_size: i32,
+    pub hidden_size: i32,
+    pub num_hidden_layers: i32,
+    pub num_attention_heads: i32,
+    pub num_key_value_heads: i32,
+    pub rms_norm_eps: f32,
+    pub tie_word_embeddings: bool,
+    pub model_type: String,    // Assuming model_type is a String
+    pub head_dim: Option<i32>, // Using Option to represent optional field
+    pub shard: Shard,          // Using the Shard struct defined above
+}
+
+// Placeholder for TransformerBlock - You'll need to implement this in Rust
+#[derive(Debug, Clone, ModuleParameters)]
+pub struct TransformerBlock {
+    // Define the layers within TransformerBlock as needed, e.g., attention, norm, etc.
+    // For now, using a linear layer as a placeholder
+    pub linear: nn::Linear,
+}
+
+impl TransformerBlock {
+    pub fn new(dims: i32, mlp_dims: i32) -> Self {
+        Self {
+            linear: nn::Linear::new(dims, mlp_dims).unwrap(), // Example, adjust params
+        }
+    }
+}
+
+impl Module<&Array> for TransformerBlock {
+    type Output = Array;
+    type Error = mlx_rs::error::Exception;
+
+    fn forward(&mut self, input: &Array) -> Result<Self::Output, Self::Error> {
+        // Implement the forward pass logic for TransformerBlock
+        // For now, just passing through the linear layer
+        self.linear.forward(input)
+    }
+
+    fn training_mode(&mut self, mode: bool) {}
+}
+
+// Define LlamaModel struct
+#[derive(Debug, Clone, ModuleParameters)]
+pub struct LlamaModelRs {
+    pub args: ModelArgsRs,
+    pub vocab_size: i32,
+    pub num_hidden_layers: i32,
+    pub embed_tokens: Option<nn::Embedding>, // Embedding layer is optional based on sharding
+    pub layers: Vec<TransformerBlock>,       // Using placeholder TransformerBlock
+    pub norm: Option<nn::RmsNorm>,           // RMSNorm layer is optional based on sharding
+}
+
+impl LlamaModelRs {
+    pub fn new(args: ModelArgsRs) -> Result<Self, mlx_rs::error::Exception> {
+        let vocab_size = args.vocab_size;
+        let num_hidden_layers = args.num_hidden_layers;
+        let mut embed_tokens = None;
+        if args.shard.is_first_layer() || (args.shard.is_last_layer() && args.tie_word_embeddings) {
+            embed_tokens = Some(nn::Embedding::new(args.vocab_size, args.hidden_size)?);
+        }
+
+        let mut layers = Vec::new();
+        for i in 0..num_hidden_layers {
+            if args.shard.start_layer <= i && i <= args.shard.end_layer {
+                // Using placeholder dimensions for TransformerBlock, adjust as needed
+                layers.push(TransformerBlock::new(
+                    args.hidden_size,
+                    args.hidden_size * 4,
+                ));
+            } else {
+                // Placeholder for IdentityBlock - you might need to create a Rust version if needed
+                // For now, just pushing a default TransformerBlock or handle differently
+                layers.push(TransformerBlock::new(
+                    args.hidden_size,
+                    args.hidden_size * 4,
+                )); // IdentityBlock() in Python seems to be a no-op
+            }
+        }
+
+        let mut norm = None;
+        if args.shard.is_last_layer() {
+            norm = Some(nn::RmsNorm::new(args.hidden_size, args.rms_norm_eps)?);
+        }
+
+        Ok(Self {
+            args,
+            vocab_size,
+            num_hidden_layers,
+            embed_tokens,
+            layers,
+            norm,
+        })
+    }
+}
+
+impl Module<&Array> for LlamaModelRs {
+    type Output = Array;
+    type Error = mlx_rs::error::Exception;
+
+    fn forward(&mut self, inputs: &Array) -> Result<Self::Output, Self::Error> {
+        let mut h;
+        if self.args.shard.is_first_layer() && self.embed_tokens.is_some() {
+            h = self.embed_tokens.as_ref().unwrap().forward(inputs)?;
+        } else {
+            h = inputs.clone(); // Assuming input is already embedded if not the first layer
+        }
+
+        // Mask creation logic would go here - needs to be implemented in Rust
+        // let mask = None;
+        // if h.ndim() > 1 && h.shape()[1] > 1 {
+        //     mask = create_attention_mask(h, cache); // Need to port create_attention_mask to Rust
+        // }
+
+        // Cache handling - needs more detailed implementation for Rust
+        // let mut cache = cache.unwrap_or_else(|| vec![None; self.layers.len()]);
+
+        for layer in &mut self.layers {
+            h = layer.forward(&h)?; // Pass mask and cache when implemented
+        }
+
+        if self.args.shard.is_last_layer() && self.norm.is_some() {
+            h = self.norm.as_ref().unwrap().forward(&h)?;
+        }
+        Ok(h)
+    }
+
+    fn training_mode(&mut self, mode: bool) {}
+}
+
+// Define Model struct
+#[derive(Debug, Clone, ModuleParameters)]
+pub struct ModelRs {
+    pub args: ModelArgsRs,
+    pub model_type: String,
+    pub model: LlamaModelRs,
+    pub lm_head: Option<nn::Linear>, // Linear layer for language model head, optional based on tie_word_embeddings
+}
+
+impl ModelRs {
+    pub fn new(args: ModelArgsRs) -> Result<Self, mlx_rs::error::Exception> {
+        let model = LlamaModelRs::new(args.clone())?; // Clone args for LlamaModel
+        let model_type = args.model_type.clone();
+        let mut lm_head = None;
+        if args.shard.is_last_layer() && !args.tie_word_embeddings {
+            lm_head = Some(nn::Linear::new(args.hidden_size, args.vocab_size)?);
+        }
+
+        Ok(Self {
+            args,
+            model_type,
+            model,
+            lm_head,
+        })
+    }
+}
+
+impl Module<&Array> for ModelRs {
+    type Output = Array;
+    type Error = mlx_rs::error::Exception;
+
+    fn forward(&mut self, inputs: &Array) -> Result<Self::Output, Self::Error> {
+        let mut out = self.model.forward(inputs)?;
+
+        if self.args.shard.is_last_layer() {
+            if self.args.tie_word_embeddings && self.model.embed_tokens.is_some() {
+                // Need to implement as_linear() equivalent in Rust or directly use embedding weights for linear transformation
+                // Placeholder - direct linear transformation using embedding weights is not directly available in mlx-rs as in python
+                if let Some(embed_tokens) = &self.model.embed_tokens {
+                    if let Ok(params) = embed_tokens.parameters() {
+                        if let Some(weight_param) = params.get("weight") {
+                            // This is a very simplified placeholder - needs proper matrix multiplication with 'out' and 'weight_param'
+                            out = weight_param.clone(); // Incorrect - replace with actual linear transformation
+                        }
+                    }
+                }
+            } else if self.lm_head.is_some() {
+                out = self.lm_head.as_ref().unwrap().forward(&out)?;
+            }
+        }
+        Ok(out)
+    }
+
+    fn training_mode(&mut self, mode: bool) {}
+}

src/llama_test.rs

@@ -0,0 +1,73 @@
+use llama_cpp_2::context::params::LlamaContextParams;
+use llama_cpp_2::llama_backend::LlamaBackend;
+use llama_cpp_2::llama_batch::LlamaBatch;
+use llama_cpp_2::model::{AddBos, LlamaModel, Special};
+use llama_cpp_2::model::params::LlamaModelParams;
+use llama_cpp_2::sampling::LlamaSampler;
+
+fn test() {
+    let model_path = std::env::args().nth(1).expect("Please specify model path");
+    let backend = LlamaBackend::init().unwrap();
+    let params = LlamaModelParams::default();
+
+    let prompt =
+        "<|im_start|>user\nHello! how are you?<|im_end|>\n<|im_start|>assistant\n".to_string();
+    LlamaContextParams::default();
+    let model =
+        LlamaModel::load_from_file(&backend, model_path, &params).expect("unable to load model");
+    let ctx_params = LlamaContextParams::default();
+    let mut ctx = model
+        .new_context(&backend, ctx_params)
+        .expect("unable to create the llama_context");
+    let tokens_list = model
+        .str_to_token(&prompt, AddBos::Always)
+        .unwrap_or_else(|_| panic!("failed to tokenize {prompt}"));
+    let n_len = 64;
+
+    // create a llama_batch with size 512
+    // we use this object to submit token data for decoding
+    let mut batch = LlamaBatch::new(512, 1);
+
+    let last_index = tokens_list.len() as i32 - 1;
+    for (i, token) in (0_i32..).zip(tokens_list.into_iter()) {
+        // llama_decode will output logits only for the last token of the prompt
+        let is_last = i == last_index;
+        batch.add(token, i, &[0], is_last).unwrap();
+    }
+    ctx.decode(&mut batch).expect("llama_decode() failed");
+
+    let mut n_cur = batch.n_tokens();
+
+    // The `Decoder`
+    let mut decoder = encoding_rs::UTF_8.new_decoder();
+    let mut sampler = LlamaSampler::greedy();
+
+    while n_cur <= n_len {
+        // sample the next token
+        {
+            let token = sampler.sample(&ctx, batch.n_tokens() - 1);
+
+            sampler.accept(token);
+
+            // is it an end of stream?
+            if token == model.token_eos() {
+                eprintln!();
+                break;
+            }
+
+            let output_bytes = model.token_to_bytes(token, Special::Tokenize).unwrap();
+            // use `Decoder.decode_to_string()` to avoid the intermediate buffer
+            let mut output_string = String::with_capacity(32);
+            let _decode_result = decoder.decode_to_string(&output_bytes, &mut output_string, false);
+            print!("{output_string}");
+            std::io::stdout().flush().unwrap();
+
+            batch.clear();
+            batch.add(token, n_cur, &[0], true).unwrap();
+        }
+
+        n_cur += 1;
+
+        ctx.decode(&mut batch).expect("failed to eval");
+    }
+}

src/main.rs

@@ -1,15 +1,20 @@
 mod device_capability_data;
 mod discovery;
+mod inference;
 mod network;
 mod orchestration;
 mod partitioning;
 mod topology;
+mod llama_test;
+mod module_loading;
+mod llama_module;
 
 use serde::{Deserialize, Serialize};
 use serde_json::Value;
 use tonic::{transport::Server, Request, Response, Status};
 
 use crate::discovery::{NodeInfo, UdpDiscovery};
+use crate::inference::InferenceEngine;
 use crate::node_service::{
     CollectTopologyRequest, Empty, ExampleRequest, HealthCheckRequest, HealthCheckResponse,
     InferenceState, Loss, PromptRequest, SendOpaqueStatusRequest, SendResultRequest, Tensor,
@@ -30,6 +35,7 @@ struct Node {
     node_info: NodeInfo,
     current_topology: Topology,
     udp_discovery: UdpDiscovery,
+    inference_engine: InferenceEngine,
 }
 
 impl Node {
@@ -45,7 +51,10 @@ impl Node {
             .current_topology
             .get_shard_for_node(base_shard, &self.node_info.node_id);
 
-        let result = self.inference_engine.infer_tensor(request_id, shard, tensor, inference_state);
+        let result: Tensor = self
+            .inference_engine
+            .infer_tensor(request_id, shard, tensor, inference_state);
+
         let result = self.process_inference_result(shard, result, request_id, inference_state);
 
         result

src/module_loading.rs

@@ -0,0 +1,52 @@
+use std::path::Path;
+
+use crate::Shard;
+
+fn load_config(
+    model_path: &Path,
+) -> serde_json::Map<String, serde_json::Value> {
+    let config_path = model_path.join("config.json");
+    let model_index = model_path.join("model_index.json");
+
+    if config_path.exists() {
+        let config = std::fs::read_to_string(config_path).unwrap();
+        serde_json::from_str(&config).unwrap()
+    } else {
+        let model_index = std::fs::read_to_string(model_index).unwrap();
+        serde_json::from_str(&model_index).unwrap()
+    }
+}
+
+async fn load_model_shard(
+    model_path: &Path,
+    shard: Shard,
+    lazy: bool,
+    model_config: serde_json::Map<String, serde_json::Value>,
+) {
+    let mut config = load_config(model_path);
+    config.extend(model_config.into_iter());
+
+    let model_name = model_path.file_name().unwrap().to_str().unwrap();
+
+    config["shard"] = serde_json::json!({
+        "model_id": model_name,
+        "start_layer": shard.start_layer,
+        "end_layer": shard.end_layer,
+        "n_layers": shard.total_layers,
+    });
+
+    let weight_files = glob::glob(model_path.join("model*.safetensors").to_str().unwrap())
+        .unwrap()
+        .collect::<Result<Vec<_>, _>>()
+        .unwrap();
+
+    let weight_files = weight_files
+        .iter()
+        .map(|path| path.file_name().unwrap().to_str().unwrap())
+        .collect::<Vec<_>>();
+
+    let weights = weight_files.iter().map(|file| {
+        
+    });
+    todo!();
+}