[graph_trainer] Add CooR precompile support for aot_fx_trace compile mode

[bobrenjc93]

Stack from ghstack (oldest at bottom):

• -> [graph_trainer] Add CooR precompile support for aot_fx_trace compile mode #2975

Introduce aot_fx_trace as an alternate precompilation path alongside
the existing AOT joint graph export. aot_fx_trace uses make_fx to
trace the full fwd+loss+bwd step as a single FX graph, then serializes
the GraphModule for loading on all ranks during training. Because
cudagraphs is a default pass, we also include cudagraph composability
related changes in this PR to get to bitwise equivalence between
aot_fx_trace with and without precompile.

Changes:

• precompile_main.py: Refactor into _common_setup() + _precompile_aot()
  • _precompile_aot_fx_trace(), route by --compile.mode. Use Python
    float for global_valid_tokens (matching dist_sum's return type) so
    make_fx bakes it as a graph constant, identical to the runtime path.
• precompile.py: PrecompiledFxTraceArtifact dataclass, save/load fns.
  Uses GraphPickler (not plain pickle) to preserve SymInt expressions
  (_runtime_compute_coordinate_on_dim) through serialization — plain
  pickle evaluates SymInts to concrete trace-time values, baking in
  rank-specific embedding vocab offsets.
• compile.py: Mode-aware validation (aot_fx_trace doesn't need
  serializable passes), logging for precompile load vs runtime trace
• trainer.py: _load_precompiled_fx_trace(), branching in
  _make_fx_forward_backward_step()
• make_fx_tracer.py: Use set_multithreading_enabled(False) so backward
  tracing runs on the calling thread, keeping the compile_on_one_rank
  ContextVar visible
• cudagraph.py: Defer _input_indices_to_copy when example_inputs is
  empty (precompile load path), accept opaque values (DeviceMesh) in
  input validation, skip non-tensors in static address checks.
  Squashed the separate cudagraph CooR PR into this one because
  cudagraph activates automatically for aot_fx_trace graphs (all
  inputs are desugared plain tensors + opaques, no HOPs), so bitwise
  equivalence testing required both fixes together.
• passes.py: Guard cudagraph_pass against non-GraphModule inputs

Test plan

• pytest torchtitan/experiments/graph_trainer/tests/test_bitwise_deterministic.py -x — 6 passed, 4 skipped (Llama3, DSv3 × {eager self-deterministic, eager vs aot_fx_trace, precompile vs trace})
• pytest torchtitan/experiments/graph_trainer/tests/test_precompile.py -x — all pass
• pytest torchtitan/experiments/graph_trainer/tests/test_trace_module.py -x — all pass
• 8-GPU integration tests via CI (ciflow/8gpu)

[SherlockNoMad]

Do we have clear instruction in claude.md for agent to use the pre-compile workflow?

I would validate by prompting ”run graph_trainer precompile training of llama_8b for me, show me the MFU, profiling result and memory snapshot as proof of work" in a /new session.

this should work in one shot.
If it doesn't, "reflect on instruction from claude.md, improve based on the experiment you just did"

[bobrenjc93]

For posterity - claude one shotted it