Fix segfaults due to bad memory view offset.

rustformers · Jul 9, 2023 · ff4bb37 · ff4bb37
1 parent ae8233a
commit ff4bb37
Show file tree

Hide file tree

Showing 4 changed files with 8 additions and 3 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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,