This document is relevant for: Trn2, Trn3
Tensor Compare Design#
Overview#
Tensor compare provides utilities for comparing intermediate tensors captured from different execution environments (HF CPU vs vLLM Neuron) to identify numerical divergence at the module level.
The core workflow is:
Capture (HF + Neuron) → Read → Align raw tensors → Reconstruct → Align reconstructed tensors → Compare → Report
Each step is pluggable: you provide model-specific reconstruction functions, and the framework handles the rest.
End-to-End Workflow#
from vllm_neuron.accuracy.tensor_io import read as tensor_io_read
from vllm_neuron.accuracy import (
align_decode_captures,
compare_captures_two_way,
compare_captures_three_way,
print_comparison_report,
print_three_way_report,
)
from vllm_neuron.accuracy.tensor_alignment_utils import (
align_and_truncate_hidden,
hf_reference_reconstruction,
)
# 1. Read captures from disk
fp32 = tensor_io_read("/tmp/captures/hf_fp32")
bf16 = tensor_io_read("/tmp/captures/hf_bf16")
neuron = tensor_io_read("/tmp/captures/neuron")
# 2. Align decode steps by position (HF and Neuron may differ)
fp32 = align_decode_captures(fp32, neuron)
bf16 = align_decode_captures(bf16, neuron)
# 3. Two-way comparison (FP32 vs Neuron, static threshold)
results = compare_captures_two_way(
fp32, neuron,
ref_reconstruction_fn=hf_reference_reconstruction,
test_reconstruction_fn=my_reconstruct,
alignment_fn=align_and_truncate_hidden,
)
print_comparison_report(results, label1="HF FP32", label2="Neuron")
# 4. Three-way comparison (dynamic threshold via BF16 baseline)
results = compare_captures_three_way(
fp32, bf16, neuron,
reference_reconstruction_fn=hf_reference_reconstruction,
target_reconstruction_fn=my_reconstruct,
alignment_fn=align_and_truncate_hidden,
)
print_three_way_report(results)
See examples/vllm_neuron/accuracy/compare_hf_vs_vllm_neuron_with_reconstruction.py for a complete runnable example with Llama-3.2-1B-Instruct.
Comparison Modes#
Two-Way Comparison#
Compares two tensors directly with a static threshold:
from vllm_neuron.accuracy import compare_tensors
result = compare_tensors(cpu_tensor, neuron_tensor, name="layer_0")
# result.linf_rel — relative L-infinity error
# result.l2_rel — relative L2 error
# result.max_abs — maximum absolute difference
When to use: Quick check with a known threshold. Good for:
CPU vs Neuron comparison (same dtype)
Comparing two Neuron runs for determinism
Module-level unit tests with known expected outputs
Limitation: Static thresholds don’t adapt to different inputs or model architectures. Use three-way comparison for robust validation.
Example output:
TENSOR COMPARISON: HF FP32 vs Neuron
================================================================================
Threshold: L-inf > 0.01 highlighted in red
=== prompt_0 ===
--- Step 0 ---
Module L-inf L2 Max Abs
--------------------------------------------------------------------------------
model_embed_tokens/token0 0.00e+00 0.00e+00 0.00e+00
model_layers_0_input_layernorm/token0 2.76e-03 1.95e-03 3.09e-02
model_layers_0_self_attn/token0 3.52e-03 2.14e-03 5.21e-02
model_layers_0_mlp/token0 4.18e-03 3.01e-03 7.83e-02
lm_head 3.78e-01 4.33e-01 4.86e+00
SUMMARY: 2 deviation(s) > 0.01, 0 shape mismatch(es)
Three-Way Comparison#
Compares both implementations against a common FP32 baseline:
Error₁ = |Neuron BF16 − FP32| (target error)
Error₂ = |CPU BF16 − FP32| (baseline error — expected precision loss)
Ratio = Error₁ / Error₂
Ratio ≈ 1.0 → Neuron matches expected BF16 quantization behavior ✓
Ratio >> 1.0 → Neuron introduces additional error beyond quantization ✗
from vllm_neuron.accuracy import compare_tensors_three_way
result = compare_tensors_three_way(
baseline=fp32_tensor, # Ground truth (FP32)
expected=bf16_cpu_tensor, # Reference quantized (BF16 CPU)
actual=bf16_neuron_tensor, # Target (BF16 Neuron)
name="layer_0"
)
# result.linf_ratio — ratio of L-inf errors
# result.l2_ratio — ratio of L2 errors
# result.bc — Bhattacharyya Coefficient (distribution similarity)
# result.passed — True if ratios < multiplier threshold
When to use: Robust validation without manual threshold tuning.
Default threshold: Ratio < 1.5x (configurable via DynamicThresholdConfig). For dense TP-only models (Llama), expect ratios close to 1.0 at the module level.
Example output:
THREE-WAY COMPARISON: HF BF16 & Neuron vs FP32
==========================================================================================
Ratio = Tgt/Base error. Red if ratio >= 1.5x
=== prompt_0 ===
--- Step 0 ---
Module L-inf Ratio L2 Ratio BC
------------------------------------------------------------------------------------------
model_embed_tokens/token0 0.0000 0.0000 1.0000
model_layers_0_input_layernorm/token0 1.0000 1.0000 1.0000
model_layers_0_self_attn/token0 1.0312 1.0099 0.9993
model_layers_0_mlp/token0 0.8904 1.0694 0.9990
lm_head 1.1419 1.2111 0.9842
SUMMARY: 0 failure(s) (ratio >= 1.5x), 0 shape mismatch(es)
Bhattacharyya Coefficient (BC)#
For multi-token or multi-prompt validation, per-element error ratios can be noisy. BC compares the full error distributions:
BC = 1.0: Identical error distributions
BC ≥ 0.99: Excellent (production ready)
BC < 0.95: Significant divergence
BC = 0.0: No overlap
BC is computed automatically in three-way comparison results.
Error Metrics#
L-infinity (Relative) — worst-case error, sensitive to outliers:
linf_rel = max|actual - reference| / max|reference|
L2 (Relative) — overall error magnitude, less sensitive to outliers:
l2_rel = norm(actual - reference) / norm(reference)
Reconstruction Functions#
The ReconstructionFn Protocol#
Reconstruction converts raw per-rank captured tensors into a single tensor comparable to the reference. The framework calls it as:
tensor = reconstruction_fn(rank_tensors, module_name, phase, positions)
All reconstruction functions must accept these four parameters (even if unused for a given model) so the compare framework can call them uniformly:
rank_tensors:List[Tensor]— per-rank tensors sorted by rank indexmodule_name:str— normalized module name (dots → underscores)phase:str—"prefill"or"decode"positions:List[int]— token positions from capture metadata
Dense TP-only Models (Llama)#
For models without sequence parallelism, all ranks produce identical outputs after the all-reduce in each TP layer. Reconstruction simply takes rank 0 and strips bucket padding:
def llama_reconstruct(rank_tensors, module_name, phase, positions):
"""All ranks identical after all-reduce. Use rank 0, strip padding."""
tensor = rank_tensors[0]
if positions:
real = sum(1 for i, p in enumerate(positions)
if i == 0 or p > positions[i - 1])
if tensor.shape[0] > real:
return tensor[:real]
return tensor
HF Reference#
HF runs full-sequence forward each decode step (recomputing all prior tokens). The hook captures [1, seq_len, hidden]. For decode, only the last token is the new result:
def hf_reference_reconstruction(rank_tensors, module_name, phase, positions):
"""Single rank; extract last token for decode."""
tensor = rank_tensors[0]
if phase != "prefill" and tensor.dim() >= 2:
return tensor[:, -1:, :] if tensor.dim() == 3 else tensor[-1:, :]
return tensor
Dense SP Models (GPT-OSS on trn2)#
GPT-OSS on trn2 uses sequence parallelism but outputs hidden states that are
already all-reduced by the time they are captured. The standard
llama_reconstruct function (rank 0, strip padding) works without modification.
Align Raw Tensors#
Before reconstruction, raw captures from HF and Neuron must be aligned so that each decode step on the ref side corresponds to the same generation position on the target side.
HF and Neuron capture decode tensors differently:
Source |
Shape per decode step |
What reconstruction does |
|---|---|---|
HF |
|
Extract last token |
Neuron |
|
Strip batch padding |
HF and Neuron may also produce different numbers of decode steps (HF captures every autoregressive step; Neuron may batch differently).
align_decode_captures(ref, target) matches them by token position:
Sort decode captures by
max(positions)Keep only ref steps whose position has a matching target step
Preserve prefill captures unchanged
from vllm_neuron.accuracy import align_decode_captures
fp32 = align_decode_captures(fp32, neuron)
bf16 = align_decode_captures(bf16, neuron)
Align Reconstructed Tensors#
After reconstruction, tensors from HF and Neuron may still differ in shape due to:
Batch dimension: HF produces
[1, seq, hidden](3D) while Neuron produces[seq, hidden](2D)Hidden dimension padding: Neuron pads hidden to a hardware-friendly size (e.g., 2048 → 2304)
align_and_truncate_hidden(baseline, expected, actual) normalizes these:
Promote 2D tensors to 3D (unsqueeze batch dim)
Truncate hidden dim to the minimum across all three tensors
Squeeze batch=1 back to 2D
Cast to float32
from vllm_neuron.accuracy.tensor_alignment_utils import align_and_truncate_hidden
# Pass as alignment_fn to compare functions:
results = compare_captures_three_way(
fp32, bf16, neuron,
alignment_fn=align_and_truncate_hidden,
...
)
The alignment_fn is called internally by the compare framework after reconstruction and before computing metrics. It returns (t1, t2, t3, shapes_match) where shapes_match is False if the tensors could not be aligned (e.g., sequence dimension mismatch).
API Reference#
# --- Tensor-level comparison ---
compare_tensors(t1, t2, name) -> ComparisonResult
compare_tensors_three_way(baseline, expected, actual, name) -> ThreeWayComparisonResult
# --- Capture-level comparison (operates on List[CapturedForwardPass]) ---
compare_captures_two_way(
ref, test,
ref_reconstruction_fn=None,
test_reconstruction_fn=None,
alignment_fn=None, # defaults to align_and_truncate_hidden
module_order=None,
phase=None, # "prefill", "decode", or None for all
) -> Dict[prompt, Dict[step, List[ComparisonResult]]]
compare_captures_three_way(
baseline, expected, actual,
reference_reconstruction_fn=None,
target_reconstruction_fn=None,
alignment_fn=None,
module_order=None,
phase=None,
) -> Dict[prompt, Dict[step, List[ThreeWayComparisonResult]]]
# --- Alignment utilities ---
align_decode_captures(ref, target) -> List[CapturedForwardPass]
align_and_truncate_hidden(t1, t2, t3) -> (t1, t2, t3, shapes_match)
# --- Reporting ---
print_comparison_report(results, threshold=0.01, label1="Ref", label2="Test")
print_three_way_report(results, label_expected="BF16", label_actual="Neuron")
Example Scripts#
examples/vllm_neuron/accuracy/compare_hf_vs_vllm_neuron_with_reconstruction.py— Full pipeline with Llama-3.2-1B-Instructexamples/vllm_neuron/accuracy/compare_hf_vs_vllm_neuron.py— Simpler rank-0 comparison (no reconstruction)
Limitations#
Same module names required: Comparison matches tensors by module name. If HF and Neuron use different module hierarchies, the reconstruction function must normalize them.
Wildcard module capture causes dynamo recompile: Using
modules: [".*"]intensor_captureconfig triggers atorch.compilerecompile guard. Use explicit module patterns instead:# Works modules = ["model.layers.0-15.input_layernorm", "model.layers.0-15.self_attn", "lm_head"] # Fails modules = [".*"]
GPT-OSS on trn3 requires custom reconstruction: trn3 shuffles the hidden dimension in attention/MLP kernels. Captured tensors must be unshuffled before comparison. Pass a custom
reconstruction_fntoTensorComparePluginthat calls_unshuffle_hidden_dimafter gathering rank tensors. GPT-OSS on trn2 works with the default reconstruction (no unshuffle needed).
Appendix: Statistical Background#
Dynamic Thresholds#
Three-way comparison uses dynamic thresholds based on the baseline error:
pass if: |actual - baseline| < multiplier × |expected - baseline|
The default multiplier is 1.5x. Why > 1.0?
Even when comparing the same precision (BF16 on both platforms), Neuron and CPU may not be numerically identical due to different accumulation orders in fused kernels. For dense TP-only models (Llama) at the module level, the ratio is typically ~1.0. Larger ratios (up to 1.5x) can appear at full-model scale due to error accumulation across layers.
Bhattacharyya Coefficient Details#
The BC measures overlap between two probability distributions. For tensor comparison, we bin the per-element errors into histograms and compute:
BC = Σ √(p_i × q_i)
where p_i and q_i are the normalized histogram bin values for the baseline and target error distributions respectively.
Why BC over simpler metrics?
Eliminates manual threshold tuning: BC ≥ 0.99 works universally across models
Captures full error behavior: Compares entire distributions, not just worst-case
Higher sensitivity: Detects subtle systematic shifts that max-error misses
Robust to outliers: A few high-error elements don’t dominate
For multi-prompt validation, use compute_aggregate_metrics() to pool errors across prompts and compute aggregate BC.
This document is relevant for: Trn2, Trn3