feat(casting): add dynamic activation range calibration in FP16 casting (#7)

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:
 

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

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