From ab04f6c52ff53fe3ab5505bf32c73872aef5095f Mon Sep 17 00:00:00 2001 From: rohith500 Date: Thu, 24 Sep 2026 09:17:38 -0400 Subject: [PATCH] feat(casting): add dynamic activation range calibration in FP16 casting (#7) --- docs/src/utils/casting.md | 2 +- src/coreai_opt/_utils/casting_utils.py | 165 +++++++++++++++++ src/coreai_opt/casting/casting.py | 32 +++- tests/casting/test_casting.py | 242 ++++++++++++++++++++++++- tests/casting/test_casting_utils.py | 176 ++++++++++++++++++ 5 files changed, 608 insertions(+), 9 deletions(-) diff --git a/docs/src/utils/casting.md b/docs/src/utils/casting.md index 124dbe9..1e9fc80 100644 --- a/docs/src/utils/casting.md +++ b/docs/src/utils/casting.md @@ -29,7 +29,7 @@ These passes mutate the `ExportedProgram` in place and may change the dtypes of The two casting passes apply different selection criteria, resulting in different levels of aggressiveness. -**FP32 → FP16.** The FP pass is relatively aggressive: it attempts to cast all floating-point state and operations to FP16 by default. The main exception is numerical safety — tensors whose values exceed the FP16 representable range (approximately ±65504) are left in FP32 to avoid overflow. Users can also provide `ignored_ops` (as a set/collection of `OpOverload` or `OpOverloadPacket` instances such as `{torch.ops.aten.exp}`, or as a predicate function `Callable[[torch.fx.Node], bool]`) to `cast_fp32_to_fp16` or `cast_to_16_bit_precision` to explicitly preserve specific operations in FP32 to avoid intermediate activation overflow. +**FP32 → FP16.** The FP pass is relatively aggressive: it attempts to cast all floating-point state and operations to FP16 by default. The main exception is numerical safety — tensors whose values exceed the FP16 representable range (approximately ±65504) are left in FP32 to avoid overflow. Users can supply `calibration_data` (an iterable yielding sample inputs, such as `[tensor]`, `[(arg1, arg2)]`, or a `DataLoader`) to `cast_fp32_to_fp16` or `cast_to_16_bit_precision` to automatically observe intermediate activation ranges and preserve operations whose activations exceed FP16 range in FP32. Users can also provide `ignored_ops` (as a set/collection of `OpOverload` or `OpOverloadPacket` instances such as `{torch.ops.aten.exp}`, or as a predicate function `Callable[[torch.fx.Node], bool]`) to explicitly preserve specific operations in FP32. **INT32/INT64 → INT16.** The INT pass is more conservative. Rather than casting all integer tensors, it targets only those that are likely to benefit without risking correctness. A tensor is skipped if any of the following apply: diff --git a/src/coreai_opt/_utils/casting_utils.py b/src/coreai_opt/_utils/casting_utils.py index 559202f..dfa32f4 100644 --- a/src/coreai_opt/_utils/casting_utils.py +++ b/src/coreai_opt/_utils/casting_utils.py @@ -9,10 +9,14 @@ passes in ``casting``. """ +import operator from collections.abc import Callable, Collection, Iterable +from typing import Any import numpy as np import torch +from torch.export.graph_signature import InputKind +from torch.fx import Interpreter # ============================================================================= # Constants @@ -870,3 +874,164 @@ def anchor_int16_reshape_inputs(graph: torch.fx.Graph, pass_inserted: set[torch. count += 1 return count + + +# ============================================================================= +# Dynamic Activation Range Calibration +# ============================================================================= +CalibrationSample = torch.Tensor | tuple[Any, ...] | list[Any] | dict[str, Any] +CalibrationDataType = Iterable[CalibrationSample] | None + + +def _iter_calibration_batches( + calibration_data: Any, + user_names: list[str], +) -> Iterable[list[Any]]: + """Yield flattened user input lists for each sample in calibration_data.""" + if isinstance(calibration_data, (torch.Tensor, dict, str, bytes)): + raise TypeError( + "calibration_data must be an iterable of samples (e.g. a list or DataLoader), " + f"not a bare {type(calibration_data).__name__}." + ) + if not isinstance(calibration_data, Iterable): + raise TypeError( + "calibration_data must be an iterable of samples, " + f"got {type(calibration_data).__name__}." + ) + + num_inputs = len(user_names) + + def _to_inputs(sample: Any) -> list[Any]: + if isinstance(sample, dict): + missing = [name for name in user_names if name not in sample] + if missing: + raise ValueError( + f"Calibration dictionary missing required input(s): {missing}. " + f"Expected inputs: {user_names}." + ) + return [sample[name] for name in user_names] + if isinstance(sample, (tuple, list)): + if len(sample) != num_inputs: + raise ValueError( + f"Calibration sample provided {len(sample)} input(s), " + f"but model expects {num_inputs} input(s): {user_names}." + ) + return list(sample) + if isinstance(sample, torch.Tensor): + if num_inputs != 1: + raise ValueError( + f"Calibration sample provided a single Tensor, " + f"but model expects {num_inputs} input(s): {user_names}." + ) + return [sample] + raise TypeError( + f"Unsupported calibration sample type: {type(sample).__name__}. " + "Expected torch.Tensor, tuple, list, or dict." + ) + + for batch in calibration_data: + yield _to_inputs(batch) + + +class _OverflowDetector(Interpreter): + """Interpreter that tracks call_function nodes producing out-of-range FP16 values.""" + + def __init__(self, ep: torch.export.ExportedProgram) -> None: + super().__init__(ep.graph_module) + self.direct_overflow_nodes: set[torch.fx.Node] = set() + + def run_node(self, n: torch.fx.Node) -> Any: + result = super().run_node(n) + if n.op == "call_function": + self._check(n, result) + return result + + def _check(self, n: torch.fx.Node, val: Any) -> None: + if isinstance(val, torch.Tensor): + if check_tensor_overflow_fp16(val): + self.direct_overflow_nodes.add(n) + elif isinstance(val, float): + abs_val = abs(val) + if abs_val > _FP16_MAX and abs_val < _FP32_INF_THRESHOLD: + self.direct_overflow_nodes.add(n) + elif isinstance(val, (tuple, list)): + for item in val: + self._check(n, item) + + +def find_overflowing_nodes( + exported_program: torch.export.ExportedProgram, + calibration_data: CalibrationDataType, +) -> set[torch.fx.Node]: + """Identify nodes that produce or consume out-of-range FP16 values during calibration. + + Runs sample inputs through the exported graph using an interpreter, checks + intermediate activations against FP16 range, and cascades overflow status + to direct consumers whose output is FP32. + + Args: + exported_program: The exported program to analyze. + calibration_data: An iterable of sample inputs (each sample a Tensor, tuple, + list, or dict) or None. + + Returns: + Set of nodes identified as producing or directly consuming overflowing activations. + """ + if calibration_data is None: + return set() + + detector = _OverflowDetector(exported_program) + specs = exported_program.graph_signature.input_specs + user_names = [s.arg.name for s in specs if s.kind == InputKind.USER_INPUT] + user_indices = [i for i, s in enumerate(specs) if s.kind == InputKind.USER_INPUT] + + base_args = [ + exported_program.state_dict.get(s.target, exported_program.constants.get(s.target)) + if s.kind != InputKind.USER_INPUT + else None + for s in specs + ] + + first_param = next( + ( + t + for store in (exported_program.state_dict, exported_program.constants) + for t in store.values() + if isinstance(t, torch.Tensor) + ), + None, + ) + model_device = first_param.device if first_param is not None else None + + with torch.no_grad(): + for user_inputs in _iter_calibration_batches(calibration_data, user_names): + graph_args = list(base_args) + for idx, user_val in zip(user_indices, user_inputs, strict=True): + if ( + model_device is not None + and isinstance(user_val, torch.Tensor) + and user_val.device != model_device + ): + user_val = user_val.to(model_device) + graph_args[idx] = user_val + detector.run(*graph_args) + detector.env.clear() + + all_overflow_nodes = set(detector.direct_overflow_nodes) + queue = list(detector.direct_overflow_nodes) + + while queue: + curr = queue.pop() + for user in curr.users: + if user.op == "call_function" and user not in all_overflow_nodes: + val = user.meta.get("val") + is_fp32 = (hasattr(val, "dtype") and val.dtype == torch.float32) or ( + isinstance(val, (tuple, list)) + and any(hasattr(v, "dtype") and v.dtype == torch.float32 for v in val) + ) + if is_fp32: + all_overflow_nodes.add(user) + if user.target == operator.getitem: + queue.append(user) + + return all_overflow_nodes diff --git a/src/coreai_opt/casting/casting.py b/src/coreai_opt/casting/casting.py index 5ed4010..13b2bac 100644 --- a/src/coreai_opt/casting/casting.py +++ b/src/coreai_opt/casting/casting.py @@ -35,6 +35,7 @@ INT16_MIN as _INT16_MIN, INT_DTYPES as _INT_DTYPES, INT_DTYPES_EXTENDED as _INT_DTYPES_EXTENDED, + CalibrationDataType as _CalibrationDataType, anchor_int16_reshape_inputs as _anchor_int16_reshape_inputs, build_castable_int16_nodes as _build_castable_int16_nodes, build_unsafe_to_cast_nodes as _build_unsafe_to_cast_nodes, @@ -44,6 +45,7 @@ check_tensor_overflow_fp16 as _check_tensor_overflow_fp16, classify_float_args as _classify_float_args, cleanup_casts as _cleanup_casts, + find_overflowing_nodes as _find_overflowing_nodes, get_node_dtype as _get_node_dtype, get_placeholder_store_and_tensor as _get_placeholder_store_and_tensor, get_to_op_dtype as _get_to_op_dtype, @@ -209,6 +211,7 @@ def __init__( | Callable[[torch.fx.Node], bool] | None ) = None, + calibration_data: _CalibrationDataType = None, ) -> None: super().__init__(exported_program) if isinstance(ignored_ops, (torch._ops.OpOverload, torch._ops.OpOverloadPacket)): @@ -220,6 +223,11 @@ def __init__( else: self._ignored_ops = ignored_ops + if calibration_data is not None: + self._overflowing_nodes = _find_overflowing_nodes(exported_program, calibration_data) + else: + self._overflowing_nodes = set() + # ------------------------------------------------------------------------- # Step 1: Convert parameters # ------------------------------------------------------------------------- @@ -328,6 +336,11 @@ def iterate_nodes_and_insert_casts(self) -> None: _maybe_update_assert_dtype(node) continue + # --- ignored ops or dynamic activation overflow: keep in FP32 --- + if node in self._overflowing_nodes or _is_ignored_op(node, self._ignored_ops): + self.handle_overflow_op(node) + continue + # --- cast ops: retarget fp32 output to fp16 --- if node.target in _CAST_OPS: target_dtype = _get_to_op_dtype(node) @@ -337,10 +350,6 @@ def iterate_nodes_and_insert_casts(self) -> None: self._pass_inserted.add(node) continue - # --- ignored ops: keep in FP32 --- - if _is_ignored_op(node, self._ignored_ops): - self.handle_overflow_op(node) - continue # --- creation ops: set dtype directly --- if node.target in _CREATION_OPS: @@ -642,6 +651,7 @@ def cast_fp32_to_fp16( | Callable[[torch.fx.Node], bool] | None ) = None, + calibration_data: _CalibrationDataType = None, ) -> torch.export.ExportedProgram: """Convert a torch exported program from FP32 to FP16 where applicable. @@ -656,11 +666,18 @@ def cast_fp32_to_fp16( a single op instance, or a predicate callable taking a ``torch.fx.Node`` and returning a boolean (e.g. ``lambda node: node.name == "exp_1"``). Ignored ops are kept in FP32 with boundary casts inserted. + calibration_data: Optional iterable yielding sample inputs (each sample a Tensor, + tuple/list of Tensors, or dict of kwargs) for dynamic activation range observation. + If provided, operations whose intermediate activations overflow the FP16 representable + range (> 65504) are automatically identified and excluded from FP16 casting, + remaining in FP32 precision with boundary casts. Returns: The modified exported program. """ - return _FP16Casting(exported_program, ignored_ops=ignored_ops)() + return _FP16Casting( + exported_program, ignored_ops=ignored_ops, calibration_data=calibration_data + )() def cast_int32_to_int16( @@ -683,6 +700,7 @@ def cast_to_16_bit_precision( | Callable[[torch.fx.Node], bool] | None ) = None, + calibration_data: _CalibrationDataType = None, ) -> torch.export.ExportedProgram: """Convert a torch exported program to 16-bit precision: FP32→FP16 and INT32/64→INT16. @@ -694,10 +712,12 @@ def cast_to_16_bit_precision( exported_program: Exported program to convert. ignored_ops: Optional operations or predicate to exclude from FP16 casting. Forwarded to :func:`cast_fp32_to_fp16`. + calibration_data: Optional iterable yielding sample inputs for dynamic activation + range observation. Forwarded to :func:`cast_fp32_to_fp16`. Returns: The modified exported program. """ - cast_fp32_to_fp16(exported_program, ignored_ops=ignored_ops) + cast_fp32_to_fp16(exported_program, ignored_ops=ignored_ops, calibration_data=calibration_data) cast_int32_to_int16(exported_program) return exported_program diff --git a/tests/casting/test_casting.py b/tests/casting/test_casting.py index 0f78f7c..6dc18d0 100644 --- a/tests/casting/test_casting.py +++ b/tests/casting/test_casting.py @@ -425,7 +425,7 @@ def test_cast_to_16_bit(self, fp_full_model_and_input): # ============================================================================= -# Selective Op Skipping tests (Issue #7) +# Selective Op Skipping tests # ============================================================================= class _ActivationOverflowModel(nn.Module): """Model that triggers FP16 activation overflow in intermediate ops. @@ -457,7 +457,7 @@ def forward(self, x): class TestSelectiveOpSkipping: - """Tests for selective op skipping in FP16 casting (Issue #7).""" + """Tests for selective op skipping in FP16 casting.""" def test_ignored_op_prevents_activation_overflow(self): """Ignoring exp and log1p keeps intermediate ops in FP32, preventing inf.""" @@ -574,3 +574,241 @@ def test_cast_to_16_bit_precision_forwards_ignored_ops(self): out = _run_ep(ep, x.half()) assert not torch.isinf(out).any() assert torch.allclose(out.float(), x, atol=1e-3) + + +# ============================================================================= +# Dynamic Activation Range Calibration tests +# ============================================================================= +class TestDynamicActivationCalibration: + """End-to-end tests for dynamic activation range calibration in FP16 casting.""" + + def test_dynamic_calibration_prevents_activation_overflow(self): + """Supplying calibration data automatically prevents FP16 overflow.""" + model = _ActivationOverflowModel().eval() + x = torch.tensor([[15.0]]) + ref = model(x) + + ep = _export(model, x) + cast_fp32_to_fp16(ep, calibration_data=[x]) + out = _run_ep(ep, x.half()) + + assert not torch.isinf(out).any(), "Expected finite output with calibration_data" + assert torch.allclose(out.float(), ref, atol=1e-3) + + # Inspect graph: overflowing ops remain in FP32 with boundary casts + exp_node = next( + n + for n in ep.graph.nodes + if n.op == "call_function" and n.target == torch.ops.aten.exp.default + ) + log1p_node = next( + n + for n in ep.graph.nodes + if n.op == "call_function" and n.target == torch.ops.aten.log1p.default + ) + assert exp_node.meta["val"].dtype == torch.float32 + assert log1p_node.meta["val"].dtype == torch.float32 + + # Check cast to FP32 before exp + cast_to_fp32 = exp_node.args[0] + assert cast_to_fp32.target == torch.ops.aten._to_copy.default + assert cast_to_fp32.kwargs.get("dtype") == torch.float32 + + # Check cast back to FP16 after log1p before graph output + output_node = next(n for n in ep.graph.nodes if n.op == "output") + cast_to_fp16 = output_node.args[0][0] + assert cast_to_fp16.target == torch.ops.aten._to_copy.default + assert cast_to_fp16.kwargs.get("dtype") == torch.float16 + + def test_dynamic_calibration_dataloader(self): + """PyTorch DataLoader yielding sample tuples works seamlessly.""" + model = _ActivationOverflowModel().eval() + x = torch.tensor([[15.0]]) + ref = model(x) + + dataset = [(x,), (x,)] + loader = torch.utils.data.DataLoader(dataset, batch_size=1) + + ep = _export(model, x) + cast_fp32_to_fp16(ep, calibration_data=loader) + out = _run_ep(ep, x.half()) + + assert not torch.isinf(out).any() + assert torch.allclose(out.float(), ref, atol=1e-3) + + exp_node = next( + n + for n in ep.graph.nodes + if n.op == "call_function" and n.target == torch.ops.aten.exp.default + ) + assert exp_node.meta["val"].dtype == torch.float32 + + def test_dynamic_calibration_no_overflow_model(self): + """Model without overflow retains >= 95% FP16 ops and expected weight reduction.""" + model = SimpleModel().eval() + x = torch.randn(1, 1, 28, 28) + ep = _export(model, x) + before_bytes = _total_param_bytes(ep) + + cast_fp32_to_fp16(ep, calibration_data=[x]) + out = _run_ep(ep, x.half()) + ref = _run_ep(_export(model, x), x) + + snr = _snr_db(ref, out) + assert snr > 30, f"SNR too low: {snr:.1f} dB" + + ratio = _fp16_ratio(_count_op_dtypes(ep)) + assert ratio >= 0.95, f"Expected >= 95% FP16 ops, got {ratio:.0%}" + + reduction = 1 - _total_param_bytes(ep) / before_bytes + assert reduction >= 0.45, f"Weight reduction only {reduction:.0%}" + + def test_dynamic_calibration_with_ignored_ops(self): + """Passing both calibration_data and ignored_ops respects both exclusions.""" + model = _ActivationOverflowModel().eval() + x = torch.tensor([[1.0]]) # x=1 does not overflow exp + ep = _export(model, x) + + # Explicitly ignore exp via ignored_ops even though calibration doesn't overflow + cast_fp32_to_fp16(ep, ignored_ops={torch.ops.aten.exp}, calibration_data=[x]) + + exp_node = next( + n + for n in ep.graph.nodes + if n.op == "call_function" and n.target == torch.ops.aten.exp.default + ) + assert exp_node.meta["val"].dtype == torch.float32 + assert exp_node.args[0].target == torch.ops.aten._to_copy.default + assert exp_node.args[0].kwargs.get("dtype") == torch.float32 + + def test_cast_to_16_bit_precision_forwards_calibration_data(self): + """cast_to_16_bit_precision forwards calibration_data to FP16 casting.""" + model = _ActivationOverflowModel().eval() + x = torch.tensor([[15.0]]) + ref = model(x) + + ep = _export(model, x) + cast_to_16_bit_precision(ep, calibration_data=[x]) + out = _run_ep(ep, x.half()) + + assert not torch.isinf(out).any() + assert torch.allclose(out.float(), ref, atol=1e-3) + + exp_node = next( + n + for n in ep.graph.nodes + if n.op == "call_function" and n.target == torch.ops.aten.exp.default + ) + assert exp_node.meta["val"].dtype == torch.float32 + + def test_attention_logit_overflow_preserved_in_fp32(self): + """Attention dot-product logits exceeding 65504 are preserved in FP32.""" + + class _AttentionBlock(nn.Module): + def __init__(self): + super().__init__() + self.q = nn.Linear(32, 32) + self.k = nn.Linear(32, 32) + self.v = nn.Linear(32, 32) + + def forward(self, x): + q = self.q(x) + k = self.k(x) + v = self.v(x) + scores = torch.matmul(q, k.transpose(-2, -1)) + attn = torch.softmax(scores, dim=-1) + return torch.matmul(attn, v) + + torch.manual_seed(42) + model = _AttentionBlock().eval() + with torch.no_grad(): + model.q.weight.mul_(150.0) + model.k.weight.mul_(150.0) + + x = torch.randn(2, 4, 32) + ref = model(x) + assert not torch.isinf(ref).any() + + ep = _export(model, x) + cast_fp32_to_fp16(ep, calibration_data=[x]) + out = _run_ep(ep, x.half()) + + assert not torch.isinf(out).any() + assert torch.allclose(out.float(), ref, atol=1e-2, rtol=1e-2) + + # The overflowing matmul and softmax should remain in FP32 + softmax_node = next( + n for n in ep.graph.nodes if n.op == "call_function" and "_softmax" in str(n.target) + ) + assert softmax_node.meta["val"].dtype == torch.float32 + + # Check cast-back to FP16 exists after softmax + softmax_cast = next( + n + for n in ep.graph.nodes + if n.op == "call_function" + and n.target == torch.ops.aten._to_copy.default + and n.args[0] is softmax_node + ) + assert softmax_cast.kwargs.get("dtype") == torch.float16 + + def test_dynamic_calibration_union_of_overflowing_nodes_across_samples(self): + """Union of overflowing nodes across distinct calibration samples stays in FP32.""" + + class _TwoBranchModel(nn.Module): + def forward(self, x): + # Branch A + branch_a = torch.log1p(torch.exp(x[:, 0])) + # Branch B + branch_b = torch.log1p(torch.exp(x[:, 1])) + return branch_a + branch_b + + model = _TwoBranchModel().eval() + # Sample A overflows branch A (exp(15) > 65504), but not branch B (exp(1) ~ 2.7) + sample_a = torch.tensor([[15.0, 1.0]]) + # Sample B overflows branch B (exp(15) > 65504), but not branch A (exp(1) ~ 2.7) + sample_b = torch.tensor([[1.0, 15.0]]) + + # Baseline: calibrating with only sample A keeps branch A in FP32 and casts branch B to FP16 + ep_single = _export(model, sample_a) + cast_fp32_to_fp16(ep_single, calibration_data=[sample_a]) + exp_single = [ + n + for n in ep_single.graph.nodes + if n.op == "call_function" and n.target == torch.ops.aten.exp.default + ] + assert [n.meta["val"].dtype for n in exp_single] == [torch.float32, torch.float16] + + # Calibrating with both samples preserves the union (both branches) in FP32 + ep = _export(model, sample_a) + cast_fp32_to_fp16(ep, calibration_data=[sample_a, sample_b]) + + # Verify both exp nodes remain in FP32 + exp_nodes = [ + n + for n in ep.graph.nodes + if n.op == "call_function" and n.target == torch.ops.aten.exp.default + ] + assert len(exp_nodes) == 2 + for node in exp_nodes: + assert node.meta["val"].dtype == torch.float32 + # Boundary cast to FP32 before exp + assert node.args[0].target == torch.ops.aten._to_copy.default + assert node.args[0].kwargs.get("dtype") == torch.float32 + + # Verify both log1p nodes remain in FP32 + log1p_nodes = [ + n + for n in ep.graph.nodes + if n.op == "call_function" and n.target == torch.ops.aten.log1p.default + ] + assert len(log1p_nodes) == 2 + for node in log1p_nodes: + assert node.meta["val"].dtype == torch.float32 + + # Verify numerical correctness for both samples + for sample in [sample_a, sample_b]: + ref = model(sample) + out = _run_ep(ep, sample.half()) + assert not torch.isinf(out).any() + assert torch.allclose(out.float(), ref, atol=1e-3) diff --git a/tests/casting/test_casting_utils.py b/tests/casting/test_casting_utils.py index 60af3ab..95eb785 100644 --- a/tests/casting/test_casting_utils.py +++ b/tests/casting/test_casting_utils.py @@ -24,6 +24,7 @@ check_tensor_overflow_fp16, classify_float_args, cleanup_casts, + find_overflowing_nodes, insert_cast_after, is_ignored_op, ) @@ -1407,3 +1408,178 @@ def test_match_predicate_function(self, exp_node, add_node): # Match by target inspection via custom predicate assert is_ignored_op(exp_node, lambda node: node.target == torch.ops.aten.exp.default) assert not is_ignored_op(add_node, lambda node: node.target == torch.ops.aten.exp.default) + + +# ============================================================================= +# find_overflowing_nodes tests +# ============================================================================= +class TestFindOverflowingNodes: + """Unit tests for find_overflowing_nodes.""" + + class _ActivationOverflowModel(torch.nn.Module): + def forward(self, x): + return torch.log1p(torch.exp(x)) + + class _ThresholdModel(torch.nn.Module): + def forward(self, x): + return x * 100.0 + + class _TupleReturnModel(torch.nn.Module): + def forward(self, x): + a = torch.exp(x) + chunks = torch.chunk(a, 2, dim=-1) + c0 = torch.log(chunks[0]) + c1 = torch.log(chunks[1]) + return torch.cat([c0, c1], dim=-1) + + class _EmptyTensorModel(torch.nn.Module): + def forward(self, x): + empty = torch.empty(0, dtype=torch.float32) + return x + 1.0, empty + + def test_empty_or_none_calibration_data(self): + """Returns empty set when calibration_data is None.""" + m = self._ActivationOverflowModel().eval() + x = torch.tensor([[15.0]]) + ep = torch.export.export(m, (x,), strict=False).run_decompositions() + + assert find_overflowing_nodes(ep, None) == set() + + def test_detects_overflow_and_consumers(self): + """Detects exp overflowing > 65504 and propagates to log1p.""" + m = self._ActivationOverflowModel().eval() + x = torch.tensor([[15.0]]) + ep = torch.export.export(m, (x,), strict=False).run_decompositions() + + overflowing = find_overflowing_nodes(ep, [x]) + targets = {n.target for n in overflowing} + assert torch.ops.aten.exp.default in targets + assert torch.ops.aten.log1p.default in targets + + def test_input_formats(self): + """Supports iterable of Tensors, tuples, dicts, or DataLoader.""" + m = self._ActivationOverflowModel().eval() + x = torch.tensor([[15.0]]) + ep = torch.export.export(m, (x,), strict=False).run_decompositions() + + # 1. Iterable of Tensors + nodes_tensor = find_overflowing_nodes(ep, [x]) + assert len(nodes_tensor) > 0 + + # 2. Iterable of tuples + nodes_tuple = find_overflowing_nodes(ep, [(x,)]) + assert nodes_tuple == nodes_tensor + + # 3. Iterable of dicts + nodes_dict = find_overflowing_nodes(ep, [{"x": x}]) + assert nodes_dict == nodes_tensor + + # 4. DataLoader + loader = torch.utils.data.DataLoader([(x,)], batch_size=1) + nodes_loader = find_overflowing_nodes(ep, loader) + assert nodes_loader == nodes_tensor + + def test_bare_tensor_raises_type_error(self): + """Passing a bare Tensor instead of an iterable of samples raises TypeError.""" + m = self._ActivationOverflowModel().eval() + x = torch.tensor([[15.0]]) + ep = torch.export.export(m, (x,), strict=False).run_decompositions() + + with pytest.raises(TypeError, match="must be an iterable of samples"): + find_overflowing_nodes(ep, x) + + def test_bare_dict_raises_type_error(self): + """Passing a bare dict instead of an iterable of samples raises TypeError.""" + m = self._ActivationOverflowModel().eval() + x = torch.tensor([[15.0]]) + ep = torch.export.export(m, (x,), strict=False).run_decompositions() + + with pytest.raises(TypeError, match="must be an iterable of samples"): + find_overflowing_nodes(ep, {"x": x}) + + def test_non_iterable_raises_type_error(self): + """Passing a non-iterable raises TypeError.""" + m = self._ActivationOverflowModel().eval() + x = torch.tensor([[15.0]]) + ep = torch.export.export(m, (x,), strict=False).run_decompositions() + + with pytest.raises(TypeError, match="must be an iterable of samples"): + find_overflowing_nodes(ep, 123) + + def test_bare_string_raises_type_error(self): + """Passing a string instead of an iterable of samples raises TypeError.""" + m = self._ActivationOverflowModel().eval() + x = torch.tensor([[15.0]]) + ep = torch.export.export(m, (x,), strict=False).run_decompositions() + + with pytest.raises(TypeError, match="must be an iterable of samples"): + find_overflowing_nodes(ep, "path/to/data") + + def test_tuple_return_propagation(self): + """Propagates through operator.getitem to downstream consumers.""" + m = self._TupleReturnModel().eval() + x = torch.tensor([[15.0, 15.0]]) + ep = torch.export.export(m, (x,), strict=False).run_decompositions() + + overflowing = find_overflowing_nodes(ep, [x]) + targets = {n.target for n in overflowing} + + assert torch.ops.aten.exp.default in targets + assert torch.ops.aten.split_with_sizes.default in targets + assert torch.ops.aten.log.default in targets + assert torch.ops.aten.cat.default not in targets + + def test_empty_tensor_no_crash(self): + """Empty tensors (numel == 0) do not cause reduction errors.""" + m = self._EmptyTensorModel().eval() + x = torch.tensor([[1.0]]) + ep = torch.export.export(m, (x,), strict=False).run_decompositions() + + overflowing = find_overflowing_nodes(ep, [x]) + assert len(overflowing) == 0 + + def test_dict_missing_input_raises_value_error(self): + """Passing a calibration dict missing required model input raises ValueError.""" + m = self._ActivationOverflowModel().eval() + x = torch.tensor([[15.0]]) + ep = torch.export.export(m, (x,), strict=False).run_decompositions() + + with pytest.raises(ValueError, match="missing required input"): + find_overflowing_nodes(ep, [{"wrong_name": x}]) + + def test_mismatched_sample_length_raises_value_error(self): + """Passing a sample with fewer or more inputs than expected raises ValueError.""" + + class _TwoInputModel(torch.nn.Module): + def forward(self, a, b): + return a + b + + m = _TwoInputModel().eval() + a = torch.tensor([[1.0]]) + ep = torch.export.export(m, (a, a), strict=False).run_decompositions() + + with pytest.raises(ValueError, match="provided 1 input"): + find_overflowing_nodes(ep, [(a,)]) + + with pytest.raises(ValueError, match="provided 3 input"): + find_overflowing_nodes(ep, [(a, a, a)]) + + def test_inf_tensor_does_not_overflow(self): + """True infinity in FP32 is allowed in FP16 and does not trigger overflow.""" + + class _InfModel(torch.nn.Module): + def forward(self, x): + return x * 2.0 + + m = _InfModel().eval() + # Input with true inf + x_inf = torch.tensor([[float("inf")]]) + ep = torch.export.export(m, (x_inf,), strict=False).run_decompositions() + + overflowing = find_overflowing_nodes(ep, [x_inf]) + assert len(overflowing) == 0 + + # Input with finite value in (65504, 1e38) does trigger overflow + x_overflow = torch.tensor([[70000.0]]) + overflowing_val = find_overflowing_nodes(ep, [x_overflow]) + assert len(overflowing_val) > 0