From 6ee0547eff652da2c8f6b5e67deaf004341efc67 Mon Sep 17 00:00:00 2001 From: Alex Cheema Date: Sat, 13 Jul 2024 15:39:31 -0700 Subject: [PATCH] fix layer calculation for sharded llama --- inference/mlx/models/sharded_llama.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/inference/mlx/models/sharded_llama.py b/inference/mlx/models/sharded_llama.py index db67607d..c2c851b7 100644 --- a/inference/mlx/models/sharded_llama.py +++ b/inference/mlx/models/sharded_llama.py @@ -166,7 +166,7 @@ class LlamaModel(nn.Module): assert self.vocab_size > 0 self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size) self.layers = [ - TransformerBlock(args=args) for _ in range(args.shard.n_layers) + TransformerBlock(args=args) for _ in range(args.shard.end_layer - args.shard.start_layer + 1) ] self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)