Skip to content

[WIP]Feat: Run checkpoint ONNX export in a detached subprocess and dedupe save_utd/best exports#261

Draft
yamsam wants to merge 1 commit into
devfrom
feat/async-onnx-export
Draft

[WIP]Feat: Run checkpoint ONNX export in a detached subprocess and dedupe save_utd/best exports#261
yamsam wants to merge 1 commit into
devfrom
feat/async-onnx-export

Conversation

@yamsam

@yamsam yamsam commented Jul 21, 2026

Copy link
Copy Markdown

Summary

  • Moved the per-checkpoint ONNX export out of the training process into a detached, CPU-only
    subprocess (launch_checkpoint_onnx_export + a new python -m diffusion_planner.utils.onnx_export
    CLI). The export already rebuilds the model from the on-disk .pth + args.json, so it needs
    nothing from the training process — rank 0 no longer stalls all ranks behind the CPU trace of 4
    graphs (the full graph contains 10 unrolled DiT evaluations)
  • Deduplicated the double export: when an epoch is both a save_utd snapshot and the new best model,
    the checkpoint is now traced once and the artifacts are copied to the second directory (previously 8
    traced graphs, now 4)
  • The child runs with CUDA_VISIBLE_DEVICES="" (never touches GPU memory), logs to
    <output_dir>/onnx_export.log, and is detached via start_new_session; the training loop reaps
    finished exports each launch and waits for stragglers before exiting so artifacts are always complete
  • Side effect: the export's global RNG seeding (random/np/torch.manual_seed(42) in
    export_checkpoint_onnx) no longer runs inside the training process, so it can no longer reset rank
    0's augmentation RNG mid-training (O-B1 in docs/onnx_export_analysis.md)
  • export_checkpoint_onnx_guarded is kept unchanged for existing callers (e.g.
    train_grpo_predictor.py); migrating GRPO can follow up separately

Test plan

  • Added diffusion_planner/tests/test_onnx_export_launch.py (5 tests): command/parser round-trip,
    artifact copying (creates missing dirs, includes external-data sidecars, excludes unrelated files),
    main() exports once then copies, launch is detached and CPU-only with a log file, launch failure
    returns None without raising
  • PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 uv run pytest diffusion_planner/tests/test_onnx_export_launch.py
    — 5 passed
  • Verified python -m diffusion_planner.utils.onnx_export --help works and diffusion_planner.train
    imports cleanly

…d it

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Isamu Yamashita <isamu.yamashita@tier4.jp>
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