| """Layer-output diagnostics; reductions are explicit local measurement choices.""" |
| from contextlib import contextmanager |
| import torch |
|
|
|
|
| @contextmanager |
| def capture_layer_outputs(layers, loop_count): |
| states, handles = [], [] |
| width = len(layers) |
| if width == 0 or loop_count <= 0: |
| raise ValueError('positive layer and loop counts required') |
| def capture(index, output): |
| if index != len(states) % width or len(states) >= width * loop_count: |
| raise ValueError('unexpected recurrent layer invocation order') |
| value = output[0] if isinstance(output, tuple) else output |
| if not isinstance(value, torch.Tensor) or value.ndim != 3: |
| raise ValueError('expected batch/token/hidden layer output') |
| states.append(value.detach().to(device='cpu', dtype=torch.float32, copy=True)) |
| try: |
| for index, layer in enumerate(layers): |
| handles.append(layer.register_forward_hook( |
| lambda module, inputs, output, index=index: capture(index, output))) |
| yield states |
| if len(states) != width * loop_count: |
| raise ValueError('incomplete recurrent layer trajectory') |
| finally: |
| for handle in handles: |
| handle.remove() |
|
|
|
|
| def compare_layer_outputs(reference, observed, *, layer_count, loop_count): |
| if layer_count <= 0 or loop_count <= 0: |
| raise ValueError('positive layer and loop counts required') |
| if len(reference) != layer_count * loop_count or len(observed) != len(reference): |
| raise ValueError('incomplete recurrent layer trajectory') |
| rows = [] |
| for index, (ref, actual) in enumerate(zip(reference, observed)): |
| if ref.shape != actual.shape or ref.ndim != 3 or ref.numel() == 0: |
| raise ValueError('layer output shapes differ or are empty') |
| ref, actual = ref.double(), actual.double() |
| if not torch.isfinite(ref).all() or not torch.isfinite(actual).all(): |
| raise ValueError('nonfinite layer output') |
| delta = actual - ref |
| norm = ref.norm().item() |
| error = delta.norm().item() |
| rows.append(dict(loop=index // layer_count, layer=index % layer_count, |
| trajectory_index=index, tokens=ref.shape[0] * ref.shape[1], |
| l2_frobenius=error, mean_token_l2=delta.norm(dim=-1).mean().item(), |
| relative_frobenius=error / norm if norm else None, |
| reference_frobenius=norm, max_abs_error=delta.abs().max().item())) |
| return rows |
|
|
