From ff4bb377d69e2d2c4cceb6efb004ffd0a9523451 Mon Sep 17 00:00:00 2001
From: Steven Weiss <steventrouble@gmail.com>
Date: Sun, 9 Jul 2023 12:34:04 -0700
Subject: [PATCH 1/3] Fix segfaults due to bad memory view offset.

---
 crates/ggml/src/context.rs       | 5 +++++
 crates/models/gpt2/src/lib.rs    | 2 +-
 crates/models/gptj/src/lib.rs    | 2 +-
 crates/models/gptneox/src/lib.rs | 2 +-
 4 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/crates/ggml/src/context.rs b/crates/ggml/src/context.rs
index 9dd7d108..e3d24017 100644
--- a/crates/ggml/src/context.rs
+++ b/crates/ggml/src/context.rs
@@ -278,6 +278,11 @@ impl Context {
 
     /// Creates a 1D view over `a`.
     pub fn op_view_1d(&self, a: &Tensor, ne0: usize, offset: usize) -> Tensor {
+        #[cfg(debug_assertions)]
+        assert!(
+            offset < a.nbytes(),
+            "Cannot create tensor view with offset larger than tensor"
+        );
         let tensor = unsafe {
             sys::ggml_view_1d(self.ptr.as_ptr(), a.ptr.as_ptr(), usize_to_i64(ne0), offset)
         };
diff --git a/crates/models/gpt2/src/lib.rs b/crates/models/gpt2/src/lib.rs
index 646d3a98..9e2102b1 100644
--- a/crates/models/gpt2/src/lib.rs
+++ b/crates/models/gpt2/src/lib.rs
@@ -102,7 +102,7 @@ impl KnownModel for Gpt2 {
     fn start_session(&self, config: InferenceSessionConfig) -> InferenceSession {
         InferenceSession::new(
             config,
-            self.hyperparameters.n_ctx,
+            self.context_size,
             self.hyperparameters.n_layer,
             self.hyperparameters.n_embd,
             self.hyperparameters.n_vocab,
diff --git a/crates/models/gptj/src/lib.rs b/crates/models/gptj/src/lib.rs
index 195f876a..14beb3e1 100644
--- a/crates/models/gptj/src/lib.rs
+++ b/crates/models/gptj/src/lib.rs
@@ -102,7 +102,7 @@ impl KnownModel for GptJ {
     fn start_session(&self, config: InferenceSessionConfig) -> InferenceSession {
         InferenceSession::new(
             config,
-            self.hyperparameters.n_ctx,
+            self.context_size,
             self.hyperparameters.n_layer,
             self.hyperparameters.n_embd,
             self.hyperparameters.n_vocab,
diff --git a/crates/models/gptneox/src/lib.rs b/crates/models/gptneox/src/lib.rs
index 5339b901..7b2e1649 100644
--- a/crates/models/gptneox/src/lib.rs
+++ b/crates/models/gptneox/src/lib.rs
@@ -115,7 +115,7 @@ impl KnownModel for GptNeoX {
     fn start_session(&self, config: InferenceSessionConfig) -> InferenceSession {
         InferenceSession::new(
             config,
-            self.hyperparameters.n_ctx,
+            self.context_size,
             self.hyperparameters.n_layer,
             self.hyperparameters.n_embd,
             self.hyperparameters.n_vocab,

From 0f841cdfd7783c5793182eb1cdec12ebe2dbb45c Mon Sep 17 00:00:00 2001
From: Steven Weiss <steventrouble@gmail.com>
Date: Sun, 9 Jul 2023 12:37:19 -0700
Subject: [PATCH 2/3] Update gpt2 to use wte if no lm_head

---
 crates/models/gpt2/src/lib.rs | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/crates/models/gpt2/src/lib.rs b/crates/models/gpt2/src/lib.rs
index 9e2102b1..2eae1ce5 100644
--- a/crates/models/gpt2/src/lib.rs
+++ b/crates/models/gpt2/src/lib.rs
@@ -31,7 +31,7 @@ pub struct Gpt2 {
     // weighted positional encodings
     wpe: Tensor,
     // language model head
-    lm_head: Tensor,
+    lm_head: Option<Tensor>,
 
     // weights for the model
     layers: Vec<Layer>,
@@ -59,7 +59,7 @@ impl KnownModel for Gpt2 {
         let ln_f_b = tl.load("model/ln_f/b")?;
         let wte = tl.load("model/wte")?;
         let wpe = tl.load("model/wpe")?;
-        let lm_head = tl.load("model/lm_head")?;
+        let lm_head = tl.load("model/lm_head").ok();
 
         let mut layers = Vec::new();
         for i in 0..hyperparameters.n_layer {
@@ -306,7 +306,8 @@ impl KnownModel for Gpt2 {
 
             let embeddings_tensor: ggml::Tensor = input_layer.share();
 
-            input_layer = ctx0.op_mul_mat(&self.lm_head, &input_layer);
+            let head = self.lm_head.as_ref().unwrap_or(&self.wte);
+            input_layer = ctx0.op_mul_mat(head, &input_layer);
 
             (
                 gf,

From c995ca8fc556560746027a4c33930260d18fce8a Mon Sep 17 00:00:00 2001
From: Steven Weiss <steventrouble@gmail.com>
Date: Mon, 10 Jul 2023 10:38:53 -0700
Subject: [PATCH 3/3] Address PR comments

---
 crates/models/gpt2/src/lib.rs | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/crates/models/gpt2/src/lib.rs b/crates/models/gpt2/src/lib.rs
index 2eae1ce5..c69832ca 100644
--- a/crates/models/gpt2/src/lib.rs
+++ b/crates/models/gpt2/src/lib.rs
@@ -31,6 +31,8 @@ pub struct Gpt2 {
     // weighted positional encodings
     wpe: Tensor,
     // language model head
+    //
+    // Optional: if not present, the `wte` tensor is used instead.
     lm_head: Option<Tensor>,
 
     // weights for the model
@@ -59,6 +61,9 @@ impl KnownModel for Gpt2 {
         let ln_f_b = tl.load("model/ln_f/b")?;
         let wte = tl.load("model/wte")?;
         let wpe = tl.load("model/wpe")?;
+
+        // GPT-2's language model head is optional; if it is not present,
+        // the `wte` tensor is used instead.
         let lm_head = tl.load("model/lm_head").ok();
 
         let mut layers = Vec::new();