"""Real pretrained-model measurements for Lesson 7. Python 3.10+, PyTorch, transformers 4.57.1. No provider API, simulated tensors, or invented GPU counters. The browser uses the same checkpoint. The notebook can update its display live via run(on_frame=callback). """ import argparse import html import json from pathlib import Path import statistics import time import torch from transformers import AutoModelForCausalLM, AutoTokenizer MODEL = 'roneneldan/TinyStories-1M' REVISION = '77f1b168e219585646439073245fe87e56b3023e' def cache_bytes(cache): """Logical payload of the tensors the implementation actually retains.""" if cache is None: return 0 if hasattr(cache, 'layers'): return sum(value.numel() * value.element_size() for layer in cache.layers for value in (getattr(layer, 'keys', None), getattr(layer, 'values', None)) if isinstance(value, torch.Tensor)) return sum(value.numel() * value.element_size() for layer in cache for value in layer) class Observer: def __init__(self, device=None, local_model=None): self.device = torch.device(device or ('cuda' if torch.cuda.is_available() else 'cpu')) if self.device.type == 'cuda' and not torch.cuda.is_available(): raise RuntimeError('CUDA is unavailable. Select a GPU runtime or explicitly use device="cpu".') if self.device.type == 'cpu': torch.set_num_threads(min(4, torch.get_num_threads())) source = local_model or MODEL options = {} if local_model else {'revision': REVISION} self.tokenizer = AutoTokenizer.from_pretrained(source, **options) self.model = AutoModelForCausalLM.from_pretrained(source, attn_implementation='eager', **options).float().to(self.device).eval() self.weight_bytes = sum(p.numel() * p.element_size() for p in self.model.parameters()) self.device_name = torch.cuda.get_device_name(self.device) if self.device.type == 'cuda' else 'CPU · PyTorch' def synchronize(self): if self.device.type == 'cuda': torch.cuda.synchronize(self.device) def memory(self, cache): result = {'weightBytes': self.weight_bytes, 'kvBytes': cache_bytes(cache)} if self.device.type == 'cuda': result.update(allocatedBytes=torch.cuda.memory_allocated(self.device), reservedBytes=torch.cuda.memory_reserved(self.device), peakAllocatedBytes=torch.cuda.max_memory_allocated(self.device)) return result def inputs(self, prompt, batch=1): ids = self.tokenizer.encode(prompt) if not ids or len(ids) > 128: raise ValueError('Use a non-empty prompt of at most 128 tokens.') return torch.tensor([ids] * batch, device=self.device) @torch.inference_mode() def run(self, prompt='Once upon a time, there was a little', tokens=16, on_frame=None): if not 1 <= tokens <= 64: raise ValueError('Choose 1–64 prediction frames.') full = self.inputs(prompt) current, cache = full, None trace = {'schema': 'vcc-observer-1', 'createdAt': time.strftime('%Y-%m-%dT%H:%M:%SZ', time.gmtime()), 'device': {'type': self.device.type, 'name': self.device_name}, 'model': {'name': MODEL, 'revision': REVISION, 'weightBytes': self.weight_bytes, 'parameterCount': sum(p.numel() for p in self.model.parameters())}, 'frames': [], 'measurement': 'Actual forward-hook values; timings include hook copying and synchronization. No hardware byte counters.'} captured, handles = {}, [] for i, block in enumerate(self.model.transformer.h): def mlp_hook(module, args, output, index=i): captured.setdefault(index, {})['mlp'] = args[0][0, -1].detach().float().cpu().tolist() handles.append(block.mlp.c_proj.register_forward_hook(mlp_hook)) def hidden_hook(module, args, output, index=i): value = output[0] if isinstance(output, tuple) else output captured.setdefault(index, {})['hidden'] = value[0, -1].detach().float().cpu().tolist() handles.append(block.register_forward_hook(hidden_hook)) pointers = [p.data_ptr() for p in self.model.parameters()] if self.device.type == 'cuda': torch.cuda.reset_peak_memory_stats(self.device) try: for step in range(tokens): captured.clear() self.synchronize() start = time.perf_counter() output = self.model(current, past_key_values=cache, use_cache=True) cache = output.past_key_values self.synchronize() elapsed = (time.perf_counter() - start) * 1000 logits = output.logits[0, -1].float() probabilities = logits.softmax(-1) best = logits.topk(8).indices.tolist() next_id = best[0] decode = lambda token: '[end]' if token == self.tokenizer.eos_token_id else self.tokenizer.decode([token]) frame = {'phase': 'prefill' if step == 0 else 'decode', 'position': full.shape[1], 'inputId': full[0, -1].item(), 'nextId': next_id, 'nextText': decode(next_id), 'context': self.tokenizer.decode(full[0].tolist()), 'elapsedMs': elapsed, 'layers': [captured[i] for i in range(self.model.config.num_layers)], 'top': [{'id': k, 'text': decode(k), 'logit': logits[k].item(), 'probability': probabilities[k].item()} for k in best], 'memory': self.memory(cache), 'weightStorageUnchanged': pointers == [p.data_ptr() for p in self.model.parameters()]} trace['frames'].append(frame) if on_frame: on_frame(frame, trace) if next_id == self.tokenizer.eos_token_id: break current = torch.tensor([[next_id]], device=self.device) full = torch.cat((full, current), dim=1) finally: for handle in handles: handle.remove() return trace @torch.inference_mode() def benchmark(self, prompt='Once upon a time, there was a little', batches=(1, 4, 8), steps=16, repeats=3): """Hook-free decode only; fresh prefill outside timing; synchronized wall time.""" if not 1 <= steps <= 64 or not 1 <= repeats <= 10: raise ValueError('Use 1–64 steps and 1–10 repeats.') rows = [] for batch in batches: if not isinstance(batch, int) or not 1 <= batch <= 16: raise ValueError('Batch sizes must be integers from 1 to 16.') prompt_ids = self.inputs(prompt, batch) times = [] for repeat in range(repeats + 1): initial = self.model(prompt_ids, use_cache=True) cache = initial.past_key_values token = initial.logits[:, -1].argmax(-1, keepdim=True) self.synchronize() start = time.perf_counter() for _ in range(steps): output = self.model(token, past_key_values=cache, use_cache=True) cache = output.past_key_values token = output.logits[:, -1].argmax(-1, keepdim=True) self.synchronize() if repeat: times.append(time.perf_counter() - start) median = statistics.median(times) rows.append({'batch': batch, 'decodeSteps': steps, 'outputTokens': batch * steps, 'medianSeconds': median, 'secondsByRepeat': times, 'msPerStep': median * 1000 / steps, 'totalTokensPerSecond': batch * steps / median}) return {'device': self.device_name, 'scope': 'Synchronized hook-free decode; prompt processing excluded; one warm-up run per batch; identical prompts; no hardware byte measurement.', 'rows': rows} @torch.inference_mode() def profile(self, path='observer-timeline.json'): """Actual operation trace; this does not measure DRAM bytes.""" activities = [torch.profiler.ProfilerActivity.CPU] if self.device.type == 'cuda': activities.append(torch.profiler.ProfilerActivity.CUDA) inputs = self.inputs('Once upon a time, there was a little') prefill = self.model(inputs, use_cache=True) token = prefill.logits[:, -1].argmax(-1, keepdim=True) with torch.profiler.profile(activities=activities, profile_memory=True, record_shapes=True) as profile: with torch.profiler.record_function('observer_decode_step'): self.model(token, past_key_values=prefill.past_key_values, use_cache=True) self.synchronize() profile.export_chrome_trace(str(path)) return str(path) @torch.inference_mode() def hardware_step(self): """Nsight Compute captures only this warmed step via CUDA profiler control.""" if self.device.type != 'cuda': raise RuntimeError('Hardware profiling requires a CUDA GPU; CPU values will not be substituted.') ids = self.inputs('Once upon a time, there was a little') # Warm up the exact decode shape without mutating the captured cache. initial = self.model(ids, use_cache=True) token = initial.logits[:, -1].argmax(-1, keepdim=True) self.model(token, past_key_values=initial.past_key_values, use_cache=True) initial = self.model(ids, use_cache=True) self.synchronize() torch.cuda.cudart().cudaProfilerStart() self.model(token, past_key_values=initial.past_key_values, use_cache=True) self.synchronize() torch.cuda.cudart().cudaProfilerStop() def render_frame(frame, trace): """Notebook display from the current run's actual values, without external JS.""" layers = frame['layers'] scale = max(abs(v) for layer in layers for v in layer['mlp']) or 1 rows = [] for i, layer in enumerate(layers): cells = [] for j, value in enumerate(layer['mlp']): amount = (abs(value) / scale) ** .5 rgb = (int(255-215*amount), int(255-155*amount), int(255-75*amount)) if value < 0 else (int(255-35*amount), int(255-155*amount), int(255-225*amount)) cells.append(f'Layer {i+1}, unit {j}: {value:.9g}') rows.append(f'
Layer {i+1}{"".join(cells)}
') m = frame['memory'] memory = f'Weights: {m["weightBytes"]:,} bytes · K/V payload: {m["kvBytes"]:,} bytes' if 'allocatedBytes' in m: memory += f'
CUDA allocated: {m["allocatedBytes"]:,} · reserved: {m["reservedBytes"]:,} bytes' else: memory += '
CPU run: GPU memory measurements unavailable.' candidates = ''.join(f'
  • {html.escape(repr(t["text"]))}: {t["probability"]:.2%}
    • ' for t in frame['top']) return f'
    LIVE {html.escape(trace["device"]["name"])} · {frame["phase"]} · {frame["position"]} tokens processed
    {html.escape(frame["context"])}

    Next selected token: {html.escape(repr(frame["nextText"]))}

    {memory}
    Hardware bytes read: not measured here. Time: {frame["elapsedMs"]:.1f} ms, including instrumentation.

    Next-token probabilities
      {candidates}

    Measured feed-forward activations after GELU. Blue is negative; orange positive. Hover for numbers. These are not physical GPU pins or semantic labels.

    {"".join(rows)}
    ' if __name__ == '__main__': parser = argparse.ArgumentParser(description=__doc__) parser.add_argument('--device', choices=['cpu', 'cuda']) parser.add_argument('--local-model', help='Optional local copy of the pinned checkpoint') parser.add_argument('--prompt', default='Once upon a time, there was a little') parser.add_argument('--tokens', type=int, default=16) parser.add_argument('--output', default='llm-observer-trace.json') parser.add_argument('--benchmark', action='store_true') parser.add_argument('--profile', action='store_true') parser.add_argument('--hardware-step', action='store_true') args = parser.parse_args() observer = Observer(device=args.device, local_model=args.local_model) if args.hardware_step: observer.hardware_step() elif args.profile: print(observer.profile(args.output)) else: result = observer.benchmark(args.prompt) if args.benchmark else observer.run(args.prompt, args.tokens) Path(args.output).write_text(json.dumps(result, allow_nan=False)) print(json.dumps({'device': observer.device_name, 'output': args.output, 'hardwareBytesMeasured': False}))