| """Dependency-free architecture arithmetic and shared initialization mapping.""" |
| from __future__ import annotations |
|
|
| MAMBA_REVISION = 'd7b1ceb3c367ec9022925e812f507bcf706937c6' |
| BASE_REVISION = '9b98f6e75713212b737d8b7bf7c3fb7913e14ee6' |
| CORPUS_SHA256 = '1786fc7cf798784edf2e0ee9400456ca93add1cbc5688cd8e83865cb230bc475' |
| CAPS = (1, 2, 4, 8, 16, 32, 64, 96, 128, 192, 256, 384, 512, 768, 1024, 1536, 2048) |
|
|
|
|
| def shared_parameter_name(name, outer_layers=4): |
| """Map new middle-block names to the original scratch tensor names. |
| |
| M4/20/M4 indices 4..23 correspond to T2/20/T2 indices 2..21. |
| The same two-for-one mapping supports M2/2/M2 tiny correctness models. |
| """ |
| if outer_layers < 2 or outer_layers % 2: |
| raise ValueError('Two native Mamba blocks replace one Transformer block') |
| if not name.startswith('layers.'): |
| return name |
| parts = name.split('.') |
| parts[1] = str(int(parts[1]) - outer_layers // 2) |
| return '.'.join(parts) |
|
|
|
|
| def parameter_counts(hidden=1024, intermediate=2816, vocab=32000, |
| middle_layers=20, outer_layers=4, state=128, |
| expand=2, headdim=64, conv=4): |
| inner = hidden * expand |
| if inner % headdim: |
| raise ValueError('Mamba inner width must divide into complete heads') |
| heads = inner // headdim |
| transformer = 4 * hidden**2 + 3 * hidden * intermediate + 2 * hidden |
| mamba = (hidden * (2 * inner + 2 * state + heads) |
| + (inner + 2 * state) * (conv + 1) |
| + 3 * heads + inner + inner * hidden + hidden) |
| shared = vocab * hidden + middle_layers * transformer + hidden |
| old = shared + outer_layers * transformer |
| new = shared + 2 * outer_layers * mamba |
| return {'transformer_block': transformer, 'mamba_block_with_prenorm': mamba, |
| 'original_model': old, 'mamba_outer_model': new, |
| 'increase_percent': 100 * (new - old) / old} |
|
|
|
|
| def theoretical_state_bytes(cap, *, batch=1, hidden=1024, middle_layers=20, |
| outer_layers=4, expand=2, state=128, conv=4, |
| kv_item_bytes=2, ssm_item_bytes=4): |
| if min(cap, batch, kv_item_bytes, ssm_item_bytes) < 1: |
| raise ValueError('Cap, batch and scalar byte sizes must be positive') |
| inner = hidden * expand |
| kv = batch * middle_layers * 2 * hidden * cap * kv_item_bytes |
| recurrent = batch * 2 * outer_layers * inner * state * ssm_item_bytes |
| convolution = batch * 2 * outer_layers * (inner + 2 * state) * conv * kv_item_bytes |
| return {'kv': kv, 'ssm': recurrent, 'convolution': convolution, |
| 'total_tensor_bytes': kv + recurrent + convolution, |
| 'excludes': ['host metadata', 'shared positional buffers', 'weights', 'temporary workspace']} |
|
|
