Skip to content

feat(aot): graceful fallback to JIT when AOT compile fails - #86

Open
cennn wants to merge 1 commit into
SandAI-org:mainfrom
cennn:feat/aot-jit-fallback
Open

cennn wants to merge 1 commit into
SandAI-org:mainfrom
cennn:feat/aot-jit-fallback

Conversation

@cennn

@cennn cennn commented Oct 1, 2026

Copy link
Copy Markdown
Collaborator

Problem

When MAGI_COMPILE_AOT=1 is set globally, models using SimpleFSDP/DTensor fail AOT compilation with:

Attempted to read undefined local variable
File ".../simple_fsdp/builder.py", line 236, in forward
    state.aot_compile(*args, **kwargs)
File ".../aot_compile.py", line 143, in aot_compile_fullgraph
    raise e.with_traceback(None) from e.__cause__

This crashes the entire warmup process, preventing AOT from being used even for compatible models (e.g., VAE which has no FSDP).

Solution

Catch exceptions in the AOT compile path and gracefully fall back to the JIT bytecode capture path. This allows MAGI_COMPILE_AOT=1 to be set globally:

  • VAE (no FSDP): AOT compiles successfully → skip Dynamo trace on verify (~14s savings)
  • Base/SR (SimpleFSDP+DTensor): AOT fails → falls back to JIT → inductor cache hits from bake tar

Key implementation details:

  1. Catch any exception from state.aot_compile()
  2. Reset compile_config.aot = False to prevent further AOT attempts for this state
  3. Reset compiled_entry and call torch._dynamo.reset() to clear AOT-only guard_filter_fn
  4. Re-run _ensure_compiled() for JIT mode and proceed with bytecode capture

Context

Part of CI cold-start optimization effort (athena PR #1112). AOT serializes Dynamo FX graphs so verify can skip Dynamo tracing + RestartAnalysis entirely (~42s savings for compatible models).

When MAGI_COMPILE_AOT=1, models using SimpleFSDP/DTensor may fail AOT
compilation (e.g. "Attempted to read undefined local variable"). Instead
of crashing, catch the exception and fall back to the JIT bytecode
capture path. This allows AOT to be enabled globally while FSDP models
auto-degrade to JIT + inductor cache hits.

This branch has not been deployed

No deployments
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