Skip to content

fix: initialize torch.compile per thread - #1495

Merged
dxqb merged 1 commit into
Nerogar:mergefrom
dxqb:fix/compiled-optimizer-thread-local
Jun 4, 2026
Merged

fix: initialize torch.compile per thread#1495
dxqb merged 1 commit into
Nerogar:mergefrom
dxqb:fix/compiled-optimizer-thread-local

Conversation

@dxqb

@dxqb dxqb commented Jun 4, 2026

Copy link
Copy Markdown
Collaborator

torch._dynamo.config overrides are thread-local. The existing call in checkpointing_util runs in the main thread and is invisible to the training thread spawned by the UI. This caused compiled optimizers (e.g. AdamW_adv with compiled_optimizer=True) to hit the default recompile_limit of 8 and abort with FailOnRecompileLimitHit when training models with more than 8 distinct parameter shapes.

Fix: call init_compile() from GenericTrainer.init, which runs in whichever thread/process owns training (UI thread, CLI main thread, or torch.multiprocessing.spawn subprocess for multi-GPU).

torch._dynamo.config overrides are thread-local. The existing call in
checkpointing_util runs in the main thread and is invisible to the
training thread spawned by the UI. This caused compiled optimizers
(e.g. AdamW_adv with compiled_optimizer=True) to hit the default
recompile_limit of 8 and abort with FailOnRecompileLimitHit when
training models with more than 8 distinct parameter shapes.

Fix: call init_compile() from GenericTrainer.__init__, which runs in
whichever thread/process owns training (UI thread, CLI main thread,
or torch.multiprocessing.spawn subprocess for multi-GPU).

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
@dxqb dxqb added the preview merged in the preview branch label Jun 4, 2026
@dxqb dxqb mentioned this pull request Jun 4, 2026
@dxqb
dxqb changed the base branch from master to merge June 4, 2026 18:40
@dxqb
dxqb merged commit 43b618c into Nerogar:merge Jun 4, 2026
1 check passed
@dxqb
dxqb deleted the fix/compiled-optimizer-thread-local branch June 4, 2026 18:40
BitcrushedHeart pushed a commit to BitcrushedHeart/OneTrainer that referenced this pull request Jun 20, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

preview merged in the preview branch

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant