Skip to content

Commit

Permalink
Rename and tidy layer printing
Browse files Browse the repository at this point in the history
  • Loading branch information
dnhkng committed Jan 12, 2024
1 parent baa1efe commit 4220642
Showing 1 changed file with 7 additions and 6 deletions.
13 changes: 7 additions & 6 deletions exllamav2/model.py
Original file line number Diff line number Diff line change
Expand Up @@ -160,7 +160,7 @@ def __init__(self, config: ExLlamaV2Config, lazy_load = False):


if hasattr(config, 'repeats'):
self.layers = []
self.layers_list = []

def listLeftIndex(alist, value):
if value == 0:
Expand All @@ -177,12 +177,13 @@ def listRightIndex(alist, value):
for interval in config.repeats:
start_idx = listLeftIndex(layer_list, interval[0])
end_idx = listRightIndex(layer_list, interval[1])
self.layers.extend(list(range(start_idx, end_idx + 1)))
self.layers.extend(list(range(listRightIndex(layer_list, config.repeats[-1][1]), len(layer_list))))
self.layers_list.extend(list(range(start_idx, end_idx + 1)))
self.layers_list.extend(list(range(listRightIndex(layer_list, config.repeats[-1][1]), len(layer_list))))

# If we have create a Frankenmerge, lets print it to verify!
for layer in self.layers:
print(layer, self.modules[layer].key)
print("Frankenstein Layers list:")
for i, layer in enumerate(self.layers_list):
print(i, self.modules[layer].key)

def set_device_map(self, allocation, embed_cpu = True):

Expand Down Expand Up @@ -639,7 +640,7 @@ def process_module(module, x, last_state):
last_state = None

if hasattr(self, 'layers'):
for i, idx in enumerate(self.layers):
for i, idx in enumerate(self.layers_list):
module = self.modules[idx]
x, last_state = process_module(module, x, last_state)
if preprocess_only and idx == self.last_kv_layer_idx:
Expand Down

0 comments on commit 4220642

Please sign in to comment.