Spaces:
Running
Running
| """Compact measurements of real decoder outputs; never synthetic activity.""" | |
| from contextlib import contextmanager | |
| from time import perf_counter | |
| import torch | |
| VERSION = 1 | |
| MAX_TOKENS = 32 | |
| MAX_NEURONS = 32 | |
| MAX_LOGITS = 64 | |
| METRIC = "hidden_state_rms" | |
| def architecture(model): | |
| config = getattr(model, "config", None) | |
| model_id = getattr(config, "_name_or_path", "") | |
| names = { | |
| "dllm-hub/Qwen3-0.6B-diffusion-bd3lm-v0.1": ("Qwen3 · 0.6B", "bd3lm"), | |
| "Efficient-Large-Model/Fast_dLLM_v2_1.5B": ("Fast-dLLM · 1.5B", "fast-dllm-v2"), | |
| } | |
| if model_id not in names: | |
| return None | |
| fields = { | |
| "layers": "num_hidden_layers", "hiddenSize": "hidden_size", | |
| "mlpWidth": "intermediate_size", "queryHeads": "num_attention_heads", | |
| "kvHeads": "num_key_value_heads", | |
| } | |
| shape = {name: getattr(config, field, None) for name, field in fields.items()} | |
| if not all(type(value) is int and value > 0 for value in shape.values()): | |
| return None | |
| head_dim = getattr(config, "head_dim", None) or shape["hiddenSize"] // shape["queryHeads"] | |
| label, family = names[model_id] | |
| return { | |
| "modelId": model_id, "modelLabel": label, | |
| "modelUrl": f"https://huggingface.co/{model_id}", "modelFamily": family, | |
| **shape, "headDim": head_dim, | |
| } | |
| def decoder_layers(model): | |
| """The deployed A2DQwen3 model does not implement output_hidden_states.""" | |
| return getattr(getattr(model, "model", None), "layers", ()) | |
| def sampled_indices(size, limit): | |
| count = min(size, limit) | |
| return ( | |
| [round(index * (size - 1) / (count - 1)) for index in range(count)] | |
| if count > 1 else [0] | |
| ) | |
| def internals_layout(layers): | |
| """Use actual projection widths, rather than assuming a checkpoint shape.""" | |
| layouts = [] | |
| for layer in layers: | |
| attention = getattr(layer, "self_attn", None) | |
| output = getattr(attention, "o_proj", None) | |
| down = getattr(getattr(layer, "mlp", None), "down_proj", None) | |
| head_dim = getattr(attention, "head_dim", None) | |
| width = getattr(output, "in_features", None) | |
| neurons = getattr(down, "in_features", None) | |
| if ( | |
| not all(isinstance(value, int) and value > 0 for value in (head_dim, width, neurons)) | |
| or width % head_dim | |
| or not callable(getattr(output, "register_forward_pre_hook", None)) | |
| or not callable(getattr(down, "register_forward_pre_hook", None)) | |
| ): | |
| return None | |
| layouts.append((width // head_dim, head_dim, neurons)) | |
| return layouts[0] if layouts and all(layout == layouts[0] for layout in layouts) else None | |
| def qkv_layout(layers, layout): | |
| """Q/K/V projections are optional; never infer their widths from a label.""" | |
| if layout is None: | |
| return None | |
| head_count, head_dim, _ = layout | |
| layouts = [] | |
| for layer in layers: | |
| attention = getattr(layer, "self_attn", None) | |
| projections = [getattr(attention, name, None) for name in ("q_proj", "k_proj", "v_proj")] | |
| widths = [getattr(projection, "out_features", None) for projection in projections] | |
| if ( | |
| not all(type(width) is int and width > 0 and width % head_dim == 0 for width in widths) | |
| or widths[0] != head_count * head_dim or widths[1] != widths[2] | |
| or widths[0] % widths[1] | |
| or not all(callable(getattr(projection, "register_forward_hook", None)) for projection in projections) | |
| ): | |
| return None | |
| layouts.append(widths[1] // head_dim) | |
| return layouts[0] if layouts and all(value == layouts[0] for value in layouts) else None | |
| def stage_modules(model): | |
| decoder = getattr(model, "model", None) | |
| candidates = { | |
| "embeddingRms": getattr(decoder, "embed_tokens", None), | |
| "finalNormRms": getattr(decoder, "norm", None), | |
| "lmHeadRms": getattr(model, "lm_head", None), | |
| } | |
| # The output width fixes a deterministic bounded vocabulary sample before | |
| # inference. Models with unconventional output heads simply omit this stage. | |
| output_width = getattr(candidates["lmHeadRms"], "out_features", None) | |
| if type(output_width) is not int or output_width < 1: | |
| candidates.pop("lmHeadRms") | |
| return { | |
| name: module for name, module in candidates.items() | |
| if callable(getattr(module, "register_forward_hook", None)) | |
| } | |
| def capability(model): | |
| layers = decoder_layers(model) | |
| layout = internals_layout(layers) | |
| return { | |
| "available": bool(len(layers)), | |
| "version": VERSION, | |
| "metric": METRIC, | |
| "layer_count": len(layers), | |
| "max_tokens": MAX_TOKENS, | |
| "token_states": True, | |
| "internals": layout is not None, | |
| "qkv": qkv_layout(layers, layout) is not None, | |
| "stages": list(stage_modules(model)), | |
| "architecture": architecture(model), | |
| } | |
| def capture_token_state( | |
| tokenizer, response_tokens, block_index, block_start, block_size, | |
| mask_before, finished, piece_cache, | |
| ): | |
| """Describe actual response tokens after this forward's commit and EOS trim. | |
| RMS hooks observe the inputs to this update. The token snapshot observes its | |
| result. Only masks that became retained non-mask tokens count as new commits; | |
| padding a discarded suffix is not a prediction. A previously committed token | |
| keeps its identity unless the sampler discards it after a settled EOS. | |
| """ | |
| mask_id = tokenizer.mask_token_id | |
| eos_id = tokenizer.eos_token_id | |
| special_ids = set(tokenizer.all_special_ids) | |
| first_eos = ( | |
| response_tokens.index(eos_id) | |
| if finished and eos_id is not None and eos_id in response_tokens else None | |
| ) | |
| tokens = [] | |
| for position, token_id in enumerate(response_tokens): | |
| if first_eos is not None and position > first_eos: | |
| state = "discarded" | |
| elif token_id == mask_id: | |
| state = "masked" | |
| else: | |
| state = "committed" | |
| local_position = position - block_start | |
| newly_committed = ( | |
| state == "committed" | |
| and 0 <= local_position < block_size | |
| and mask_before[local_position] | |
| ) | |
| if token_id not in piece_cache: | |
| piece_cache[token_id] = tokenizer.decode( | |
| [token_id], skip_special_tokens=False, | |
| clean_up_tokenization_spaces=False, | |
| ) | |
| tokens.append({ | |
| "position": position, | |
| "tokenId": token_id, | |
| "piece": piece_cache[token_id], | |
| "state": state, | |
| "newlyCommitted": newly_committed, | |
| "special": token_id in special_ids, | |
| }) | |
| return { | |
| "version": VERSION, | |
| "blockIndex": block_index, | |
| "blockStart": block_start, | |
| "blockSize": block_size, | |
| "tokens": tokens, | |
| } | |
| def capture_forward(model, token_offset, token_count, query_indices=None, query_count=None): | |
| """Register hooks only around one conditional denoising forward pass. | |
| The caller holds the model inference lock. We reduce each decoder's residual | |
| output [batch, tokens, hidden] immediately, keeping no full hidden tensors. | |
| Layer indices are zero-based; token positions are relative to the response. | |
| Optional query indices map response positions to actual tensor rows, so a | |
| first block containing prompt tokens can exclude those rows completely. | |
| """ | |
| layers = decoder_layers(model) | |
| if not len(layers): | |
| raise ValueError("This model does not expose decoder layers for telemetry") | |
| if token_count < 1: | |
| raise ValueError("Telemetry needs an active token block") | |
| query_count = token_count if query_count is None else query_count | |
| query_indices = list(range(token_count)) if query_indices is None else list(query_indices) | |
| if ( | |
| type(query_count) is not int or query_count < 1 | |
| or len(query_indices) != token_count | |
| or any(type(index) is not int or not 0 <= index < query_count for index in query_indices) | |
| or any(left >= right for left, right in zip(query_indices, query_indices[1:])) | |
| ): | |
| raise ValueError("Invalid response-to-query mapping for telemetry") | |
| sample_indices = sampled_indices(token_count, MAX_TOKENS) | |
| sampled_queries = [query_indices[index] for index in sample_indices] | |
| frame = { | |
| "version": VERSION, | |
| "metric": METRIC, | |
| "normalization": "rms/(1+rms)", | |
| "activations": [], | |
| "rawRms": [], | |
| "layerIndices": list(range(len(layers))), | |
| "tokenPositions": [token_offset + index for index in sample_indices], | |
| "architecture": architecture(model), | |
| "phase": "denoise", | |
| } | |
| measurements = {} | |
| attention_measurements = {} | |
| neuron_measurements = {} | |
| qkv_measurements = {name: {} for name in ("queryHeadRms", "keyHeadRms", "valueHeadRms")} | |
| stages = stage_modules(model) | |
| stage_measurements = {} | |
| logit_indices = ( | |
| sampled_indices(stages["lmHeadRms"].out_features, MAX_LOGITS) | |
| if "lmHeadRms" in stages else None | |
| ) | |
| layout = internals_layout(layers) | |
| kv_head_count = qkv_layout(layers, layout) | |
| if layout is not None: | |
| head_count, head_dim, neuron_count = layout | |
| neuron_indices = sampled_indices(neuron_count, MAX_NEURONS) | |
| handles = [] | |
| def make_hook(layer_index): | |
| def measure(_module, _inputs, output): | |
| hidden = output[0] if isinstance(output, (tuple, list)) else output | |
| if not isinstance(hidden, torch.Tensor) or hidden.ndim != 3 or hidden.shape[1] != query_count: | |
| raise ValueError("Unexpected decoder output shape for telemetry") | |
| values = hidden.detach()[0, sampled_queries, :].float() | |
| rms = values.square().mean(dim=-1).sqrt() | |
| if not bool(torch.isfinite(rms).all()): | |
| raise ValueError("Non-finite decoder activations") | |
| measurements[layer_index] = ( | |
| rms.cpu().tolist(), (rms / (1.0 + rms)).cpu().tolist() | |
| ) | |
| return measure | |
| def make_stage_hook(name): | |
| def measure(_module, _inputs, output): | |
| if ( | |
| not isinstance(output, torch.Tensor) or output.ndim != 3 | |
| or output.shape[0] < 1 or output.shape[1] != query_count or output.shape[2] < 1 | |
| ): | |
| raise ValueError("Unexpected architecture stage output shape for telemetry") | |
| output = output.detach() | |
| if name == "lmHeadRms": | |
| if output.shape[2] != stages[name].out_features: | |
| raise ValueError("Unexpected language-model head width for telemetry") | |
| # Index rows and vocabulary columns together. Selecting rows | |
| # first would copy the full 151,936-wide logits unnecessarily. | |
| rows = torch.tensor(sampled_queries, device=output.device).unsqueeze(1) | |
| columns = torch.tensor(logit_indices, device=output.device) | |
| values = output[0, rows, columns].float() | |
| else: | |
| values = output[0, sampled_queries, :].float() | |
| rms = values.square().mean(dim=-1).sqrt() | |
| if not bool(torch.isfinite(rms).all()): | |
| raise ValueError("Non-finite architecture stage activations") | |
| stage_measurements[name] = rms.cpu().tolist() | |
| return measure | |
| def make_qkv_hook(layer_index, name, count): | |
| def measure(_module, _inputs, output): | |
| if ( | |
| not isinstance(output, torch.Tensor) or output.ndim != 3 | |
| or output.shape[0] < 1 or output.shape[1] != query_count | |
| or output.shape[2] != count * head_dim | |
| ): | |
| raise ValueError("Unexpected Q/K/V projection output shape for telemetry") | |
| # These are projection outputs before Q/K norm, RoPE, KV caching, | |
| # head repetition, or attention. Only current response rows count. | |
| values = output.detach()[0, sampled_queries, :].float() | |
| values = values.reshape(len(sample_indices), count, head_dim) | |
| rms = values.square().mean(dim=(0, 2)).sqrt() | |
| if not bool(torch.isfinite(rms).all()): | |
| raise ValueError("Non-finite Q/K/V activations") | |
| qkv_measurements[name][layer_index] = rms.cpu().tolist() | |
| return measure | |
| def make_projection_hook(layer_index, attention): | |
| def measure(_module, inputs): | |
| values = inputs[0] if inputs else None | |
| width = head_count * head_dim if attention else neuron_count | |
| if ( | |
| not isinstance(values, torch.Tensor) or values.ndim != 3 | |
| or values.shape[1] != query_count or values.shape[2] != width | |
| ): | |
| raise ValueError("Unexpected projection input shape for telemetry") | |
| values = values.detach()[0, sampled_queries, :] | |
| if attention: | |
| # o_proj receives concatenated, attention-weighted value heads. | |
| # These are head outputs, not query/key scores or attention weights. | |
| values = values.float().reshape(len(sample_indices), head_count, head_dim) | |
| rms = values.square().mean(dim=(0, 2)).sqrt() | |
| target = attention_measurements | |
| else: | |
| # down_proj receives actual SwiGLU channel outputs. Preserve the | |
| # sampled channel identities and aggregate only over active tokens. | |
| values = values[:, neuron_indices].float() | |
| rms = values.square().mean(dim=0).sqrt() | |
| target = neuron_measurements | |
| if not bool(torch.isfinite(rms).all()): | |
| raise ValueError("Non-finite internal activations") | |
| target[layer_index] = rms.cpu().tolist() | |
| return measure | |
| started = perf_counter() | |
| try: | |
| for name, module in stages.items(): | |
| handles.append(module.register_forward_hook(make_stage_hook(name))) | |
| for index, layer in enumerate(layers): | |
| handles.append(layer.register_forward_hook(make_hook(index))) | |
| if layout is not None: | |
| handles.append(layer.self_attn.o_proj.register_forward_pre_hook(make_projection_hook(index, True))) | |
| handles.append(layer.mlp.down_proj.register_forward_pre_hook(make_projection_hook(index, False))) | |
| if kv_head_count is not None: | |
| for projection, name, count in ( | |
| ("q_proj", "queryHeadRms", head_count), | |
| ("k_proj", "keyHeadRms", kv_head_count), | |
| ("v_proj", "valueHeadRms", kv_head_count), | |
| ): | |
| handles.append(getattr(layer.self_attn, projection).register_forward_hook(make_qkv_hook(index, name, count))) | |
| yield frame | |
| if len(measurements) != len(layers): | |
| raise ValueError("Incomplete decoder telemetry") | |
| if len(stage_measurements) != len(stages): | |
| raise ValueError("Incomplete architecture stage telemetry") | |
| result = { | |
| "rawRms": [measurements[index][0] for index in range(len(layers))], | |
| "activations": [measurements[index][1] for index in range(len(layers))], | |
| } | |
| if layout is not None: | |
| if len(attention_measurements) != len(layers) or len(neuron_measurements) != len(layers): | |
| raise ValueError("Incomplete internal telemetry") | |
| result["internals"] = { | |
| "version": VERSION, | |
| "headCount": head_count, | |
| "headDim": head_dim, | |
| "neuronIndices": neuron_indices, | |
| "attentionHeadRms": [attention_measurements[index] for index in range(len(layers))], | |
| "mlpNeuronRms": [neuron_measurements[index] for index in range(len(layers))], | |
| } | |
| if kv_head_count is not None: | |
| if any(len(values) != len(layers) for values in qkv_measurements.values()): | |
| raise ValueError("Incomplete Q/K/V telemetry") | |
| result["internals"].update({ | |
| "kvHeadCount": kv_head_count, | |
| **{name: [values[index] for index in range(len(layers))] for name, values in qkv_measurements.items()}, | |
| }) | |
| if stages: | |
| result["stages"] = stage_measurements | |
| if logit_indices is not None: | |
| result["stages"]["lmHeadTokenIndices"] = logit_indices | |
| result["forwardMs"] = round((perf_counter() - started) * 1000, 2) | |
| # Publish only after every installed hook has supplied a valid result. | |
| frame.update(result) | |
| finally: | |
| for handle in handles: | |
| handle.remove() | |