Skip to content

Add pwrite-based single-file ONNX export + fix GPT-OSS dynamo export - #1232

Draft
smedhe wants to merge 3 commits into
quic:mainfrom
smedhe:nit/gpt_oss_0727
Draft

Add pwrite-based single-file ONNX export + fix GPT-OSS dynamo export#1232
smedhe wants to merge 3 commits into
quic:mainfrom
smedhe:nit/gpt_oss_0727

Conversation

@smedhe

@smedhe smedhe commented Aug 3, 2026

Copy link
Copy Markdown
Contributor

Summary

  • Replaces onnx_program.save() with a new save_dynamo_onnx() helper (QEfficient/utils/onnx_save_utils.py) that writes all initializer tensors in parallel via os.pwrite into a single pre-allocated .data file, then serializes a small graph-only .onnx
    proto referencing them via ExternalTensor offsets. This avoids the ~17M small write() syscalls and per-file inode overhead of the default save path, which is especially costly on large models on HDD-backed storage. Includes disk-space precheck, atomic
    rename via .data.tmp, per-tensor byte-count verification, RAM-headroom warning (advisory, psutil optional), and an opt-in sync flag for fdatasync durability.
  • Wires the new save path into QEFFBaseModel's dynamo export flow (modeling_qeff.py), and explicitly del onnx_program; gc.collect() afterward to release ExportedProgram weakrefs so PT weight offload can proceed.
  • Fixes GPT-OSS dynamo export/compile issues in modeling_gpt_oss.py (rotate_half rewritten using torch.chunk) and cache_utils.py (QEffHybridCacheForGPTOSS — drops redundant torch._dynamo.is_compiling() int-cast branching for
    max_cache_len/sliding_window_len).
  • Adds a dynamo: bool = False parameter to QEFFAutoModelForCausalLM.compile()/export plumbing (modeling_auto.py) to opt into the dynamo-based export path, and stores self.onnx_path on the base model for later reference.

Why

The default onnx_program.save() serializes weights through a BufferedWriter with an 8KB buffer, resulting in millions of small writes — this becomes a major bottleneck for large models (e.g. 70B+) on spinning-disk-backed storage. The single-file pwrite
approach lets the OS page cache coalesce writes into large sequential flushes, cutting save time from ~45 min to ~2 min for a 131GB model in local testing. The GPT-OSS fixes address failures encountered when exporting that model through the dynamo path.

smedhe added 2 commits August 3, 2026 11:52
Signed-off-by: smedhe <smedhe@qti.qualcomm.com>
Signed-off-by: smedhe <smedhe@qti.qualcomm.com>
@ochougul
ochougul marked this pull request as draft August 4, 2026 16:39
Signed-off-by: smedhe <smedhe@qti.qualcomm.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant