amachine.am_transformers.am_module_resolver

View Source

 1import torch
 2from typing import cast
 3
 4AM_LAYER_PATHS = [
 5    ("model", "layers"),
 6    ("transformer", "h"),
 7    ("gpt_neox", "layers"),
 8    ("model", "decoder", "layers"),
 9    ("layers",),
10]
11
12AM_WRAPPER_ATTRS = (
13    "module",                 # DDP / DataParallel
14    "_orig_mod",              # torch.compile
15    "_fsdp_wrapped_module",   # FSDP fallback
16)
17
18def find_decoder_layers(model) -> torch.nn.ModuleList :
19    """Find decoder layers, descending through common wrappers if needed."""
20    pending, visited = [model], set()
21
22    while pending:
23        current = pending.pop()
24        if id(current) in visited:
25            continue
26        visited.add(id(current))
27
28        for path in AM_LAYER_PATHS:
29            obj = current
30            for name in path:
31                obj = getattr(obj, name, None)
32                if obj is None:
33                    break
34            if obj is not None and hasattr(obj, "__len__"):
35                return cast(torch.nn.ModuleList, obj)
36
37        for attr in AM_WRAPPER_ATTRS:
38            wrapped = getattr(current, attr, None)
39            if isinstance(wrapped, torch.nn.Module) and wrapped is not current:
40                pending.append(wrapped)
41
42    raise RuntimeError("Could not locate decoder layers.")

AM_LAYER_PATHS = [('model', 'layers'), ('transformer', 'h'), ('gpt_neox', 'layers'), ('model', 'decoder', 'layers'), ('layers',)]

AM_WRAPPER_ATTRS = ('module', '_orig_mod', '_fsdp_wrapped_module')

def find_decoder_layers(model) -> torch.nn.modules.container.ModuleList: View Source

19def find_decoder_layers(model) -> torch.nn.ModuleList :
20    """Find decoder layers, descending through common wrappers if needed."""
21    pending, visited = [model], set()
22
23    while pending:
24        current = pending.pop()
25        if id(current) in visited:
26            continue
27        visited.add(id(current))
28
29        for path in AM_LAYER_PATHS:
30            obj = current
31            for name in path:
32                obj = getattr(obj, name, None)
33                if obj is None:
34                    break
35            if obj is not None and hasattr(obj, "__len__"):
36                return cast(torch.nn.ModuleList, obj)
37
38        for attr in AM_WRAPPER_ATTRS:
39            wrapped = getattr(current, attr, None)
40            if isinstance(wrapped, torch.nn.Module) and wrapped is not current:
41                pending.append(wrapped)
42
43    raise RuntimeError("Could not locate decoder layers.")

Find decoder layers, descending through common wrappers if needed.