[acceleration] Model-agnostic acceleration plugin layer (torch.compile / attention backend / feature caching)

[Jayce-Ping]

Summary

Adds a model-agnostic, registry-based acceleration plugin layer (src/flow_factory/acceleration/) that speeds up training without touching trainer or model math, gated for train-inference correctness by the trainer's RL paradigm.

• Two ordered config slots under acceleration: — shared (lossless/stage=both, applied to both rollout and the training forward via setup()) and rollout (Stage-3 only via rollout_context(), may be lossy).
• Accelerators: attention_backend (folds in the removed model.attn_backend knob), torch_compile, diffusers_cache, cache_dit.
• acceleration/validator.py enforces the safety contract (constraints [Feat] Support Pairwise and Groupwise reward computation #7/[Feat] Extend EMA schedules #26): a lossy rollout accelerator is rejected on coupled trainers; a lossy stage=both accelerator (torch_compile) is allowed everywhere but warns on coupled that the on-policy ratio is ≈1, not bit-exact.
• model.attn_backend is removed (fails fast with a migration message → use the attention_backend accelerator).

Correctness — GRPO on-policy ratio (zero tolerance: must be exactly 1, atol=0)

Verified on 16×H20 (2-node) under both DeepSpeed ZeRO-2 and FSDP2 HYBRID_SHARD, per-element assert at the first on-policy step on all 16 ranks.

method	backend / mode	max|ratio-1| (worst of 16)	verdict
eager	SDPA	0.000e+00	bit-exact
attention_backend	_flash_3_hub (FA3)	0.000e+00 (ZeRO-2 and FSDP2)	bit-exact
torch_compile	full	1.13e-5 (2/16; FSDP2 run 0/16)	not bit-exact (intermittent)
attention_backend+torch_compile	_flash_3_hub + full	1.31e-5	not bit-exact (compile)
torch_compile (Qwen-Image)	regional	1.69e-5 (1/16)	not bit-exact
torch_compile (no-CFG)	full	7.15e-7 (1/16)	not bit-exact

Conclusion: eager and attention_backend are strictly bit-exact. torch_compile (any mode/model, CFG or not) leaves an intermittent ~1e-5 residual on a minority of samples — numerically on-policy (≪ clip_range=1e-4) but not bit-exact. Root cause: Inductor compiles a separate graph for grad vs no-grad mode; the grad-forced rollout fix removes the dominant divergence but rollout and training remain distinct kernel invocations and bf16 accumulation is non-associative. Hence torch_compile is now classified lossy (validator warns on coupled).

Note: regional compile is unavailable on SD3.5 (SD3Transformer2DModel has no _repeated_blocks) → clean fail-fast (works on Qwen-Image).

Speedup — lossless (GRPO / SD3.5 LoRA, 16-GPU ZeRO-2, CFG g=4.5, 8 steps; rank0 steady-state median)

method	rollout (s)	rollout ×	train-step (s)	train ×	peak VRAM
eager (baseline)	13.65	1.00×	15.38	1.00×	14.49 GB
attention_backend(_flash_3_hub)	13.03	1.05×	15.40	1.00×	14.49 GB
torch_compile(full)	11.44	1.19×	13.04	1.18×	12.00 GB (−17%)
_flash_3_hub+torch_compile(full)	10.38	1.32×	11.01	1.40×	12.00 GB

• FA3 is ~neutral on H20 (bandwidth-bound Hopper; attention is not the bottleneck).
• torch.compile is the real win (~1.2× + −17% VRAM; combined with FA3 ~1.3–1.4×) — at the cost of strict bit-exactness.

Speedup — lossy rollout caching (DiffusionNFT / Qwen-Image LoRA, 16-GPU ZeRO-2, 28 steps; same seed; rank0 steady-state)

Fresh post-removal measurement (rank0, median of n=6 steady-state epochs, 28 inference steps):

method	rollout (s)	rollout ×	train-step (s)	reward drift vs baseline
cache off (baseline)	27.02	1.00×	19.52	—
diffusers_cache (first_block, thr 0.08)	22.01	1.23×	19.77	negligible
cache_dit (default / aggressive)	25.5-26.2	~1.0×	—	0 (bit-identical — no-op; accelerator removed)

• Caching is rollout-only (train-step unaffected) and validator-gated to decoupled/distillation.
• Fixed (commit 247a64a): periodic eval + diffusers_cache previously crashed with a diffusers FirstBlockCache "No context is set" error. Root cause: diffusers' HookRegistry._get_child_registries() caches its child list on first use; during eval (cache disabled) the adapter still opens cache_context(...), caching an empty child list before the block cache-hooks exist, so the later training _set_context skipped them. The accelerator now invalidates that stale _child_registries_cache after enable_cache. Verified on 16xH20 Qwen-Image NFT (eval_freq=1, rc=0, zero errors).
• Verified: torch_compile (regional) + diffusers_cache coexist — both apply, run completes (rc=0), and per-epoch rollout time is stable (no per-epoch Dynamo recompile churn).
• diffusers_cache works: ~1.2× rollout (1.23× here on n=6; 1.18-1.26× across runs) with negligible reward drift. At 8 inference steps it never triggers (no gain, bit-identical output) — it needs enough denoising steps to engage.
• cache_dit accelerator was REMOVED (commit 0170953). diffusers_cache is now the single lossy rollout backend; a config still naming cache_dit fails fast with an actionable error pointing to diffusers_cache. Full investigation that led to removal:
  • Root cause of the no-op: CacheDitAccelerator called cache_dit.enable_cache(adapter.pipeline), which sets persistent_context=False for a real pipeline — caching then only activates inside pipeline.__call__, but Flow-Factory's rollout calls the transformer directly (adapter.inference()), so cache_dit's per-call context is never created and every step logs "Cache context not exist … skip cache" (bit-identical output, default and aggressive configs alike).
  • Engagement fix (validated): passing the bare base transformer instead (cache_dit wraps it in a FakeDiffusionPipeline → persistent_context=True) makes caching actually engage on Qwen-Image (epoch 0 ran with no "skip cache").
  • Remaining blockers (why it's not shipped): (1) cache_dit's enable_cache/disable_cache is not idempotent across the per-epoch rollout scoping — re-enabling on epoch 2 asserts hasattr(pipe, "_context_manager"); (2) it conflicts with the adapter's native diffusers cache_context (_diffusers_hook is not None) — cache_dit is a competing caching system to the diffusers-native one that diffusers_cache already uses; (3) on SD3.5 the bare transformer class SD3Transformer2DModel isn't registered (only StableDiffusion3Pipeline is), so cache_dit's bare-module path is unsupported there.
  • Conclusion: a robust cache_dit integration would need to coexist with / replace the adapters' diffusers cache_context and use cache_dit's context-manager + refresh_context for per-rollout scoping, built on private dev-build internals — substantial work for no gain over diffusers-native caching. So the CacheDitAccelerator was removed; diffusers_cache is the lossy backend (clean per-epoch lifecycle, ~1.2× rollout on Qwen, negligible reward drift).

Test plan

• GRPO on-policy ratio bit-exact assert (eager/FA3 pass; compile ~1e-5) — ZeRO-2 + FSDP2, 16 GPU.
• Lossless speedup (FA3 / compile / combo) vs eager — SD3.5, 16 GPU.
• Lossy NFT caching speedup + reward drift — Qwen-Image, 16 GPU (diffusers_cache ~1.2×; cache_dit found to be a no-op — bug noted).
• Paradigm-gated fail-fast (lossy-rollout on coupled rejected; lossy-shared on coupled warns).