Update gpt2 to use wte if no lm_head

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)
  };

crates/models/gpt2/src/lib.rs

@@ -31,7 +31,9 @@ pub struct Gpt2 {
  // weighted positional encodings
  wpe: Tensor,
  // language model head
- lm_head: Tensor,
+ //
+ // Optional: if not present, the `wte` tensor is used instead.
+ lm_head: Option<Tensor>,
 
  // weights for the model
  layers: Vec<Layer>,
@@ -59,7 +61,10 @@ 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")?;
+
+ // 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();
  for i in 0..hyperparameters.n_layer {
@@ -102,7 +107,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,
@@ -306,7 +311,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,

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,

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,