Skip to content

Commit 48690b1

Browse files
ssjiamanuelcandales
authored andcommitted
[ET-VK] Fix fp16 inference: clamp constants and sync view_convert shader dtype combos
Two fixes for force_fp16 inference mode: 1. Clamp constant tensors to the representable range of the effective dtype after conversion in VkGraphBuilder. When force_fp16 converts fp32 constants to fp16, values like -100000 (used as attention mask fill values) overflow to -inf. A subsequent `0 * (-inf)` produces NaN per IEEE 754, which propagates through the entire model. Clamping to finfo(fp16).min/max (-65504/65504) prevents the overflow. 2. Sync dtype combos across view_convert_buffer.yaml and view_convert_texture.yaml so both define the same set of conversion variants. Added float<->half, half<->int32 to the buffer YAML, and float<->half, half<->float, uint8<->float/half/int32 to the texture YAML. Both files now cover all 9 dtype combos: int32<->float, int32<->half, uint8<->float/half/int32, float<->half, half<->float, half<->int32. The float<->half variants are needed by the _to_dim_order_copy operator when running fp16 models (e.g. SceneX). Verified on edgeTAM first frame: mask IoU improved from 0.18 (NaN-corrupted) to 0.9991, matching the fp32 baseline range of 0.9926-0.9995. Also verified SceneX fp16 model (scenex_v9_512_vulkan_fp16_v3.et) no longer crashes with missing view_convert_buffer_float_half shader. Differential Revision: [D96036512](https://our.internmc.facebook.com/intern/diff/D96036512/) ghstack-source-id: 350221317 Pull Request resolved: pytorch#18076
1 parent 76dfb19 commit 48690b1

3 files changed

Lines changed: 18 additions & 0 deletions

File tree

‎backends/vulkan/runtime/graph/ops/glsl/view_convert_buffer.yaml‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,5 +19,8 @@ view_convert_buffer:
1919
- parameter_values: [uint8, half]
2020
- parameter_values: [uint8, int32]
2121
- parameter_values: [float, int32]
22+
- parameter_values: [float, half]
23+
- parameter_values: [half, float]
24+
- parameter_values: [half, int32]
2225
shader_variants:
2326
- NAME: view_convert_buffer

‎backends/vulkan/runtime/graph/ops/glsl/view_convert_texture.yaml‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,13 @@ view_convert_texture:
1414
parameter_names: [IN_DTYPE, OUT_DTYPE]
1515
combos:
1616
- parameter_values: [int32, float]
17+
- parameter_values: [int32, half]
18+
- parameter_values: [uint8, float]
19+
- parameter_values: [uint8, half]
20+
- parameter_values: [uint8, int32]
1721
- parameter_values: [float, int32]
22+
- parameter_values: [float, half]
23+
- parameter_values: [half, float]
24+
- parameter_values: [half, int32]
1825
shader_variants:
1926
- NAME: view_convert_texture

‎backends/vulkan/serialization/vulkan_graph_builder.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -118,6 +118,14 @@ def maybe_add_constant_tensor(self, node: Node) -> int:
118118

119119
# Convert the tensor dtype if needed
120120
if tensor.dtype != effective_dtype:
121+
# Clamp float tensors to the representable range of the effective
122+
# dtype to avoid infinities. This is needed when force_fp16 converts
123+
# fp32 constants (e.g. -100000 attention mask values) to fp16 where
124+
# the max representable value is ~65504.
125+
if effective_dtype.is_floating_point:
126+
dtype_info = torch.finfo(effective_dtype)
127+
tensor = tensor.clamp(min=dtype_info.min, max=dtype_info.max)
128+
121129
tensor = tensor.to(effective_dtype)
122130

123131
# Serialize tensor data to bytes

0 commit comments

Comments
 (0)