[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
Draft
[WIP]Feat: Run checkpoint ONNX export in a detached subprocess and dedupe save_utd/best exports#261yamsam wants to merge 1 commit into
yamsam wants to merge 1 commit into
Conversation
…d it Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Isamu Yamashita <isamu.yamashita@tier4.jp>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
subprocess (
launch_checkpoint_onnx_export+ a newpython -m diffusion_planner.utils.onnx_exportCLI). The export already rebuilds the model from the on-disk
.pth+args.json, so it needsnothing 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)
save_utdsnapshot 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)
CUDA_VISIBLE_DEVICES=""(never touches GPU memory), logs to<output_dir>/onnx_export.log, and is detached viastart_new_session; the training loop reapsfinished exports each launch and waits for stragglers before exiting so artifacts are always complete
random/np/torch.manual_seed(42)inexport_checkpoint_onnx) no longer runs inside the training process, so it can no longer reset rank0's augmentation RNG mid-training (O-B1 in
docs/onnx_export_analysis.md)export_checkpoint_onnx_guardedis kept unchanged for existing callers (e.g.train_grpo_predictor.py); migrating GRPO can follow up separatelyTest plan
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 failurereturns
Nonewithout raisingPYTEST_DISABLE_PLUGIN_AUTOLOAD=1 uv run pytest diffusion_planner/tests/test_onnx_export_launch.py— 5 passed
python -m diffusion_planner.utils.onnx_export --helpworks anddiffusion_planner.trainimports cleanly