Skip to content

Commit 86cb354

Browse files
Ronald Tseronaldtse
authored andcommitted
fix: deterministically release ORT sessions in run_parity
Six interleaved session lifetimes (3 precisions x parity+margins) plus the torch reference OOMed the 32GB container at the int8 stage; del + gc in a finally block instead of waiting on GC.
1 parent a120384 commit 86cb354

1 file changed

Lines changed: 19 additions & 10 deletions

File tree

‎src/imf/parity.py‎

Lines changed: 19 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -166,6 +166,8 @@ def run_parity(
166166
zip_path = Path(zip_path)
167167
enc, kv = _sessions_from_zip(zip_path)
168168

169+
import gc
170+
169171
import yaml
170172

171173
with zipfile.ZipFile(zip_path) as zf:
@@ -175,16 +177,23 @@ def run_parity(
175177
mismatches = 0
176178
cer_ref_sum = 0.0
177179
cer_onnx_sum = 0.0
178-
for i, (source, gold) in enumerate(pairs):
179-
n += 1
180-
ref = reference[i] if reference is not None else _torch_greedy_tokens(
181-
model, source, max_len
182-
)
183-
got = onnx_greedy_kv(enc, kv, source, max_len)
184-
if ref != got:
185-
mismatches += 1
186-
cer_ref_sum += char_error_rate(_decode_tokens(ref), gold)
187-
cer_onnx_sum += char_error_rate(_decode_tokens(got), gold)
180+
try:
181+
for i, (source, gold) in enumerate(pairs):
182+
n += 1
183+
ref = reference[i] if reference is not None else _torch_greedy_tokens(
184+
model, source, max_len
185+
)
186+
got = onnx_greedy_kv(enc, kv, source, max_len)
187+
if ref != got:
188+
mismatches += 1
189+
cer_ref_sum += char_error_rate(_decode_tokens(ref), gold)
190+
cer_onnx_sum += char_error_rate(_decode_tokens(got), gold)
191+
finally:
192+
# ORT arenas accumulate across precision stages (six session
193+
# lifetimes + the torch reference OOMed the container at int8);
194+
# release deterministically instead of waiting on GC
195+
del enc, kv
196+
gc.collect()
188197

189198
cer_ref = 100.0 * cer_ref_sum / max(n, 1)
190199
cer_onnx = 100.0 * cer_onnx_sum / max(n, 1)

0 commit comments

Comments
 (0)