[graph_trainer] Replace trace_module/run_traced_module with minimal_fx_tracer API

Authored-by: Claude

Redesign the graph trainer's tracing API.

Key changes:

make_fx_tracer.py:
- Rename trace_module -> minimal_fx_tracer. Takes a callable and args where
  args[0] must be the nn.Module. Its params/buffers are lifted as graph
  inputs. fn must be a plain callable (not nn.Module) so trace and execute
  have the same calling convention — no hidden arg prepending.
- Delete run_traced_module. TracedResult is now directly callable — pass the
  same positional args (with the live module first) to execute the graph.
  Fresh params are read from the module automatically on each call.
- Store and restore output pytree spec so TracedResult.__call__ returns the
  same pytree structure as the original function (e.g. single tensor, list,
  tuple, dict), not a flat list.
- Store param/buffer FQNs at trace time, validate on first execution to catch
  module structure mismatches early.
- Install TracingContext before make_fx so invoke_subgraph deduplication works.
- Validate that all pytree leaves in args are tensors or primitives
  (int/float/bool/str). Callables like loss_fn must be captured in fn's
  closure, not passed as args.
- Handle tensor subclass (e.g. DTensor) unwrapping/rewrapping via dict-based
  layouts — only subclass positions get entries, plain tensors have no entry.

trainer.py:
- Replace FwdBwdStepModule (nn.Module wrapper that only existed because the
  old trace_module required nn.Module as fn) with make_fwd_bwd_step, a plain
  function factory. The model is now passed as the first arg, loss_fn is
  captured in the closure.
- Remove manual params_and_buffers dict construction — TracedResult.__call__
  reads fresh params from the live module automatically.

test_trace_module.py:
- Replace TrainStepModule with make_train_step plain function factory.
- Remove _get_params_and_buffers helper (no longer needed).
- Update all callsites: trace_module -> minimal_fx_tracer, run_traced_module
  -> direct TracedResult.__call__.
- Register BlockMask as pytree node at module level so flex_attention tests
  pass the leaf validation.
- Add test_mismatched_module_raises: FQN validation catches wrong module.
- Add test_non_tensor_leaf_raises: callable leaf in args raises ValueError.

All 7 model tests pass (llama3, llama4, qwen3, qwen3_moe, deepseek_v3,
gpt_oss, flex_attention_annotations).

ghstack-source-id: 3fa5e5c
Pull Request resolved: #2753

Author: tugsbayasgalan