From c6925546b6cb0e539912990409e347c2f18b4a72 Mon Sep 17 00:00:00 2001 From: Jayce-Ping <315229706@qq.com> Date: Sat, 27 Jun 2026 15:47:36 +0800 Subject: [PATCH 01/15] [docs] docs: add acceleration plugin-layer roadmap & phase-1 plan --- .../plans/acceleration_plugin_layer.plan.md | 216 ++++++++++++++++++ 1 file changed, 216 insertions(+) create mode 100644 .cursor/plans/acceleration_plugin_layer.plan.md diff --git a/.cursor/plans/acceleration_plugin_layer.plan.md b/.cursor/plans/acceleration_plugin_layer.plan.md new file mode 100644 index 00000000..46e0c0ac --- /dev/null +++ b/.cursor/plans/acceleration_plugin_layer.plan.md @@ -0,0 +1,216 @@ +# Flow-Factory Acceleration Roadmap & Phase-1 Plan + +## Context + +Flow-Factory is an **online RL fine-tuning framework** for diffusion/flow-matching models. Its +dominant runtime cost is **Stage 3 — Trajectory Generation (rollout)**: for every epoch it runs +`num_inference_steps` denoising steps × `group_size (K)` samples per prompt × `num_batches_per_epoch`, +all under `torch.no_grad()`. This is *exactly* the inference workload that cache-dit / lightx2v / SGLang +accelerate. Stage 6 (policy optimization, forward+backward) is the secondary cost. + +The framework already has a solid **lossless** optimization base (FSDP/FSDP2/DeepSpeed, grad-checkpoint, +LoRA, bf16/fp16, CPU-offload + async H2D prefetch, Qwen CFG-merge, comm fusion, `StackedSampleBatch`, +single-root `ModelBundle`). What is missing is the **inference-acceleration layer** that the referenced +projects specialize in: feature caching (cache-dit DBCache/TaylorSeer/FBCache), whole-/region +`torch.compile`, long-sequence context parallelism (lightx2v / xDiT / SGLang style), and +quantization of frozen modules. + +This plan proposes a **model-agnostic, registry-based acceleration plugin layer** that respects the +framework's hard algorithm/model decoupling, plus a phased roadmap. Goal of Phase 1: cut rollout wall-time +on image models (FLUX/Qwen-Image) with **zero changes to trainers or model math**, gated for correctness. + +--- + +## The Decoupling-Critical Insight (read first — shapes the whole design) + +Acceleration techniques split into two correctness classes, and the split maps **exactly** onto +Constraint #7 (coupled vs decoupled) + Philosophy #1 / Constraint #20a (train-inference consistency): + +| Class | Examples | Where safe | Why | +|---|---|---|---| +| **Lossy** (changes `noise_pred`) | feature caching, step-reduction, int8/Sage attention | **rollout-only, decoupled/distillation-only** | Caching skips block compute; it cannot be replicated in the Stage-6 training forward (that needs full gradient through every block, every step). For **coupled** algos (GRPO/GRPO-Guard/DPPO) the rollout's per-step `log_prob` becomes the PPO "old log-prob"; if rollout is approximated but the training forward is exact, the importance ratio is biased → **silently wrong gradients**. For **decoupled** algos (NFT/AWM/DGPO/DPO/CRD) rollout runs with `compute_log_prob=False`; only the final image enters the reward → caching merely shifts the sample distribution, which is acceptable/tunable. | +| **Lossless** (numerically ~identical, applied to the *shared* module) | `torch.compile`, exact attention backends (FA2/FA3/xformers), context/sequence parallelism, quantizing **frozen-only** modules, comm/overlap/offload | **everywhere** | Because both `inference()` and `forward()` call the *same* `adapter.transformer`, a transform applied to that module is consistent across rollout and training by construction. | + +**Design rule enforced by the plan:** every accelerator declares `safety ∈ {lossless, lossy}` and +`stage ∈ {rollout, train, both}`. A fail-fast validator (Constraint #26) rejects `lossy` accel on a +coupled trainer, and rejects `lossy` accel on the `train` stage for any trainer. This is the mechanism +that lets us add aggressive rollout speedups *without* touching the train-inference consistency contract. + +--- + +## Current State (verified) + +**Already present / reusable seams:** +- `attn_backend` config (`hparams/model_args.py:110`) → `BaseAdapter._set_attention_backend()` + (`models/abc.py:852-866`) routes to diffusers' attention dispatcher: `native / flash / flash_hub / + _flash_3 / _flash_3_hub / sage / xformers`. **Attention backends are already wired framework-wide** — + they need benchmarking/per-stage exposure, not invention. +- Transformers are diffusers `CacheMixin` models: adapters already open `transformer.cache_context(...)` + (Qwen, Wan, LTX2, FLUX2-Klein). This is the exact hook diffusers-native & cache-dit caching plug into. +- Pluggable **scheduler registry** (`scheduler/registry.py`): FlowMatchEuler + UniPC multistep — solver + swaps / step schedules are already extensible. +- `RoutedComponentProxy` (`models/model_bundle.py:93-121`) delegates attribute access (incl. + `enable_cache`, `cache_context`, `compile`) to the inner module → an accelerator can call + `adapter.transformer.enable_cache(...)` model-agnostically, post-`prepare`. +- Rollout entry points: `BaseTrainer.generate_samples()` / `sample_batch()` (`trainers/abc.py:550-798`). + +**Absent (opportunities):** feature caching of any kind; whole-/region `torch.compile` (only +`bagel/.../qwen2_navit.py:43` compiles `flex_attention`); context/sequence parallelism for video; +quantization integration (`bitsandbytes` is an optional dep but unused); rollout throughput/VRAM/quality +benchmark harness. + +--- + +## Roadmap (phased, by ROI × safety × effort) + +| Phase | Track | What | Safety | Primary win | Effort | +|---|---|---|---|---|---| +| **0** | Foundation | Benchmark harness (rollout samples/s, step-time, peak VRAM, **reward-distribution regression**) + the `acceleration/` plugin layer + paradigm-gated validator | infra | enables safe tuning | S | +| **1** | Rollout (image-first) | **Feature caching** (cache-dit / diffusers-native) gated decoupled-rollout-only; **`torch.compile`** of shared transformer (lossless, everywhere); finish **attn_backend** benchmark + per-stage selection | mixed | **2–4× rollout on NFT/AWM/DGPO/CRD; ~1.2–1.8× compile everywhere** | M | +| **2** | Video long-seq | **Context/sequence parallelism** (Ulysses/Ring attention) for the DiT during rollout *and* training (Wan2, LTX2) — lossless | lossless | enables/scales long video; near-linear on seq dim | L | +| **3** | Memory→throughput | Quantize **frozen** modules only (text encoders, VAE, reference/EMA, reward models) via torchao fp8 / bnb nf4; FSDP2 selective activation checkpointing; smarter offload scheduling rollout vs train | lossless | bigger batch / larger models | M | +| **4** | Pipeline overlap | Async reward (Stage 4) + advantage (Stage 5) overlapped with next rollout; extend existing comm fusion | lossless | hide reward/comm latency | M | + +Phases 1 and 2 are the user-prioritized fronts (rollout throughput, then video parallelism). 3–4 are +follow-ons. Everything hangs off the Phase-0 plugin layer so each later track is a registry entry, not a +trainer/adapter edit. + +**Execution order requested by user:** implement all **lossless** accelerators first (torch.compile + +per-stage attention backend within the plugin layer), then the **lossy** feature-caching accelerators. + +--- + +## Proposed Architecture: `src/flow_factory/acceleration/` (model-agnostic plugin layer) + +Mirror the existing `rewards/` module shape (registry + abc + concrete impls + hparams), so it inherits +the framework's decoupling guarantees and lazy-import conventions (Constraints #1–#4). + +``` +src/flow_factory/acceleration/ + __init__.py + abc.py # BaseAccelerator: declares safety + stage; setup()/rollout_context() + registry.py # _ACCELERATOR_REGISTRY {id -> path} + get_accelerator_class() w/ direct-path fallback + validator.py # paradigm-gated safety check (fail-fast, Constraint #26) + torch_compile.py # CompileAccelerator (lossless; applied to shared transformer at init) + attention_backend.py # AttentionBackendAccelerator (lossless exact / lossy sage, per-stage) + cache_dit.py # CacheDitAccelerator (lossy; wraps cache_dit.enable_cache / disable_cache) + diffusers_cache.py # DiffusersCacheAccelerator (lossy; FirstBlockCache/FasterCache/PAB) +``` + +**`BaseAccelerator` (abc.py)** — minimal contract: +```python +class BaseAccelerator(ABC): + safety: ClassVar[Literal["lossless", "lossy"]] + stage: ClassVar[Literal["rollout", "train", "both"]] + + @abstractmethod + def setup(self, adapter: "BaseAdapter") -> None: ... # one-time (e.g. torch.compile) + @contextmanager + def rollout_context(self, adapter: "BaseAdapter"): ... # enable cache → yield → disable+reset +``` + +**Config (`hparams/acceleration_args.py`)** — new `AccelerationArguments`, aggregated into top-level +`Arguments` (Constraint #15). Per-stage so lossy stays rollout-only: +```yaml +acceleration: + rollout: { name: cache_dit, params: { policy: DBCache, rdt: 0.08 } } # lossy, rollout-only + shared: { name: torch_compile, params: { mode: default, dynamic: true } } # lossless, both stages +``` + +**Trainer paradigm tag.** Add `paradigm: ClassVar[Literal["coupled","decoupled","distillation"]]` to each +trainer (values already fixed by Constraint #7). `validator.py` reads it and the chosen accelerators' +`safety/stage` to enforce the table above before training starts. + +**Integration points (only two, both model-agnostic, no per-adapter math change):** +1. `BaseTrainer._initialization()` builds `self.rollout_accelerator` / `self.shared_accelerator` from + config via the registry, runs `validator.validate(paradigm, accelerators)`, and calls + `shared_accelerator.setup(adapter)` (e.g. compile) after `accelerator.prepare()`. +2. `BaseTrainer.generate_samples()` (`trainers/abc.py:701`) wraps the batch loop: + `with self.rollout_accelerator.rollout_context(self.adapter): ...`. The context calls + `adapter.transformer.enable_cache(...)` (reachable through `RoutedComponentProxy`) and tears it down + after, so cache state never leaks into the Stage-6 forward. + +Trainers and adapters stay unchanged except: (a) the one-line `paradigm` tag, (b) the two +`BaseTrainer` hooks above. No `optimize()` / `forward()` / `inference()` signature changes +(Constraint #12), no scheduler-to-trainer coupling. + +--- + +## Phase 1 — Detailed Implementation Spec + +**Objective:** lossless `torch.compile` + per-stage attention path usable by all algos, then 2–4× +rollout speedup on decoupled image trainers via feature caching, with a benchmark that proves reward +isn't degraded. + +### 1. Phase-0 prerequisites bundled in +- **Benchmark harness** under `.scratch/bench/` (scratch per Constraint #28) + a small reusable profiler + in `utils/`: measure rollout samples/s, mean step time, peak VRAM, and **mean/std of the reward + buffer** before vs after accel. The reward-regression check is mandatory for the lossy path. +- Build `acceleration/{abc,registry,validator}.py` and `hparams/acceleration_args.py`; wire + `AccelerationArguments` into `hparams/args.py` `Arguments`. + +### 2. Lossless first — `torch.compile` (both stages) +- **`CompileAccelerator`** (`acceleration/torch_compile.py`): in `setup(adapter)` apply + `torch.compile` to the shared transformer — prefer diffusers `transformer.compile_repeated_blocks()` + (region compile → fast warmup, robust to the variable image/seq-len shapes set by + `calculate_shift`/`set_scheduler_timesteps`). Because the same compiled module backs both + `inference()` and `forward()`, it is consistent for coupled algos too. Handle LoRA (compile after PEFT + wrap) and `dynamic=True` for variable resolution. + +### 3. Lossless — attention backend (per-stage) +- No new mechanism — extend `attn_backend` to be **per-stage** through the same accel config, and add a + benchmark sweep (FA2/FA3/xformers exact = lossless/both; Sage int8 = lossy/rollout-only). Document + results in `guidance/`. + +### 4. Lossy — feature caching (rollout-only, decoupled-only) +- **`DiffusersCacheAccelerator`** (`acceleration/diffusers_cache.py`): zero-extra-dep, uses diffusers' + built-in `transformer.enable_cache(FirstBlockCacheConfig | FasterCacheConfig | PyramidAttentionBroadcastConfig)` / `disable_cache()`. **Default lossy backend.** +- **`CacheDitAccelerator`** (`acceleration/cache_dit.py`): optional dep (`cache-dit`), imported under + `try/except ImportError` per Constraint #22(a). `rollout_context` calls + `cache_dit.enable_cache(adapter.transformer, ...)` on enter, `cache_dit.disable_cache(...)` on exit. +- **Adapter cache-readiness:** caching hooks fire only inside a `cache_context`. Adapters that already + wrap their transformer call in `transformer.cache_context(...)` (Qwen-Image `qwen_image.py:598`, Wan, + LTX2, FLUX2-Klein) are ready. Adapters that call the transformer bare (e.g. FLUX.1 `flux1.py:323`) + need the existing call wrapped in a `cache_context` — a localized, behavior-preserving change, not a + decoupling violation. **Validate first on Qwen-Image (already ready) + an NFT/AWM/DGPO config.** +- New optional dep: add `cache-dit` to `[project.optional-dependencies]` as `acceleration = [...]`. + +### 5. Files touched (Phase 1) +- **New:** `src/flow_factory/acceleration/*` (7 files), `src/flow_factory/hparams/acceleration_args.py`. +- **Edited (small):** `hparams/args.py` (aggregate args); `trainers/abc.py` (`_initialization` build+validate+setup, `generate_samples` context wrap); each trainer class (one `paradigm` ClassVar); `flux1.py`-style adapters lacking `cache_context` (wrap existing transformer call); `pyproject.toml` (optional dep); `examples/` configs gain an optional `acceleration:` block (Constraint #15/#17); `guidance/` + `.agents/knowledge/architecture.md` doc updates (registry table + new module). + +### 6. Reuse (do not reinvent) +- Registry pattern: copy `rewards/registry.py` resolution + direct-path fallback. +- Hook seam: `RoutedComponentProxy.__getattr__` already exposes `enable_cache`/`cache_context`/`compile`. +- Rollout loop already isolated in `generate_samples()`/`sample_batch()` — single wrap point. +- diffusers native caches need no new dependency. + +--- + +## Verification + +1. **Correctness gating (unit):** assert `validator` raises (Constraint #26) for `lossy` accel on a + coupled trainer and for `lossy` on the `train` stage; passes for decoupled-rollout. +2. **Numerical (lossless):** with `torch.compile` only, a short GRPO run's per-step `log_prob` and loss + match the eager baseline within tolerance (confirms coupled-safety of the shared-module path). +3. **Rollout speed + quality (lossy):** run an NFT (or AWM/DGPO) image config on Qwen-Image with + caching vs baseline for N epochs; report rollout samples/s, peak VRAM, and **reward mean/std drift** + from the harness. Accept only if speedup ≥ target and reward drift within budget. +4. **End-to-end:** `ff-train examples//lora/qwen_image/default.yaml` with and without the + `acceleration:` block trains and checkpoints identically in shape; `black --check src/ && isort + --check src/` clean (Constraint #21); run `/ff-review` before commit. +5. **Video smoke (Phase 2 entry):** confirm the same plugin enables on Wan2/LTX2 (already use + `cache_context`) before building context parallelism. + +--- + +## Open Risks / Decisions +- **cache-dit vs diffusers-native** as the default lossy backend: cache-dit has richer policies + (DBCache/TaylorSeer) but adds a dep; diffusers-native (FBCache/FasterCache/PAB) is dep-free. Plan + ships **both** behind the registry; default = diffusers-native, cache-dit opt-in. +- **FLUX.1 cache_context wrap** is the only per-adapter code change in Phase 1 — keep it behavior-identical + when caching is off (empty context). +- **Reward-drift budget** for the lossy path must be set with the user (per reward model) — the harness + measures it but the accept/reject threshold is a product decision. +- Phase 2 context parallelism interacts with FSDP/`ModelBundle` sharding and the SDE scheduler's + per-token noise (`scheduler/flow_match_euler_discrete.py`) — needs its own design pass (separate plan). From 5a72cba781e610bad78b4a7929ed6cab019194ff Mon Sep 17 00:00:00 2001 From: Jayce-Ping <315229706@qq.com> Date: Sat, 27 Jun 2026 16:04:00 +0800 Subject: [PATCH 02/15] [acceleration,trainers,hparams] feat: add model-agnostic acceleration plugin layer (lossless) Introduce a registry-based acceleration plugin layer that respects the algorithm/model decoupling, plus the first lossless accelerators. - acceleration/: BaseAccelerator (safety/stage contract), registry with direct-path fallback, paradigm-gated validator, CompileAccelerator (torch.compile, regional/full), AttentionBackendAccelerator (exact backends). - hparams: AccelerationArguments with shared/rollout slots; wired into Arguments (field + nested_map) and exported. Off by default (backward compatible). - trainers: BaseTrainer builds and validates accelerators after prepare, applies the shared accelerator via setup() and wraps the Stage-3 rollout loop with the rollout accelerator context; per-trainer paradigm tags (coupled/decoupled/ distillation) drive the lossy-safety gate (constraints.md #7, #20a, #26). --- src/flow_factory/acceleration/__init__.py | 36 ++++++ src/flow_factory/acceleration/abc.py | 114 ++++++++++++++++++ .../acceleration/attention_backend.py | 86 +++++++++++++ src/flow_factory/acceleration/registry.py | 103 ++++++++++++++++ .../acceleration/torch_compile.py | 78 ++++++++++++ src/flow_factory/acceleration/validator.py | 114 ++++++++++++++++++ src/flow_factory/hparams/__init__.py | 2 + src/flow_factory/hparams/acceleration_args.py | 80 ++++++++++++ src/flow_factory/hparams/args.py | 6 + src/flow_factory/trainers/abc.py | 79 +++++++++++- src/flow_factory/trainers/awm.py | 3 + src/flow_factory/trainers/crd.py | 3 + src/flow_factory/trainers/dgpo.py | 3 + src/flow_factory/trainers/dpo.py | 3 + src/flow_factory/trainers/grpo.py | 6 +- src/flow_factory/trainers/nft.py | 4 + src/flow_factory/trainers/opd/trainer.py | 4 + 17 files changed, 721 insertions(+), 3 deletions(-) create mode 100644 src/flow_factory/acceleration/__init__.py create mode 100644 src/flow_factory/acceleration/abc.py create mode 100644 src/flow_factory/acceleration/attention_backend.py create mode 100644 src/flow_factory/acceleration/registry.py create mode 100644 src/flow_factory/acceleration/torch_compile.py create mode 100644 src/flow_factory/acceleration/validator.py create mode 100644 src/flow_factory/hparams/acceleration_args.py diff --git a/src/flow_factory/acceleration/__init__.py b/src/flow_factory/acceleration/__init__.py new file mode 100644 index 00000000..0065acdf --- /dev/null +++ b/src/flow_factory/acceleration/__init__.py @@ -0,0 +1,36 @@ +# Copyright 2026 Jayce-Ping +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# src/flow_factory/acceleration/__init__.py +"""Model-agnostic acceleration plugin layer. + +Concrete accelerators are resolved lazily through the registry (so optional +dependencies such as ``cache-dit`` are only imported when actually requested). +""" + +from .abc import BaseAccelerator +from .registry import ( + build_accelerator, + get_accelerator_class, + list_registered_accelerators, +) +from .validator import validate_accelerator + +__all__ = [ + "BaseAccelerator", + "build_accelerator", + "get_accelerator_class", + "list_registered_accelerators", + "validate_accelerator", +] diff --git a/src/flow_factory/acceleration/abc.py b/src/flow_factory/acceleration/abc.py new file mode 100644 index 00000000..08cc1e47 --- /dev/null +++ b/src/flow_factory/acceleration/abc.py @@ -0,0 +1,114 @@ +# Copyright 2026 Jayce-Ping +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# src/flow_factory/acceleration/abc.py +"""Abstract base class for the model-agnostic acceleration plugin layer. + +An accelerator is a pluggable speedup applied to a model adapter's transformer(s) +without touching trainer or model math. Each accelerator declares two invariants +that the validator (:mod:`flow_factory.acceleration.validator`) enforces against +the trainer's RL paradigm: + +* ``safety`` — ``"lossless"`` (numerically ~identical; safe for any algorithm and + any stage because the same module backs both rollout ``inference()`` and the + training ``forward()``) or ``"lossy"`` (changes ``noise_pred``; only safe during + rollout for **decoupled / distillation** algorithms — see ``constraints.md`` #7 + and #20a). +* ``stage`` — ``"both"`` (a one-time mutation applied via :meth:`setup`, e.g. + ``torch.compile``) or ``"rollout"`` (a per-epoch context applied via + :meth:`rollout_context`, e.g. feature caching). + +Subclasses implement only what they need: ``setup`` defaults to a no-op, +``rollout_context`` defaults to yielding without modification. +""" + +from __future__ import annotations + +from abc import ABC +from contextlib import contextmanager +from typing import TYPE_CHECKING, Any, ClassVar, Iterator, Literal + +if TYPE_CHECKING: + from ..models.abc import BaseAdapter + + +class BaseAccelerator(ABC): + """Base class for all acceleration plugins. + + Subclasses MUST define the ``safety`` and ``stage`` class attributes. They + are validated at construction time by ``__init_subclass__`` so a misdeclared + accelerator fails fast at import rather than mid-training. + + Args: + **params: Accelerator-specific parameters forwarded verbatim from the + ``acceleration`` config block (e.g. ``mode`` for the compile + accelerator). Stored on ``self.params``. + """ + + safety: ClassVar[Literal["lossless", "lossy"]] + stage: ClassVar[Literal["rollout", "both"]] + + def __init_subclass__(cls, **kwargs: Any) -> None: + super().__init_subclass__(**kwargs) + # Skip intermediate ABCs that intentionally leave the markers unset. + if getattr(cls, "__abstractmethods__", None): + return + for attr in ("safety", "stage"): + if not hasattr(cls, attr): + raise TypeError( + f"Accelerator '{cls.__name__}' must define the class attribute " + f"'{attr}'. Declare it on the class body (e.g. `safety = 'lossless'`)." + ) + if cls.safety not in ("lossless", "lossy"): + raise TypeError( + f"Accelerator '{cls.__name__}' has invalid safety={cls.safety!r}; " + "expected 'lossless' or 'lossy'." + ) + if cls.stage not in ("rollout", "both"): + raise TypeError( + f"Accelerator '{cls.__name__}' has invalid stage={cls.stage!r}; " + "expected 'rollout' or 'both'." + ) + + def __init__(self, **params: Any) -> None: + self.params = params + + def setup(self, adapter: "BaseAdapter") -> None: + """Apply a one-time mutation to the adapter's prepared transformer(s). + + Called once from ``BaseTrainer._initialization`` after + ``accelerator.prepare()`` and after the routing proxies are installed, so + attribute access (``compile``, ``set_attention_backend``, ...) reaches the + inner module while forwards still route through the prepared bundle root. + + Args: + adapter: The model adapter whose transformer(s) to accelerate. + """ + return None + + @contextmanager + def rollout_context(self, adapter: "BaseAdapter") -> Iterator[None]: + """Wrap one epoch of rollout (Stage 3) with stage-scoped acceleration. + + The default implementation is a no-op. Stateful accelerators (feature + caching) enable their state on enter and tear it down on exit so nothing + leaks into the Stage-6 training forward. + + Args: + adapter: The model adapter being sampled from. + + Yields: + ``None``; the rollout loop runs inside the ``with`` block. + """ + yield diff --git a/src/flow_factory/acceleration/attention_backend.py b/src/flow_factory/acceleration/attention_backend.py new file mode 100644 index 00000000..243f37d5 --- /dev/null +++ b/src/flow_factory/acceleration/attention_backend.py @@ -0,0 +1,86 @@ +# Copyright 2026 Jayce-Ping +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# src/flow_factory/acceleration/attention_backend.py +"""Lossless attention-backend accelerator (unified config knob). + +Selects an *exact* attention backend (FlashAttention 2/3, xformers, native SDPA) +for every transformer through the same ``acceleration`` block that drives compile +and caching, complementing the lower-level ``model_args.attn_backend`` field. + +Approximate backends (e.g. ``sage``, which quantizes attention to int8) are +**lossy** and intentionally rejected here — they belong to a rollout-only lossy +accelerator so the paradigm validator can gate them. +""" + +from typing import TYPE_CHECKING + +from .abc import BaseAccelerator +from ..utils.logger_utils import setup_logger + +if TYPE_CHECKING: + from ..models.abc import BaseAdapter + +logger = setup_logger(__name__) + +# Exact, numerically-faithful diffusers attention backends. Kept conservative on +# purpose; lossy/approximate backends are excluded so this stays a lossless knob. +_EXACT_BACKENDS = frozenset( + {"native", "flash", "flash_hub", "_flash_3", "_flash_3_hub", "xformers"} +) + + +class AttentionBackendAccelerator(BaseAccelerator): + """Set an exact attention backend on every transformer component. + + Lossless and stage-``both``. Reuses diffusers' ``set_attention_backend`` (the + same mechanism behind ``model_args.attn_backend``) but routes the choice + through the unified ``acceleration`` config. + + Parameters (from ``acceleration.shared_params``): + backend: One of ``native`` / ``flash`` / ``flash_hub`` / ``_flash_3`` / + ``_flash_3_hub`` / ``xformers``. + """ + + safety = "lossless" + stage = "both" + + def setup(self, adapter: "BaseAdapter") -> None: + backend = self.params.get("backend") + if backend is None: + raise ValueError( + "AttentionBackendAccelerator requires a `backend` parameter " + f"(one of {sorted(_EXACT_BACKENDS)})." + ) + if backend not in _EXACT_BACKENDS: + raise ValueError( + f"AttentionBackendAccelerator: backend={backend!r} is not an exact (lossless) " + f"backend. Allowed: {sorted(_EXACT_BACKENDS)}. Approximate backends such as " + "'sage' are lossy and must be configured as a rollout-only accelerator." + ) + + applied = False + for name in adapter.transformer_names: + transformer = adapter.get_component(name) + if hasattr(transformer, "set_attention_backend"): + transformer.set_attention_backend(backend) + applied = True + logger.info( + "AttentionBackendAccelerator: set backend '%s' for '%s'.", backend, name + ) + if not applied: + raise ValueError( + "AttentionBackendAccelerator: no transformer component supports " + "`set_attention_backend`; check the diffusers version or adapter." + ) diff --git a/src/flow_factory/acceleration/registry.py b/src/flow_factory/acceleration/registry.py new file mode 100644 index 00000000..fd1a7f4c --- /dev/null +++ b/src/flow_factory/acceleration/registry.py @@ -0,0 +1,103 @@ +# Copyright 2026 Jayce-Ping +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# src/flow_factory/acceleration/registry.py +"""Accelerator registry: maps string identifiers to accelerator class paths. + +Mirrors the trainer / model / reward registries (see ``constraints.md`` #1-#3): +case-insensitive keys, lazy ``importlib`` resolution, and a direct-python-path +fallback so users can plug in a custom accelerator without editing this file. +""" + +from typing import Any, Dict, Type +import importlib + +from .abc import BaseAccelerator +from ..utils.logger_utils import setup_logger + +logger = setup_logger(__name__) + + +# Accelerator Registry Storage +_ACCELERATOR_REGISTRY: Dict[str, str] = { + # Lossless (safe for any algorithm, applied to the shared transformer). + 'torch_compile': 'flow_factory.acceleration.torch_compile.CompileAccelerator', + 'attention_backend': 'flow_factory.acceleration.attention_backend.AttentionBackendAccelerator', +} +_ACCELERATOR_REGISTRY = {k.lower(): v for k, v in _ACCELERATOR_REGISTRY.items()} + + +def get_accelerator_class(identifier: str) -> Type[BaseAccelerator]: + """Resolve and import an accelerator class from the registry or a python path. + + Supports two modes: + 1. Registry lookup: ``'torch_compile'`` -> ``CompileAccelerator``. + 2. Direct import: ``'my_pkg.accel.CustomAccelerator'`` -> ``CustomAccelerator``. + + Args: + identifier: Accelerator name or fully qualified class path. + + Returns: + The accelerator class (a ``BaseAccelerator`` subclass). + + Raises: + ImportError: If the accelerator cannot be loaded. + """ + identifier_lower = identifier.lower() + if identifier_lower in _ACCELERATOR_REGISTRY: + class_path = _ACCELERATOR_REGISTRY[identifier_lower] + else: + class_path = identifier + + try: + module_path, class_name = class_path.rsplit('.', 1) + module = importlib.import_module(module_path) + accelerator_class = getattr(module, class_name) + logger.debug(f"Loaded accelerator: {identifier} -> {class_name}") + return accelerator_class + except (ImportError, AttributeError, ValueError) as e: + raise ImportError( + f"Could not load accelerator '{identifier}'. " + f"Ensure it is either:\n" + f" 1. A registered accelerator: {list(_ACCELERATOR_REGISTRY.keys())}\n" + f" 2. A valid python path (e.g., 'my_package.accel.CustomAccelerator')\n" + f"Error: {e}" + ) from e + + +def build_accelerator(identifier: str, params: Dict[str, Any]) -> BaseAccelerator: + """Instantiate an accelerator from its identifier and parameters. + + Args: + identifier: Accelerator name or fully qualified class path. + params: Keyword parameters forwarded to the accelerator constructor. + + Returns: + A constructed ``BaseAccelerator`` instance. + + Raises: + TypeError: If the resolved class is not a ``BaseAccelerator`` subclass. + """ + accelerator_class = get_accelerator_class(identifier) + if not (isinstance(accelerator_class, type) and issubclass(accelerator_class, BaseAccelerator)): + raise TypeError( + f"Accelerator '{identifier}' resolved to {accelerator_class!r}, which is not a " + "BaseAccelerator subclass." + ) + return accelerator_class(**(params or {})) + + +def list_registered_accelerators() -> Dict[str, str]: + """Return a copy of the accelerator name -> class-path mapping.""" + return _ACCELERATOR_REGISTRY.copy() diff --git a/src/flow_factory/acceleration/torch_compile.py b/src/flow_factory/acceleration/torch_compile.py new file mode 100644 index 00000000..009cc892 --- /dev/null +++ b/src/flow_factory/acceleration/torch_compile.py @@ -0,0 +1,78 @@ +# Copyright 2026 Jayce-Ping +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# src/flow_factory/acceleration/torch_compile.py +"""Lossless ``torch.compile`` accelerator for the shared transformer(s).""" + +from typing import TYPE_CHECKING, Any, Dict + +from .abc import BaseAccelerator +from ..utils.logger_utils import setup_logger + +if TYPE_CHECKING: + from ..models.abc import BaseAdapter + +logger = setup_logger(__name__) + + +class CompileAccelerator(BaseAccelerator): + """Apply ``torch.compile`` to every transformer the adapter exposes. + + Lossless and stage-``both``: the compiled module backs both rollout + ``inference()`` and the training ``forward()`` (they share + ``adapter.transformer``), so numerical behavior stays consistent across the + two — safe even for coupled algorithms. + + Parameters (from ``acceleration.shared_params``): + mode: ``"regional"`` (default) compiles only the repeated transformer + blocks via diffusers' ``compile_repeated_blocks`` — fast warmup and + robust to the variable image/sequence lengths set per resolution. + ``"full"`` compiles the whole module in place. + compile_kwargs: Extra kwargs forwarded to the underlying compile call + (e.g. ``{"mode": "max-autotune", "dynamic": true}``). + """ + + safety = "lossless" + stage = "both" + + def setup(self, adapter: "BaseAdapter") -> None: + mode = self.params.get("mode", "regional") + compile_kwargs: Dict[str, Any] = self.params.get("compile_kwargs", {}) + + if mode not in ("regional", "full"): + raise ValueError( + f"CompileAccelerator: unknown mode={mode!r}; expected 'regional' or 'full'." + ) + + transformer_names = adapter.transformer_names + if not transformer_names: + raise ValueError( + "CompileAccelerator: adapter exposes no transformer components to compile." + ) + + for name in transformer_names: + module = adapter.get_component(name) + if mode == "regional": + if not hasattr(module, "compile_repeated_blocks"): + raise ValueError( + f"CompileAccelerator: component '{name}' " + f"({type(adapter._unwrap(module)).__name__}) has no " + "`compile_repeated_blocks`; use `mode: full` for whole-module compilation." + ) + module.compile_repeated_blocks(**compile_kwargs) + else: + # nn.Module.compile compiles the module's forward in place, so the + # routed bundle call hits the compiled path on subsequent forwards. + module.compile(**compile_kwargs) + logger.info("CompileAccelerator: compiled '%s' (mode=%s).", name, mode) diff --git a/src/flow_factory/acceleration/validator.py b/src/flow_factory/acceleration/validator.py new file mode 100644 index 00000000..b89e48b2 --- /dev/null +++ b/src/flow_factory/acceleration/validator.py @@ -0,0 +1,114 @@ +# Copyright 2026 Jayce-Ping +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# src/flow_factory/acceleration/validator.py +"""Paradigm-gated safety validation for accelerators (fail-fast, ``constraints.md`` #26). + +The correctness contract (``constraints.md`` #7 + #20a): + +* A ``lossy`` accelerator changes ``noise_pred`` and cannot be replicated in the + Stage-6 training forward (which needs full gradient through every block). It is + therefore only valid in the ``rollout`` slot, and only for **decoupled** / + **distillation** algorithms — for **coupled** algorithms (GRPO / GRPO-Guard / + DPPO) the rollout log-prob becomes the PPO "old log-prob" and an approximated + rollout would bias the importance ratio -> silently wrong gradients. +* A ``lossless`` accelerator is numerically ~identical and, because it mutates the + shared transformer used by both ``inference()`` and ``forward()``, is safe in + any slot for any algorithm. +""" + +from typing import Optional + +from .abc import BaseAccelerator +from ..utils.logger_utils import setup_logger + +logger = setup_logger(__name__) + +# RL paradigms that tolerate a lossy (distribution-shifting) rollout: the rollout +# trajectory's per-step log-prob does not enter the training loss. +_LOSSY_SAFE_PARADIGMS = frozenset({"decoupled", "distillation"}) + + +def validate_accelerator( + accelerator: BaseAccelerator, + *, + slot: str, + paradigm: Optional[str], + trainer_name: str, +) -> None: + """Validate one accelerator against its config slot and the trainer paradigm. + + Args: + accelerator: The constructed accelerator instance. + slot: The config slot it was placed in — ``"shared"`` (applied to both + rollout and training via :meth:`~BaseAccelerator.setup`) or + ``"rollout"`` (applied only during Stage 3 via + :meth:`~BaseAccelerator.rollout_context`). + paradigm: The trainer's RL paradigm + (``"coupled"`` / ``"decoupled"`` / ``"distillation"``), or ``None`` if + the trainer did not declare one. + trainer_name: Trainer class name, for error messages. + + Raises: + ValueError: If the accelerator is unsafe for the given slot / paradigm. + """ + name = type(accelerator).__name__ + + if slot not in ("shared", "rollout"): + raise ValueError( + f"Unknown acceleration slot {slot!r} for '{name}'; expected 'shared' or 'rollout'." + ) + + # The `shared` slot applies to BOTH rollout and the training forward, so only + # lossless accelerators may live there. + if slot == "shared": + if accelerator.safety != "lossless": + raise ValueError( + f"Accelerator '{name}' (safety='{accelerator.safety}') is configured under " + "`acceleration.shared_accelerator`, but the shared slot is applied to BOTH " + "rollout and the training forward — only lossless accelerators are allowed there. " + "Move a lossy accelerator to `acceleration.rollout_accelerator`." + ) + if accelerator.stage != "both": + raise ValueError( + f"Accelerator '{name}' (stage='{accelerator.stage}') cannot occupy the shared " + "slot; the shared slot requires a stage='both' accelerator." + ) + return + + # slot == "rollout": lossless is always fine; lossy is gated on paradigm. + if accelerator.safety == "lossy": + if paradigm is None: + raise ValueError( + f"Trainer '{trainer_name}' did not declare a `paradigm`, so a lossy rollout " + f"accelerator ('{name}') cannot be validated for correctness. Set the trainer's " + "`paradigm` class attribute to one of 'coupled' / 'decoupled' / 'distillation'." + ) + if paradigm not in _LOSSY_SAFE_PARADIGMS: + raise ValueError( + f"Lossy rollout accelerator '{name}' is unsafe for the '{paradigm}' trainer " + f"'{trainer_name}'. Lossy acceleration changes `noise_pred` during rollout but " + "cannot be replicated in the training forward; for coupled algorithms " + "(GRPO / GRPO-Guard / DPPO) this biases the PPO importance ratio and silently " + f"corrupts gradients. Allowed paradigms: {sorted(_LOSSY_SAFE_PARADIGMS)}. Use a " + "lossless accelerator (e.g. 'torch_compile') instead, or switch to a decoupled " + "algorithm (NFT / AWM / DGPO / DPO / CRD)." + ) + logger.warning( + "Lossy rollout accelerator '%s' enabled for '%s' (paradigm='%s'). This shifts the " + "generated-sample distribution; monitor the reward mean/std for regression.", + name, + trainer_name, + paradigm, + ) diff --git a/src/flow_factory/hparams/__init__.py b/src/flow_factory/hparams/__init__.py index 99ef7a35..ab21973f 100644 --- a/src/flow_factory/hparams/__init__.py +++ b/src/flow_factory/hparams/__init__.py @@ -33,6 +33,7 @@ get_training_args_class, ) from .reward_args import RewardArguments, MultiRewardArguments +from .acceleration_args import AccelerationArguments from .dataset_args import DatasetArguments, DatasetTrainSpec, DatasetEvalSpec from .log_args import LogArguments @@ -55,6 +56,7 @@ "get_training_args_class", "RewardArguments", "MultiRewardArguments", + "AccelerationArguments", "DatasetArguments", "DatasetTrainSpec", "DatasetEvalSpec", diff --git a/src/flow_factory/hparams/acceleration_args.py b/src/flow_factory/hparams/acceleration_args.py new file mode 100644 index 00000000..433ea5b9 --- /dev/null +++ b/src/flow_factory/hparams/acceleration_args.py @@ -0,0 +1,80 @@ +# Copyright 2026 Jayce-Ping +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# src/flow_factory/hparams/acceleration_args.py +import yaml +from dataclasses import dataclass, field +from typing import Any, Dict, Optional + +from .abc import ArgABC + + +@dataclass +class AccelerationArguments(ArgABC): + r"""Arguments for the model-agnostic acceleration plugin layer. + + Two independent slots, both off by default (so existing configs are + unaffected): + + * ``shared_*`` — a lossless accelerator applied to BOTH rollout and the + training forward (e.g. ``torch_compile``, ``attention_backend``). + * ``rollout_*`` — an accelerator applied only during Stage-3 rollout. May be + lossy (e.g. feature caching), in which case the trainer paradigm validator + restricts it to decoupled / distillation algorithms. + + Example YAML:: + + acceleration: + shared_accelerator: torch_compile + shared_params: { mode: regional } + rollout_accelerator: diffusers_cache + rollout_params: { policy: first_block, threshold: 0.08 } + """ + + shared_accelerator: Optional[str] = field( + default=None, + metadata={ + "help": ( + "Lossless accelerator id (or python path) applied to both rollout and training. " + "Options: 'torch_compile', 'attention_backend'. None disables it." + ) + }, + ) + shared_params: Dict[str, Any] = field( + default_factory=dict, + metadata={"help": "Keyword parameters forwarded to the shared accelerator constructor."}, + ) + rollout_accelerator: Optional[str] = field( + default=None, + metadata={ + "help": ( + "Accelerator id (or python path) applied only during Stage-3 rollout. " + "Options: 'diffusers_cache', 'cache_dit' (lossy, decoupled/distillation only), or any " + "lossless accelerator. None disables it." + ) + }, + ) + rollout_params: Dict[str, Any] = field( + default_factory=dict, + metadata={"help": "Keyword parameters forwarded to the rollout accelerator constructor."}, + ) + + def to_dict(self) -> dict[str, Any]: + return super().to_dict() + + def __str__(self) -> str: + return yaml.dump(self.to_dict(), default_flow_style=False, sort_keys=False, indent=2) + + def __repr__(self) -> str: + return self.__str__() diff --git a/src/flow_factory/hparams/args.py b/src/flow_factory/hparams/args.py index 34f64ae8..fee261d6 100644 --- a/src/flow_factory/hparams/args.py +++ b/src/flow_factory/hparams/args.py @@ -39,6 +39,7 @@ get_training_args_class, ) from .reward_args import RewardArguments, MultiRewardArguments +from .acceleration_args import AccelerationArguments from .log_args import LogArguments from ..utils.logger_utils import setup_logger from ..utils.dist import get_world_size @@ -148,6 +149,10 @@ class Arguments(ArgABC): default_factory=LogArguments, metadata={"help": "Arguments for logging configuration."}, ) + acceleration_args: AccelerationArguments = field( + default_factory=AccelerationArguments, + metadata={"help": "Arguments for the acceleration plugin layer."}, + ) reward_args: MultiRewardArguments = field( default_factory=MultiRewardArguments, metadata={"help": "Arguments for multiple reward configurations."}, @@ -1000,6 +1005,7 @@ def from_dict(cls, args_dict: dict[str, Any]) -> Arguments: 'train': ('training_args', training_args_cls), 'eval': ('eval_args', EvaluationArguments), 'log': ('log_args', LogArguments), + 'acceleration': ('acceleration_args', AccelerationArguments), 'rewards': ('reward_args', MultiRewardArguments), 'eval_rewards': ('eval_reward_args', MultiRewardArguments), } diff --git a/src/flow_factory/trainers/abc.py b/src/flow_factory/trainers/abc.py index 91bc1890..de6bfad8 100644 --- a/src/flow_factory/trainers/abc.py +++ b/src/flow_factory/trainers/abc.py @@ -16,7 +16,8 @@ import json import os from abc import ABC, abstractmethod -from typing import Dict, Any, Optional, Tuple, List, Union, Literal, Iterator +from contextlib import AbstractContextManager, nullcontext +from typing import Dict, Any, ClassVar, Optional, Tuple, List, Union, Literal, Iterator from functools import partial import numpy as np import torch @@ -40,6 +41,7 @@ ) from ..rewards import load_reward_model, BaseRewardModel, MultiRewardLoader, RewardProcessor, RewardBuffer from ..advantage import AdvantageProcessor +from ..acceleration import BaseAccelerator, build_accelerator, validate_accelerator from ..logger import load_logger, LogFormatter from ..samples import BaseSample, StackedSampleBatch from ..utils.logger_utils import setup_logger @@ -61,6 +63,13 @@ class BaseTrainer(ABC): """ Abstract Base Class for Flow-Factory trainers. """ + + # RL paradigm of this algorithm (``constraints.md`` #7). Read by the + # acceleration validator to gate lossy rollout accelerators: only + # 'decoupled' / 'distillation' trainers may use them. Concrete trainers + # MUST override this; leaving it None disables lossy acceleration. + paradigm: ClassVar[Optional[Literal["coupled", "decoupled", "distillation"]]] = None + def __init__( self, accelerator: Accelerator, @@ -367,9 +376,73 @@ def _initialization(self): # Load inference modules, excluding all bundle members (already prepared). self._load_inference_components(bundle_names) + # Build acceleration plugins now that the prepared transformer proxies are + # installed (shared accelerators like torch.compile mutate them in place). + self._init_acceleration() + # Initialize reward model self._init_reward_model() + def _init_acceleration(self): + """Build and install acceleration plugins from ``config.acceleration_args``. + + Two independent slots (both off by default): + + * ``shared_accelerator`` — a lossless accelerator applied here via + :meth:`~BaseAccelerator.setup` (e.g. ``torch.compile``); it affects both + rollout and the training forward. + * ``rollout_accelerator`` — applied per-epoch in :meth:`generate_samples` + via :meth:`~BaseAccelerator.rollout_context`; may be lossy. + + Each accelerator is validated against this trainer's ``paradigm`` before + use (fail-fast, ``constraints.md`` #26). + """ + accel_args = getattr(self.config, "acceleration_args", None) + self.shared_accelerator: Optional[BaseAccelerator] = None + self.rollout_accelerator: Optional[BaseAccelerator] = None + if accel_args is None: + return + + trainer_name = type(self).__name__ + paradigm = type(self).paradigm + + if accel_args.shared_accelerator: + accelerator = build_accelerator( + accel_args.shared_accelerator, accel_args.shared_params + ) + validate_accelerator( + accelerator, slot="shared", paradigm=paradigm, trainer_name=trainer_name + ) + accelerator.setup(self.adapter) + self.shared_accelerator = accelerator + if self.accelerator.is_main_process: + logger.info( + "Acceleration: shared accelerator '%s' (safety=%s) applied.", + accel_args.shared_accelerator, + accelerator.safety, + ) + + if accel_args.rollout_accelerator: + accelerator = build_accelerator( + accel_args.rollout_accelerator, accel_args.rollout_params + ) + validate_accelerator( + accelerator, slot="rollout", paradigm=paradigm, trainer_name=trainer_name + ) + self.rollout_accelerator = accelerator + if self.accelerator.is_main_process: + logger.info( + "Acceleration: rollout accelerator '%s' (safety=%s) enabled.", + accel_args.rollout_accelerator, + accelerator.safety, + ) + + def _rollout_acceleration(self) -> AbstractContextManager: + """Return the rollout accelerator's context, or a no-op when disabled.""" + if self.rollout_accelerator is not None: + return self.rollout_accelerator.rollout_context(self.adapter) + return nullcontext() + def _synchronize_frozen_components(self): if self.accelerator.num_processes <= 1: return @@ -758,7 +831,9 @@ def generate_samples( samples: List[BaseSample] = [] data_iter = iter(self.dataloader) - with torch.no_grad(), self.autocast(): + # Stage-3-only acceleration (e.g. feature caching) is scoped to this loop + # so its state never leaks into the Stage-6 training forward. + with self._rollout_acceleration(), torch.no_grad(), self.autocast(): for _ in tqdm( range(self.training_args.num_batches_per_epoch), desc=f'Epoch {self.epoch} Sampling', diff --git a/src/flow_factory/trainers/awm.py b/src/flow_factory/trainers/awm.py index 87a0bd3d..803d59c9 100644 --- a/src/flow_factory/trainers/awm.py +++ b/src/flow_factory/trainers/awm.py @@ -54,6 +54,9 @@ class AWMTrainer(BaseTrainer): - https://arxiv.org/pdf/2509.25050 """ + # Decoupled paradigm: lossy rollout acceleration is permitted (constraints.md #7). + paradigm = "decoupled" + def __init__(self, **kwargs): super().__init__(**kwargs) diff --git a/src/flow_factory/trainers/crd.py b/src/flow_factory/trainers/crd.py index 9135bbb9..7517192c 100644 --- a/src/flow_factory/trainers/crd.py +++ b/src/flow_factory/trainers/crd.py @@ -145,6 +145,9 @@ class CRDTrainer(BaseTrainer): Reference: https://arxiv.org/abs/2603.14128 """ + # Decoupled paradigm: lossy rollout acceleration is permitted (constraints.md #7). + paradigm = "decoupled" + _OLD_PARAMS_NAME = '_crd_old' _SAMPLING_PARAMS_NAME = '_crd_sampling' diff --git a/src/flow_factory/trainers/dgpo.py b/src/flow_factory/trainers/dgpo.py index bc544be2..3aee5f55 100644 --- a/src/flow_factory/trainers/dgpo.py +++ b/src/flow_factory/trainers/dgpo.py @@ -122,6 +122,9 @@ class DGPOTrainer(BaseTrainer): Reference: [1] DGPO: Reinforcing Diffusion Models by Direct Group Preference Optimization (ICLR 2026). """ + # Decoupled paradigm: lossy rollout acceleration is permitted (constraints.md #7). + paradigm = "decoupled" + def __init__(self, **kwargs): super().__init__(**kwargs) diff --git a/src/flow_factory/trainers/dpo.py b/src/flow_factory/trainers/dpo.py index 13deef81..6cdd0874 100644 --- a/src/flow_factory/trainers/dpo.py +++ b/src/flow_factory/trainers/dpo.py @@ -75,6 +75,9 @@ class DPOTrainer(BaseTrainer): - https://arxiv.org/abs/2311.12908 """ + # Decoupled paradigm: lossy rollout acceleration is permitted (constraints.md #7). + paradigm = "decoupled" + def __init__(self, **kwargs): super().__init__(**kwargs) self.training_args: DPOTrainingArguments diff --git a/src/flow_factory/trainers/grpo.py b/src/flow_factory/trainers/grpo.py index c87fbc21..f08a2af9 100644 --- a/src/flow_factory/trainers/grpo.py +++ b/src/flow_factory/trainers/grpo.py @@ -46,7 +46,11 @@ class GRPOTrainer(BaseTrainer): [1] Flow-GRPO: Training Flow Matching Models via Online RL - https://arxiv.org/abs/2505.05470 """ - + + # Coupled paradigm: rollout log-probs feed the PPO ratio, so lossy rollout + # acceleration is disallowed (constraints.md #7). Inherited by GRPOGuard/DPPO. + paradigm = "coupled" + def __init__(self, **kwargs): super().__init__(**kwargs) self.training_args : GRPOTrainingArguments diff --git a/src/flow_factory/trainers/nft.py b/src/flow_factory/trainers/nft.py index edad5928..543dc050 100644 --- a/src/flow_factory/trainers/nft.py +++ b/src/flow_factory/trainers/nft.py @@ -50,6 +50,10 @@ class DiffusionNFTTrainer(BaseTrainer): Reference: https://arxiv.org/abs/2509.16117 """ + # Decoupled paradigm: rollout trajectory log-probs do not enter the loss, + # so lossy rollout acceleration is permitted (constraints.md #7). + paradigm = "decoupled" + def __init__(self, **kwargs): super().__init__(**kwargs) diff --git a/src/flow_factory/trainers/opd/trainer.py b/src/flow_factory/trainers/opd/trainer.py index 2c0b60a3..df11b5fe 100644 --- a/src/flow_factory/trainers/opd/trainer.py +++ b/src/flow_factory/trainers/opd/trainer.py @@ -72,6 +72,10 @@ class DiffusionOPDTrainer(BaseTrainer): """Multi-teacher on-policy distillation trainer (ODE + SDE).""" + # Distillation paradigm: no reward/advantage stage and rollout log-probs do not + # enter the loss, so lossy rollout acceleration is permitted (constraints.md #7). + paradigm = "distillation" + def __init__(self, **kwargs: Any) -> None: super().__init__(**kwargs) self.training_args: DiffusionOPDTrainingArguments From 21e21c060a2a218b7eb575e7e999823018ceeb34 Mon Sep 17 00:00:00 2001 From: Jayce-Ping <315229706@qq.com> Date: Sat, 27 Jun 2026 16:07:20 +0800 Subject: [PATCH 03/15] [acceleration] feat: add lossy rollout feature-caching accelerators Rollout-only (Stage 3) feature caching, gated by the paradigm validator to decoupled/distillation trainers so train-inference consistency is preserved. - DiffusersCacheAccelerator: zero-extra-dep, wraps diffusers CacheMixin enable_cache/disable_cache with policies first_block (default) / faster / pyramid / taylorseer / magcache; enabled per rollout epoch, torn down on exit. - CacheDitAccelerator: optional cache-dit backend (richer DBCache/TaylorSeer), imported defensively with an install hint when absent (constraints.md #22a). - registry: register both lossy accelerators. - pyproject: add optional 'acceleration' extra (cache-dit) and include in 'all'. --- pyproject.toml | 9 +- src/flow_factory/acceleration/cache_dit.py | 77 +++++++++++ .../acceleration/diffusers_cache.py | 123 ++++++++++++++++++ src/flow_factory/acceleration/registry.py | 3 + 4 files changed, 211 insertions(+), 1 deletion(-) create mode 100644 src/flow_factory/acceleration/cache_dit.py create mode 100644 src/flow_factory/acceleration/diffusers_cache.py diff --git a/pyproject.toml b/pyproject.toml index f4fc067f..f183a1e0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -64,6 +64,13 @@ quantization = [ "bitsandbytes>=0.45.3", ] +# Rollout acceleration: richer feature-caching policies (DBCache / TaylorSeer). +# Diffusers-native caching (the `diffusers_cache` accelerator) needs nothing extra; +# this extra is only for the `cache_dit` accelerator. +acceleration = [ + "cache-dit>=0.2.0", +] + # Bagel adapter runtime (flash-attn varlen kernels + opencv transforms) bagel = [ "flash-attn>=2.5.8", @@ -98,7 +105,7 @@ geneval2 = [ # flow-factory: uv pip install hpsv2 --no-deps (runtime works with protobuf 6+). all = [ - "flow-factory[deepspeed,quantization]", + "flow-factory[deepspeed,quantization,acceleration]", ] [project.scripts] diff --git a/src/flow_factory/acceleration/cache_dit.py b/src/flow_factory/acceleration/cache_dit.py new file mode 100644 index 00000000..1b55d57b --- /dev/null +++ b/src/flow_factory/acceleration/cache_dit.py @@ -0,0 +1,77 @@ +# Copyright 2026 Jayce-Ping +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# src/flow_factory/acceleration/cache_dit.py +"""Lossy rollout-only feature caching via the optional ``cache-dit`` library. + +`cache-dit `_ offers richer cache policies +(DBCache, TaylorSeer, ...) than diffusers' built-ins. It is an optional dependency +(``pip install flow-factory[acceleration]``); when absent, this accelerator raises +a clear install hint instead of failing at import (``constraints.md`` #22a). + +Only valid in the rollout slot of a decoupled / distillation trainer — the +paradigm validator enforces this (``constraints.md`` #7). +""" + +from contextlib import contextmanager +from typing import TYPE_CHECKING, Iterator + +from .abc import BaseAccelerator +from ..utils.logger_utils import setup_logger + +try: + import cache_dit +except ImportError: + cache_dit = None + +if TYPE_CHECKING: + from ..models.abc import BaseAdapter + +logger = setup_logger(__name__) + + +class CacheDitAccelerator(BaseAccelerator): + """Enable cache-dit caching on the adapter pipeline during rollout. + + Lossy and rollout-scoped. All params are forwarded verbatim to + ``cache_dit.enable_cache`` (see the cache-dit docs for policy options such as + ``cache_type``/``Fn_compute_blocks``/``rdt``), and cleared on context exit so + the Stage-6 training forward stays exact. + + cache-dit operates on ``adapter.pipeline``, whose ``transformer`` attribute is + the unwrapped inner module; rollout forwards still route through it via the + prepared bundle, so the cache hooks fire as expected. + """ + + safety = "lossy" + stage = "rollout" + + def __init__(self, **params) -> None: + super().__init__(**params) + if cache_dit is None: + raise ImportError( + "CacheDitAccelerator requires the optional 'cache-dit' package. " + "Install it with `pip install flow-factory[acceleration]` (or " + "`pip install cache-dit`)." + ) + + @contextmanager + def rollout_context(self, adapter: "BaseAdapter") -> Iterator[None]: + pipeline = adapter.pipeline + cache_dit.enable_cache(pipeline, **self.params) + logger.info("CacheDitAccelerator: cache-dit enabled on the rollout pipeline.") + try: + yield + finally: + cache_dit.disable_cache(pipeline) diff --git a/src/flow_factory/acceleration/diffusers_cache.py b/src/flow_factory/acceleration/diffusers_cache.py new file mode 100644 index 00000000..47993f83 --- /dev/null +++ b/src/flow_factory/acceleration/diffusers_cache.py @@ -0,0 +1,123 @@ +# Copyright 2026 Jayce-Ping +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# src/flow_factory/acceleration/diffusers_cache.py +"""Lossy rollout-only feature caching via diffusers' native ``CacheMixin``. + +Zero extra dependency: diffusers transformers already inherit ``enable_cache`` / +``disable_cache`` (they are ``CacheMixin`` models — the same machinery behind the +``cache_context(...)`` calls Flow-Factory adapters already use). This accelerator +enables a cache policy for the duration of one rollout epoch and tears it down on +exit so the Stage-6 training forward stays exact. + +Only valid in the rollout slot of a decoupled / distillation trainer — the +paradigm validator enforces this (``constraints.md`` #7). +""" + +from contextlib import contextmanager +from typing import TYPE_CHECKING, Iterator, List + +from diffusers.hooks import ( + FasterCacheConfig, + FirstBlockCacheConfig, + MagCacheConfig, + PyramidAttentionBroadcastConfig, + TaylorSeerCacheConfig, +) + +from .abc import BaseAccelerator +from ..utils.logger_utils import setup_logger + +if TYPE_CHECKING: + import torch + + from ..models.abc import BaseAdapter + +logger = setup_logger(__name__) + +# Map the user-facing ``policy`` string to its diffusers config class. +_POLICY_CONFIGS = { + "first_block": FirstBlockCacheConfig, + "faster": FasterCacheConfig, + "pyramid": PyramidAttentionBroadcastConfig, + "taylorseer": TaylorSeerCacheConfig, + "magcache": MagCacheConfig, +} + + +class DiffusersCacheAccelerator(BaseAccelerator): + """Enable a diffusers cache policy on every transformer during rollout. + + Lossy and rollout-scoped. The default policy is ``first_block`` (FirstBlockCache, + aka FBCache) which is robust across models and needs only a single ``threshold``. + + Parameters (from ``acceleration.rollout_params``): + policy: One of ``first_block`` / ``faster`` / ``pyramid`` / ``taylorseer`` / + ``magcache``. Defaults to ``first_block``. + : All remaining params are forwarded verbatim to the selected + diffusers config class (e.g. ``threshold`` for FirstBlockCache). + + Note: + Caching reuses block outputs across denoising steps, so the per-step + ``cache_context`` the adapter opens around its transformer call must persist + cache state across the loop. Adapters that already wrap their forward in + ``transformer.cache_context(...)`` (Qwen-Image, Wan2, LTX2, FLUX.2-Klein) + are cache-ready; verify the reward distribution before/after enabling. + """ + + safety = "lossy" + stage = "rollout" + + def _build_config(self): + params = dict(self.params) + policy = params.pop("policy", "first_block") + if policy not in _POLICY_CONFIGS: + raise ValueError( + f"DiffusersCacheAccelerator: unknown policy={policy!r}; " + f"expected one of {sorted(_POLICY_CONFIGS)}." + ) + config_cls = _POLICY_CONFIGS[policy] + try: + return config_cls(**params) + except TypeError as e: + raise ValueError( + f"DiffusersCacheAccelerator: invalid parameters {sorted(params)} for policy " + f"'{policy}' ({config_cls.__name__}): {e}" + ) from e + + @contextmanager + def rollout_context(self, adapter: "BaseAdapter") -> Iterator[None]: + transformer_names = adapter.transformer_names + if not transformer_names: + raise ValueError("DiffusersCacheAccelerator: adapter exposes no transformer to cache.") + + enabled: List["torch.nn.Module"] = [] + try: + for name in transformer_names: + transformer = adapter.get_component(name) + if not hasattr(transformer, "enable_cache"): + raise ValueError( + f"DiffusersCacheAccelerator: component '{name}' is not a diffusers " + "CacheMixin (no `enable_cache`); use a different accelerator." + ) + # Defensive: clear any stale cache left enabled by a prior epoch. + if getattr(transformer, "is_cache_enabled", False): + transformer.disable_cache() + transformer.enable_cache(self._build_config()) + enabled.append(transformer) + logger.info("DiffusersCacheAccelerator: cache enabled for '%s'.", name) + yield + finally: + for transformer in enabled: + transformer.disable_cache() diff --git a/src/flow_factory/acceleration/registry.py b/src/flow_factory/acceleration/registry.py index fd1a7f4c..46d88532 100644 --- a/src/flow_factory/acceleration/registry.py +++ b/src/flow_factory/acceleration/registry.py @@ -34,6 +34,9 @@ # Lossless (safe for any algorithm, applied to the shared transformer). 'torch_compile': 'flow_factory.acceleration.torch_compile.CompileAccelerator', 'attention_backend': 'flow_factory.acceleration.attention_backend.AttentionBackendAccelerator', + # Lossy (rollout-only; validator restricts to decoupled / distillation algos). + 'diffusers_cache': 'flow_factory.acceleration.diffusers_cache.DiffusersCacheAccelerator', + 'cache_dit': 'flow_factory.acceleration.cache_dit.CacheDitAccelerator', } _ACCELERATOR_REGISTRY = {k.lower(): v for k, v in _ACCELERATOR_REGISTRY.items()} From ca6908f57c842c83698c5d416f9248f5d6304e8f Mon Sep 17 00:00:00 2001 From: Jayce-Ping <315229706@qq.com> Date: Sat, 27 Jun 2026 16:09:45 +0800 Subject: [PATCH 04/15] [docs,examples] docs: document acceleration plugin layer - guidance/acceleration.md: safety model, config schema, accelerator table, model cache-readiness, and how to add a new accelerator. - architecture.md: register the acceleration registry (now four), accelerator table, validator/paradigm gating note, and the new-accelerator extension point. - AGENTS.md: link the new guide. - examples: commented acceleration block in the NFT/Qwen-Image (decoupled) config. --- .agents/knowledge/architecture.md | 13 +++- AGENTS.md | 1 + .../lora/qwen_image/rational_rewards_t2i.yaml | 8 ++ guidance/acceleration.md | 74 +++++++++++++++++++ 4 files changed, 95 insertions(+), 1 deletion(-) create mode 100644 guidance/acceleration.md diff --git a/.agents/knowledge/architecture.md b/.agents/knowledge/architecture.md index 1507f835..567847fa 100644 --- a/.agents/knowledge/architecture.md +++ b/.agents/knowledge/architecture.md @@ -92,7 +92,7 @@ Stage 6: Policy Optimization ## Registry System -All three registries map string keys → lazy import paths. Resolution: registry lookup → fallback to direct Python path → dynamic import. See `trainers/registry.py`, `models/registry.py`, `rewards/registry.py` for implementation. +All four registries map string keys → lazy import paths. Resolution: registry lookup → fallback to direct Python path → dynamic import. See `trainers/registry.py`, `models/registry.py`, `rewards/registry.py`, `acceleration/registry.py` for implementation. ### Registered Components @@ -147,6 +147,16 @@ All three registries map string keys → lazy import paths. Resolution: registry | `hpsv2` | `HPSv2RewardModel` | Pointwise | | `qwen_image_bench` | `QwenImageBenchRewardModel` | Pointwise | +**Accelerators** (`acceleration/registry.py`): +| Key | Class | Safety | Stage | Notes | +|-----|-------|--------|-------|-------| +| `torch_compile` | `CompileAccelerator` | lossless | both | `torch.compile` of the shared transformer (regional/full). | +| `attention_backend` | `AttentionBackendAccelerator` | lossless | both | Sets an exact diffusers attention backend (FA2/FA3/xformers/native). | +| `diffusers_cache` | `DiffusersCacheAccelerator` | lossy | rollout | Diffusers `CacheMixin` feature caching (first_block/faster/pyramid/taylorseer/magcache). | +| `cache_dit` | `CacheDitAccelerator` | lossy | rollout | Optional `cache-dit` backend (DBCache/TaylorSeer). | + +Configured via the `acceleration:` block (`hparams/acceleration_args.py`): `shared_accelerator` (lossless, both stages) and `rollout_accelerator` (Stage-3 only). The `acceleration/validator.py` enforces that lossy accelerators are rollout-only and only on `decoupled`/`distillation` trainers (each trainer declares a `paradigm`), preserving train-inference consistency (constraints #7, #20a). Off by default. + --- ## Extension Points @@ -154,6 +164,7 @@ All three registries map string keys → lazy import paths. Resolution: registry - **New model adapter**: `guidance/new_model.md`, skill `/ff-new-model`, conventions `topics/adapter_conventions.md` - **New reward model**: `guidance/rewards.md`, skill `/ff-new-reward` - **New algorithm**: `guidance/algorithms.md`, skill `/ff-new-algorithm` +- **New accelerator**: subclass `acceleration/abc.py::BaseAccelerator` (declare `safety`/`stage`), register in `acceleration/registry.py` --- diff --git a/AGENTS.md b/AGENTS.md index 9a56e323..88437123 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -63,6 +63,7 @@ See `.agents/knowledge/architecture.md` "Module Dependency Graph" for full detai | `guidance/algorithms.md` | All 9 algorithms (GRPO, GRPO-Guard, DPPO, DPO, DGPO, DiffusionNFT, AWM, CRD, DiffusionOPD) deep dive | | `guidance/rewards.md` | Reward system design, custom model creation | | `guidance/new_model.md` | Step-by-step model adapter integration | +| `guidance/acceleration.md` | Acceleration plugin layer (compile, attention backend, feature caching) | ## Available Skills diff --git a/examples/nft/lora/qwen_image/rational_rewards_t2i.yaml b/examples/nft/lora/qwen_image/rational_rewards_t2i.yaml index bfffb2c1..0c8f7d6a 100644 --- a/examples/nft/lora/qwen_image/rational_rewards_t2i.yaml +++ b/examples/nft/lora/qwen_image/rational_rewards_t2i.yaml @@ -82,6 +82,14 @@ train: scheduler: dynamics_type: "ODE" +# Optional rollout/throughput acceleration (off by default). See guidance/acceleration.md. +# NFT is a decoupled trainer, so lossy rollout feature-caching is allowed here. +# acceleration: +# shared_accelerator: "torch_compile" # lossless, both stages. Options: torch_compile, attention_backend +# shared_params: { mode: "regional" } # regional (compile_repeated_blocks) | full +# rollout_accelerator: "diffusers_cache" # lossy, rollout-only. Options: diffusers_cache, cache_dit +# rollout_params: { policy: "first_block", threshold: 0.08 } + eval: resolution: 512 condition_image_size: 512 diff --git a/guidance/acceleration.md b/guidance/acceleration.md new file mode 100644 index 00000000..a22079b2 --- /dev/null +++ b/guidance/acceleration.md @@ -0,0 +1,74 @@ +# Acceleration + +Flow-Factory's acceleration layer is a **model-agnostic, registry-based plugin system** +that speeds up training without touching trainer or model math. It lives in +`src/flow_factory/acceleration/` and mirrors the reward/model/trainer registries. + +The dominant cost of online RL fine-tuning is **Stage 3 — rollout** (multi-step denoising +under `torch.no_grad()`), so that is where most accelerators apply. + +## Safety model (read this first) + +Every accelerator declares two markers, and a validator (`acceleration/validator.py`) +enforces them against the trainer's `paradigm` before training starts (fail-fast): + +| `safety` | Meaning | Where allowed | +|----------|---------|---------------| +| `lossless` | Numerically ~identical; applied to the transformer shared by both rollout `inference()` and training `forward()`. | Any algorithm, any stage. | +| `lossy` | Changes `noise_pred` (e.g. feature caching). Cannot be replicated in the training forward. | **Rollout only**, and only on `decoupled` / `distillation` trainers. | + +Why the restriction: for **coupled** algorithms (GRPO, GRPO-Guard, DPPO) the rollout's +per-step log-prob becomes the PPO "old log-prob". Approximating the rollout while the +training forward stays exact biases the importance ratio and silently corrupts gradients +(`.agents/knowledge/constraints.md` #7, #20a). For **decoupled** algorithms (NFT, AWM, +DGPO, DPO, CRD) and **distillation** (diffusion-opd), the rollout trajectory's log-prob +does not enter the loss, so lossy caching only shifts the generated-sample distribution — +acceptable and tunable. Monitor the reward mean/std when enabling a lossy accelerator. + +## Configuration + +Add an optional `acceleration:` block to any config. Two independent slots, both off by +default: + +```yaml +acceleration: + # Lossless, applied to BOTH rollout and the training forward. + shared_accelerator: "torch_compile" # torch_compile | attention_backend + shared_params: { mode: "regional" } # regional (compile_repeated_blocks) | full + + # Rollout-only (Stage 3). May be lossy (paradigm-gated). + rollout_accelerator: "diffusers_cache" # diffusers_cache | cache_dit | + rollout_params: { policy: "first_block", threshold: 0.08 } +``` + +Either slot may be omitted. A direct python path (e.g. `my_pkg.accel.MyAccelerator`) is +accepted in place of a registered id. + +## Available accelerators + +| id | safety | stage | Notes | +|----|--------|-------|-------| +| `torch_compile` | lossless | both | `torch.compile` of the shared transformer. `mode: regional` uses diffusers' `compile_repeated_blocks` (fast warmup, robust to variable resolution); `mode: full` compiles the whole module. Extra `compile_kwargs` forwarded to the compile call. | +| `attention_backend` | lossless | both | Sets an exact diffusers attention backend on every transformer. `backend`: `native` / `flash` / `flash_hub` / `_flash_3` / `_flash_3_hub` / `xformers`. Complements `model.attn_backend`. | +| `diffusers_cache` | lossy | rollout | Diffusers-native feature caching (no extra dependency). `policy`: `first_block` (default) / `faster` / `pyramid` / `taylorseer` / `magcache`; remaining params forwarded to the policy's diffusers config (e.g. `threshold`). | +| `cache_dit` | lossy | rollout | [cache-dit](https://github.com/vipshop/cache-dit) backend (DBCache/TaylorSeer). Requires `pip install flow-factory[acceleration]`. All params forwarded to `cache_dit.enable_cache`. | + +## Model cache-readiness (lossy caching) + +Feature caching reuses block outputs across denoising steps via the transformer's +`cache_context(...)`. Adapters that already wrap their transformer call in +`transformer.cache_context(...)` — Qwen-Image, Qwen-Image-Edit-Plus, Wan2, LTX2, +FLUX.2-Klein — are cache-ready. Adapters that call the transformer bare (e.g. FLUX.1) +need their forward wrapped in a `cache_context` first. Validate the reward distribution +before/after enabling on a new model. + +`torch_compile` and `attention_backend` are model-agnostic and apply to every adapter. + +## Adding a new accelerator + +1. Subclass `acceleration/abc.py::BaseAccelerator`; set the `safety` and `stage` class + attributes. Implement `setup()` (one-time mutation, e.g. compile) and/or + `rollout_context()` (per-epoch context, e.g. caching). +2. Register the id → class path in `acceleration/registry.py`. +3. Optional dependencies must be imported defensively (`try/except ImportError`) per + constraint #22(a). From 1bb70f5c582f4b9c0262a2f9077a4fe02af4ae95 Mon Sep 17 00:00:00 2001 From: Jayce-Ping <315229706@qq.com> Date: Sat, 27 Jun 2026 16:28:10 +0800 Subject: [PATCH 05/15] [acceleration,trainers] fix: remove duplicate attention-backend accelerator; compile after post_init Two correctness/cleanup fixes from review: 1. Remove AttentionBackendAccelerator: it duplicated model.attn_backend (BaseAdapter._set_attention_backend, run in adapter __init__) with no added capability, and re-set the backend a second time after prepare. Attention backend stays a single canonical knob (model.attn_backend). Drop the registry entry, file, and docs/example references. 2. Apply the shared (lossless) accelerator AFTER adapter.post_init() instead of inside _initialization(). post_init() performs state-checkpoint resume, _init_ema(), and _init_ref_parameters(); compiling first was fragile. Now _init_acceleration() only builds+validates; _apply_shared_acceleration() runs torch.compile after weights are final. In-place compile (nn.Module.compile / compile_repeated_blocks) keeps state_dict keys and parameter identity stable, so save/load checkpoint, LoRA, and copy_-based EMA/ref/named-param swaps remain correct. Also fix the regional-mode guard to check _repeated_blocks (the method always exists on ModelMixin) and document the LoRA disable_adapter recompile. --- .agents/knowledge/architecture.md | 5 +- .../lora/qwen_image/rational_rewards_t2i.yaml | 2 +- guidance/acceleration.md | 12 ++- .../acceleration/attention_backend.py | 86 ------------------- src/flow_factory/acceleration/registry.py | 4 +- .../acceleration/torch_compile.py | 29 +++++-- src/flow_factory/hparams/acceleration_args.py | 5 +- src/flow_factory/trainers/abc.py | 46 +++++++--- 8 files changed, 74 insertions(+), 115 deletions(-) delete mode 100644 src/flow_factory/acceleration/attention_backend.py diff --git a/.agents/knowledge/architecture.md b/.agents/knowledge/architecture.md index 567847fa..a05798dd 100644 --- a/.agents/knowledge/architecture.md +++ b/.agents/knowledge/architecture.md @@ -150,12 +150,11 @@ All four registries map string keys → lazy import paths. Resolution: registry **Accelerators** (`acceleration/registry.py`): | Key | Class | Safety | Stage | Notes | |-----|-------|--------|-------|-------| -| `torch_compile` | `CompileAccelerator` | lossless | both | `torch.compile` of the shared transformer (regional/full). | -| `attention_backend` | `AttentionBackendAccelerator` | lossless | both | Sets an exact diffusers attention backend (FA2/FA3/xformers/native). | +| `torch_compile` | `CompileAccelerator` | lossless | both | `torch.compile` of the shared transformer (regional/full); applied in-place after `post_init` so checkpoint keys / param identity stay stable. | | `diffusers_cache` | `DiffusersCacheAccelerator` | lossy | rollout | Diffusers `CacheMixin` feature caching (first_block/faster/pyramid/taylorseer/magcache). | | `cache_dit` | `CacheDitAccelerator` | lossy | rollout | Optional `cache-dit` backend (DBCache/TaylorSeer). | -Configured via the `acceleration:` block (`hparams/acceleration_args.py`): `shared_accelerator` (lossless, both stages) and `rollout_accelerator` (Stage-3 only). The `acceleration/validator.py` enforces that lossy accelerators are rollout-only and only on `decoupled`/`distillation` trainers (each trainer declares a `paradigm`), preserving train-inference consistency (constraints #7, #20a). Off by default. +Configured via the `acceleration:` block (`hparams/acceleration_args.py`): `shared_accelerator` (lossless, both stages) and `rollout_accelerator` (Stage-3 only). The `acceleration/validator.py` enforces that lossy accelerators are rollout-only and only on `decoupled`/`distillation` trainers (each trainer declares a `paradigm`), preserving train-inference consistency (constraints #7, #20a). Off by default. Attention-backend selection is **not** an accelerator — use `model.attn_backend` (`BaseAdapter._set_attention_backend`). --- diff --git a/examples/nft/lora/qwen_image/rational_rewards_t2i.yaml b/examples/nft/lora/qwen_image/rational_rewards_t2i.yaml index 0c8f7d6a..a6433712 100644 --- a/examples/nft/lora/qwen_image/rational_rewards_t2i.yaml +++ b/examples/nft/lora/qwen_image/rational_rewards_t2i.yaml @@ -85,7 +85,7 @@ scheduler: # Optional rollout/throughput acceleration (off by default). See guidance/acceleration.md. # NFT is a decoupled trainer, so lossy rollout feature-caching is allowed here. # acceleration: -# shared_accelerator: "torch_compile" # lossless, both stages. Options: torch_compile, attention_backend +# shared_accelerator: "torch_compile" # lossless, both stages (attention backend -> model.attn_backend) # shared_params: { mode: "regional" } # regional (compile_repeated_blocks) | full # rollout_accelerator: "diffusers_cache" # lossy, rollout-only. Options: diffusers_cache, cache_dit # rollout_params: { policy: "first_block", threshold: 0.08 } diff --git a/guidance/acceleration.md b/guidance/acceleration.md index a22079b2..6c5da3eb 100644 --- a/guidance/acceleration.md +++ b/guidance/acceleration.md @@ -33,7 +33,7 @@ default: ```yaml acceleration: # Lossless, applied to BOTH rollout and the training forward. - shared_accelerator: "torch_compile" # torch_compile | attention_backend + shared_accelerator: "torch_compile" # torch_compile shared_params: { mode: "regional" } # regional (compile_repeated_blocks) | full # Rollout-only (Stage 3). May be lossy (paradigm-gated). @@ -48,11 +48,15 @@ accepted in place of a registered id. | id | safety | stage | Notes | |----|--------|-------|-------| -| `torch_compile` | lossless | both | `torch.compile` of the shared transformer. `mode: regional` uses diffusers' `compile_repeated_blocks` (fast warmup, robust to variable resolution); `mode: full` compiles the whole module. Extra `compile_kwargs` forwarded to the compile call. | -| `attention_backend` | lossless | both | Sets an exact diffusers attention backend on every transformer. `backend`: `native` / `flash` / `flash_hub` / `_flash_3` / `_flash_3_hub` / `xformers`. Complements `model.attn_backend`. | +| `torch_compile` | lossless | both | `torch.compile` of the shared transformer. `mode: regional` uses diffusers' `compile_repeated_blocks` (fast warmup, robust to variable resolution); `mode: full` compiles the whole module. Extra `compile_kwargs` forwarded to the compile call. Compiles in place (checkpoint- and EMA/ref-safe), applied after `post_init`. | | `diffusers_cache` | lossy | rollout | Diffusers-native feature caching (no extra dependency). `policy`: `first_block` (default) / `faster` / `pyramid` / `taylorseer` / `magcache`; remaining params forwarded to the policy's diffusers config (e.g. `threshold`). | | `cache_dit` | lossy | rollout | [cache-dit](https://github.com/vipshop/cache-dit) backend (DBCache/TaylorSeer). Requires `pip install flow-factory[acceleration]`. All params forwarded to `cache_dit.enable_cache`. | +> Attention-backend selection (FlashAttention 2/3, xformers, native SDPA) is **not** an +> accelerator — set it once via `model.attn_backend` (handled by +> `BaseAdapter._set_attention_backend`), which already applies to every transformer +> framework-wide and is lossless. + ## Model cache-readiness (lossy caching) Feature caching reuses block outputs across denoising steps via the transformer's @@ -62,7 +66,7 @@ FLUX.2-Klein — are cache-ready. Adapters that call the transformer bare (e.g. need their forward wrapped in a `cache_context` first. Validate the reward distribution before/after enabling on a new model. -`torch_compile` and `attention_backend` are model-agnostic and apply to every adapter. +`torch_compile` is model-agnostic and applies to every adapter. ## Adding a new accelerator diff --git a/src/flow_factory/acceleration/attention_backend.py b/src/flow_factory/acceleration/attention_backend.py deleted file mode 100644 index 243f37d5..00000000 --- a/src/flow_factory/acceleration/attention_backend.py +++ /dev/null @@ -1,86 +0,0 @@ -# Copyright 2026 Jayce-Ping -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -# src/flow_factory/acceleration/attention_backend.py -"""Lossless attention-backend accelerator (unified config knob). - -Selects an *exact* attention backend (FlashAttention 2/3, xformers, native SDPA) -for every transformer through the same ``acceleration`` block that drives compile -and caching, complementing the lower-level ``model_args.attn_backend`` field. - -Approximate backends (e.g. ``sage``, which quantizes attention to int8) are -**lossy** and intentionally rejected here — they belong to a rollout-only lossy -accelerator so the paradigm validator can gate them. -""" - -from typing import TYPE_CHECKING - -from .abc import BaseAccelerator -from ..utils.logger_utils import setup_logger - -if TYPE_CHECKING: - from ..models.abc import BaseAdapter - -logger = setup_logger(__name__) - -# Exact, numerically-faithful diffusers attention backends. Kept conservative on -# purpose; lossy/approximate backends are excluded so this stays a lossless knob. -_EXACT_BACKENDS = frozenset( - {"native", "flash", "flash_hub", "_flash_3", "_flash_3_hub", "xformers"} -) - - -class AttentionBackendAccelerator(BaseAccelerator): - """Set an exact attention backend on every transformer component. - - Lossless and stage-``both``. Reuses diffusers' ``set_attention_backend`` (the - same mechanism behind ``model_args.attn_backend``) but routes the choice - through the unified ``acceleration`` config. - - Parameters (from ``acceleration.shared_params``): - backend: One of ``native`` / ``flash`` / ``flash_hub`` / ``_flash_3`` / - ``_flash_3_hub`` / ``xformers``. - """ - - safety = "lossless" - stage = "both" - - def setup(self, adapter: "BaseAdapter") -> None: - backend = self.params.get("backend") - if backend is None: - raise ValueError( - "AttentionBackendAccelerator requires a `backend` parameter " - f"(one of {sorted(_EXACT_BACKENDS)})." - ) - if backend not in _EXACT_BACKENDS: - raise ValueError( - f"AttentionBackendAccelerator: backend={backend!r} is not an exact (lossless) " - f"backend. Allowed: {sorted(_EXACT_BACKENDS)}. Approximate backends such as " - "'sage' are lossy and must be configured as a rollout-only accelerator." - ) - - applied = False - for name in adapter.transformer_names: - transformer = adapter.get_component(name) - if hasattr(transformer, "set_attention_backend"): - transformer.set_attention_backend(backend) - applied = True - logger.info( - "AttentionBackendAccelerator: set backend '%s' for '%s'.", backend, name - ) - if not applied: - raise ValueError( - "AttentionBackendAccelerator: no transformer component supports " - "`set_attention_backend`; check the diffusers version or adapter." - ) diff --git a/src/flow_factory/acceleration/registry.py b/src/flow_factory/acceleration/registry.py index 46d88532..bf537fd5 100644 --- a/src/flow_factory/acceleration/registry.py +++ b/src/flow_factory/acceleration/registry.py @@ -32,8 +32,10 @@ # Accelerator Registry Storage _ACCELERATOR_REGISTRY: Dict[str, str] = { # Lossless (safe for any algorithm, applied to the shared transformer). + # NOTE: attention-backend selection is handled by `model.attn_backend` + # (BaseAdapter._set_attention_backend); it is intentionally NOT an accelerator + # here to avoid a second, redundant set after prepare. 'torch_compile': 'flow_factory.acceleration.torch_compile.CompileAccelerator', - 'attention_backend': 'flow_factory.acceleration.attention_backend.AttentionBackendAccelerator', # Lossy (rollout-only; validator restricts to decoupled / distillation algos). 'diffusers_cache': 'flow_factory.acceleration.diffusers_cache.DiffusersCacheAccelerator', 'cache_dit': 'flow_factory.acceleration.cache_dit.CacheDitAccelerator', diff --git a/src/flow_factory/acceleration/torch_compile.py b/src/flow_factory/acceleration/torch_compile.py index 009cc892..3cde20bb 100644 --- a/src/flow_factory/acceleration/torch_compile.py +++ b/src/flow_factory/acceleration/torch_compile.py @@ -34,11 +34,26 @@ class CompileAccelerator(BaseAccelerator): ``adapter.transformer``), so numerical behavior stays consistent across the two — safe even for coupled algorithms. + Both modes compile **in place** (``nn.Module.compile`` / + ``compile_repeated_blocks``), which preserves parameter identity and leaves + ``state_dict`` keys unchanged (no ``_orig_mod.`` prefix). Consequences: + + * **Checkpointing** stays correct — save/load operate on the same keys. + * **EMA / reference / named-parameter swaps** stay correct — they mutate + ``param.data`` in place (``copy_``), which the compiled graph reads. + * **LoRA reference forwards** (``use_ref_parameters`` -> PEFT + ``disable_adapter``) toggle control flow, so Dynamo recompiles the + adapter-disabled path once; this is correct but adds a one-time recompile. + + Applied after ``post_init`` (see ``BaseTrainer._apply_shared_acceleration``), + so it wraps the final, fully-loaded weights. + Parameters (from ``acceleration.shared_params``): mode: ``"regional"`` (default) compiles only the repeated transformer blocks via diffusers' ``compile_repeated_blocks`` — fast warmup and robust to the variable image/sequence lengths set per resolution. - ``"full"`` compiles the whole module in place. + Requires the transformer to declare ``_repeated_blocks``. ``"full"`` + compiles the whole module in place. compile_kwargs: Extra kwargs forwarded to the underlying compile call (e.g. ``{"mode": "max-autotune", "dynamic": true}``). """ @@ -64,11 +79,15 @@ def setup(self, adapter: "BaseAdapter") -> None: for name in transformer_names: module = adapter.get_component(name) if mode == "regional": - if not hasattr(module, "compile_repeated_blocks"): + # `compile_repeated_blocks` exists on every diffusers ModelMixin but + # only works when the model declares `_repeated_blocks`; check the + # unwrapped module so the error is actionable. + inner = adapter._unwrap(module) + if not getattr(inner, "_repeated_blocks", None): raise ValueError( - f"CompileAccelerator: component '{name}' " - f"({type(adapter._unwrap(module)).__name__}) has no " - "`compile_repeated_blocks`; use `mode: full` for whole-module compilation." + f"CompileAccelerator: component '{name}' ({type(inner).__name__}) does " + "not declare `_repeated_blocks`, so regional compilation is unavailable. " + "Use `mode: full` for whole-module compilation." ) module.compile_repeated_blocks(**compile_kwargs) else: diff --git a/src/flow_factory/hparams/acceleration_args.py b/src/flow_factory/hparams/acceleration_args.py index 433ea5b9..4888675f 100644 --- a/src/flow_factory/hparams/acceleration_args.py +++ b/src/flow_factory/hparams/acceleration_args.py @@ -28,7 +28,7 @@ class AccelerationArguments(ArgABC): unaffected): * ``shared_*`` — a lossless accelerator applied to BOTH rollout and the - training forward (e.g. ``torch_compile``, ``attention_backend``). + training forward (e.g. ``torch_compile``). * ``rollout_*`` — an accelerator applied only during Stage-3 rollout. May be lossy (e.g. feature caching), in which case the trainer paradigm validator restricts it to decoupled / distillation algorithms. @@ -47,7 +47,8 @@ class AccelerationArguments(ArgABC): metadata={ "help": ( "Lossless accelerator id (or python path) applied to both rollout and training. " - "Options: 'torch_compile', 'attention_backend'. None disables it." + "Options: 'torch_compile'. None disables it. " + "(Attention backend is set via model.attn_backend, not here.)" ) }, ) diff --git a/src/flow_factory/trainers/abc.py b/src/flow_factory/trainers/abc.py index de6bfad8..daac3729 100644 --- a/src/flow_factory/trainers/abc.py +++ b/src/flow_factory/trainers/abc.py @@ -93,6 +93,10 @@ def __init__( self._initialization() self.adapter.post_init() + # Apply the shared (lossless) accelerator last: after prepare, state-resume, + # EMA, and reference-parameter setup, so e.g. torch.compile wraps the final + # weights and keeps state_dict keys / parameter identity stable. + self._apply_shared_acceleration() self._init_logging_backend() self._patch_deepspeed_autocast(accelerator) @@ -376,21 +380,24 @@ def _initialization(self): # Load inference modules, excluding all bundle members (already prepared). self._load_inference_components(bundle_names) - # Build acceleration plugins now that the prepared transformer proxies are - # installed (shared accelerators like torch.compile mutate them in place). + # Build + validate acceleration plugins. The shared (lossless) accelerator + # is *applied* later via _apply_shared_acceleration(), after post_init() + # finishes any state-resume / EMA / reference setup. self._init_acceleration() # Initialize reward model self._init_reward_model() def _init_acceleration(self): - """Build and install acceleration plugins from ``config.acceleration_args``. + """Build and validate acceleration plugins from ``config.acceleration_args``. Two independent slots (both off by default): - * ``shared_accelerator`` — a lossless accelerator applied here via - :meth:`~BaseAccelerator.setup` (e.g. ``torch.compile``); it affects both - rollout and the training forward. + * ``shared_accelerator`` — a lossless accelerator (e.g. ``torch.compile``) + that affects both rollout and the training forward. Only built/validated + here; it is *applied* later by :meth:`_apply_shared_acceleration` (after + ``post_init`` finishes state-resume / EMA / reference setup), so it + transforms the final weights. * ``rollout_accelerator`` — applied per-epoch in :meth:`generate_samples` via :meth:`~BaseAccelerator.rollout_context`; may be lossy. @@ -413,14 +420,7 @@ def _init_acceleration(self): validate_accelerator( accelerator, slot="shared", paradigm=paradigm, trainer_name=trainer_name ) - accelerator.setup(self.adapter) self.shared_accelerator = accelerator - if self.accelerator.is_main_process: - logger.info( - "Acceleration: shared accelerator '%s' (safety=%s) applied.", - accel_args.shared_accelerator, - accelerator.safety, - ) if accel_args.rollout_accelerator: accelerator = build_accelerator( @@ -437,6 +437,26 @@ def _init_acceleration(self): accelerator.safety, ) + def _apply_shared_acceleration(self) -> None: + """Apply the shared (lossless) accelerator to the adapter. + + Called from ``__init__`` AFTER ``adapter.post_init()`` so that compilation + happens once all weights are final — i.e. after ``accelerator.prepare``, + any ``state`` checkpoint resume, and EMA / reference-parameter snapshotting. + In-place compilation (``nn.Module.compile`` / ``compile_repeated_blocks``) + preserves parameter identity and ``state_dict`` keys, so checkpointing and + the ``copy_``-based EMA / ref / named-parameter swaps stay correct. + """ + accelerator = getattr(self, "shared_accelerator", None) + if accelerator is None: + return + accelerator.setup(self.adapter) + if self.accelerator.is_main_process: + logger.info( + "Acceleration: shared accelerator (safety=%s) applied to adapter.", + accelerator.safety, + ) + def _rollout_acceleration(self) -> AbstractContextManager: """Return the rollout accelerator's context, or a no-op when disabled.""" if self.rollout_accelerator is not None: From 566ea67224609e7160bda1dada88743946fbde4a Mon Sep 17 00:00:00 2001 From: Jayce-Ping <315229706@qq.com> Date: Sat, 27 Jun 2026 16:38:31 +0800 Subject: [PATCH 06/15] [acceleration,models,trainers] refactor: make attention backend accelerator the single path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per review preference, invert the previous fix: instead of removing the AttentionBackendAccelerator and keeping BaseAdapter._set_attention_backend, remove the in-adapter call and route attention-backend selection exclusively through the accelerator, for one consistent transformer-acceleration mechanism. - Remove BaseAdapter._set_attention_backend() call and method from models/abc.py. - Restore AttentionBackendAccelerator: reads backend from model.attn_backend (with optional 'backend' param override) and forwards it verbatim to diffusers' set_attention_backend (incl. approximate backends like 'sage' — matches old behavior; no whitelist). Re-register in the accelerator registry. - BaseTrainer._apply_shared_acceleration now applies attention backend (from model.attn_backend) first, then the shared accelerator (e.g. torch.compile), so the compiled graph captures the chosen backend. Runs after prepare()/post_init(). - model.attn_backend stays the config knob (used in examples); no config breakage. Standalone inference uses raw diffusers pipelines, not the adapter, so dropping the constructor call has no effect there. - Docs (architecture.md, guidance/acceleration.md, acceleration_args help) updated. --- .agents/knowledge/architecture.md | 3 +- guidance/acceleration.md | 19 ++++- .../acceleration/attention_backend.py | 82 +++++++++++++++++++ src/flow_factory/acceleration/registry.py | 6 +- src/flow_factory/hparams/acceleration_args.py | 3 +- src/flow_factory/models/abc.py | 25 +----- src/flow_factory/trainers/abc.py | 39 ++++++--- 7 files changed, 135 insertions(+), 42 deletions(-) create mode 100644 src/flow_factory/acceleration/attention_backend.py diff --git a/.agents/knowledge/architecture.md b/.agents/knowledge/architecture.md index a05798dd..808df979 100644 --- a/.agents/knowledge/architecture.md +++ b/.agents/knowledge/architecture.md @@ -150,11 +150,12 @@ All four registries map string keys → lazy import paths. Resolution: registry **Accelerators** (`acceleration/registry.py`): | Key | Class | Safety | Stage | Notes | |-----|-------|--------|-------|-------| +| `attention_backend` | `AttentionBackendAccelerator` | lossless | both | Sets the diffusers attention backend on every transformer. Auto-applied from `model.attn_backend` by the trainer (before compile); this is the single code path for backend selection (the old `BaseAdapter._set_attention_backend` was removed). | | `torch_compile` | `CompileAccelerator` | lossless | both | `torch.compile` of the shared transformer (regional/full); applied in-place after `post_init` so checkpoint keys / param identity stay stable. | | `diffusers_cache` | `DiffusersCacheAccelerator` | lossy | rollout | Diffusers `CacheMixin` feature caching (first_block/faster/pyramid/taylorseer/magcache). | | `cache_dit` | `CacheDitAccelerator` | lossy | rollout | Optional `cache-dit` backend (DBCache/TaylorSeer). | -Configured via the `acceleration:` block (`hparams/acceleration_args.py`): `shared_accelerator` (lossless, both stages) and `rollout_accelerator` (Stage-3 only). The `acceleration/validator.py` enforces that lossy accelerators are rollout-only and only on `decoupled`/`distillation` trainers (each trainer declares a `paradigm`), preserving train-inference consistency (constraints #7, #20a). Off by default. Attention-backend selection is **not** an accelerator — use `model.attn_backend` (`BaseAdapter._set_attention_backend`). +Configured via the `acceleration:` block (`hparams/acceleration_args.py`): `shared_accelerator` (lossless, both stages) and `rollout_accelerator` (Stage-3 only). Attention backend keeps its own knob, `model.attn_backend`, and is applied by the trainer's acceleration step (`BaseTrainer._apply_shared_acceleration`) — attention backend first, then the configured shared accelerator (e.g. compile). The `acceleration/validator.py` enforces that lossy accelerators are rollout-only and only on `decoupled`/`distillation` trainers (each trainer declares a `paradigm`), preserving train-inference consistency (constraints #7, #20a). Off by default. --- diff --git a/guidance/acceleration.md b/guidance/acceleration.md index 6c5da3eb..dbed1195 100644 --- a/guidance/acceleration.md +++ b/guidance/acceleration.md @@ -48,14 +48,25 @@ accepted in place of a registered id. | id | safety | stage | Notes | |----|--------|-------|-------| +| `attention_backend` | lossless | both | Sets the diffusers attention backend on every transformer. Configured via `model.attn_backend` (see below); `backend` param can override. Forwards any backend (`native` / `flash` / `_flash_3` / `_flash_3_hub` / `sage` / `xformers`) to `set_attention_backend`. | | `torch_compile` | lossless | both | `torch.compile` of the shared transformer. `mode: regional` uses diffusers' `compile_repeated_blocks` (fast warmup, robust to variable resolution); `mode: full` compiles the whole module. Extra `compile_kwargs` forwarded to the compile call. Compiles in place (checkpoint- and EMA/ref-safe), applied after `post_init`. | | `diffusers_cache` | lossy | rollout | Diffusers-native feature caching (no extra dependency). `policy`: `first_block` (default) / `faster` / `pyramid` / `taylorseer` / `magcache`; remaining params forwarded to the policy's diffusers config (e.g. `threshold`). | | `cache_dit` | lossy | rollout | [cache-dit](https://github.com/vipshop/cache-dit) backend (DBCache/TaylorSeer). Requires `pip install flow-factory[acceleration]`. All params forwarded to `cache_dit.enable_cache`. | -> Attention-backend selection (FlashAttention 2/3, xformers, native SDPA) is **not** an -> accelerator — set it once via `model.attn_backend` (handled by -> `BaseAdapter._set_attention_backend`), which already applies to every transformer -> framework-wide and is lossless. +### Attention backend + +Attention-backend selection keeps its dedicated config knob, `model.attn_backend`: + +```yaml +model: + attn_backend: "_flash_3_hub" # native | flash | flash_hub | _flash_3 | _flash_3_hub | sage | xformers +``` + +It is applied through `AttentionBackendAccelerator` by the trainer +(`BaseTrainer._apply_shared_acceleration`) — after `accelerator.prepare` / `post_init` +and **before** compile, so the compiled graph captures the chosen backend. This is the +single code path for backend selection (the old `BaseAdapter._set_attention_backend` was +removed); it applies whether or not an `acceleration:` block is present. ## Model cache-readiness (lossy caching) diff --git a/src/flow_factory/acceleration/attention_backend.py b/src/flow_factory/acceleration/attention_backend.py new file mode 100644 index 00000000..66b9f295 --- /dev/null +++ b/src/flow_factory/acceleration/attention_backend.py @@ -0,0 +1,82 @@ +# Copyright 2026 Jayce-Ping +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# src/flow_factory/acceleration/attention_backend.py +"""Attention-backend accelerator — the single code path that selects the +diffusers attention backend for every transformer. + +This replaces the old ``BaseAdapter._set_attention_backend`` call: the backend is +applied here (after ``accelerator.prepare`` / ``post_init`` and before compile) +instead of in the adapter constructor, so all transformer-level acceleration flows +through the same plugin mechanism. + +The backend is read from ``model.attn_backend`` by default (so existing configs +keep working) and may be overridden by an explicit ``backend`` param. Whatever +string is given is forwarded to diffusers' ``set_attention_backend`` verbatim — +including approximate backends like ``sage`` — matching the previous behavior. + +Marked ``stage='both'`` / ``safety='lossless'``: a backend is applied to the +transformer shared by rollout ``inference()`` and training ``forward()``, so the +two stay consistent (the property the validator's lossless category guarantees) +even for approximate kernels and coupled algorithms. +""" + +from typing import TYPE_CHECKING + +from .abc import BaseAccelerator +from ..utils.logger_utils import setup_logger + +if TYPE_CHECKING: + from ..models.abc import BaseAdapter + +logger = setup_logger(__name__) + + +class AttentionBackendAccelerator(BaseAccelerator): + """Set the diffusers attention backend on every transformer component. + + Parameters (from ``acceleration.shared_params``, all optional): + backend: Backend name forwarded to ``transformer.set_attention_backend`` + (e.g. ``native`` / ``flash`` / ``_flash_3`` / ``_flash_3_hub`` / + ``sage`` / ``xformers``). Defaults to ``model.attn_backend``. + + See https://huggingface.co/docs/diffusers/main/en/optimization/attention_backends + for the full list of supported backends. + """ + + safety = "lossless" + stage = "both" + + def setup(self, adapter: "BaseAdapter") -> None: + backend = self.params.get("backend") or adapter.model_args.attn_backend + if backend is None: + # Nothing requested (no `backend` param and `model.attn_backend` unset). + return + + applied = False + for name in adapter.transformer_names: + transformer = adapter.get_component(name) + if hasattr(transformer, "set_attention_backend"): + transformer.set_attention_backend(backend) + applied = True + if adapter.accelerator.is_main_process: + logger.info( + "AttentionBackendAccelerator: set backend '%s' for '%s'.", backend, name + ) + if not applied: + logger.warning( + "AttentionBackendAccelerator: backend '%s' requested but no transformer component " + "supports `set_attention_backend`; leaving the diffusers default.", + backend, + ) diff --git a/src/flow_factory/acceleration/registry.py b/src/flow_factory/acceleration/registry.py index bf537fd5..1565735d 100644 --- a/src/flow_factory/acceleration/registry.py +++ b/src/flow_factory/acceleration/registry.py @@ -32,9 +32,9 @@ # Accelerator Registry Storage _ACCELERATOR_REGISTRY: Dict[str, str] = { # Lossless (safe for any algorithm, applied to the shared transformer). - # NOTE: attention-backend selection is handled by `model.attn_backend` - # (BaseAdapter._set_attention_backend); it is intentionally NOT an accelerator - # here to avoid a second, redundant set after prepare. + # `attention_backend` is the single code path for attention-backend selection; + # it is auto-applied from `model.attn_backend` by the trainer (before compile). + 'attention_backend': 'flow_factory.acceleration.attention_backend.AttentionBackendAccelerator', 'torch_compile': 'flow_factory.acceleration.torch_compile.CompileAccelerator', # Lossy (rollout-only; validator restricts to decoupled / distillation algos). 'diffusers_cache': 'flow_factory.acceleration.diffusers_cache.DiffusersCacheAccelerator', diff --git a/src/flow_factory/hparams/acceleration_args.py b/src/flow_factory/hparams/acceleration_args.py index 4888675f..4512614e 100644 --- a/src/flow_factory/hparams/acceleration_args.py +++ b/src/flow_factory/hparams/acceleration_args.py @@ -48,7 +48,8 @@ class AccelerationArguments(ArgABC): "help": ( "Lossless accelerator id (or python path) applied to both rollout and training. " "Options: 'torch_compile'. None disables it. " - "(Attention backend is set via model.attn_backend, not here.)" + "(Attention backend has its own knob, model.attn_backend, applied " + "automatically before this.)" ) }, ) diff --git a/src/flow_factory/models/abc.py b/src/flow_factory/models/abc.py index bbe7586e..5dbdc6d9 100644 --- a/src/flow_factory/models/abc.py +++ b/src/flow_factory/models/abc.py @@ -194,8 +194,10 @@ def __init__(self, config: Arguments, accelerator : Accelerator): # Set precision self._mix_precision() - # Set attention backend for all transformers - self._set_attention_backend() + # NOTE: attention-backend selection is applied later by the trainer's + # acceleration step (AttentionBackendAccelerator, from `model.attn_backend`), + # after accelerator.prepare()/post_init() and before torch.compile — see + # BaseTrainer._apply_shared_acceleration(). It is intentionally NOT set here. # Enable gradient checkpointing if needed if self.training_args.enable_gradient_checkpointing: @@ -848,25 +850,6 @@ def enable_gradient_checkpointing(self): else: logger.warning(f"{comp_name} does not support gradient checkpointing") - # ============================== Attention Backend ============================== - def _set_attention_backend(self) -> None: - """ - Set attention backend for all transformer components. - - Refer to https://huggingface.co/docs/diffusers/main/en/optimization/attention_backends#available-backends - to see supported backends. - """ - backend = self.model_args.attn_backend - if backend is None: - return - - for transformer_name in self.transformer_names: - transformer = self.get_component(transformer_name) - if hasattr(transformer, 'set_attention_backend'): - transformer.set_attention_backend(backend) - if self.accelerator.is_main_process: - logger.info(f"Set attention backend '{backend}' for {transformer_name}") - # ============================== Precision Management ============================== def _cast_module_mixed_precision( self, diff --git a/src/flow_factory/trainers/abc.py b/src/flow_factory/trainers/abc.py index daac3729..6dec9253 100644 --- a/src/flow_factory/trainers/abc.py +++ b/src/flow_factory/trainers/abc.py @@ -438,24 +438,39 @@ def _init_acceleration(self): ) def _apply_shared_acceleration(self) -> None: - """Apply the shared (lossless) accelerator to the adapter. + """Apply transformer-level (shared) acceleration to the adapter. + + Called from ``__init__`` AFTER ``adapter.post_init()`` so transforms wrap the + final weights — i.e. after ``accelerator.prepare``, any ``state`` checkpoint + resume, and EMA / reference-parameter snapshotting. Applied in order: + + 1. **Attention backend** from ``model.attn_backend`` — the single code path + for backend selection (replaces ``BaseAdapter._set_attention_backend``), + run here so it sits after ``prepare`` and before compile. + 2. The configured **shared accelerator** (e.g. ``torch.compile``), applied + last so it captures the chosen attention backend. - Called from ``__init__`` AFTER ``adapter.post_init()`` so that compilation - happens once all weights are final — i.e. after ``accelerator.prepare``, - any ``state`` checkpoint resume, and EMA / reference-parameter snapshotting. In-place compilation (``nn.Module.compile`` / ``compile_repeated_blocks``) preserves parameter identity and ``state_dict`` keys, so checkpointing and the ``copy_``-based EMA / ref / named-parameter swaps stay correct. """ + # 1. Attention backend (lossless / both stages), sourced from model.attn_backend. + # Always safe (applied to the shared transformer), so it bypasses the + # paradigm validator; AttentionBackendAccelerator.setup is a no-op when + # `model.attn_backend` is unset. + if self.model_args.attn_backend: + build_accelerator("attention_backend", {}).setup(self.adapter) + + # 2. Configured shared accelerator (e.g. torch.compile), applied after the + # backend so the compiled graph captures it. accelerator = getattr(self, "shared_accelerator", None) - if accelerator is None: - return - accelerator.setup(self.adapter) - if self.accelerator.is_main_process: - logger.info( - "Acceleration: shared accelerator (safety=%s) applied to adapter.", - accelerator.safety, - ) + if accelerator is not None: + accelerator.setup(self.adapter) + if self.accelerator.is_main_process: + logger.info( + "Acceleration: shared accelerator (safety=%s) applied to adapter.", + accelerator.safety, + ) def _rollout_acceleration(self) -> AbstractContextManager: """Return the rollout accelerator's context, or a no-op when disabled.""" From 433685c28cf79cec0c5ec3db09ae8be25ad1a27d Mon Sep 17 00:00:00 2001 From: Jayce-Ping <315229706@qq.com> Date: Sat, 27 Jun 2026 16:49:02 +0800 Subject: [PATCH 07/15] [acceleration] refactor: gate shared slot on stage (symmetric), not numerical safety Answers the review questions on attention-backend timing and lossy backends: - Correctness axis is SYMMETRIC application (stage), not bit-exactness. A stage='both' transform runs in both rollout inference() and training forward() on the shared module, so they stay train-inference consistent even when the transform is numerically approximate (e.g. Sage int8 attention). We therefore do NOT reject 'sage' and do not need per-instance 'two safety states'. - Validator: the shared slot now requires stage=='both' (drop the safety=='lossless' requirement); the rollout slot requires stage=='rollout' and still gates safety=='lossy' to decoupled/distillation paradigms. safety is only consulted for rollout-stage accelerators. - Clarify BaseAccelerator safety/stage semantics and the validator/guidance docs. Timing (confirmed, no code change): set_attention_backend only sets processor._attention_backend attributes + a global active backend (touches no params), so it is safe AFTER accelerator.prepare() and is applied BEFORE compile (the dispatch is read at forward time); applying after prepare is also required for context-parallel backend validation. _apply_shared_acceleration runs after post_init (final weights), attention backend first, then compile. --- guidance/acceleration.md | 32 ++++++----- src/flow_factory/acceleration/abc.py | 36 ++++++++----- src/flow_factory/acceleration/validator.py | 63 ++++++++++++++-------- 3 files changed, 85 insertions(+), 46 deletions(-) diff --git a/guidance/acceleration.md b/guidance/acceleration.md index dbed1195..060b58cb 100644 --- a/guidance/acceleration.md +++ b/guidance/acceleration.md @@ -9,21 +9,29 @@ under `torch.no_grad()`), so that is where most accelerators apply. ## Safety model (read this first) +The correctness axis is **symmetric application** (`stage`), not numerical bit-exactness. Every accelerator declares two markers, and a validator (`acceleration/validator.py`) enforces them against the trainer's `paradigm` before training starts (fail-fast): -| `safety` | Meaning | Where allowed | -|----------|---------|---------------| -| `lossless` | Numerically ~identical; applied to the transformer shared by both rollout `inference()` and training `forward()`. | Any algorithm, any stage. | -| `lossy` | Changes `noise_pred` (e.g. feature caching). Cannot be replicated in the training forward. | **Rollout only**, and only on `decoupled` / `distillation` trainers. | - -Why the restriction: for **coupled** algorithms (GRPO, GRPO-Guard, DPPO) the rollout's -per-step log-prob becomes the PPO "old log-prob". Approximating the rollout while the -training forward stays exact biases the importance ratio and silently corrupts gradients -(`.agents/knowledge/constraints.md` #7, #20a). For **decoupled** algorithms (NFT, AWM, -DGPO, DPO, CRD) and **distillation** (diffusion-opd), the rollout trajectory's log-prob -does not enter the loss, so lossy caching only shifts the generated-sample distribution — -acceptable and tunable. Monitor the reward mean/std when enabling a lossy accelerator. +| marker | values | meaning | +|--------|--------|---------| +| `stage` | `both` / `rollout` | `both`: persistent transform applied to the transformer shared by rollout `inference()` and training `forward()` → **consistent by construction, safe for any algorithm** (the `shared` slot). `rollout`: per-epoch context, torn down before training (the `rollout` slot). | +| `safety` | `lossless` / `lossy` | Only consulted for `stage='rollout'`. `lossy` = changes rollout outputs → diverges from training. `lossless` = bit-identical. | + +The single restriction: a **`lossy` rollout** accelerator is allowed **only on +`decoupled` / `distillation`** trainers. Why: for **coupled** algorithms (GRPO, +GRPO-Guard, DPPO) the rollout's per-step log-prob becomes the PPO "old log-prob"; +changing the rollout while the training forward stays exact biases the importance ratio +and silently corrupts gradients (`.agents/knowledge/constraints.md` #7, #20a). For +**decoupled** (NFT, AWM, DGPO, DPO, CRD) and **distillation** (diffusion-opd), the rollout +log-prob does not enter the loss, so a lossy rollout only shifts the generated-sample +distribution — acceptable and tunable. Monitor the reward mean/std when enabling it. + +A `stage='both'` accelerator is **always safe** — even a numerically-approximate one. For +example, Sage int8 attention used as the attention backend runs in *both* rollout and +training, so the two stay consistent; there is no need to reject it or give it a special +"lossy" status. Numerical exactness only matters when a transform is applied to one stage +but not the other, which is exactly what `stage='rollout'` + `safety='lossy'` captures. ## Configuration diff --git a/src/flow_factory/acceleration/abc.py b/src/flow_factory/acceleration/abc.py index 08cc1e47..fd43c35d 100644 --- a/src/flow_factory/acceleration/abc.py +++ b/src/flow_factory/acceleration/abc.py @@ -16,18 +16,30 @@ """Abstract base class for the model-agnostic acceleration plugin layer. An accelerator is a pluggable speedup applied to a model adapter's transformer(s) -without touching trainer or model math. Each accelerator declares two invariants -that the validator (:mod:`flow_factory.acceleration.validator`) enforces against -the trainer's RL paradigm: - -* ``safety`` — ``"lossless"`` (numerically ~identical; safe for any algorithm and - any stage because the same module backs both rollout ``inference()`` and the - training ``forward()``) or ``"lossy"`` (changes ``noise_pred``; only safe during - rollout for **decoupled / distillation** algorithms — see ``constraints.md`` #7 - and #20a). -* ``stage`` — ``"both"`` (a one-time mutation applied via :meth:`setup`, e.g. - ``torch.compile``) or ``"rollout"`` (a per-epoch context applied via - :meth:`rollout_context`, e.g. feature caching). +without touching trainer or model math. Each accelerator declares two markers that +the validator (:mod:`flow_factory.acceleration.validator`) uses to preserve +train-inference consistency against the trainer's RL paradigm: + +* ``stage`` — *where/how* it is applied, and which config slot it belongs to: + + - ``"both"``: a persistent, one-time mutation via :meth:`setup` (e.g. + ``torch.compile``, attention backend). Because it transforms the module shared + by rollout ``inference()`` and training ``forward()``, the two stay CONSISTENT + by construction — safe for any algorithm. Belongs in the ``shared`` slot. + - ``"rollout"``: a per-epoch context via :meth:`rollout_context` (e.g. feature + caching), torn down before the training forward. Belongs in the ``rollout`` slot. + +* ``safety`` — the train-inference consistency class, **only consulted for + ``stage='rollout'`` accelerators**: + + - ``"lossless"``: bit-identical outputs (rollout unchanged) — safe for any paradigm. + - ``"lossy"``: changes outputs, so rollout diverges from the (un-accelerated) + training forward — only safe when the rollout log-prob never feeds the loss, i.e. + **decoupled / distillation** algorithms (see ``constraints.md`` #7, #20a). + + For ``stage='both'`` accelerators ``safety`` is informational: a symmetric + transform is consistent regardless of numerical exactness, so even an approximate + attention backend (e.g. Sage int8) used in *both* stages is safe. Subclasses implement only what they need: ``setup`` defaults to a no-op, ``rollout_context`` defaults to yielding without modification. diff --git a/src/flow_factory/acceleration/validator.py b/src/flow_factory/acceleration/validator.py index b89e48b2..f25c42fa 100644 --- a/src/flow_factory/acceleration/validator.py +++ b/src/flow_factory/acceleration/validator.py @@ -15,17 +15,23 @@ # src/flow_factory/acceleration/validator.py """Paradigm-gated safety validation for accelerators (fail-fast, ``constraints.md`` #26). -The correctness contract (``constraints.md`` #7 + #20a): +The correctness contract (``constraints.md`` #7 + #20a) hinges on **symmetric +application**, encoded by ``stage``, not on numerical bit-exactness: -* A ``lossy`` accelerator changes ``noise_pred`` and cannot be replicated in the - Stage-6 training forward (which needs full gradient through every block). It is - therefore only valid in the ``rollout`` slot, and only for **decoupled** / - **distillation** algorithms — for **coupled** algorithms (GRPO / GRPO-Guard / - DPPO) the rollout log-prob becomes the PPO "old log-prob" and an approximated - rollout would bias the importance ratio -> silently wrong gradients. -* A ``lossless`` accelerator is numerically ~identical and, because it mutates the - shared transformer used by both ``inference()`` and ``forward()``, is safe in - any slot for any algorithm. +* ``stage='both'`` accelerators mutate the transformer persistently, so the exact + same transform is in both rollout ``inference()`` and training ``forward()``. + Rollout and training therefore stay CONSISTENT — even for numerically-approximate + transforms (e.g. Sage int8 attention) — so they are safe for any algorithm. They + belong in the ``shared`` slot; ``safety`` is not consulted for them. +* ``stage='rollout'`` accelerators run only during Stage-3 rollout. If such an + accelerator changes outputs (``safety='lossy'``, e.g. feature caching), rollout + diverges from the training forward, which it cannot be replicated in (that needs + full gradient through every block). That is only safe when the rollout + trajectory's log-prob never feeds the loss — i.e. **decoupled** / **distillation** + algorithms. For **coupled** algorithms (GRPO / GRPO-Guard / DPPO) the rollout + log-prob becomes the PPO "old log-prob", so a divergent rollout biases the + importance ratio -> silently wrong gradients. A bit-identical rollout accelerator + (``safety='lossless'``) is safe for any paradigm. """ from typing import Optional @@ -70,24 +76,37 @@ def validate_accelerator( f"Unknown acceleration slot {slot!r} for '{name}'; expected 'shared' or 'rollout'." ) - # The `shared` slot applies to BOTH rollout and the training forward, so only - # lossless accelerators may live there. + # The slot must match the accelerator's `stage`: `shared` accelerators mutate the + # transformer persistently (applied to both rollout and training via `setup`); + # `rollout` accelerators are a per-epoch context (`rollout_context`). A mismatch + # would silently no-op (e.g. a stage='both' accelerator in the rollout slot has + # no rollout_context), so reject it. if slot == "shared": - if accelerator.safety != "lossless": - raise ValueError( - f"Accelerator '{name}' (safety='{accelerator.safety}') is configured under " - "`acceleration.shared_accelerator`, but the shared slot is applied to BOTH " - "rollout and the training forward — only lossless accelerators are allowed there. " - "Move a lossy accelerator to `acceleration.rollout_accelerator`." - ) if accelerator.stage != "both": raise ValueError( f"Accelerator '{name}' (stage='{accelerator.stage}') cannot occupy the shared " - "slot; the shared slot requires a stage='both' accelerator." + "slot, which applies a persistent transform to both rollout and training. Put a " + "stage='both' accelerator here, or move this one to `acceleration.rollout_accelerator`." ) + # No safety gate here: a stage='both' transform is applied identically to + # rollout `inference()` and training `forward()` (the same shared module), so + # rollout and training stay CONSISTENT even for numerically-approximate + # backends (e.g. Sage int8 attention). Train-inference consistency depends on + # symmetric application, not on bit-exactness. return - # slot == "rollout": lossless is always fine; lossy is gated on paradigm. + # slot == "rollout". + if accelerator.stage != "rollout": + raise ValueError( + f"Accelerator '{name}' (stage='{accelerator.stage}') cannot occupy the rollout slot, " + "which only runs a per-epoch `rollout_context`. Put a stage='both' accelerator under " + "`acceleration.shared_accelerator` instead." + ) + + # A rollout-only accelerator that changes outputs (`lossy`) makes rollout diverge + # from the training forward, so it is only safe when the rollout trajectory's + # log-prob never feeds the loss (decoupled / distillation). A `lossless` + # rollout accelerator (bit-identical) is safe for any paradigm. if accelerator.safety == "lossy": if paradigm is None: raise ValueError( @@ -102,7 +121,7 @@ def validate_accelerator( "cannot be replicated in the training forward; for coupled algorithms " "(GRPO / GRPO-Guard / DPPO) this biases the PPO importance ratio and silently " f"corrupts gradients. Allowed paradigms: {sorted(_LOSSY_SAFE_PARADIGMS)}. Use a " - "lossless accelerator (e.g. 'torch_compile') instead, or switch to a decoupled " + "stage='both' accelerator (e.g. 'torch_compile') instead, or switch to a decoupled " "algorithm (NFT / AWM / DGPO / DPO / CRD)." ) logger.warning( From 25fd6e7ddc6b0a5eedd162918a27dba238d2f2e3 Mon Sep 17 00:00:00 2001 From: Jayce-Ping <315229706@qq.com> Date: Sun, 28 Jun 2026 08:37:44 +0800 Subject: [PATCH 08/15] [acceleration,hparams,models,trainers,examples,docs] refactor: fold attn_backend into the acceleration layer + ordered multi-accelerator Remove the standalone `model.attn_backend` knob and express attention-backend selection only as an `attention_backend` entry in the acceleration layer. Each acceleration slot (`shared` / `rollout`) becomes an ordered list of {name, params} entries (list order = application order), mirroring MultiRewardArguments, so multiple accelerators can be combined deterministically (e.g. attention_backend before torch_compile so the compiled graph captures the backend). - hparams/acceleration_args.py: add AccelerationSpec; AccelerationArguments.shared and .rollout are List[AccelerationSpec]; from_dict accepts list or single-entry shorthand; to_dict round-trips. - hparams/model_args.py: drop attn_backend field; __post_init__ fail-fast migration error if a stale `model.attn_backend` is present. - acceleration/attention_backend.py: require an explicit `backend` param; raise when no transformer supports set_attention_backend (e.g. Bagel forces fa2 at load). - trainers/abc.py: build/validate accelerator lists; _apply_shared_acceleration runs setup() in order; _rollout_acceleration nests rollout contexts via ExitStack. - examples: replace commented attn_backend lines with a commented acceleration block (coupled => shared-only; decoupled => shared+rollout). Bagel examples drop the inert attn_backend line and document the forced flash_attention_2 path. - docs: guidance/acceleration.md, architecture.md, README.md updated to the list form. Co-Authored-By: Claude Opus 4.8 (1M context) --- .agents/knowledge/architecture.md | 4 +- README.md | 10 +- examples/awm/lora/flux1/default.yaml | 13 +- .../awm/lora/flux2_klein_base/default.yaml | 13 +- examples/awm/lora/sd3_5/default.yaml | 13 +- examples/crd/lora/sd3_5/default.yaml | 13 +- examples/grpo/full/qwen_image/default.yaml | 10 +- .../full/qwen_image_edit_plus/default.yaml | 10 +- examples/grpo/lora/bagel/i2i.yaml | 3 +- examples/grpo/lora/ltx2/t2av.yaml | 10 +- examples/grpo/lora/ltx2/t2av_pickscore.yaml | 10 +- examples/grpo/lora/qwen_image/default.yaml | 10 +- .../lora/qwen_image_edit_plus/default.yaml | 10 +- examples/grpo/lora/sd3_5/default.yaml | 10 +- examples/grpo/lora/sd3_5/geneval.yaml | 10 +- examples/grpo/lora/sd3_5/nocfg.yaml | 10 +- examples/nft/lora/bagel/default.yaml | 3 +- examples/nft/lora/bagel/i2i.yaml | 3 +- examples/nft/lora/flux1/default.yaml | 13 +- .../lora/qwen_image/rational_rewards_t2i.yaml | 17 +- examples/nft/lora/sd3_5/default.yaml | 13 +- examples/nft/lora/wan21/i2v.yaml | 13 +- examples/nft/lora/wan21/t2v.yaml | 13 +- examples/nft/lora/wan22/t2v.yaml | 13 +- examples/template/sd3_5/async_reward.yaml | 10 +- guidance/acceleration.md | 56 ++++--- multinode_examples/train.yaml | 11 +- .../acceleration/attention_backend.py | 40 ++--- .../acceleration/diffusers_cache.py | 2 +- src/flow_factory/acceleration/registry.py | 2 +- .../acceleration/torch_compile.py | 2 +- src/flow_factory/acceleration/validator.py | 6 +- src/flow_factory/hparams/__init__.py | 3 +- src/flow_factory/hparams/acceleration_args.py | 145 +++++++++++++----- src/flow_factory/hparams/model_args.py | 22 +-- src/flow_factory/models/abc.py | 7 +- src/flow_factory/trainers/abc.py | 91 +++++------ 37 files changed, 460 insertions(+), 184 deletions(-) diff --git a/.agents/knowledge/architecture.md b/.agents/knowledge/architecture.md index 808df979..9c7b5d02 100644 --- a/.agents/knowledge/architecture.md +++ b/.agents/knowledge/architecture.md @@ -150,12 +150,12 @@ All four registries map string keys → lazy import paths. Resolution: registry **Accelerators** (`acceleration/registry.py`): | Key | Class | Safety | Stage | Notes | |-----|-------|--------|-------|-------| -| `attention_backend` | `AttentionBackendAccelerator` | lossless | both | Sets the diffusers attention backend on every transformer. Auto-applied from `model.attn_backend` by the trainer (before compile); this is the single code path for backend selection (the old `BaseAdapter._set_attention_backend` was removed). | +| `attention_backend` | `AttentionBackendAccelerator` | lossless | both | Sets the diffusers attention backend on every transformer (requires a `backend` param). Listed as a `shared` entry (before `torch_compile`); this is the single code path for backend selection (the old `BaseAdapter._set_attention_backend` and the `model.attn_backend` knob were both removed — a config still setting `model.attn_backend` fails fast). Bagel forces flash_attention_2 at load and does not use it. | | `torch_compile` | `CompileAccelerator` | lossless | both | `torch.compile` of the shared transformer (regional/full); applied in-place after `post_init` so checkpoint keys / param identity stay stable. | | `diffusers_cache` | `DiffusersCacheAccelerator` | lossy | rollout | Diffusers `CacheMixin` feature caching (first_block/faster/pyramid/taylorseer/magcache). | | `cache_dit` | `CacheDitAccelerator` | lossy | rollout | Optional `cache-dit` backend (DBCache/TaylorSeer). | -Configured via the `acceleration:` block (`hparams/acceleration_args.py`): `shared_accelerator` (lossless, both stages) and `rollout_accelerator` (Stage-3 only). Attention backend keeps its own knob, `model.attn_backend`, and is applied by the trainer's acceleration step (`BaseTrainer._apply_shared_acceleration`) — attention backend first, then the configured shared accelerator (e.g. compile). The `acceleration/validator.py` enforces that lossy accelerators are rollout-only and only on `decoupled`/`distillation` trainers (each trainer declares a `paradigm`), preserving train-inference consistency (constraints #7, #20a). Off by default. +Configured via the `acceleration:` block (`hparams/acceleration_args.py`): two ordered lists of `{name, params}` entries — `shared` (lossless, both stages) and `rollout` (Stage-3 only). **List order is application order**: `shared` entries run their `setup()` in order (so `attention_backend` must precede `torch_compile`), and `rollout` entries nest their `rollout_context()` in order. The `acceleration/validator.py` enforces that lossy accelerators are rollout-only and only on `decoupled`/`distillation` trainers (each trainer declares a `paradigm`), preserving train-inference consistency (constraints #7, #20a). Off by default. --- diff --git a/README.md b/README.md index 76158b1e..872959a1 100644 --- a/README.md +++ b/README.md @@ -15,12 +15,14 @@ git submodule update --init pip install -e ./diffusers ``` -* **[2026-02-01]** Support for multiple **Attention Backends**! You can now optimize memory and speed by setting the `attn_backend` parameter in your config: +* **[2026-02-01]** Support for multiple **Attention Backends**! Attention-backend selection now lives in the unified `acceleration:` block (the old `model.attn_backend` knob was removed), where it can be combined with `torch.compile` and feature caching — applied in list order: ```yaml - model: - attn_backend: "flash" # Options: "native", "xformers", "flash_hub", "_flash_3_hub", "_flash_3_varlen_hub" + acceleration: + shared: + - name: attention_backend + params: { backend: "flash" } # Options: "native", "xformers", "flash_hub", "_flash_3_hub", "_flash_3_varlen_hub" ``` -This experimental feature leverages `diffusers`'s `transformer.set_attention_backend`. Check the [official diffusers documentation](https://huggingface.co/docs/diffusers/main/en/optimization/attention_backends#available-backends) for all available options. +This experimental feature leverages `diffusers`'s `transformer.set_attention_backend`. Check the [official diffusers documentation](https://huggingface.co/docs/diffusers/main/en/optimization/attention_backends#available-backends) for all available options. See [`guidance/acceleration.md`](guidance/acceleration.md) for the full acceleration layer. > We recommend installing the `kernels` package (`pip install kernels`) and using `flash_hub`, `flash_varlen_hub`, `_flash_3_hub`, or `_flash_3_varlen_hub` to avoid the complexity and potential incompatibility of installing Flash-Attention directly. # 📕 Table of Contents diff --git a/examples/awm/lora/flux1/default.yaml b/examples/awm/lora/flux1/default.yaml index fb2d4161..7ee7b4cc 100644 --- a/examples/awm/lora/flux1/default.yaml +++ b/examples/awm/lora/flux1/default.yaml @@ -30,7 +30,18 @@ model: model_type: "flux1" resume_path: null # Local path or HF repo id (e.g. 'owner/repo[/subdir][@rev]') for previous checkpoint/lora adapter resume_type: null # Options: lora, full, state. Null to auto-detect based on `finetune_type` - # attn_backend: '_flash_3_hub' +# Optional acceleration plugins (off by default). Applied in list order. +# See guidance/acceleration.md. This is a decoupled trainer, so lossy rollout +# feature-caching is allowed. +# acceleration: +# shared: # lossless, applied to BOTH rollout and training (in order) +# - name: attention_backend # set the attention backend first... +# params: { backend: _flash_3_hub } +# - name: torch_compile # ...so the compiled graph captures it +# params: { mode: regional } # regional (compile_repeated_blocks) | full +# rollout: # rollout-only (Stage 3); lossy -> decoupled/distillation only +# - name: diffusers_cache # Options: diffusers_cache, cache_dit +# params: { policy: first_block, threshold: 0.08 } log: run_name: null # Run name (auto: {model_type}_{finetune_type}_{trainer_type}_{timestamp}) diff --git a/examples/awm/lora/flux2_klein_base/default.yaml b/examples/awm/lora/flux2_klein_base/default.yaml index 15f55684..891680ab 100644 --- a/examples/awm/lora/flux2_klein_base/default.yaml +++ b/examples/awm/lora/flux2_klein_base/default.yaml @@ -30,7 +30,18 @@ model: model_type: "flux2-klein" resume_path: null # Local path or HF repo id (e.g. 'owner/repo[/subdir][@rev]') for previous checkpoint/lora adapter resume_type: null # Options: lora, full, state. Null to auto-detect based on `finetune_type` - # attn_backend: '_flash_3_hub' # Attention backend for training. +# Optional acceleration plugins (off by default). Applied in list order. +# See guidance/acceleration.md. This is a decoupled trainer, so lossy rollout +# feature-caching is allowed. +# acceleration: +# shared: # lossless, applied to BOTH rollout and training (in order) +# - name: attention_backend # set the attention backend first... +# params: { backend: _flash_3_hub } +# - name: torch_compile # ...so the compiled graph captures it +# params: { mode: regional } # regional (compile_repeated_blocks) | full +# rollout: # rollout-only (Stage 3); lossy -> decoupled/distillation only +# - name: diffusers_cache # Options: diffusers_cache, cache_dit +# params: { policy: first_block, threshold: 0.08 } log: run_name: null # Run name (auto: {model_type}_{finetune_type}_{trainer_type}_{timestamp}) diff --git a/examples/awm/lora/sd3_5/default.yaml b/examples/awm/lora/sd3_5/default.yaml index 3e771dd5..ea2a6829 100644 --- a/examples/awm/lora/sd3_5/default.yaml +++ b/examples/awm/lora/sd3_5/default.yaml @@ -31,7 +31,18 @@ model: model_type: "sd3-5" resume_path: null # Local path or HF repo id (e.g. 'owner/repo[/subdir][@rev]') for previous checkpoint/lora adapter resume_type: null # Options: lora, full, state. Null to auto-detect based on `finetune_type` - # attn_backend: '_flash_3_hub' # Attention backend for training. +# Optional acceleration plugins (off by default). Applied in list order. +# See guidance/acceleration.md. This is a decoupled trainer, so lossy rollout +# feature-caching is allowed. +# acceleration: +# shared: # lossless, applied to BOTH rollout and training (in order) +# - name: attention_backend # set the attention backend first... +# params: { backend: _flash_3_hub } +# - name: torch_compile # ...so the compiled graph captures it +# params: { mode: regional } # regional (compile_repeated_blocks) | full +# rollout: # rollout-only (Stage 3); lossy -> decoupled/distillation only +# - name: diffusers_cache # Options: diffusers_cache, cache_dit +# params: { policy: first_block, threshold: 0.08 } log: run_name: null # Run name (auto: {model_type}_{finetune_type}_{trainer_type}_{timestamp}) diff --git a/examples/crd/lora/sd3_5/default.yaml b/examples/crd/lora/sd3_5/default.yaml index 34fc1a6a..1a9515c0 100644 --- a/examples/crd/lora/sd3_5/default.yaml +++ b/examples/crd/lora/sd3_5/default.yaml @@ -32,7 +32,18 @@ model: model_type: "sd3-5" resume_path: null # Local path or HF repo id (e.g. 'owner/repo[/subdir][@rev]') for previous checkpoint/lora adapter resume_type: null # Options: lora, full, state. Null to auto-detect based on `finetune_type` - # attn_backend: '_flash_3_hub' # Attention backend for training. +# Optional acceleration plugins (off by default). Applied in list order. +# See guidance/acceleration.md. This is a decoupled trainer, so lossy rollout +# feature-caching is allowed. +# acceleration: +# shared: # lossless, applied to BOTH rollout and training (in order) +# - name: attention_backend # set the attention backend first... +# params: { backend: _flash_3_hub } +# - name: torch_compile # ...so the compiled graph captures it +# params: { mode: regional } # regional (compile_repeated_blocks) | full +# rollout: # rollout-only (Stage 3); lossy -> decoupled/distillation only +# - name: diffusers_cache # Options: diffusers_cache, cache_dit +# params: { policy: first_block, threshold: 0.08 } log: run_name: null # Run name (auto: {model_type}_{finetune_type}_{timestamp}) diff --git a/examples/grpo/full/qwen_image/default.yaml b/examples/grpo/full/qwen_image/default.yaml index 5347cca6..dc982f6e 100644 --- a/examples/grpo/full/qwen_image/default.yaml +++ b/examples/grpo/full/qwen_image/default.yaml @@ -29,7 +29,15 @@ model: model_type: "qwen-image" # Options: flux1, flux1-kontext, flux2, qwenimage, qwenimage-edit resume_path: null # Local path or HF repo id (e.g. 'owner/repo[/subdir][@rev]') for previous checkpoint/lora adapter resume_type: null # Options: lora, full, state. Null to auto-detect based on `finetune_type` - # attn_backend: '_flash_3_varlen_hub' # Attention backend for Qwen-Image Series, which uses masked attention with variable sequence length. +# Optional acceleration plugins (off by default). Applied in list order. +# See guidance/acceleration.md. GRPO is a coupled trainer, so only lossless +# `shared` accelerators are allowed (lossy rollout caching would bias the PPO ratio). +# acceleration: +# shared: # lossless, applied to BOTH rollout and training (in order) +# - name: attention_backend # set the attention backend first... +# params: { backend: _flash_3_varlen_hub } +# - name: torch_compile # ...so the compiled graph captures it +# params: { mode: regional } # regional (compile_repeated_blocks) | full log: run_name: null # Run name (auto: {model_type}_{finetune_type}_{trainer_type}_{timestamp}) diff --git a/examples/grpo/full/qwen_image_edit_plus/default.yaml b/examples/grpo/full/qwen_image_edit_plus/default.yaml index d2bae8a8..4e28bd9f 100644 --- a/examples/grpo/full/qwen_image_edit_plus/default.yaml +++ b/examples/grpo/full/qwen_image_edit_plus/default.yaml @@ -28,7 +28,15 @@ model: model_type: "qwen-image-edit-plus" resume_path: null # Local path or HF repo id (e.g. 'owner/repo[/subdir][@rev]') for previous checkpoint/lora adapter resume_type: null # Options: lora, full, state. Null to auto-detect based on `finetune_type` - # attn_backend: '_flash_3_varlen_hub' # Attention backend for Qwen-Image Series, which uses masked attention with variable sequence length. +# Optional acceleration plugins (off by default). Applied in list order. +# See guidance/acceleration.md. GRPO is a coupled trainer, so only lossless +# `shared` accelerators are allowed (lossy rollout caching would bias the PPO ratio). +# acceleration: +# shared: # lossless, applied to BOTH rollout and training (in order) +# - name: attention_backend # set the attention backend first... +# params: { backend: _flash_3_varlen_hub } +# - name: torch_compile # ...so the compiled graph captures it +# params: { mode: regional } # regional (compile_repeated_blocks) | full log: run_name: null # Run name (auto: {model_type}_{finetune_type}_{trainer_type}_{timestamp}) diff --git a/examples/grpo/lora/bagel/i2i.yaml b/examples/grpo/lora/bagel/i2i.yaml index 6b56b443..b83ab83d 100644 --- a/examples/grpo/lora/bagel/i2i.yaml +++ b/examples/grpo/lora/bagel/i2i.yaml @@ -38,7 +38,8 @@ model: model_type: "bagel" # Model adapter key; see src/flow_factory/models/registry.py for options resume_path: null # Path to load previous checkpoint/lora adapter resume_type: null # Options: lora, full, state. Null to auto-detect based on `finetune_type` - attn_backend: 'flash' # Options: auto, flash, sdpa + # Bagel forces flash_attention_2 at load (pip install -e ".[bagel]"); it has no diffusers + # attention backend, so it takes no acceleration `attention_backend` entry. # Training Configuration train: diff --git a/examples/grpo/lora/ltx2/t2av.yaml b/examples/grpo/lora/ltx2/t2av.yaml index ef9a7df4..15ac742a 100644 --- a/examples/grpo/lora/ltx2/t2av.yaml +++ b/examples/grpo/lora/ltx2/t2av.yaml @@ -33,7 +33,15 @@ model: model_type: "ltx2_t2av" resume_path: null # Local path or HF repo id (e.g. 'owner/repo[/subdir][@rev]') for previous checkpoint/lora adapter resume_type: null # Options: lora, full, state. Null to auto-detect based on `finetune_type` - # attn_backend: '_flash_3_hub' # Attention backend. +# Optional acceleration plugins (off by default). Applied in list order. +# See guidance/acceleration.md. GRPO is a coupled trainer, so only lossless +# `shared` accelerators are allowed (lossy rollout caching would bias the PPO ratio). +# acceleration: +# shared: # lossless, applied to BOTH rollout and training (in order) +# - name: attention_backend # set the attention backend first... +# params: { backend: _flash_3_hub } +# - name: torch_compile # ...so the compiled graph captures it +# params: { mode: regional } # regional (compile_repeated_blocks) | full log: run_name: null # Run name (auto: {model_type}_{finetune_type}_{trainer_type}_{timestamp}) diff --git a/examples/grpo/lora/ltx2/t2av_pickscore.yaml b/examples/grpo/lora/ltx2/t2av_pickscore.yaml index 86972357..21038e92 100644 --- a/examples/grpo/lora/ltx2/t2av_pickscore.yaml +++ b/examples/grpo/lora/ltx2/t2av_pickscore.yaml @@ -31,7 +31,15 @@ model: model_type: "ltx2_t2av" resume_path: null # Local path or HF repo id (e.g. 'owner/repo[/subdir][@rev]') for previous checkpoint/lora adapter resume_type: null # Options: lora, full, state. Null to auto-detect based on `finetune_type` - # attn_backend: '_flash_3_hub' # Attention backend. +# Optional acceleration plugins (off by default). Applied in list order. +# See guidance/acceleration.md. GRPO is a coupled trainer, so only lossless +# `shared` accelerators are allowed (lossy rollout caching would bias the PPO ratio). +# acceleration: +# shared: # lossless, applied to BOTH rollout and training (in order) +# - name: attention_backend # set the attention backend first... +# params: { backend: _flash_3_hub } +# - name: torch_compile # ...so the compiled graph captures it +# params: { mode: regional } # regional (compile_repeated_blocks) | full log: run_name: null # Run name (auto: {model_type}_{finetune_type}_{trainer_type}_{timestamp}) diff --git a/examples/grpo/lora/qwen_image/default.yaml b/examples/grpo/lora/qwen_image/default.yaml index fb39330f..dd6747ea 100644 --- a/examples/grpo/lora/qwen_image/default.yaml +++ b/examples/grpo/lora/qwen_image/default.yaml @@ -30,7 +30,15 @@ model: model_type: "qwen-image" resume_path: null # Local path or HF repo id (e.g. 'owner/repo[/subdir][@rev]') for previous checkpoint/lora adapter resume_type: null # Options: lora, full, state. Null to auto-detect based on `finetune_type` - # attn_backend: '_flash_3_varlen_hub' # Attention backend for Qwen-Image Series, which uses masked attention with variable sequence length. +# Optional acceleration plugins (off by default). Applied in list order. +# See guidance/acceleration.md. GRPO is a coupled trainer, so only lossless +# `shared` accelerators are allowed (lossy rollout caching would bias the PPO ratio). +# acceleration: +# shared: # lossless, applied to BOTH rollout and training (in order) +# - name: attention_backend # set the attention backend first... +# params: { backend: _flash_3_varlen_hub } +# - name: torch_compile # ...so the compiled graph captures it +# params: { mode: regional } # regional (compile_repeated_blocks) | full log: run_name: null # Run name (auto: {model_type}_{finetune_type}_{trainer_type}_{timestamp}) diff --git a/examples/grpo/lora/qwen_image_edit_plus/default.yaml b/examples/grpo/lora/qwen_image_edit_plus/default.yaml index 274b1270..bfdd6669 100644 --- a/examples/grpo/lora/qwen_image_edit_plus/default.yaml +++ b/examples/grpo/lora/qwen_image_edit_plus/default.yaml @@ -31,7 +31,15 @@ model: model_type: "qwen-image-edit-plus" resume_path: null # Local path or HF repo id (e.g. 'owner/repo[/subdir][@rev]') for previous checkpoint/lora adapter resume_type: null # Options: lora, full, state. Null to auto-detect based on `finetune_type` - # attn_backend: '_flash_3_varlen_hub' # Attention backend for Qwen-Image Series, which uses masked attention with variable sequence length. +# Optional acceleration plugins (off by default). Applied in list order. +# See guidance/acceleration.md. GRPO is a coupled trainer, so only lossless +# `shared` accelerators are allowed (lossy rollout caching would bias the PPO ratio). +# acceleration: +# shared: # lossless, applied to BOTH rollout and training (in order) +# - name: attention_backend # set the attention backend first... +# params: { backend: _flash_3_varlen_hub } +# - name: torch_compile # ...so the compiled graph captures it +# params: { mode: regional } # regional (compile_repeated_blocks) | full log: run_name: null # Run name (auto: {model_type}_{finetune_type}_{trainer_type}_{timestamp}) diff --git a/examples/grpo/lora/sd3_5/default.yaml b/examples/grpo/lora/sd3_5/default.yaml index d4766a2f..6399ce89 100644 --- a/examples/grpo/lora/sd3_5/default.yaml +++ b/examples/grpo/lora/sd3_5/default.yaml @@ -33,7 +33,15 @@ model: model_type: "sd3-5" resume_path: null # Local path or HF repo id (e.g. 'owner/repo[/subdir][@rev]') for previous checkpoint/lora adapter resume_type: null # Options: lora, full, state. Null to auto-detect based on `finetune_type` - # attn_backend: '_flash_3_hub' # Attention backend for training. +# Optional acceleration plugins (off by default). Applied in list order. +# See guidance/acceleration.md. GRPO is a coupled trainer, so only lossless +# `shared` accelerators are allowed (lossy rollout caching would bias the PPO ratio). +# acceleration: +# shared: # lossless, applied to BOTH rollout and training (in order) +# - name: attention_backend # set the attention backend first... +# params: { backend: _flash_3_hub } +# - name: torch_compile # ...so the compiled graph captures it +# params: { mode: regional } # regional (compile_repeated_blocks) | full log: run_name: null # Run name (auto: {model_type}_{finetune_type}_{trainer_type}_{timestamp}) diff --git a/examples/grpo/lora/sd3_5/geneval.yaml b/examples/grpo/lora/sd3_5/geneval.yaml index f0f6d319..3116654b 100644 --- a/examples/grpo/lora/sd3_5/geneval.yaml +++ b/examples/grpo/lora/sd3_5/geneval.yaml @@ -30,7 +30,15 @@ model: model_type: "sd3-5" resume_path: null # Local path or HF repo id (e.g. 'owner/repo[/subdir][@rev]') for previous checkpoint/lora adapter resume_type: null # Options: lora, full, state. Null to auto-detect based on `finetune_type` - # attn_backend: '_flash_3_hub' # Attention backend for training. +# Optional acceleration plugins (off by default). Applied in list order. +# See guidance/acceleration.md. GRPO is a coupled trainer, so only lossless +# `shared` accelerators are allowed (lossy rollout caching would bias the PPO ratio). +# acceleration: +# shared: # lossless, applied to BOTH rollout and training (in order) +# - name: attention_backend # set the attention backend first... +# params: { backend: _flash_3_hub } +# - name: torch_compile # ...so the compiled graph captures it +# params: { mode: regional } # regional (compile_repeated_blocks) | full log: run_name: null # Run name (auto: {model_type}_{finetune_type}_{trainer_type}_{timestamp}) diff --git a/examples/grpo/lora/sd3_5/nocfg.yaml b/examples/grpo/lora/sd3_5/nocfg.yaml index 25618f09..dd7a8be5 100644 --- a/examples/grpo/lora/sd3_5/nocfg.yaml +++ b/examples/grpo/lora/sd3_5/nocfg.yaml @@ -33,7 +33,15 @@ model: model_type: "sd3-5" resume_path: null # Local path or HF repo id (e.g. 'owner/repo[/subdir][@rev]') for previous checkpoint/lora adapter resume_type: null # Options: lora, full, state. Null to auto-detect based on `finetune_type` - # attn_backend: '_flash_3_hub' # Attention backend for training. +# Optional acceleration plugins (off by default). Applied in list order. +# See guidance/acceleration.md. GRPO is a coupled trainer, so only lossless +# `shared` accelerators are allowed (lossy rollout caching would bias the PPO ratio). +# acceleration: +# shared: # lossless, applied to BOTH rollout and training (in order) +# - name: attention_backend # set the attention backend first... +# params: { backend: _flash_3_hub } +# - name: torch_compile # ...so the compiled graph captures it +# params: { mode: regional } # regional (compile_repeated_blocks) | full log: run_name: null # Run name (auto: {model_type}_{finetune_type}_{trainer_type}_{timestamp}) diff --git a/examples/nft/lora/bagel/default.yaml b/examples/nft/lora/bagel/default.yaml index 29dcb510..35af97ff 100644 --- a/examples/nft/lora/bagel/default.yaml +++ b/examples/nft/lora/bagel/default.yaml @@ -37,7 +37,8 @@ model: model_type: "bagel" # Model adapter key; see src/flow_factory/models/registry.py for options resume_path: null # Path to load previous checkpoint/lora adapter resume_type: null # Options: lora, full, state. Null to auto-detect based on `finetune_type` - attn_backend: 'flash' # Options: auto, flash, sdpa + # Bagel forces flash_attention_2 at load (pip install -e ".[bagel]"); it has no diffusers + # attention backend, so it takes no acceleration `attention_backend` entry. # Training Configuration train: diff --git a/examples/nft/lora/bagel/i2i.yaml b/examples/nft/lora/bagel/i2i.yaml index 36c17c7a..23660ae5 100644 --- a/examples/nft/lora/bagel/i2i.yaml +++ b/examples/nft/lora/bagel/i2i.yaml @@ -37,7 +37,8 @@ model: model_type: "bagel" # Model adapter key; see src/flow_factory/models/registry.py for options resume_path: null # Path to load previous checkpoint/lora adapter resume_type: null # Options: lora, full, state. Null to auto-detect based on `finetune_type` - attn_backend: 'flash' # Options: auto, flash, sdpa + # Bagel forces flash_attention_2 at load (pip install -e ".[bagel]"); it has no diffusers + # attention backend, so it takes no acceleration `attention_backend` entry. # Training Configuration train: diff --git a/examples/nft/lora/flux1/default.yaml b/examples/nft/lora/flux1/default.yaml index 886427b3..d221a91e 100644 --- a/examples/nft/lora/flux1/default.yaml +++ b/examples/nft/lora/flux1/default.yaml @@ -30,7 +30,18 @@ model: model_type: "flux1" resume_path: null # Local path or HF repo id (e.g. 'owner/repo[/subdir][@rev]') for previous checkpoint/lora adapter resume_type: null # Options: lora, full, state. Null to auto-detect based on `finetune_type` - # attn_backend: '_flash_3_hub' # Use flash attention 3 backend. +# Optional acceleration plugins (off by default). Applied in list order. +# See guidance/acceleration.md. This is a decoupled trainer, so lossy rollout +# feature-caching is allowed. +# acceleration: +# shared: # lossless, applied to BOTH rollout and training (in order) +# - name: attention_backend # set the attention backend first... +# params: { backend: _flash_3_hub } +# - name: torch_compile # ...so the compiled graph captures it +# params: { mode: regional } # regional (compile_repeated_blocks) | full +# rollout: # rollout-only (Stage 3); lossy -> decoupled/distillation only +# - name: diffusers_cache # Options: diffusers_cache, cache_dit +# params: { policy: first_block, threshold: 0.08 } log: run_name: null # Run name (auto: {model_type}_{finetune_type}_{trainer_type}_{timestamp}) diff --git a/examples/nft/lora/qwen_image/rational_rewards_t2i.yaml b/examples/nft/lora/qwen_image/rational_rewards_t2i.yaml index a6433712..586029b0 100644 --- a/examples/nft/lora/qwen_image/rational_rewards_t2i.yaml +++ b/examples/nft/lora/qwen_image/rational_rewards_t2i.yaml @@ -82,13 +82,18 @@ train: scheduler: dynamics_type: "ODE" -# Optional rollout/throughput acceleration (off by default). See guidance/acceleration.md. -# NFT is a decoupled trainer, so lossy rollout feature-caching is allowed here. +# Optional acceleration plugins (off by default). Applied in list order. +# See guidance/acceleration.md. NFT is a decoupled trainer, so lossy rollout +# feature-caching is allowed. # acceleration: -# shared_accelerator: "torch_compile" # lossless, both stages (attention backend -> model.attn_backend) -# shared_params: { mode: "regional" } # regional (compile_repeated_blocks) | full -# rollout_accelerator: "diffusers_cache" # lossy, rollout-only. Options: diffusers_cache, cache_dit -# rollout_params: { policy: "first_block", threshold: 0.08 } +# shared: # lossless, applied to BOTH rollout and training (in order) +# - name: attention_backend # set the attention backend first... +# params: { backend: _flash_3_varlen_hub } +# - name: torch_compile # ...so the compiled graph captures it +# params: { mode: regional } # regional (compile_repeated_blocks) | full +# rollout: # rollout-only (Stage 3); lossy -> decoupled/distillation only +# - name: diffusers_cache # Options: diffusers_cache, cache_dit +# params: { policy: first_block, threshold: 0.08 } eval: resolution: 512 diff --git a/examples/nft/lora/sd3_5/default.yaml b/examples/nft/lora/sd3_5/default.yaml index bd4fc8bd..46a72dc0 100644 --- a/examples/nft/lora/sd3_5/default.yaml +++ b/examples/nft/lora/sd3_5/default.yaml @@ -31,7 +31,18 @@ model: model_type: "sd3-5" resume_path: null # Local path or HF repo id (e.g. 'owner/repo[/subdir][@rev]') for previous checkpoint/lora adapter resume_type: null # Options: lora, full, state. Null to auto-detect based on `finetune_type` - # attn_backend: '_flash_3_hub' # Attention backend for training. +# Optional acceleration plugins (off by default). Applied in list order. +# See guidance/acceleration.md. This is a decoupled trainer, so lossy rollout +# feature-caching is allowed. +# acceleration: +# shared: # lossless, applied to BOTH rollout and training (in order) +# - name: attention_backend # set the attention backend first... +# params: { backend: _flash_3_hub } +# - name: torch_compile # ...so the compiled graph captures it +# params: { mode: regional } # regional (compile_repeated_blocks) | full +# rollout: # rollout-only (Stage 3); lossy -> decoupled/distillation only +# - name: diffusers_cache # Options: diffusers_cache, cache_dit +# params: { policy: first_block, threshold: 0.08 } log: run_name: null # Run name (auto: {model_type}_{finetune_type}_{timestamp}) diff --git a/examples/nft/lora/wan21/i2v.yaml b/examples/nft/lora/wan21/i2v.yaml index 2e96026e..e54c2e6e 100644 --- a/examples/nft/lora/wan21/i2v.yaml +++ b/examples/nft/lora/wan21/i2v.yaml @@ -31,7 +31,18 @@ model: model_type: "wan2_i2v" # wan2_t2v, wan2_i2v, wan2_v2v resume_path: null # Local path or HF repo id (e.g. 'owner/repo[/subdir][@rev]') for previous checkpoint/lora adapter resume_type: null # Options: lora, full, state. Null to auto-detect based on `finetune_type` - # attn_backend: '_flash_3_hub' # Use flash attention 3 backend. +# Optional acceleration plugins (off by default). Applied in list order. +# See guidance/acceleration.md. This is a decoupled trainer, so lossy rollout +# feature-caching is allowed. +# acceleration: +# shared: # lossless, applied to BOTH rollout and training (in order) +# - name: attention_backend # set the attention backend first... +# params: { backend: _flash_3_hub } +# - name: torch_compile # ...so the compiled graph captures it +# params: { mode: regional } # regional (compile_repeated_blocks) | full +# rollout: # rollout-only (Stage 3); lossy -> decoupled/distillation only +# - name: diffusers_cache # Options: diffusers_cache, cache_dit +# params: { policy: first_block, threshold: 0.08 } log: run_name: null # Run name (auto: {model_type}_{finetune_type}_{trainer_type}_{timestamp}) diff --git a/examples/nft/lora/wan21/t2v.yaml b/examples/nft/lora/wan21/t2v.yaml index c536fa6f..72106e22 100644 --- a/examples/nft/lora/wan21/t2v.yaml +++ b/examples/nft/lora/wan21/t2v.yaml @@ -31,7 +31,18 @@ model: model_type: "wan2_t2v" # wan2_t2v, wan2_i2v, wan2_v2v resume_path: null # Local path or HF repo id (e.g. 'owner/repo[/subdir][@rev]') for previous checkpoint/lora adapter resume_type: null # Options: lora, full, state. Null to auto-detect based on `finetune_type` - # attn_backend: '_flash_3_hub' # Use flash attention 3 backend. +# Optional acceleration plugins (off by default). Applied in list order. +# See guidance/acceleration.md. This is a decoupled trainer, so lossy rollout +# feature-caching is allowed. +# acceleration: +# shared: # lossless, applied to BOTH rollout and training (in order) +# - name: attention_backend # set the attention backend first... +# params: { backend: _flash_3_hub } +# - name: torch_compile # ...so the compiled graph captures it +# params: { mode: regional } # regional (compile_repeated_blocks) | full +# rollout: # rollout-only (Stage 3); lossy -> decoupled/distillation only +# - name: diffusers_cache # Options: diffusers_cache, cache_dit +# params: { policy: first_block, threshold: 0.08 } log: run_name: null # Run name (auto: {model_type}_{finetune_type}_{trainer_type}_{timestamp}) diff --git a/examples/nft/lora/wan22/t2v.yaml b/examples/nft/lora/wan22/t2v.yaml index 8a77b03e..a1bcd7c7 100644 --- a/examples/nft/lora/wan22/t2v.yaml +++ b/examples/nft/lora/wan22/t2v.yaml @@ -31,7 +31,18 @@ model: model_type: "wan2_t2v" # wan2_t2v, wan2_i2v, wan2_v2v resume_path: null # Local path or HF repo id (e.g. 'owner/repo[/subdir][@rev]') for previous checkpoint/lora adapter resume_type: null # Options: lora, full, state. Null to auto-detect based on `finetune_type` - # attn_backend: '_flash_3_hub' # Use flash attention 3 backend. +# Optional acceleration plugins (off by default). Applied in list order. +# See guidance/acceleration.md. This is a decoupled trainer, so lossy rollout +# feature-caching is allowed. +# acceleration: +# shared: # lossless, applied to BOTH rollout and training (in order) +# - name: attention_backend # set the attention backend first... +# params: { backend: _flash_3_hub } +# - name: torch_compile # ...so the compiled graph captures it +# params: { mode: regional } # regional (compile_repeated_blocks) | full +# rollout: # rollout-only (Stage 3); lossy -> decoupled/distillation only +# - name: diffusers_cache # Options: diffusers_cache, cache_dit +# params: { policy: first_block, threshold: 0.08 } log: run_name: null # Run name (auto: {model_type}_{finetune_type}_{trainer_type}_{timestamp}) diff --git a/examples/template/sd3_5/async_reward.yaml b/examples/template/sd3_5/async_reward.yaml index 1fa18625..336b14d6 100644 --- a/examples/template/sd3_5/async_reward.yaml +++ b/examples/template/sd3_5/async_reward.yaml @@ -30,7 +30,15 @@ model: model_type: "sd3-5" resume_path: null # Local path or HF repo id (e.g. 'owner/repo[/subdir][@rev]') for previous checkpoint/lora adapter resume_type: null # Options: lora, full, state. Null to auto-detect based on `finetune_type` - # attn_backend: '_flash_3_hub' # Attention backend for training. +# Optional acceleration plugins (off by default). Applied in list order. +# See guidance/acceleration.md. GRPO is a coupled trainer, so only lossless +# `shared` accelerators are allowed (lossy rollout caching would bias the PPO ratio). +# acceleration: +# shared: # lossless, applied to BOTH rollout and training (in order) +# - name: attention_backend # set the attention backend first... +# params: { backend: _flash_3_hub } +# - name: torch_compile # ...so the compiled graph captures it +# params: { mode: regional } # regional (compile_repeated_blocks) | full log: run_name: null # Run name (auto: {model_type}_{finetune_type}_{trainer_type}_{timestamp}) diff --git a/guidance/acceleration.md b/guidance/acceleration.md index 060b58cb..62964ac5 100644 --- a/guidance/acceleration.md +++ b/guidance/acceleration.md @@ -35,46 +35,64 @@ but not the other, which is exactly what `stage='rollout'` + `safety='lossy'` ca ## Configuration -Add an optional `acceleration:` block to any config. Two independent slots, both off by -default: +Add an optional `acceleration:` block to any config. Two independent slots, each an +**ordered list** of `{name, params}` entries (both empty by default). **List order is the +application order**: ```yaml acceleration: - # Lossless, applied to BOTH rollout and the training forward. - shared_accelerator: "torch_compile" # torch_compile - shared_params: { mode: "regional" } # regional (compile_repeated_blocks) | full - - # Rollout-only (Stage 3). May be lossy (paradigm-gated). - rollout_accelerator: "diffusers_cache" # diffusers_cache | cache_dit | - rollout_params: { policy: "first_block", threshold: 0.08 } + # Lossless, applied to BOTH rollout and the training forward, in list order. + shared: + - name: attention_backend # set the diffusers backend first... + params: { backend: _flash_3_hub } + - name: torch_compile # ...so the compiled graph captures it + params: { mode: regional } # regional (compile_repeated_blocks) | full + + # Rollout-only (Stage 3), nested in list order. May be lossy (paradigm-gated). + rollout: + - name: diffusers_cache # diffusers_cache | cache_dit | + params: { policy: first_block, threshold: 0.08 } ``` -Either slot may be omitted. A direct python path (e.g. `my_pkg.accel.MyAccelerator`) is -accepted in place of a registered id. +Either slot may be omitted or left empty. A single entry dict (without the list dashes) is +accepted as shorthand for a one-element list. A direct python path (e.g. +`my_pkg.accel.MyAccelerator`) is accepted in place of a registered id. ## Available accelerators | id | safety | stage | Notes | |----|--------|-------|-------| -| `attention_backend` | lossless | both | Sets the diffusers attention backend on every transformer. Configured via `model.attn_backend` (see below); `backend` param can override. Forwards any backend (`native` / `flash` / `_flash_3` / `_flash_3_hub` / `sage` / `xformers`) to `set_attention_backend`. | +| `attention_backend` | lossless | both | Sets the diffusers attention backend on every transformer. Requires a `backend` param. Forwards any backend (`native` / `flash` / `_flash_3` / `_flash_3_hub` / `sage` / `xformers`) to `set_attention_backend`. List it in `shared` **before** `torch_compile` so the compiled graph captures the backend. | | `torch_compile` | lossless | both | `torch.compile` of the shared transformer. `mode: regional` uses diffusers' `compile_repeated_blocks` (fast warmup, robust to variable resolution); `mode: full` compiles the whole module. Extra `compile_kwargs` forwarded to the compile call. Compiles in place (checkpoint- and EMA/ref-safe), applied after `post_init`. | | `diffusers_cache` | lossy | rollout | Diffusers-native feature caching (no extra dependency). `policy`: `first_block` (default) / `faster` / `pyramid` / `taylorseer` / `magcache`; remaining params forwarded to the policy's diffusers config (e.g. `threshold`). | | `cache_dit` | lossy | rollout | [cache-dit](https://github.com/vipshop/cache-dit) backend (DBCache/TaylorSeer). Requires `pip install flow-factory[acceleration]`. All params forwarded to `cache_dit.enable_cache`. | ### Attention backend -Attention-backend selection keeps its dedicated config knob, `model.attn_backend`: +Attention-backend selection is a `shared` accelerator (it transforms the module shared by +rollout and training, so it is consistent for any algorithm — even an approximate kernel +like Sage int8). It used to live under the dedicated `model.attn_backend` knob; that knob +was **removed** and folded into the acceleration layer: ```yaml -model: - attn_backend: "_flash_3_hub" # native | flash | flash_hub | _flash_3 | _flash_3_hub | sage | xformers +acceleration: + shared: + - name: attention_backend + params: { backend: _flash_3_hub } # native | flash | flash_hub | _flash_3 | _flash_3_hub | sage | xformers + - name: torch_compile # optional; if present, list it AFTER attention_backend + params: { mode: regional } ``` It is applied through `AttentionBackendAccelerator` by the trainer -(`BaseTrainer._apply_shared_acceleration`) — after `accelerator.prepare` / `post_init` -and **before** compile, so the compiled graph captures the chosen backend. This is the -single code path for backend selection (the old `BaseAdapter._set_attention_backend` was -removed); it applies whether or not an `acceleration:` block is present. +(`BaseTrainer._apply_shared_acceleration`) — after `accelerator.prepare` / `post_init`, in +list order. Place it **before** `torch_compile` so the compiled graph captures the chosen +backend. This is the single code path for backend selection (the old +`BaseAdapter._set_attention_backend` and `model.attn_backend` were both removed). A config +that still sets `model.attn_backend` fails fast with a migration error. + +> **Bagel** forces `flash_attention_2` at model load (requires `pip install -e ".[bagel]"`) +> and its custom transformer has no `set_attention_backend`, so it does **not** take an +> `attention_backend` entry — omit it (the accelerator raises if applied to bagel). ## Model cache-readiness (lossy caching) diff --git a/multinode_examples/train.yaml b/multinode_examples/train.yaml index 538a994d..d7cfe81a 100644 --- a/multinode_examples/train.yaml +++ b/multinode_examples/train.yaml @@ -30,7 +30,16 @@ model: model_type: "wan2_t2v" # wan2_t2v, wan2_i2v, wan2_v2v resume_path: null # Path to load previous checkpoint/lora adapter resume_type: null # Options: lora, full, state. Null to auto-detect based on `finetune_type` - attn_backend: '_flash_3_hub' # Attention backend for training. + +# Optional acceleration plugins (off by default). Applied in list order. +# See guidance/acceleration.md. GRPO is a coupled trainer, so only lossless +# `shared` accelerators are allowed (lossy rollout caching would bias the PPO ratio). +# acceleration: +# shared: # lossless, applied to BOTH rollout and training (in order) +# - name: attention_backend # set the attention backend first... +# params: { backend: _flash_3_hub } +# - name: torch_compile # ...so the compiled graph captures it +# params: { mode: regional } # regional (compile_repeated_blocks) | full # Training Configuration train: diff --git a/src/flow_factory/acceleration/attention_backend.py b/src/flow_factory/acceleration/attention_backend.py index 66b9f295..6996b98b 100644 --- a/src/flow_factory/acceleration/attention_backend.py +++ b/src/flow_factory/acceleration/attention_backend.py @@ -16,15 +16,15 @@ """Attention-backend accelerator — the single code path that selects the diffusers attention backend for every transformer. -This replaces the old ``BaseAdapter._set_attention_backend`` call: the backend is -applied here (after ``accelerator.prepare`` / ``post_init`` and before compile) -instead of in the adapter constructor, so all transformer-level acceleration flows -through the same plugin mechanism. +This replaces the old ``BaseAdapter._set_attention_backend`` call and the +``model.attn_backend`` knob: the backend is now requested as an ``attention_backend`` +entry in the acceleration ``shared`` list and applied here (after +``accelerator.prepare`` / ``post_init`` and before compile), so all transformer-level +acceleration flows through the same plugin mechanism. -The backend is read from ``model.attn_backend`` by default (so existing configs -keep working) and may be overridden by an explicit ``backend`` param. Whatever -string is given is forwarded to diffusers' ``set_attention_backend`` verbatim — -including approximate backends like ``sage`` — matching the previous behavior. +The backend name is taken from the required ``backend`` param and forwarded to +diffusers' ``set_attention_backend`` verbatim — including approximate backends like +``sage`` — matching the previous behavior. Marked ``stage='both'`` / ``safety='lossless'``: a backend is applied to the transformer shared by rollout ``inference()`` and training ``forward()``, so the @@ -46,10 +46,10 @@ class AttentionBackendAccelerator(BaseAccelerator): """Set the diffusers attention backend on every transformer component. - Parameters (from ``acceleration.shared_params``, all optional): + Parameters (from the entry's ``params``): backend: Backend name forwarded to ``transformer.set_attention_backend`` (e.g. ``native`` / ``flash`` / ``_flash_3`` / ``_flash_3_hub`` / - ``sage`` / ``xformers``). Defaults to ``model.attn_backend``. + ``sage`` / ``xformers``). Required. See https://huggingface.co/docs/diffusers/main/en/optimization/attention_backends for the full list of supported backends. @@ -59,10 +59,12 @@ class AttentionBackendAccelerator(BaseAccelerator): stage = "both" def setup(self, adapter: "BaseAdapter") -> None: - backend = self.params.get("backend") or adapter.model_args.attn_backend - if backend is None: - # Nothing requested (no `backend` param and `model.attn_backend` unset). - return + backend = self.params.get("backend") + if not backend: + raise ValueError( + "AttentionBackendAccelerator requires a `backend` param, e.g. " + "`{ name: attention_backend, params: { backend: _flash_3_hub } }`." + ) applied = False for name in adapter.transformer_names: @@ -75,8 +77,10 @@ def setup(self, adapter: "BaseAdapter") -> None: "AttentionBackendAccelerator: set backend '%s' for '%s'.", backend, name ) if not applied: - logger.warning( - "AttentionBackendAccelerator: backend '%s' requested but no transformer component " - "supports `set_attention_backend`; leaving the diffusers default.", - backend, + raise ValueError( + f"AttentionBackendAccelerator: backend '{backend}' requested but none of the " + f"adapter's transformer components {adapter.transformer_names} support " + "`set_attention_backend`. Models with a custom attention implementation (e.g. " + "Bagel, which forces flash_attention_2 at load) must not use this accelerator; " + "remove the `attention_backend` entry from the acceleration config." ) diff --git a/src/flow_factory/acceleration/diffusers_cache.py b/src/flow_factory/acceleration/diffusers_cache.py index 47993f83..e56b2120 100644 --- a/src/flow_factory/acceleration/diffusers_cache.py +++ b/src/flow_factory/acceleration/diffusers_cache.py @@ -62,7 +62,7 @@ class DiffusersCacheAccelerator(BaseAccelerator): Lossy and rollout-scoped. The default policy is ``first_block`` (FirstBlockCache, aka FBCache) which is robust across models and needs only a single ``threshold``. - Parameters (from ``acceleration.rollout_params``): + Parameters (from the entry's ``params``): policy: One of ``first_block`` / ``faster`` / ``pyramid`` / ``taylorseer`` / ``magcache``. Defaults to ``first_block``. : All remaining params are forwarded verbatim to the selected diff --git a/src/flow_factory/acceleration/registry.py b/src/flow_factory/acceleration/registry.py index 1565735d..0fa8b388 100644 --- a/src/flow_factory/acceleration/registry.py +++ b/src/flow_factory/acceleration/registry.py @@ -33,7 +33,7 @@ _ACCELERATOR_REGISTRY: Dict[str, str] = { # Lossless (safe for any algorithm, applied to the shared transformer). # `attention_backend` is the single code path for attention-backend selection; - # it is auto-applied from `model.attn_backend` by the trainer (before compile). + # add it as a `shared` entry (before `torch_compile`) in the acceleration block. 'attention_backend': 'flow_factory.acceleration.attention_backend.AttentionBackendAccelerator', 'torch_compile': 'flow_factory.acceleration.torch_compile.CompileAccelerator', # Lossy (rollout-only; validator restricts to decoupled / distillation algos). diff --git a/src/flow_factory/acceleration/torch_compile.py b/src/flow_factory/acceleration/torch_compile.py index 3cde20bb..4c93f24c 100644 --- a/src/flow_factory/acceleration/torch_compile.py +++ b/src/flow_factory/acceleration/torch_compile.py @@ -48,7 +48,7 @@ class CompileAccelerator(BaseAccelerator): Applied after ``post_init`` (see ``BaseTrainer._apply_shared_acceleration``), so it wraps the final, fully-loaded weights. - Parameters (from ``acceleration.shared_params``): + Parameters (from the entry's ``params``): mode: ``"regional"`` (default) compiles only the repeated transformer blocks via diffusers' ``compile_repeated_blocks`` — fast warmup and robust to the variable image/sequence lengths set per resolution. diff --git a/src/flow_factory/acceleration/validator.py b/src/flow_factory/acceleration/validator.py index f25c42fa..850f3c63 100644 --- a/src/flow_factory/acceleration/validator.py +++ b/src/flow_factory/acceleration/validator.py @@ -86,7 +86,7 @@ def validate_accelerator( raise ValueError( f"Accelerator '{name}' (stage='{accelerator.stage}') cannot occupy the shared " "slot, which applies a persistent transform to both rollout and training. Put a " - "stage='both' accelerator here, or move this one to `acceleration.rollout_accelerator`." + "stage='both' accelerator here, or move this one to the `acceleration.rollout` list." ) # No safety gate here: a stage='both' transform is applied identically to # rollout `inference()` and training `forward()` (the same shared module), so @@ -99,8 +99,8 @@ def validate_accelerator( if accelerator.stage != "rollout": raise ValueError( f"Accelerator '{name}' (stage='{accelerator.stage}') cannot occupy the rollout slot, " - "which only runs a per-epoch `rollout_context`. Put a stage='both' accelerator under " - "`acceleration.shared_accelerator` instead." + "which only runs a per-epoch `rollout_context`. Put a stage='both' accelerator in the " + "`acceleration.shared` list instead." ) # A rollout-only accelerator that changes outputs (`lossy`) makes rollout diverge diff --git a/src/flow_factory/hparams/__init__.py b/src/flow_factory/hparams/__init__.py index ab21973f..f7f29d2c 100644 --- a/src/flow_factory/hparams/__init__.py +++ b/src/flow_factory/hparams/__init__.py @@ -33,7 +33,7 @@ get_training_args_class, ) from .reward_args import RewardArguments, MultiRewardArguments -from .acceleration_args import AccelerationArguments +from .acceleration_args import AccelerationArguments, AccelerationSpec from .dataset_args import DatasetArguments, DatasetTrainSpec, DatasetEvalSpec from .log_args import LogArguments @@ -57,6 +57,7 @@ "RewardArguments", "MultiRewardArguments", "AccelerationArguments", + "AccelerationSpec", "DatasetArguments", "DatasetTrainSpec", "DatasetEvalSpec", diff --git a/src/flow_factory/hparams/acceleration_args.py b/src/flow_factory/hparams/acceleration_args.py index 4512614e..85031eff 100644 --- a/src/flow_factory/hparams/acceleration_args.py +++ b/src/flow_factory/hparams/acceleration_args.py @@ -15,65 +15,126 @@ # src/flow_factory/hparams/acceleration_args.py import yaml from dataclasses import dataclass, field -from typing import Any, Dict, Optional +from typing import Any, Dict, List, Union from .abc import ArgABC +@dataclass +class AccelerationSpec(ArgABC): + r"""One accelerator entry: an id (or python path) plus its constructor params. + + Mirrors the ``{name, params}`` shape of a reward entry (see + :class:`~flow_factory.hparams.reward_args.RewardArguments`). The ``name`` + resolves through the accelerator registry + (:mod:`flow_factory.acceleration.registry`); ``params`` is forwarded verbatim + to the accelerator constructor. + + Example YAML:: + + - name: attention_backend + params: { backend: _flash_3_hub } + """ + + name: str = "" + params: Dict[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + if not self.name: + raise ValueError( + "Each acceleration entry must declare a non-empty `name` " + "(e.g. `attention_backend`, `torch_compile`, `diffusers_cache`)." + ) + + def to_dict(self) -> dict[str, Any]: + return {"name": self.name, "params": dict(self.params)} + + +def _parse_specs( + value: Union[Dict[str, Any], List[Dict[str, Any]], None], +) -> List[AccelerationSpec]: + """Parse a slot value into an ordered list of :class:`AccelerationSpec`. + + Accepts a list of entries (canonical) or a single entry dict (shorthand for a + one-element list), mirroring ``MultiRewardArguments.from_dict``. List order is + the application order. + + Args: + value: The raw ``shared`` / ``rollout`` value from the config. + + Returns: + Ordered list of specs (empty when ``value`` is ``None``). + + Raises: + ValueError: If ``value`` is neither a list nor a dict. + """ + if value is None: + return [] + if isinstance(value, dict): + return [AccelerationSpec.from_dict(value)] + if isinstance(value, list): + return [AccelerationSpec.from_dict(entry) for entry in value] + raise ValueError( + f"Acceleration slot must be a list of {{name, params}} entries or a single " + f"entry dict; got {type(value).__name__}." + ) + + @dataclass class AccelerationArguments(ArgABC): r"""Arguments for the model-agnostic acceleration plugin layer. - Two independent slots, both off by default (so existing configs are - unaffected): + Two independent slots, each an **ordered list** of accelerator entries (both + empty by default, so existing configs are unaffected). List order is the + application order: - * ``shared_*`` — a lossless accelerator applied to BOTH rollout and the - training forward (e.g. ``torch_compile``). - * ``rollout_*`` — an accelerator applied only during Stage-3 rollout. May be - lossy (e.g. feature caching), in which case the trainer paradigm validator - restricts it to decoupled / distillation algorithms. + * ``shared`` — lossless accelerators applied to BOTH rollout and the training + forward, in order, via each accelerator's ``setup()`` (e.g. + ``attention_backend`` then ``torch_compile`` — backend first so the compiled + graph captures it). + * ``rollout`` — accelerators applied only during Stage-3 rollout, nested in + order via each accelerator's ``rollout_context()``. May be lossy (e.g. + feature caching), in which case the trainer-paradigm validator restricts them + to decoupled / distillation algorithms. Example YAML:: acceleration: - shared_accelerator: torch_compile - shared_params: { mode: regional } - rollout_accelerator: diffusers_cache - rollout_params: { policy: first_block, threshold: 0.08 } + shared: + - name: attention_backend + params: { backend: _flash_3_hub } + - name: torch_compile + params: { mode: regional } + rollout: + - name: diffusers_cache + params: { policy: first_block, threshold: 0.08 } """ - shared_accelerator: Optional[str] = field( - default=None, - metadata={ - "help": ( - "Lossless accelerator id (or python path) applied to both rollout and training. " - "Options: 'torch_compile'. None disables it. " - "(Attention backend has its own knob, model.attn_backend, applied " - "automatically before this.)" - ) - }, - ) - shared_params: Dict[str, Any] = field( - default_factory=dict, - metadata={"help": "Keyword parameters forwarded to the shared accelerator constructor."}, - ) - rollout_accelerator: Optional[str] = field( - default=None, - metadata={ - "help": ( - "Accelerator id (or python path) applied only during Stage-3 rollout. " - "Options: 'diffusers_cache', 'cache_dit' (lossy, decoupled/distillation only), or any " - "lossless accelerator. None disables it." - ) - }, - ) - rollout_params: Dict[str, Any] = field( - default_factory=dict, - metadata={"help": "Keyword parameters forwarded to the rollout accelerator constructor."}, - ) + shared: List[AccelerationSpec] = field(default_factory=list) + rollout: List[AccelerationSpec] = field(default_factory=list) + + @classmethod + def from_dict(cls, args_dict: Dict[str, Any]) -> "AccelerationArguments": + """Build from a config dict, parsing each slot into an ordered spec list. + + Args: + args_dict: The ``acceleration`` config block. + + Returns: + An ``AccelerationArguments`` with parsed ``shared`` / ``rollout`` lists. + """ + args_dict = dict(args_dict or {}) + return cls( + shared=_parse_specs(args_dict.pop("shared", None)), + rollout=_parse_specs(args_dict.pop("rollout", None)), + extra_kwargs=args_dict, + ) def to_dict(self) -> dict[str, Any]: - return super().to_dict() + return { + "shared": [spec.to_dict() for spec in self.shared], + "rollout": [spec.to_dict() for spec in self.rollout], + } def __str__(self) -> str: return yaml.dump(self.to_dict(), default_flow_style=False, sort_keys=False, indent=2) diff --git a/src/flow_factory/hparams/model_args.py b/src/flow_factory/hparams/model_args.py index e70bccc1..267343f2 100644 --- a/src/flow_factory/hparams/model_args.py +++ b/src/flow_factory/hparams/model_args.py @@ -107,17 +107,19 @@ class ModelArguments(ArgABC): } ) - attn_backend: Optional[str] = field( - default=None, - metadata={ - "help": "Attention backend for transformers. " - "Options: 'native', 'flash', 'flash_hub', '_flash_3', '_flash_3_hub', 'sage', 'xformers'. " - "None means use diffusers default." - "See https://huggingface.co/docs/diffusers/main/en/optimization/attention_backends for all details." - }, - ) - def __post_init__(self): + if "attn_backend" in self.extra_kwargs: + raise ValueError( + "`model.attn_backend` has been removed. Attention-backend selection now lives in " + "the acceleration layer as a `shared` accelerator. Replace it with:\n" + " acceleration:\n" + " shared:\n" + " - name: attention_backend\n" + f" params: {{ backend: {self.extra_kwargs['attn_backend']!r} }}\n" + "See guidance/acceleration.md. (Bagel forces flash_attention_2 at load and ignores " + "this knob — just drop the line.)" + ) + if isinstance(self.master_weight_dtype, str): self.master_weight_dtype = dtype_map[self.master_weight_dtype] diff --git a/src/flow_factory/models/abc.py b/src/flow_factory/models/abc.py index 5dbdc6d9..c96d1e14 100644 --- a/src/flow_factory/models/abc.py +++ b/src/flow_factory/models/abc.py @@ -195,9 +195,10 @@ def __init__(self, config: Arguments, accelerator : Accelerator): self._mix_precision() # NOTE: attention-backend selection is applied later by the trainer's - # acceleration step (AttentionBackendAccelerator, from `model.attn_backend`), - # after accelerator.prepare()/post_init() and before torch.compile — see - # BaseTrainer._apply_shared_acceleration(). It is intentionally NOT set here. + # acceleration step (AttentionBackendAccelerator, configured as a `shared` + # entry in the acceleration block), after accelerator.prepare()/post_init() + # and before torch.compile — see BaseTrainer._apply_shared_acceleration(). + # It is intentionally NOT set here. # Enable gradient checkpointing if needed if self.training_args.enable_gradient_checkpointing: diff --git a/src/flow_factory/trainers/abc.py b/src/flow_factory/trainers/abc.py index 6dec9253..2a388b2a 100644 --- a/src/flow_factory/trainers/abc.py +++ b/src/flow_factory/trainers/abc.py @@ -16,7 +16,7 @@ import json import os from abc import ABC, abstractmethod -from contextlib import AbstractContextManager, nullcontext +from contextlib import ExitStack, contextmanager from typing import Dict, Any, ClassVar, Optional, Tuple, List, Union, Literal, Iterator from functools import partial import numpy as np @@ -391,92 +391,83 @@ def _initialization(self): def _init_acceleration(self): """Build and validate acceleration plugins from ``config.acceleration_args``. - Two independent slots (both off by default): + Two independent slots, each an **ordered list** (both empty by default). + List order is the application order: - * ``shared_accelerator`` — a lossless accelerator (e.g. ``torch.compile``) - that affects both rollout and the training forward. Only built/validated - here; it is *applied* later by :meth:`_apply_shared_acceleration` (after - ``post_init`` finishes state-resume / EMA / reference setup), so it - transforms the final weights. - * ``rollout_accelerator`` — applied per-epoch in :meth:`generate_samples` + * ``shared`` — lossless accelerators (e.g. ``attention_backend`` then + ``torch_compile``) that affect both rollout and the training forward. Only + built/validated here; they are *applied* later by + :meth:`_apply_shared_acceleration` (after ``post_init`` finishes + state-resume / EMA / reference setup), so they transform the final weights. + * ``rollout`` — accelerators applied per-epoch in :meth:`generate_samples` via :meth:`~BaseAccelerator.rollout_context`; may be lossy. Each accelerator is validated against this trainer's ``paradigm`` before use (fail-fast, ``constraints.md`` #26). """ accel_args = getattr(self.config, "acceleration_args", None) - self.shared_accelerator: Optional[BaseAccelerator] = None - self.rollout_accelerator: Optional[BaseAccelerator] = None + self.shared_accelerators: List[BaseAccelerator] = [] + self.rollout_accelerators: List[BaseAccelerator] = [] if accel_args is None: return trainer_name = type(self).__name__ paradigm = type(self).paradigm - if accel_args.shared_accelerator: - accelerator = build_accelerator( - accel_args.shared_accelerator, accel_args.shared_params - ) + for spec in accel_args.shared: + accelerator = build_accelerator(spec.name, spec.params) validate_accelerator( accelerator, slot="shared", paradigm=paradigm, trainer_name=trainer_name ) - self.shared_accelerator = accelerator + self.shared_accelerators.append(accelerator) - if accel_args.rollout_accelerator: - accelerator = build_accelerator( - accel_args.rollout_accelerator, accel_args.rollout_params - ) + for spec in accel_args.rollout: + accelerator = build_accelerator(spec.name, spec.params) validate_accelerator( accelerator, slot="rollout", paradigm=paradigm, trainer_name=trainer_name ) - self.rollout_accelerator = accelerator + self.rollout_accelerators.append(accelerator) if self.accelerator.is_main_process: logger.info( "Acceleration: rollout accelerator '%s' (safety=%s) enabled.", - accel_args.rollout_accelerator, + spec.name, accelerator.safety, ) def _apply_shared_acceleration(self) -> None: - """Apply transformer-level (shared) acceleration to the adapter. + """Apply the shared (lossless) accelerators to the adapter, in config order. Called from ``__init__`` AFTER ``adapter.post_init()`` so transforms wrap the final weights — i.e. after ``accelerator.prepare``, any ``state`` checkpoint - resume, and EMA / reference-parameter snapshotting. Applied in order: - - 1. **Attention backend** from ``model.attn_backend`` — the single code path - for backend selection (replaces ``BaseAdapter._set_attention_backend``), - run here so it sits after ``prepare`` and before compile. - 2. The configured **shared accelerator** (e.g. ``torch.compile``), applied - last so it captures the chosen attention backend. - - In-place compilation (``nn.Module.compile`` / ``compile_repeated_blocks``) - preserves parameter identity and ``state_dict`` keys, so checkpointing and - the ``copy_``-based EMA / ref / named-parameter swaps stay correct. + resume, and EMA / reference-parameter snapshotting. + + Each entry's ``setup`` runs in list order, so a config that lists + ``attention_backend`` before ``torch_compile`` sets the backend first and + then compiles the graph capturing it. In-place compilation + (``nn.Module.compile`` / ``compile_repeated_blocks``) preserves parameter + identity and ``state_dict`` keys, so checkpointing and the ``copy_``-based + EMA / ref / named-parameter swaps stay correct. """ - # 1. Attention backend (lossless / both stages), sourced from model.attn_backend. - # Always safe (applied to the shared transformer), so it bypasses the - # paradigm validator; AttentionBackendAccelerator.setup is a no-op when - # `model.attn_backend` is unset. - if self.model_args.attn_backend: - build_accelerator("attention_backend", {}).setup(self.adapter) - - # 2. Configured shared accelerator (e.g. torch.compile), applied after the - # backend so the compiled graph captures it. - accelerator = getattr(self, "shared_accelerator", None) - if accelerator is not None: + for accelerator in self.shared_accelerators: accelerator.setup(self.adapter) if self.accelerator.is_main_process: logger.info( - "Acceleration: shared accelerator (safety=%s) applied to adapter.", + "Acceleration: shared accelerator '%s' (safety=%s) applied to adapter.", + type(accelerator).__name__, accelerator.safety, ) - def _rollout_acceleration(self) -> AbstractContextManager: - """Return the rollout accelerator's context, or a no-op when disabled.""" - if self.rollout_accelerator is not None: - return self.rollout_accelerator.rollout_context(self.adapter) - return nullcontext() + @contextmanager + def _rollout_acceleration(self) -> Iterator[None]: + """Nest every rollout accelerator's context (first in list = outermost). + + A no-op when no rollout accelerator is configured. + """ + with ExitStack() as stack: + for accelerator in self.rollout_accelerators: + stack.enter_context(accelerator.rollout_context(self.adapter)) + yield + def _synchronize_frozen_components(self): if self.accelerator.num_processes <= 1: From becc159c3ba68406115c46fff84c7359ed1c2906 Mon Sep 17 00:00:00 2001 From: Jayce-Ping <315229706@qq.com> Date: Mon, 29 Jun 2026 16:28:28 +0800 Subject: [PATCH 09/15] [acceleration,trainers,models,docs] fix: bit-exact torch.compile train-inference consistency for coupled algorithms MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit torch.compile (Inductor) compiles a separate, numerically non-identical graph for grad vs no-grad mode (Dynamo guards on grad_mode). Rollout ran under no_grad while the training forward runs with grad, so the recomputed on-policy log-prob diverged from the rollout log-prob (~1.5e-5) and broke the coupled PPO ratio==1 invariant (GRPO/GRPO-Guard/DPPO). This is the ONLY source — CFG is applied identically in both stages (batch 2N each, verified) and is not a separate cause. Fix (CompileAccelerator + trainer + base adapter): - requires_grad_rollout flag: the trainer's _rollout_grad_context runs the rollout with grad enabled and sets adapter._rollout_detach when a compile accelerator is active (otherwise unchanged: no_grad). - The compiled transformer's call entry points (forward + _compiled_call_impl) are wrapped to force torch.enable_grad(), overriding the @torch.no_grad() on inference(), so rollout and training execute the same grad-mode graph. - The wrapper returns the grad-carrying output directly and does NOT detach it: an inner detach lets Inductor pick a divergent inference-optimized kernel and silently re-introduces the drift. Instead the per-step latent feedback is detached in BaseAdapter.cast_latents (gated by _rollout_detach), which breaks the autograd graph chain across denoising steps so rollout memory stays bounded. - Compile/wrap the BASE transformer under PEFT/LoRA (_peel_peft -> get_base_model), not the PeftModel: compiling the wrapper drifts (~2.0) and hides _repeated_blocks (so regional compile now also works under LoRA). LoRA submodules stay inside the compiled graph and still train. Result (verified, SD3.5, first on-policy step): max|ratio-1| = 0.000e+00 (bit-exact) for coupled training with AND without CFG (old_lp == new_lp exactly; rollout and training both call the transformer at the identical 2N shape under CFG). A train-vs-train recompute is exactly 0.0, confirming compile is deterministic (no RNG/autotune nondeterminism). Non-compile runs are unchanged. Full analysis: .scratch/torch_compile_consistency_report.md; guidance/acceleration.md updated. Co-Authored-By: Claude Opus 4.8 (1M context) --- guidance/acceleration.md | 28 +++++ src/flow_factory/acceleration/abc.py | 15 +++ .../acceleration/torch_compile.py | 106 +++++++++++++++++- src/flow_factory/models/abc.py | 13 ++- src/flow_factory/trainers/abc.py | 45 +++++++- 5 files changed, 202 insertions(+), 5 deletions(-) diff --git a/guidance/acceleration.md b/guidance/acceleration.md index 62964ac5..d0fab993 100644 --- a/guidance/acceleration.md +++ b/guidance/acceleration.md @@ -94,6 +94,34 @@ that still sets `model.attn_backend` fails fast with a migration error. > and its custom transformer has no `set_attention_backend`, so it does **not** take an > `attention_backend` entry — omit it (the accelerator raises if applied to bagel). +### torch.compile train-inference consistency (coupled algorithms) + +For coupled algorithms (GRPO / GRPO-Guard / DPPO) the on-policy PPO ratio on the first inner +step must be **1.0**. `torch.compile` (Inductor) threatens this because it compiles a +**separate, numerically non-identical graph for grad vs no-grad mode** (Dynamo guards on +`grad_mode`): rollout normally runs the transformer under `torch.no_grad()` and the training +forward under grad, so a naive compiled rollout would diverge from training (~1e-5) and bias +the ratio. + +`CompileAccelerator` handles this automatically (no user action needed): it declares +`requires_grad_rollout = True`, so the trainer runs the rollout transformer under +`torch.enable_grad()` (overriding the `@torch.no_grad()` on `inference()`) and detaches only +the per-step latent feedback (in `cast_latents`) to keep memory bounded. It also compiles the +**base** transformer under any PEFT/LoRA wrapper (not the wrapper itself). With this, the +on-policy ratio is **bit-exact** (`max|ratio-1| = 0.000e+00`, verified on SD3.5) for coupled +training — **both with and without CFG** (CFG concatenates `[uncond, cond]` identically in +rollout and training, so it aligns exactly). + +Two implementation notes that matter for correctness: +- The grad-force wrapper must **not** detach its own output — an inner detach lets Inductor + pick a divergent inference-optimized kernel and re-introduces ~1e-5 drift. The detach lives + at the latent-feedback chokepoint (`cast_latents`) instead. +- Determinism knobs (`cudnn.deterministic`, `fallback_random`) are irrelevant here — the cause + was the grad-vs-no-grad graph split, not RNG (a `train-vs-train` recompute is exactly 0.0). + +See `.scratch/torch_compile_consistency_report.md` for the full analysis. + + ## Model cache-readiness (lossy caching) Feature caching reuses block outputs across denoising steps via the transformer's diff --git a/src/flow_factory/acceleration/abc.py b/src/flow_factory/acceleration/abc.py index fd43c35d..4ac802a8 100644 --- a/src/flow_factory/acceleration/abc.py +++ b/src/flow_factory/acceleration/abc.py @@ -71,6 +71,21 @@ class BaseAccelerator(ABC): safety: ClassVar[Literal["lossless", "lossy"]] stage: ClassVar[Literal["rollout", "both"]] + # Whether this accelerator requires the Stage-3 rollout to run with autograd + # ENABLED (instead of the default ``torch.no_grad()``) so the transformer + # executes the *same* graph in rollout and the training forward. Only + # ``torch_compile`` needs this: Inductor compiles a separate, numerically + # non-identical graph for grad vs no-grad mode (Dynamo guards on grad_mode), + # so a no-grad rollout would diverge from the grad training forward and break + # the on-policy PPO ratio==1 invariant for coupled algorithms. When set, + # ``CompileAccelerator`` wraps the compiled transformer to force grad (returning + # the grad-carrying output directly — an inner detach would let Inductor pick a + # divergent inference kernel), and the trainer flags the rollout via + # ``_rollout_grad_context`` so the latent feedback is detached in ``cast_latents``. + # Result: bit-exact on-policy ratio (max|ratio-1| == 0) for coupled training, with + # or without CFG. See ``guidance/acceleration.md``. + requires_grad_rollout: ClassVar[bool] = False + def __init_subclass__(cls, **kwargs: Any) -> None: super().__init_subclass__(**kwargs) # Skip intermediate ABCs that intentionally leave the markers unset. diff --git a/src/flow_factory/acceleration/torch_compile.py b/src/flow_factory/acceleration/torch_compile.py index 4c93f24c..58621a2c 100644 --- a/src/flow_factory/acceleration/torch_compile.py +++ b/src/flow_factory/acceleration/torch_compile.py @@ -15,8 +15,11 @@ # src/flow_factory/acceleration/torch_compile.py """Lossless ``torch.compile`` accelerator for the shared transformer(s).""" +import functools from typing import TYPE_CHECKING, Any, Dict +import torch + from .abc import BaseAccelerator from ..utils.logger_utils import setup_logger @@ -60,6 +63,11 @@ class CompileAccelerator(BaseAccelerator): safety = "lossless" stage = "both" + # Inductor compiles a separate graph for grad vs no-grad mode whose fused + # kernels are NOT bit-identical. To keep rollout (Stage 3) and the training + # forward (Stage 6) on the exact same graph — required for the coupled + # on-policy ratio==1 invariant — the rollout must run with grad enabled. + requires_grad_rollout = True def setup(self, adapter: "BaseAdapter") -> None: mode = self.params.get("mode", "regional") @@ -78,20 +86,112 @@ def setup(self, adapter: "BaseAdapter") -> None: for name in transformer_names: module = adapter.get_component(name) + # The routed call chain is proxy(...) -> bundle(name, ...) -> + # members[name](...) -> PeftModel -> LoraModel -> the BASE transformer's + # __call__. So compilation and the grad-consistency wrap must target the + # unwrapped *base* transformer: + # * `adapter._unwrap` peels the RoutedComponentProxy + DDP/FSDP/DeepSpeed + # wrapper, but NOT PEFT — it returns the PeftModel. + # * compiling the PeftModel is wrong: Dynamo specializes the LoRA wrapper + # on grad mode in a way the outer grad-force cannot unify (noise_pred + # drifts ~2.0 between rollout/train), AND the PeftModel has no + # `_repeated_blocks` so regional compile is unavailable under LoRA. + # * the base transformer keeps its LoRA submodules inside the compiled + # graph (they still train), exposes `_repeated_blocks`, and the + # grad-force wrap makes rollout/train bit-identical. + inner = self._peel_peft(adapter._unwrap(module)) if mode == "regional": # `compile_repeated_blocks` exists on every diffusers ModelMixin but # only works when the model declares `_repeated_blocks`; check the # unwrapped module so the error is actionable. - inner = adapter._unwrap(module) if not getattr(inner, "_repeated_blocks", None): raise ValueError( f"CompileAccelerator: component '{name}' ({type(inner).__name__}) does " "not declare `_repeated_blocks`, so regional compilation is unavailable. " "Use `mode: full` for whole-module compilation." ) - module.compile_repeated_blocks(**compile_kwargs) + inner.compile_repeated_blocks(**compile_kwargs) else: # nn.Module.compile compiles the module's forward in place, so the # routed bundle call hits the compiled path on subsequent forwards. - module.compile(**compile_kwargs) + inner.compile(**compile_kwargs) + # Force every forward of the compiled transformer to run under + # torch.enable_grad() (overriding the @torch.no_grad() on + # adapter.inference()), and detach the output during rollout. This keeps + # rollout (Stage 3) and the training forward (Stage 6) on the SAME + # compiled graph — Inductor emits numerically different kernels for the + # grad vs no-grad graph, so a no-grad rollout would break the coupled + # on-policy ratio==1 invariant. Detaching during rollout (gated by the + # trainer's `adapter._rollout_detach` flag) stops autograd from chaining + # across denoising steps, so rollout memory stays bounded. + self._wrap_forward_grad_consistent(adapter, inner) logger.info("CompileAccelerator: compiled '%s' (mode=%s).", name, mode) + + @staticmethod + def _peel_peft(module): + """Return the base transformer under a PEFT/LoRA wrapper (else ``module``). + + ``adapter._unwrap`` peels the routing proxy + DDP/FSDP/DeepSpeed wrapper but + leaves a ``PeftModel`` in place. Compiling the ``PeftModel`` is incorrect: + Dynamo specializes the LoRA wrapper on grad mode in a way the grad-force wrap + cannot unify, and the wrapper hides the base model's ``_repeated_blocks``. The + base transformer keeps its LoRA submodules (they still train) while being the + actual compiled call target, so we compile/wrap it instead. + """ + get_base = getattr(module, "get_base_model", None) + if callable(get_base): + try: + return get_base() + except Exception: + return module + return module + + @staticmethod + def _wrap_forward_grad_consistent(adapter: "BaseAdapter", module) -> None: + """Force the compiled transformer to run grad-consistently across stages. + + Wraps the module's call entry point(s) so every forward runs under + ``torch.enable_grad()`` — overriding the ``@torch.no_grad()`` on + ``adapter.inference()``. This pins rollout (Stage 3) and the training forward + (Stage 6) to the *same* Inductor graph: Inductor compiles a separate, + numerically non-identical graph for grad vs no-grad mode (Dynamo guards on + ``grad_mode``), so a no-grad rollout would diverge from the grad training + forward and break the coupled on-policy PPO ratio==1 invariant. + + Crucially the output is **NOT detached here**: detaching inside the wrapper + lets Inductor see the result is unused-for-grad and pick a different + (inference-optimized) kernel, re-introducing the divergence. Instead the + rollout's per-step latent feedback is detached in + :meth:`BaseAdapter.cast_latents` (gated by ``adapter._rollout_detach``), which + breaks the autograd graph chain across denoising steps so rollout memory stays + bounded while every transformer call still executes the identical grad-mode + graph (bit-exact noise_pred → bit-exact on-policy ratio). + + Two entry points must be covered because the two compile modes dispatch + differently: + * ``mode='regional'`` (``compile_repeated_blocks``): the top ``forward`` + stays eager and calls the compiled blocks — wrapping ``forward`` suffices. + * ``mode='full'`` (``nn.Module.compile``): ``module(...)`` routes through + ``module._compiled_call_impl``, bypassing a reassigned ``forward`` — so + that attribute must be wrapped too. + + Idempotent per attribute (``_ff_grad_consistent`` marker). + """ + + def _wrap(fn): + if fn is None or getattr(fn, "_ff_grad_consistent", False): + return fn + + @functools.wraps(fn) + def wrapped(*args, **kwargs): + with torch.enable_grad(): + return fn(*args, **kwargs) + + wrapped._ff_grad_consistent = True + return wrapped + + # Regional + eager fallback: top-level forward. + module.forward = _wrap(module.forward) + # Full mode: nn.Module.compile() dispatches through _compiled_call_impl. + if getattr(module, "_compiled_call_impl", None) is not None: + module._compiled_call_impl = _wrap(module._compiled_call_impl) diff --git a/src/flow_factory/models/abc.py b/src/flow_factory/models/abc.py index c96d1e14..80342d93 100644 --- a/src/flow_factory/models/abc.py +++ b/src/flow_factory/models/abc.py @@ -222,7 +222,18 @@ def latent_storage_dtype(self) -> Optional[torch.dtype]: return self._DTYPE_MAP.get(val) if val else None def cast_latents(self, latents: torch.Tensor, default_dtype: Optional[torch.dtype] = None) -> torch.Tensor: - """Cast latents to storage dtype with float16 overflow protection.""" + """Cast latents to storage dtype with float16 overflow protection. + + During a compile-accelerated rollout the trainer sets ``self._rollout_detach`` + (see ``BaseTrainer._rollout_grad_context``): the compiled transformer is forced + to run under ``torch.enable_grad()`` for graph-consistency with training, so the + per-step latent feedback would otherwise chain an autograd graph across the + whole denoising loop. Detaching here — the single feedback chokepoint every + adapter routes through — breaks that chain so rollout memory stays bounded while + each transformer call still executes the identical grad-mode graph. + """ + if getattr(self, "_rollout_detach", False) and latents.requires_grad: + latents = latents.detach() target = self.latent_storage_dtype or default_dtype if target is None or latents.dtype == target: return latents diff --git a/src/flow_factory/trainers/abc.py b/src/flow_factory/trainers/abc.py index 2a388b2a..b21f3ca6 100644 --- a/src/flow_factory/trainers/abc.py +++ b/src/flow_factory/trainers/abc.py @@ -468,6 +468,46 @@ def _rollout_acceleration(self) -> Iterator[None]: stack.enter_context(accelerator.rollout_context(self.adapter)) yield + @contextmanager + def _rollout_grad_context(self) -> Iterator[None]: + """Grad context for the Stage-3 rollout loop. + + Default: ``torch.no_grad()`` (rollout needs no gradients — saves memory). + + Exception: if any active **shared** accelerator declares + ``requires_grad_rollout`` (only ``torch_compile`` does), rollout instead runs + with grad enabled AND sets ``adapter._rollout_detach``. Reason: + ``torch.compile`` (Inductor) emits a *separate, numerically non-identical* + graph for grad vs no-grad mode (Dynamo guards on ``grad_mode``), so a no-grad + rollout would diverge from the grad-mode training forward and break the + coupled on-policy PPO ratio==1 invariant. The compiled transformer is wrapped + (see :class:`CompileAccelerator`) so its forward always executes under + ``torch.enable_grad()`` — overriding the ``@torch.no_grad()`` on + ``adapter.inference()``. The output is NOT detached inside that wrapper (an + inner detach lets Inductor pick a divergent inference kernel); instead the + per-step latent feedback is detached in ``BaseAdapter.cast_latents`` (gated by + this flag), so rollout and training share the identical grad-mode graph + (bit-exact on-policy ratio, with or without CFG — see + ``guidance/acceleration.md``) while the autograd graph never chains across + denoising steps (memory stays bounded). + """ + needs_grad = any( + getattr(acc, "requires_grad_rollout", False) + for acc in getattr(self, "shared_accelerators", []) + ) + if not needs_grad: + with torch.no_grad(): + yield + return + # Compile active: the compiled-transformer wrapper forces grad locally and + # detaches its output while this flag is set, so rollout matches the training + # graph bit-for-bit without chaining autograd across steps. + prev = getattr(self.adapter, "_rollout_detach", False) + self.adapter._rollout_detach = True + try: + yield + finally: + self.adapter._rollout_detach = prev def _synchronize_frozen_components(self): if self.accelerator.num_processes <= 1: @@ -859,7 +899,10 @@ def generate_samples( # Stage-3-only acceleration (e.g. feature caching) is scoped to this loop # so its state never leaks into the Stage-6 training forward. - with self._rollout_acceleration(), torch.no_grad(), self.autocast(): + # Grad context is normally no_grad, but flips to enable_grad when a + # compile accelerator is active (see _rollout_grad_context) so rollout and + # training share the same compiled graph (bit-exact on-policy ratio). + with self._rollout_acceleration(), self._rollout_grad_context(), self.autocast(): for _ in tqdm( range(self.training_args.num_batches_per_epoch), desc=f'Epoch {self.epoch} Sampling', From b97d398650d91d8551078aecd287652b923ce153 Mon Sep 17 00:00:00 2001 From: Jayce-Ping <315229706@qq.com> Date: Mon, 29 Jun 2026 20:15:38 +0800 Subject: [PATCH 10/15] [acceleration,models,trainers,docs] refactor: mark torch_compile lossy (not bit-exact across rollout/train) + fail-fast tidy torch.compile is stage='both' (applied to the shared transformer) but Inductor compiles a separate graph for grad vs no-grad mode; even with the grad-forced rollout fix an intermittent ~1e-5 on-policy ratio residual remains on a minority of samples (verified 16xH20, 2-node ZeRO-2/FSDP2, SD3.5 + Qwen-Image, regional/full, CFG and no-CFG). So it is NOT bit-exact across stages. - Re-class CompileAccelerator safety='lossy' (lossless now means bit-exact across stages, not merely "applied to the shared module"); update abc/validator/torch_compile docstrings. - validator: WARN (do not reject) when a stage='both' lossy accelerator runs on a coupled trainer -- it stays within clip_range so it remains allowed, but the on-policy ratio is ~1, not exactly 1. eager + attention_backend stay lossless/bit-exact. - fail-fast: drop the swallowing try/except in CompileAccelerator._peel_peft. - init BaseAdapter._rollout_detach=False and read it directly (drop getattr defaults); remove the dead acceleration_args None-branch in _init_acceleration. - correct guidance/acceleration.md + architecture.md tables/notes that overclaimed bit-exact. --- .agents/knowledge/architecture.md | 4 +- guidance/acceleration.md | 14 ++++-- src/flow_factory/acceleration/abc.py | 36 +++++++++----- .../acceleration/torch_compile.py | 47 +++++++++++++++---- src/flow_factory/acceleration/validator.py | 38 +++++++++++---- src/flow_factory/models/abc.py | 6 ++- src/flow_factory/trainers/abc.py | 11 ++--- 7 files changed, 108 insertions(+), 48 deletions(-) diff --git a/.agents/knowledge/architecture.md b/.agents/knowledge/architecture.md index 9c7b5d02..2963b4ba 100644 --- a/.agents/knowledge/architecture.md +++ b/.agents/knowledge/architecture.md @@ -151,11 +151,11 @@ All four registries map string keys → lazy import paths. Resolution: registry | Key | Class | Safety | Stage | Notes | |-----|-------|--------|-------|-------| | `attention_backend` | `AttentionBackendAccelerator` | lossless | both | Sets the diffusers attention backend on every transformer (requires a `backend` param). Listed as a `shared` entry (before `torch_compile`); this is the single code path for backend selection (the old `BaseAdapter._set_attention_backend` and the `model.attn_backend` knob were both removed — a config still setting `model.attn_backend` fails fast). Bagel forces flash_attention_2 at load and does not use it. | -| `torch_compile` | `CompileAccelerator` | lossless | both | `torch.compile` of the shared transformer (regional/full); applied in-place after `post_init` so checkpoint keys / param identity stay stable. | +| `torch_compile` | `CompileAccelerator` | lossy | both | `torch.compile` of the shared transformer (regional/full); applied in-place after `post_init` so checkpoint keys / param identity stay stable. Marked `lossy` because it is applied symmetrically but is **not bit-exact across rollout vs training** (Inductor's grad/no-grad graph split → intermittent ~1e-5 on-policy residual, within `clip_range`); still allowed on coupled algos, validator warns. | | `diffusers_cache` | `DiffusersCacheAccelerator` | lossy | rollout | Diffusers `CacheMixin` feature caching (first_block/faster/pyramid/taylorseer/magcache). | | `cache_dit` | `CacheDitAccelerator` | lossy | rollout | Optional `cache-dit` backend (DBCache/TaylorSeer). | -Configured via the `acceleration:` block (`hparams/acceleration_args.py`): two ordered lists of `{name, params}` entries — `shared` (lossless, both stages) and `rollout` (Stage-3 only). **List order is application order**: `shared` entries run their `setup()` in order (so `attention_backend` must precede `torch_compile`), and `rollout` entries nest their `rollout_context()` in order. The `acceleration/validator.py` enforces that lossy accelerators are rollout-only and only on `decoupled`/`distillation` trainers (each trainer declares a `paradigm`), preserving train-inference consistency (constraints #7, #20a). Off by default. +Configured via the `acceleration:` block (`hparams/acceleration_args.py`): two ordered lists of `{name, params}` entries — `shared` (lossless, both stages) and `rollout` (Stage-3 only). **List order is application order**: `shared` entries run their `setup()` in order (so `attention_backend` must precede `torch_compile`), and `rollout` entries nest their `rollout_context()` in order. The `acceleration/validator.py` enforces that a **lossy `rollout`** accelerator runs only on `decoupled`/`distillation` trainers (each trainer declares a `paradigm`), preserving train-inference consistency (constraints #7, #20a). A **lossy `shared` (stage=both)** accelerator (e.g. `torch_compile`, applied symmetrically but not bit-exact across stages) is allowed on any paradigm but the validator **warns** on coupled trainers that the on-policy ratio will be ≈1, not exactly 1. Off by default. --- diff --git a/guidance/acceleration.md b/guidance/acceleration.md index d0fab993..f3d33745 100644 --- a/guidance/acceleration.md +++ b/guidance/acceleration.md @@ -63,7 +63,7 @@ accepted as shorthand for a one-element list. A direct python path (e.g. | id | safety | stage | Notes | |----|--------|-------|-------| | `attention_backend` | lossless | both | Sets the diffusers attention backend on every transformer. Requires a `backend` param. Forwards any backend (`native` / `flash` / `_flash_3` / `_flash_3_hub` / `sage` / `xformers`) to `set_attention_backend`. List it in `shared` **before** `torch_compile` so the compiled graph captures the backend. | -| `torch_compile` | lossless | both | `torch.compile` of the shared transformer. `mode: regional` uses diffusers' `compile_repeated_blocks` (fast warmup, robust to variable resolution); `mode: full` compiles the whole module. Extra `compile_kwargs` forwarded to the compile call. Compiles in place (checkpoint- and EMA/ref-safe), applied after `post_init`. | +| `torch_compile` | lossy | both | `torch.compile` of the shared transformer. `mode: regional` uses diffusers' `compile_repeated_blocks` (fast warmup, robust to variable resolution); `mode: full` compiles the whole module. Extra `compile_kwargs` forwarded to the compile call. Compiles in place (checkpoint- and EMA/ref-safe), applied after `post_init`. Marked **lossy** because it is applied symmetrically but is **not bit-exact across rollout vs training** (grad/no-grad graph split → intermittent ~1e-5 on-policy residual, within `clip_range`); allowed on coupled algos, but the validator warns. | | `diffusers_cache` | lossy | rollout | Diffusers-native feature caching (no extra dependency). `policy`: `first_block` (default) / `faster` / `pyramid` / `taylorseer` / `magcache`; remaining params forwarded to the policy's diffusers config (e.g. `threshold`). | | `cache_dit` | lossy | rollout | [cache-dit](https://github.com/vipshop/cache-dit) backend (DBCache/TaylorSeer). Requires `pip install flow-factory[acceleration]`. All params forwarded to `cache_dit.enable_cache`. | @@ -108,9 +108,15 @@ the ratio. `torch.enable_grad()` (overriding the `@torch.no_grad()` on `inference()`) and detaches only the per-step latent feedback (in `cast_latents`) to keep memory bounded. It also compiles the **base** transformer under any PEFT/LoRA wrapper (not the wrapper itself). With this, the -on-policy ratio is **bit-exact** (`max|ratio-1| = 0.000e+00`, verified on SD3.5) for coupled -training — **both with and without CFG** (CFG concatenates `[uncond, cond]` identically in -rollout and training, so it aligns exactly). +on-policy ratio is driven to **≈1, well within `clip_range` (1e-4)** — but **not strictly +bit-exact**: an intermittent **~1e-5** residual remains on a minority of samples/ranks +(16×H20, ZeRO-2/FSDP2, SD3.5+Qwen-Image, regional/full, CFG and no-CFG). Forcing grad removes +the *dominant* grad-vs-no-grad graph split, but rollout and training are still distinct Inductor +kernel invocations (different latent stride/contiguity + autograd-graph context), and bf16's +non-associative accumulation surfaces a last-bit difference for some inputs. So compile is +**numerically on-policy, not bit-exact** — if you need a strictly bit-exact ratio, use eager or +the `attention_backend` accelerator (a backend's forward is grad-mode-independent and stays +exactly 0). See `CompileAccelerator._wrap_forward_grad_consistent` for details. Two implementation notes that matter for correctness: - The grad-force wrapper must **not** detach its own output — an inner detach lets Inductor diff --git a/src/flow_factory/acceleration/abc.py b/src/flow_factory/acceleration/abc.py index 4ac802a8..bbbafe19 100644 --- a/src/flow_factory/acceleration/abc.py +++ b/src/flow_factory/acceleration/abc.py @@ -29,17 +29,24 @@ - ``"rollout"``: a per-epoch context via :meth:`rollout_context` (e.g. feature caching), torn down before the training forward. Belongs in the ``rollout`` slot. -* ``safety`` — the train-inference consistency class, **only consulted for - ``stage='rollout'`` accelerators**: - - - ``"lossless"``: bit-identical outputs (rollout unchanged) — safe for any paradigm. - - ``"lossy"``: changes outputs, so rollout diverges from the (un-accelerated) - training forward — only safe when the rollout log-prob never feeds the loss, i.e. - **decoupled / distillation** algorithms (see ``constraints.md`` #7, #20a). - - For ``stage='both'`` accelerators ``safety`` is informational: a symmetric - transform is consistent regardless of numerical exactness, so even an approximate - attention backend (e.g. Sage int8) used in *both* stages is safe. +* ``safety`` — the train-inference consistency class: + + - ``"lossless"``: **bit-exact across rollout and training**. For ``stage='rollout'`` + this means the rollout output is bit-identical to the un-accelerated forward; for + ``stage='both'`` it means the shared transform produces identical results in both + stages (an exact transform, or a symmetric-approximate one such as Sage int8 + attention whose same int8 kernel runs in both stages). Safe for any paradigm. + - ``"lossy"``: **NOT bit-exact across the two stages.** + + - ``stage='rollout'`` + lossy (e.g. feature caching): rollout diverges from the + training forward, which cannot replicate it — only safe when the rollout + log-prob never feeds the loss, i.e. **decoupled / distillation** algorithms + (validator *rejects* it on coupled; see ``constraints.md`` #7, #20a). + - ``stage='both'`` + lossy (e.g. ``torch.compile``, whose grad/no-grad + compiled-graph split leaves a ~1e-5 residual): applied symmetrically and stays + within ``clip_range``, so it is *allowed* on any paradigm — but the validator + *warns* on a coupled trainer that the on-policy PPO ratio will be ~1, not + bit-exact. Subclasses implement only what they need: ``setup`` defaults to a no-op, ``rollout_context`` defaults to yielding without modification. @@ -82,8 +89,11 @@ class BaseAccelerator(ABC): # the grad-carrying output directly — an inner detach would let Inductor pick a # divergent inference kernel), and the trainer flags the rollout via # ``_rollout_grad_context`` so the latent feedback is detached in ``cast_latents``. - # Result: bit-exact on-policy ratio (max|ratio-1| == 0) for coupled training, with - # or without CFG. See ``guidance/acceleration.md``. + # Result: removes the dominant grad/no-grad divergence so the on-policy ratio is + # ~1 (well within ``clip_range``), but NOT strictly bit-exact — an intermittent + # ~1e-5 residual remains on a minority of samples (different Inductor kernel + # invocations + bf16 non-associativity). See + # ``CompileAccelerator._wrap_forward_grad_consistent`` for the full reason. requires_grad_rollout: ClassVar[bool] = False def __init_subclass__(cls, **kwargs: Any) -> None: diff --git a/src/flow_factory/acceleration/torch_compile.py b/src/flow_factory/acceleration/torch_compile.py index 58621a2c..a9cae4bc 100644 --- a/src/flow_factory/acceleration/torch_compile.py +++ b/src/flow_factory/acceleration/torch_compile.py @@ -32,10 +32,21 @@ class CompileAccelerator(BaseAccelerator): """Apply ``torch.compile`` to every transformer the adapter exposes. - Lossless and stage-``both``: the compiled module backs both rollout - ``inference()`` and the training ``forward()`` (they share - ``adapter.transformer``), so numerical behavior stays consistent across the - two — safe even for coupled algorithms. + Stage-``both`` but ``safety='lossy'``. The compiled module backs both rollout + ``inference()`` and the training ``forward()`` (they share ``adapter.transformer``), + so it is applied *symmetrically* — but, unlike an exact or symmetric-approximate + transform, ``torch.compile`` is NOT numerically identical across the two stages: + Inductor compiles a separate graph for grad vs no-grad mode, and even with the + grad-forced rollout fix below an intermittent ~1e-5 on-policy ratio residual remains + on a minority of samples (see :meth:`_wrap_forward_grad_consistent`). That is why it + is marked ``lossy`` rather than ``lossless`` (which here means *bit-exact across + stages*, not merely "applied to the shared module"). + + It stays well within ``clip_range`` — numerically on-policy — and is the main + compute speedup on real hardware, so it remains a ``stage='both'`` accelerator + allowed on coupled algorithms. The validator does not reject it, but it WARNS on a + coupled trainer that the on-policy PPO ratio will be ~1, not bit-exact. Use eager or + the ``attention_backend`` accelerator if a strictly bit-exact ratio is required. Both modes compile **in place** (``nn.Module.compile`` / ``compile_repeated_blocks``), which preserves parameter identity and leaves @@ -61,7 +72,11 @@ class CompileAccelerator(BaseAccelerator): (e.g. ``{"mode": "max-autotune", "dynamic": true}``). """ - safety = "lossless" + # `lossy` not because it is asymmetric (it is stage='both', applied to the shared + # module) but because it is NOT bit-exact across rollout vs training — see the class + # docstring and `_wrap_forward_grad_consistent`. The validator warns (does not + # reject) when this runs on a coupled trainer. + safety = "lossy" stage = "both" # Inductor compiles a separate graph for grad vs no-grad mode whose fused # kernels are NOT bit-identical. To keep rollout (Stage 3) and the training @@ -140,10 +155,7 @@ def _peel_peft(module): """ get_base = getattr(module, "get_base_model", None) if callable(get_base): - try: - return get_base() - except Exception: - return module + return get_base() return module @staticmethod @@ -165,7 +177,22 @@ def _wrap_forward_grad_consistent(adapter: "BaseAdapter", module) -> None: :meth:`BaseAdapter.cast_latents` (gated by ``adapter._rollout_detach``), which breaks the autograd graph chain across denoising steps so rollout memory stays bounded while every transformer call still executes the identical grad-mode - graph (bit-exact noise_pred → bit-exact on-policy ratio). + graph (near-exact noise_pred → on-policy ratio ≈ 1). + + Known residual (NOT strictly bit-exact): forcing grad removes the *dominant* + divergence (the grad-vs-no-grad graph split), but it does not make the rollout + and training forwards a literally identical Inductor kernel invocation. They are + different call sites: the stored-then-reloaded latents may differ in + stride/contiguity, and the surrounding autograd graphs differ (rollout detaches + per step and discards the graph; training keeps it and backwards). For some + inputs Inductor's shape/layout-specialized, autotuned kernels then take a + different reduction order, and bf16's non-associative accumulation turns that + into an *intermittent* ~1e-5 on-policy ratio drift on a minority of + samples/ranks (value-dependent — some runs/shardings show exactly 0). This is + well within ``clip_range`` (1e-4), i.e. numerically on-policy, but not + bit-exact. Eager and the attention-backend accelerator stay exactly 0 + (an attention kernel's forward is grad-mode-independent). Measured on 16×H20 + (2-node ZeRO-2/FSDP2), SD3.5 + Qwen-Image, regional + full, CFG and no-CFG. Two entry points must be covered because the two compile modes dispatch differently: diff --git a/src/flow_factory/acceleration/validator.py b/src/flow_factory/acceleration/validator.py index 850f3c63..56218aa0 100644 --- a/src/flow_factory/acceleration/validator.py +++ b/src/flow_factory/acceleration/validator.py @@ -18,11 +18,15 @@ The correctness contract (``constraints.md`` #7 + #20a) hinges on **symmetric application**, encoded by ``stage``, not on numerical bit-exactness: -* ``stage='both'`` accelerators mutate the transformer persistently, so the exact - same transform is in both rollout ``inference()`` and training ``forward()``. - Rollout and training therefore stay CONSISTENT — even for numerically-approximate - transforms (e.g. Sage int8 attention) — so they are safe for any algorithm. They - belong in the ``shared`` slot; ``safety`` is not consulted for them. +* ``stage='both'`` accelerators mutate the transformer persistently, so the same + transform runs in both rollout ``inference()`` and training ``forward()``. When that + transform is identical across the two stages (exact, or symmetric-approximate like + Sage int8 attention) rollout and training stay CONSISTENT — safe for any algorithm. + They belong in the ``shared`` slot and are never rejected. ``safety`` is only used to + *warn*: a ``stage='both'`` + ``lossy`` accelerator (e.g. ``torch.compile``, which is + applied symmetrically but is not bit-exact across stages due to its grad/no-grad + graph split) is still allowed, but on a **coupled** trainer the on-policy PPO ratio + will be ≈1, not exactly 1, so the validator logs a warning. * ``stage='rollout'`` accelerators run only during Stage-3 rollout. If such an accelerator changes outputs (``safety='lossy'``, e.g. feature caching), rollout diverges from the training forward, which it cannot be replicated in (that needs @@ -88,11 +92,25 @@ def validate_accelerator( "slot, which applies a persistent transform to both rollout and training. Put a " "stage='both' accelerator here, or move this one to the `acceleration.rollout` list." ) - # No safety gate here: a stage='both' transform is applied identically to - # rollout `inference()` and training `forward()` (the same shared module), so - # rollout and training stay CONSISTENT even for numerically-approximate - # backends (e.g. Sage int8 attention). Train-inference consistency depends on - # symmetric application, not on bit-exactness. + # A stage='both' transform is applied to the SAME module in rollout + # `inference()` and training `forward()`. For a transform that is identical + # across the two stages (exact, or symmetric-approximate like Sage int8) this is + # consistent by construction — safe for any paradigm, no gate. A `lossy` + # stage='both' accelerator (e.g. torch.compile, whose grad/no-grad compiled-graph + # split leaves a ~1e-5 residual) is applied symmetrically but is NOT bit-exact + # across stages; it stays within clip_range so it is allowed, but on a coupled + # trainer the on-policy ratio will be ~1, not exactly 1 — warn so the user can + # pick eager / an exact attention backend if strict ratio==1 is required. + if accelerator.safety == "lossy" and paradigm == "coupled": + logger.warning( + "Accelerator '%s' (stage='both', safety='lossy') is applied symmetrically " + "but is not bit-exact across rollout and training (e.g. torch.compile's " + "grad/no-grad graph split). On the coupled trainer '%s' the on-policy PPO " + "ratio will be ~1 but NOT exactly 1 (within clip_range). Use eager or an " + "exact attention backend if a strictly bit-exact ratio is required.", + name, + trainer_name, + ) return # slot == "rollout". diff --git a/src/flow_factory/models/abc.py b/src/flow_factory/models/abc.py index 80342d93..f64f917a 100644 --- a/src/flow_factory/models/abc.py +++ b/src/flow_factory/models/abc.py @@ -154,6 +154,10 @@ def __init__(self, config: Arguments, accelerator : Accelerator): self.training_args = config.training_args self.eval_args = config.eval_args self._mode : str = 'train' # ['train', 'eval', 'rollout'] + # Set by the trainer's `_rollout_grad_context` only when a compile accelerator + # forces a grad-enabled rollout: detaches per-step latent feedback in + # `cast_latents` to keep rollout memory bounded. Off otherwise. + self._rollout_detach: bool = False self._named_parameters : Dict[str, NamedParametersInfo] = {} # Load pipeline and scheduler (delegated to subclasses) @@ -232,7 +236,7 @@ def cast_latents(self, latents: torch.Tensor, default_dtype: Optional[torch.dtyp adapter routes through — breaks that chain so rollout memory stays bounded while each transformer call still executes the identical grad-mode graph. """ - if getattr(self, "_rollout_detach", False) and latents.requires_grad: + if self._rollout_detach and latents.requires_grad: latents = latents.detach() target = self.latent_storage_dtype or default_dtype if target is None or latents.dtype == target: diff --git a/src/flow_factory/trainers/abc.py b/src/flow_factory/trainers/abc.py index b21f3ca6..16c5b606 100644 --- a/src/flow_factory/trainers/abc.py +++ b/src/flow_factory/trainers/abc.py @@ -405,11 +405,9 @@ def _init_acceleration(self): Each accelerator is validated against this trainer's ``paradigm`` before use (fail-fast, ``constraints.md`` #26). """ - accel_args = getattr(self.config, "acceleration_args", None) + accel_args = self.config.acceleration_args self.shared_accelerators: List[BaseAccelerator] = [] self.rollout_accelerators: List[BaseAccelerator] = [] - if accel_args is None: - return trainer_name = type(self).__name__ paradigm = type(self).paradigm @@ -491,10 +489,7 @@ def _rollout_grad_context(self) -> Iterator[None]: ``guidance/acceleration.md``) while the autograd graph never chains across denoising steps (memory stays bounded). """ - needs_grad = any( - getattr(acc, "requires_grad_rollout", False) - for acc in getattr(self, "shared_accelerators", []) - ) + needs_grad = any(acc.requires_grad_rollout for acc in self.shared_accelerators) if not needs_grad: with torch.no_grad(): yield @@ -502,7 +497,7 @@ def _rollout_grad_context(self) -> Iterator[None]: # Compile active: the compiled-transformer wrapper forces grad locally and # detaches its output while this flag is set, so rollout matches the training # graph bit-for-bit without chaining autograd across steps. - prev = getattr(self.adapter, "_rollout_detach", False) + prev = self.adapter._rollout_detach self.adapter._rollout_detach = True try: yield From 01709533124e99cef1476e9e0590da884c75c018 Mon Sep 17 00:00:00 2001 From: Jayce-Ping <315229706@qq.com> Date: Tue, 30 Jun 2026 08:46:30 +0800 Subject: [PATCH 11/15] [acceleration,docs] refactor: remove the cache_dit accelerator; diffusers_cache is the single lossy backend cache-dit only caches inside the pipeline `__call__` it monkeypatches, but Flow-Factory's rollout drives the transformer directly (adapter.inference()), so CacheDitAccelerator was a silent no-op (every step logged "Cache context not exist, skip cache"). Its transformer-only path (FakeDiffusionPipeline -> persistent_context) does engage but assumes enable-once: the class-level _is_cached flag is not reset on disable, so the per-epoch rollout_context enable/disable cycle crashes on epoch 2; it also conflicts with the cache-ready adapters' own diffusers `cache_context` (two caching systems on the same blocks). It adds nothing over the diffusers-native `diffusers_cache` (FBCache/TaylorSeer/FasterCache, ~1.2x rollout, clean per-epoch lifecycle), so it is removed rather than force-fit on private dev-build internals. - Delete acceleration/cache_dit.py; drop it from the registry. - Add _REMOVED_ACCELERATORS guard: a config still naming `cache_dit` fails fast with an actionable message pointing to `diffusers_cache`. - Drop the `acceleration`/cache-dit optional-dependency group (diffusers caching needs none). - Update guidance/acceleration.md, architecture.md, and example-config comments. --- .agents/knowledge/architecture.md | 2 - examples/awm/lora/flux1/default.yaml | 2 +- .../awm/lora/flux2_klein_base/default.yaml | 2 +- examples/awm/lora/sd3_5/default.yaml | 2 +- examples/crd/lora/sd3_5/default.yaml | 2 +- examples/nft/lora/flux1/default.yaml | 2 +- .../lora/qwen_image/rational_rewards_t2i.yaml | 2 +- examples/nft/lora/sd3_5/default.yaml | 2 +- examples/nft/lora/wan21/i2v.yaml | 2 +- examples/nft/lora/wan21/t2v.yaml | 2 +- examples/nft/lora/wan22/t2v.yaml | 2 +- guidance/acceleration.md | 5 +- pyproject.toml | 10 +-- src/flow_factory/acceleration/__init__.py | 4 +- src/flow_factory/acceleration/cache_dit.py | 77 ------------------- src/flow_factory/acceleration/registry.py | 19 ++++- 16 files changed, 35 insertions(+), 102 deletions(-) delete mode 100644 src/flow_factory/acceleration/cache_dit.py diff --git a/.agents/knowledge/architecture.md b/.agents/knowledge/architecture.md index 2963b4ba..0600b6ca 100644 --- a/.agents/knowledge/architecture.md +++ b/.agents/knowledge/architecture.md @@ -153,8 +153,6 @@ All four registries map string keys → lazy import paths. Resolution: registry | `attention_backend` | `AttentionBackendAccelerator` | lossless | both | Sets the diffusers attention backend on every transformer (requires a `backend` param). Listed as a `shared` entry (before `torch_compile`); this is the single code path for backend selection (the old `BaseAdapter._set_attention_backend` and the `model.attn_backend` knob were both removed — a config still setting `model.attn_backend` fails fast). Bagel forces flash_attention_2 at load and does not use it. | | `torch_compile` | `CompileAccelerator` | lossy | both | `torch.compile` of the shared transformer (regional/full); applied in-place after `post_init` so checkpoint keys / param identity stay stable. Marked `lossy` because it is applied symmetrically but is **not bit-exact across rollout vs training** (Inductor's grad/no-grad graph split → intermittent ~1e-5 on-policy residual, within `clip_range`); still allowed on coupled algos, validator warns. | | `diffusers_cache` | `DiffusersCacheAccelerator` | lossy | rollout | Diffusers `CacheMixin` feature caching (first_block/faster/pyramid/taylorseer/magcache). | -| `cache_dit` | `CacheDitAccelerator` | lossy | rollout | Optional `cache-dit` backend (DBCache/TaylorSeer). | - Configured via the `acceleration:` block (`hparams/acceleration_args.py`): two ordered lists of `{name, params}` entries — `shared` (lossless, both stages) and `rollout` (Stage-3 only). **List order is application order**: `shared` entries run their `setup()` in order (so `attention_backend` must precede `torch_compile`), and `rollout` entries nest their `rollout_context()` in order. The `acceleration/validator.py` enforces that a **lossy `rollout`** accelerator runs only on `decoupled`/`distillation` trainers (each trainer declares a `paradigm`), preserving train-inference consistency (constraints #7, #20a). A **lossy `shared` (stage=both)** accelerator (e.g. `torch_compile`, applied symmetrically but not bit-exact across stages) is allowed on any paradigm but the validator **warns** on coupled trainers that the on-policy ratio will be ≈1, not exactly 1. Off by default. --- diff --git a/examples/awm/lora/flux1/default.yaml b/examples/awm/lora/flux1/default.yaml index 7ee7b4cc..00c2ee3a 100644 --- a/examples/awm/lora/flux1/default.yaml +++ b/examples/awm/lora/flux1/default.yaml @@ -40,7 +40,7 @@ model: # - name: torch_compile # ...so the compiled graph captures it # params: { mode: regional } # regional (compile_repeated_blocks) | full # rollout: # rollout-only (Stage 3); lossy -> decoupled/distillation only -# - name: diffusers_cache # Options: diffusers_cache, cache_dit +# - name: diffusers_cache # Options: diffusers_cache # params: { policy: first_block, threshold: 0.08 } log: diff --git a/examples/awm/lora/flux2_klein_base/default.yaml b/examples/awm/lora/flux2_klein_base/default.yaml index 891680ab..aa470d80 100644 --- a/examples/awm/lora/flux2_klein_base/default.yaml +++ b/examples/awm/lora/flux2_klein_base/default.yaml @@ -40,7 +40,7 @@ model: # - name: torch_compile # ...so the compiled graph captures it # params: { mode: regional } # regional (compile_repeated_blocks) | full # rollout: # rollout-only (Stage 3); lossy -> decoupled/distillation only -# - name: diffusers_cache # Options: diffusers_cache, cache_dit +# - name: diffusers_cache # Options: diffusers_cache # params: { policy: first_block, threshold: 0.08 } log: diff --git a/examples/awm/lora/sd3_5/default.yaml b/examples/awm/lora/sd3_5/default.yaml index ea2a6829..f7eaa4ed 100644 --- a/examples/awm/lora/sd3_5/default.yaml +++ b/examples/awm/lora/sd3_5/default.yaml @@ -41,7 +41,7 @@ model: # - name: torch_compile # ...so the compiled graph captures it # params: { mode: regional } # regional (compile_repeated_blocks) | full # rollout: # rollout-only (Stage 3); lossy -> decoupled/distillation only -# - name: diffusers_cache # Options: diffusers_cache, cache_dit +# - name: diffusers_cache # Options: diffusers_cache # params: { policy: first_block, threshold: 0.08 } log: diff --git a/examples/crd/lora/sd3_5/default.yaml b/examples/crd/lora/sd3_5/default.yaml index 1a9515c0..7f490059 100644 --- a/examples/crd/lora/sd3_5/default.yaml +++ b/examples/crd/lora/sd3_5/default.yaml @@ -42,7 +42,7 @@ model: # - name: torch_compile # ...so the compiled graph captures it # params: { mode: regional } # regional (compile_repeated_blocks) | full # rollout: # rollout-only (Stage 3); lossy -> decoupled/distillation only -# - name: diffusers_cache # Options: diffusers_cache, cache_dit +# - name: diffusers_cache # Options: diffusers_cache # params: { policy: first_block, threshold: 0.08 } log: diff --git a/examples/nft/lora/flux1/default.yaml b/examples/nft/lora/flux1/default.yaml index d221a91e..6c53e19e 100644 --- a/examples/nft/lora/flux1/default.yaml +++ b/examples/nft/lora/flux1/default.yaml @@ -40,7 +40,7 @@ model: # - name: torch_compile # ...so the compiled graph captures it # params: { mode: regional } # regional (compile_repeated_blocks) | full # rollout: # rollout-only (Stage 3); lossy -> decoupled/distillation only -# - name: diffusers_cache # Options: diffusers_cache, cache_dit +# - name: diffusers_cache # Options: diffusers_cache # params: { policy: first_block, threshold: 0.08 } log: diff --git a/examples/nft/lora/qwen_image/rational_rewards_t2i.yaml b/examples/nft/lora/qwen_image/rational_rewards_t2i.yaml index 586029b0..875c264b 100644 --- a/examples/nft/lora/qwen_image/rational_rewards_t2i.yaml +++ b/examples/nft/lora/qwen_image/rational_rewards_t2i.yaml @@ -92,7 +92,7 @@ scheduler: # - name: torch_compile # ...so the compiled graph captures it # params: { mode: regional } # regional (compile_repeated_blocks) | full # rollout: # rollout-only (Stage 3); lossy -> decoupled/distillation only -# - name: diffusers_cache # Options: diffusers_cache, cache_dit +# - name: diffusers_cache # Options: diffusers_cache # params: { policy: first_block, threshold: 0.08 } eval: diff --git a/examples/nft/lora/sd3_5/default.yaml b/examples/nft/lora/sd3_5/default.yaml index 46a72dc0..ff4b7e95 100644 --- a/examples/nft/lora/sd3_5/default.yaml +++ b/examples/nft/lora/sd3_5/default.yaml @@ -41,7 +41,7 @@ model: # - name: torch_compile # ...so the compiled graph captures it # params: { mode: regional } # regional (compile_repeated_blocks) | full # rollout: # rollout-only (Stage 3); lossy -> decoupled/distillation only -# - name: diffusers_cache # Options: diffusers_cache, cache_dit +# - name: diffusers_cache # Options: diffusers_cache # params: { policy: first_block, threshold: 0.08 } log: diff --git a/examples/nft/lora/wan21/i2v.yaml b/examples/nft/lora/wan21/i2v.yaml index e54c2e6e..973f4e9b 100644 --- a/examples/nft/lora/wan21/i2v.yaml +++ b/examples/nft/lora/wan21/i2v.yaml @@ -41,7 +41,7 @@ model: # - name: torch_compile # ...so the compiled graph captures it # params: { mode: regional } # regional (compile_repeated_blocks) | full # rollout: # rollout-only (Stage 3); lossy -> decoupled/distillation only -# - name: diffusers_cache # Options: diffusers_cache, cache_dit +# - name: diffusers_cache # Options: diffusers_cache # params: { policy: first_block, threshold: 0.08 } log: diff --git a/examples/nft/lora/wan21/t2v.yaml b/examples/nft/lora/wan21/t2v.yaml index 72106e22..0026a9a8 100644 --- a/examples/nft/lora/wan21/t2v.yaml +++ b/examples/nft/lora/wan21/t2v.yaml @@ -41,7 +41,7 @@ model: # - name: torch_compile # ...so the compiled graph captures it # params: { mode: regional } # regional (compile_repeated_blocks) | full # rollout: # rollout-only (Stage 3); lossy -> decoupled/distillation only -# - name: diffusers_cache # Options: diffusers_cache, cache_dit +# - name: diffusers_cache # Options: diffusers_cache # params: { policy: first_block, threshold: 0.08 } log: diff --git a/examples/nft/lora/wan22/t2v.yaml b/examples/nft/lora/wan22/t2v.yaml index a1bcd7c7..8a12fe92 100644 --- a/examples/nft/lora/wan22/t2v.yaml +++ b/examples/nft/lora/wan22/t2v.yaml @@ -41,7 +41,7 @@ model: # - name: torch_compile # ...so the compiled graph captures it # params: { mode: regional } # regional (compile_repeated_blocks) | full # rollout: # rollout-only (Stage 3); lossy -> decoupled/distillation only -# - name: diffusers_cache # Options: diffusers_cache, cache_dit +# - name: diffusers_cache # Options: diffusers_cache # params: { policy: first_block, threshold: 0.08 } log: diff --git a/guidance/acceleration.md b/guidance/acceleration.md index f3d33745..570dc15d 100644 --- a/guidance/acceleration.md +++ b/guidance/acceleration.md @@ -50,7 +50,7 @@ acceleration: # Rollout-only (Stage 3), nested in list order. May be lossy (paradigm-gated). rollout: - - name: diffusers_cache # diffusers_cache | cache_dit | + - name: diffusers_cache # diffusers_cache | params: { policy: first_block, threshold: 0.08 } ``` @@ -64,8 +64,7 @@ accepted as shorthand for a one-element list. A direct python path (e.g. |----|--------|-------|-------| | `attention_backend` | lossless | both | Sets the diffusers attention backend on every transformer. Requires a `backend` param. Forwards any backend (`native` / `flash` / `_flash_3` / `_flash_3_hub` / `sage` / `xformers`) to `set_attention_backend`. List it in `shared` **before** `torch_compile` so the compiled graph captures the backend. | | `torch_compile` | lossy | both | `torch.compile` of the shared transformer. `mode: regional` uses diffusers' `compile_repeated_blocks` (fast warmup, robust to variable resolution); `mode: full` compiles the whole module. Extra `compile_kwargs` forwarded to the compile call. Compiles in place (checkpoint- and EMA/ref-safe), applied after `post_init`. Marked **lossy** because it is applied symmetrically but is **not bit-exact across rollout vs training** (grad/no-grad graph split → intermittent ~1e-5 on-policy residual, within `clip_range`); allowed on coupled algos, but the validator warns. | -| `diffusers_cache` | lossy | rollout | Diffusers-native feature caching (no extra dependency). `policy`: `first_block` (default) / `faster` / `pyramid` / `taylorseer` / `magcache`; remaining params forwarded to the policy's diffusers config (e.g. `threshold`). | -| `cache_dit` | lossy | rollout | [cache-dit](https://github.com/vipshop/cache-dit) backend (DBCache/TaylorSeer). Requires `pip install flow-factory[acceleration]`. All params forwarded to `cache_dit.enable_cache`. | +| `diffusers_cache` | lossy | rollout | Diffusers-native feature caching (no extra dependency). `policy`: `first_block` (default) / `faster` / `pyramid` / `taylorseer` / `magcache`; remaining params forwarded to the policy's diffusers config (e.g. `threshold`). The single lossy rollout backend. | ### Attention backend diff --git a/pyproject.toml b/pyproject.toml index f183a1e0..e4e4c8b5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -64,12 +64,8 @@ quantization = [ "bitsandbytes>=0.45.3", ] -# Rollout acceleration: richer feature-caching policies (DBCache / TaylorSeer). -# Diffusers-native caching (the `diffusers_cache` accelerator) needs nothing extra; -# this extra is only for the `cache_dit` accelerator. -acceleration = [ - "cache-dit>=0.2.0", -] +# Rollout feature-caching (the `diffusers_cache` accelerator) is diffusers-native and +# needs no extra dependency, so there is no `acceleration` optional-dependency group. # Bagel adapter runtime (flash-attn varlen kernels + opencv transforms) bagel = [ @@ -105,7 +101,7 @@ geneval2 = [ # flow-factory: uv pip install hpsv2 --no-deps (runtime works with protobuf 6+). all = [ - "flow-factory[deepspeed,quantization,acceleration]", + "flow-factory[deepspeed,quantization]", ] [project.scripts] diff --git a/src/flow_factory/acceleration/__init__.py b/src/flow_factory/acceleration/__init__.py index 0065acdf..62e22b00 100644 --- a/src/flow_factory/acceleration/__init__.py +++ b/src/flow_factory/acceleration/__init__.py @@ -15,8 +15,8 @@ # src/flow_factory/acceleration/__init__.py """Model-agnostic acceleration plugin layer. -Concrete accelerators are resolved lazily through the registry (so optional -dependencies such as ``cache-dit`` are only imported when actually requested). +Concrete accelerators are resolved lazily through the registry, so an accelerator +module (and any heavy import it needs) is only loaded when actually requested. """ from .abc import BaseAccelerator diff --git a/src/flow_factory/acceleration/cache_dit.py b/src/flow_factory/acceleration/cache_dit.py deleted file mode 100644 index 1b55d57b..00000000 --- a/src/flow_factory/acceleration/cache_dit.py +++ /dev/null @@ -1,77 +0,0 @@ -# Copyright 2026 Jayce-Ping -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -# src/flow_factory/acceleration/cache_dit.py -"""Lossy rollout-only feature caching via the optional ``cache-dit`` library. - -`cache-dit `_ offers richer cache policies -(DBCache, TaylorSeer, ...) than diffusers' built-ins. It is an optional dependency -(``pip install flow-factory[acceleration]``); when absent, this accelerator raises -a clear install hint instead of failing at import (``constraints.md`` #22a). - -Only valid in the rollout slot of a decoupled / distillation trainer — the -paradigm validator enforces this (``constraints.md`` #7). -""" - -from contextlib import contextmanager -from typing import TYPE_CHECKING, Iterator - -from .abc import BaseAccelerator -from ..utils.logger_utils import setup_logger - -try: - import cache_dit -except ImportError: - cache_dit = None - -if TYPE_CHECKING: - from ..models.abc import BaseAdapter - -logger = setup_logger(__name__) - - -class CacheDitAccelerator(BaseAccelerator): - """Enable cache-dit caching on the adapter pipeline during rollout. - - Lossy and rollout-scoped. All params are forwarded verbatim to - ``cache_dit.enable_cache`` (see the cache-dit docs for policy options such as - ``cache_type``/``Fn_compute_blocks``/``rdt``), and cleared on context exit so - the Stage-6 training forward stays exact. - - cache-dit operates on ``adapter.pipeline``, whose ``transformer`` attribute is - the unwrapped inner module; rollout forwards still route through it via the - prepared bundle, so the cache hooks fire as expected. - """ - - safety = "lossy" - stage = "rollout" - - def __init__(self, **params) -> None: - super().__init__(**params) - if cache_dit is None: - raise ImportError( - "CacheDitAccelerator requires the optional 'cache-dit' package. " - "Install it with `pip install flow-factory[acceleration]` (or " - "`pip install cache-dit`)." - ) - - @contextmanager - def rollout_context(self, adapter: "BaseAdapter") -> Iterator[None]: - pipeline = adapter.pipeline - cache_dit.enable_cache(pipeline, **self.params) - logger.info("CacheDitAccelerator: cache-dit enabled on the rollout pipeline.") - try: - yield - finally: - cache_dit.disable_cache(pipeline) diff --git a/src/flow_factory/acceleration/registry.py b/src/flow_factory/acceleration/registry.py index 0fa8b388..1e88fa18 100644 --- a/src/flow_factory/acceleration/registry.py +++ b/src/flow_factory/acceleration/registry.py @@ -38,10 +38,25 @@ 'torch_compile': 'flow_factory.acceleration.torch_compile.CompileAccelerator', # Lossy (rollout-only; validator restricts to decoupled / distillation algos). 'diffusers_cache': 'flow_factory.acceleration.diffusers_cache.DiffusersCacheAccelerator', - 'cache_dit': 'flow_factory.acceleration.cache_dit.CacheDitAccelerator', } _ACCELERATOR_REGISTRY = {k.lower(): v for k, v in _ACCELERATOR_REGISTRY.items()} +# Accelerators that were removed; mapped to an actionable migration message so a +# config still naming them fails fast with guidance instead of a generic error. +_REMOVED_ACCELERATORS: Dict[str, str] = { + 'cache_dit': ( + "The 'cache_dit' accelerator was removed. cache-dit only caches inside a " + "pipeline `__call__` it patches, but Flow-Factory's rollout drives the " + "transformer directly (adapter.inference()), so it was a silent no-op; its " + "transformer-only path also assumes enable-once (incompatible with the " + "per-epoch rollout lifecycle) and conflicts with the adapters' native " + "diffusers `cache_context`. Use the `diffusers_cache` accelerator instead " + "(diffusers-native FirstBlockCache / TaylorSeer / FasterCache; no extra " + "dependency): `{ name: diffusers_cache, params: { policy: first_block, " + "threshold: 0.08 } }`." + ), +} + def get_accelerator_class(identifier: str) -> Type[BaseAccelerator]: """Resolve and import an accelerator class from the registry or a python path. @@ -60,6 +75,8 @@ def get_accelerator_class(identifier: str) -> Type[BaseAccelerator]: ImportError: If the accelerator cannot be loaded. """ identifier_lower = identifier.lower() + if identifier_lower in _REMOVED_ACCELERATORS: + raise ValueError(_REMOVED_ACCELERATORS[identifier_lower]) if identifier_lower in _ACCELERATOR_REGISTRY: class_path = _ACCELERATOR_REGISTRY[identifier_lower] else: From 247a64a0ae42b584d1f190bbd160c0e4bc59e5f2 Mon Sep 17 00:00:00 2001 From: Jayce-Ping <315229706@qq.com> Date: Tue, 30 Jun 2026 10:06:18 +0800 Subject: [PATCH 12/15] [acceleration] fix: diffusers_cache + periodic eval crash ("No context is set") A decoupled trainer with a lossy `diffusers_cache` rollout accelerator AND periodic eval (eval_freq>0) crashed on the first cached training rollout after an eval: `ValueError: No context is set` from diffusers FirstBlockCache. Root cause (diffusers HookRegistry): `_get_child_registries()` caches the child-registry list on first use. During eval the cache is disabled, but the adapter still opens `transformer.cache_context(...)`, which creates the top HookRegistry and caches an EMPTY child list (no block cache-hooks exist yet). When training then enables the cache, the FBC block hooks are added, but the stale empty cache makes `_set_context` skip them, so the block forward reads an unset context and raises. Fix: after enable_cache, invalidate the stale `_child_registries_cache` on the unwrapped transformer (the exact object the adapter's `cache_context` targets) so the next context build rediscovers the freshly added block hooks. No-op on the common no-eval path. Verified on 16xH20 (2-node) Qwen-Image NFT + diffusers_cache + eval_freq=1: eval->train interleave runs clean (rc=0, zero "No context" errors). Also smoke-verified torch_compile (regional) + diffusers_cache coexist with stable per-epoch rollout time (no recompile churn). --- src/flow_factory/acceleration/diffusers_cache.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/src/flow_factory/acceleration/diffusers_cache.py b/src/flow_factory/acceleration/diffusers_cache.py index e56b2120..a1fe8f04 100644 --- a/src/flow_factory/acceleration/diffusers_cache.py +++ b/src/flow_factory/acceleration/diffusers_cache.py @@ -116,6 +116,20 @@ def rollout_context(self, adapter: "BaseAdapter") -> Iterator[None]: transformer.disable_cache() transformer.enable_cache(self._build_config()) enabled.append(transformer) + # diffusers' HookRegistry caches its child-registry list the first time + # a `cache_context` sets a context. If a `cache_context` ran while the + # cache was DISABLED (e.g. during eval, where the adapter still opens + # `transformer.cache_context(...)`), that cache was populated EMPTY -- + # before the per-block cache hooks above existed -- so a later + # `_set_context` never reaches the freshly added block hooks and the + # block forward raises "No context is set". Invalidate the stale cache on + # the unwrapped module (the exact object the adapter's `cache_context` + # targets) so the next context build rediscovers the new hooks. Safe/no-op + # when no such registry exists yet (the common no-eval path). + unwrapped = adapter.get_component_unwrapped(name) + cache_hook = getattr(unwrapped, "_diffusers_hook", None) + if cache_hook is not None: + cache_hook._child_registries_cache = None logger.info("DiffusersCacheAccelerator: cache enabled for '%s'.", name) yield finally: From 65afd3547e966bc60a592aed61a131a4f864874b Mon Sep 17 00:00:00 2001 From: Jayce-Ping <315229706@qq.com> Date: Tue, 30 Jun 2026 11:39:00 +0800 Subject: [PATCH 13/15] [repo] chore: untrack .cursor/plans internal roadmap doc acceleration_plugin_layer.plan.md is an internal design/roadmap note, not part of the shipped feature; remove it from the PR and gitignore .cursor/plans/ so scratch plans don't re-enter the tree. The tracked .cursor/rules/*.mdc project rules are unaffected. --- .../plans/acceleration_plugin_layer.plan.md | 216 ------------------ .gitignore | 5 + 2 files changed, 5 insertions(+), 216 deletions(-) delete mode 100644 .cursor/plans/acceleration_plugin_layer.plan.md diff --git a/.cursor/plans/acceleration_plugin_layer.plan.md b/.cursor/plans/acceleration_plugin_layer.plan.md deleted file mode 100644 index 46e0c0ac..00000000 --- a/.cursor/plans/acceleration_plugin_layer.plan.md +++ /dev/null @@ -1,216 +0,0 @@ -# Flow-Factory Acceleration Roadmap & Phase-1 Plan - -## Context - -Flow-Factory is an **online RL fine-tuning framework** for diffusion/flow-matching models. Its -dominant runtime cost is **Stage 3 — Trajectory Generation (rollout)**: for every epoch it runs -`num_inference_steps` denoising steps × `group_size (K)` samples per prompt × `num_batches_per_epoch`, -all under `torch.no_grad()`. This is *exactly* the inference workload that cache-dit / lightx2v / SGLang -accelerate. Stage 6 (policy optimization, forward+backward) is the secondary cost. - -The framework already has a solid **lossless** optimization base (FSDP/FSDP2/DeepSpeed, grad-checkpoint, -LoRA, bf16/fp16, CPU-offload + async H2D prefetch, Qwen CFG-merge, comm fusion, `StackedSampleBatch`, -single-root `ModelBundle`). What is missing is the **inference-acceleration layer** that the referenced -projects specialize in: feature caching (cache-dit DBCache/TaylorSeer/FBCache), whole-/region -`torch.compile`, long-sequence context parallelism (lightx2v / xDiT / SGLang style), and -quantization of frozen modules. - -This plan proposes a **model-agnostic, registry-based acceleration plugin layer** that respects the -framework's hard algorithm/model decoupling, plus a phased roadmap. Goal of Phase 1: cut rollout wall-time -on image models (FLUX/Qwen-Image) with **zero changes to trainers or model math**, gated for correctness. - ---- - -## The Decoupling-Critical Insight (read first — shapes the whole design) - -Acceleration techniques split into two correctness classes, and the split maps **exactly** onto -Constraint #7 (coupled vs decoupled) + Philosophy #1 / Constraint #20a (train-inference consistency): - -| Class | Examples | Where safe | Why | -|---|---|---|---| -| **Lossy** (changes `noise_pred`) | feature caching, step-reduction, int8/Sage attention | **rollout-only, decoupled/distillation-only** | Caching skips block compute; it cannot be replicated in the Stage-6 training forward (that needs full gradient through every block, every step). For **coupled** algos (GRPO/GRPO-Guard/DPPO) the rollout's per-step `log_prob` becomes the PPO "old log-prob"; if rollout is approximated but the training forward is exact, the importance ratio is biased → **silently wrong gradients**. For **decoupled** algos (NFT/AWM/DGPO/DPO/CRD) rollout runs with `compute_log_prob=False`; only the final image enters the reward → caching merely shifts the sample distribution, which is acceptable/tunable. | -| **Lossless** (numerically ~identical, applied to the *shared* module) | `torch.compile`, exact attention backends (FA2/FA3/xformers), context/sequence parallelism, quantizing **frozen-only** modules, comm/overlap/offload | **everywhere** | Because both `inference()` and `forward()` call the *same* `adapter.transformer`, a transform applied to that module is consistent across rollout and training by construction. | - -**Design rule enforced by the plan:** every accelerator declares `safety ∈ {lossless, lossy}` and -`stage ∈ {rollout, train, both}`. A fail-fast validator (Constraint #26) rejects `lossy` accel on a -coupled trainer, and rejects `lossy` accel on the `train` stage for any trainer. This is the mechanism -that lets us add aggressive rollout speedups *without* touching the train-inference consistency contract. - ---- - -## Current State (verified) - -**Already present / reusable seams:** -- `attn_backend` config (`hparams/model_args.py:110`) → `BaseAdapter._set_attention_backend()` - (`models/abc.py:852-866`) routes to diffusers' attention dispatcher: `native / flash / flash_hub / - _flash_3 / _flash_3_hub / sage / xformers`. **Attention backends are already wired framework-wide** — - they need benchmarking/per-stage exposure, not invention. -- Transformers are diffusers `CacheMixin` models: adapters already open `transformer.cache_context(...)` - (Qwen, Wan, LTX2, FLUX2-Klein). This is the exact hook diffusers-native & cache-dit caching plug into. -- Pluggable **scheduler registry** (`scheduler/registry.py`): FlowMatchEuler + UniPC multistep — solver - swaps / step schedules are already extensible. -- `RoutedComponentProxy` (`models/model_bundle.py:93-121`) delegates attribute access (incl. - `enable_cache`, `cache_context`, `compile`) to the inner module → an accelerator can call - `adapter.transformer.enable_cache(...)` model-agnostically, post-`prepare`. -- Rollout entry points: `BaseTrainer.generate_samples()` / `sample_batch()` (`trainers/abc.py:550-798`). - -**Absent (opportunities):** feature caching of any kind; whole-/region `torch.compile` (only -`bagel/.../qwen2_navit.py:43` compiles `flex_attention`); context/sequence parallelism for video; -quantization integration (`bitsandbytes` is an optional dep but unused); rollout throughput/VRAM/quality -benchmark harness. - ---- - -## Roadmap (phased, by ROI × safety × effort) - -| Phase | Track | What | Safety | Primary win | Effort | -|---|---|---|---|---|---| -| **0** | Foundation | Benchmark harness (rollout samples/s, step-time, peak VRAM, **reward-distribution regression**) + the `acceleration/` plugin layer + paradigm-gated validator | infra | enables safe tuning | S | -| **1** | Rollout (image-first) | **Feature caching** (cache-dit / diffusers-native) gated decoupled-rollout-only; **`torch.compile`** of shared transformer (lossless, everywhere); finish **attn_backend** benchmark + per-stage selection | mixed | **2–4× rollout on NFT/AWM/DGPO/CRD; ~1.2–1.8× compile everywhere** | M | -| **2** | Video long-seq | **Context/sequence parallelism** (Ulysses/Ring attention) for the DiT during rollout *and* training (Wan2, LTX2) — lossless | lossless | enables/scales long video; near-linear on seq dim | L | -| **3** | Memory→throughput | Quantize **frozen** modules only (text encoders, VAE, reference/EMA, reward models) via torchao fp8 / bnb nf4; FSDP2 selective activation checkpointing; smarter offload scheduling rollout vs train | lossless | bigger batch / larger models | M | -| **4** | Pipeline overlap | Async reward (Stage 4) + advantage (Stage 5) overlapped with next rollout; extend existing comm fusion | lossless | hide reward/comm latency | M | - -Phases 1 and 2 are the user-prioritized fronts (rollout throughput, then video parallelism). 3–4 are -follow-ons. Everything hangs off the Phase-0 plugin layer so each later track is a registry entry, not a -trainer/adapter edit. - -**Execution order requested by user:** implement all **lossless** accelerators first (torch.compile + -per-stage attention backend within the plugin layer), then the **lossy** feature-caching accelerators. - ---- - -## Proposed Architecture: `src/flow_factory/acceleration/` (model-agnostic plugin layer) - -Mirror the existing `rewards/` module shape (registry + abc + concrete impls + hparams), so it inherits -the framework's decoupling guarantees and lazy-import conventions (Constraints #1–#4). - -``` -src/flow_factory/acceleration/ - __init__.py - abc.py # BaseAccelerator: declares safety + stage; setup()/rollout_context() - registry.py # _ACCELERATOR_REGISTRY {id -> path} + get_accelerator_class() w/ direct-path fallback - validator.py # paradigm-gated safety check (fail-fast, Constraint #26) - torch_compile.py # CompileAccelerator (lossless; applied to shared transformer at init) - attention_backend.py # AttentionBackendAccelerator (lossless exact / lossy sage, per-stage) - cache_dit.py # CacheDitAccelerator (lossy; wraps cache_dit.enable_cache / disable_cache) - diffusers_cache.py # DiffusersCacheAccelerator (lossy; FirstBlockCache/FasterCache/PAB) -``` - -**`BaseAccelerator` (abc.py)** — minimal contract: -```python -class BaseAccelerator(ABC): - safety: ClassVar[Literal["lossless", "lossy"]] - stage: ClassVar[Literal["rollout", "train", "both"]] - - @abstractmethod - def setup(self, adapter: "BaseAdapter") -> None: ... # one-time (e.g. torch.compile) - @contextmanager - def rollout_context(self, adapter: "BaseAdapter"): ... # enable cache → yield → disable+reset -``` - -**Config (`hparams/acceleration_args.py`)** — new `AccelerationArguments`, aggregated into top-level -`Arguments` (Constraint #15). Per-stage so lossy stays rollout-only: -```yaml -acceleration: - rollout: { name: cache_dit, params: { policy: DBCache, rdt: 0.08 } } # lossy, rollout-only - shared: { name: torch_compile, params: { mode: default, dynamic: true } } # lossless, both stages -``` - -**Trainer paradigm tag.** Add `paradigm: ClassVar[Literal["coupled","decoupled","distillation"]]` to each -trainer (values already fixed by Constraint #7). `validator.py` reads it and the chosen accelerators' -`safety/stage` to enforce the table above before training starts. - -**Integration points (only two, both model-agnostic, no per-adapter math change):** -1. `BaseTrainer._initialization()` builds `self.rollout_accelerator` / `self.shared_accelerator` from - config via the registry, runs `validator.validate(paradigm, accelerators)`, and calls - `shared_accelerator.setup(adapter)` (e.g. compile) after `accelerator.prepare()`. -2. `BaseTrainer.generate_samples()` (`trainers/abc.py:701`) wraps the batch loop: - `with self.rollout_accelerator.rollout_context(self.adapter): ...`. The context calls - `adapter.transformer.enable_cache(...)` (reachable through `RoutedComponentProxy`) and tears it down - after, so cache state never leaks into the Stage-6 forward. - -Trainers and adapters stay unchanged except: (a) the one-line `paradigm` tag, (b) the two -`BaseTrainer` hooks above. No `optimize()` / `forward()` / `inference()` signature changes -(Constraint #12), no scheduler-to-trainer coupling. - ---- - -## Phase 1 — Detailed Implementation Spec - -**Objective:** lossless `torch.compile` + per-stage attention path usable by all algos, then 2–4× -rollout speedup on decoupled image trainers via feature caching, with a benchmark that proves reward -isn't degraded. - -### 1. Phase-0 prerequisites bundled in -- **Benchmark harness** under `.scratch/bench/` (scratch per Constraint #28) + a small reusable profiler - in `utils/`: measure rollout samples/s, mean step time, peak VRAM, and **mean/std of the reward - buffer** before vs after accel. The reward-regression check is mandatory for the lossy path. -- Build `acceleration/{abc,registry,validator}.py` and `hparams/acceleration_args.py`; wire - `AccelerationArguments` into `hparams/args.py` `Arguments`. - -### 2. Lossless first — `torch.compile` (both stages) -- **`CompileAccelerator`** (`acceleration/torch_compile.py`): in `setup(adapter)` apply - `torch.compile` to the shared transformer — prefer diffusers `transformer.compile_repeated_blocks()` - (region compile → fast warmup, robust to the variable image/seq-len shapes set by - `calculate_shift`/`set_scheduler_timesteps`). Because the same compiled module backs both - `inference()` and `forward()`, it is consistent for coupled algos too. Handle LoRA (compile after PEFT - wrap) and `dynamic=True` for variable resolution. - -### 3. Lossless — attention backend (per-stage) -- No new mechanism — extend `attn_backend` to be **per-stage** through the same accel config, and add a - benchmark sweep (FA2/FA3/xformers exact = lossless/both; Sage int8 = lossy/rollout-only). Document - results in `guidance/`. - -### 4. Lossy — feature caching (rollout-only, decoupled-only) -- **`DiffusersCacheAccelerator`** (`acceleration/diffusers_cache.py`): zero-extra-dep, uses diffusers' - built-in `transformer.enable_cache(FirstBlockCacheConfig | FasterCacheConfig | PyramidAttentionBroadcastConfig)` / `disable_cache()`. **Default lossy backend.** -- **`CacheDitAccelerator`** (`acceleration/cache_dit.py`): optional dep (`cache-dit`), imported under - `try/except ImportError` per Constraint #22(a). `rollout_context` calls - `cache_dit.enable_cache(adapter.transformer, ...)` on enter, `cache_dit.disable_cache(...)` on exit. -- **Adapter cache-readiness:** caching hooks fire only inside a `cache_context`. Adapters that already - wrap their transformer call in `transformer.cache_context(...)` (Qwen-Image `qwen_image.py:598`, Wan, - LTX2, FLUX2-Klein) are ready. Adapters that call the transformer bare (e.g. FLUX.1 `flux1.py:323`) - need the existing call wrapped in a `cache_context` — a localized, behavior-preserving change, not a - decoupling violation. **Validate first on Qwen-Image (already ready) + an NFT/AWM/DGPO config.** -- New optional dep: add `cache-dit` to `[project.optional-dependencies]` as `acceleration = [...]`. - -### 5. Files touched (Phase 1) -- **New:** `src/flow_factory/acceleration/*` (7 files), `src/flow_factory/hparams/acceleration_args.py`. -- **Edited (small):** `hparams/args.py` (aggregate args); `trainers/abc.py` (`_initialization` build+validate+setup, `generate_samples` context wrap); each trainer class (one `paradigm` ClassVar); `flux1.py`-style adapters lacking `cache_context` (wrap existing transformer call); `pyproject.toml` (optional dep); `examples/` configs gain an optional `acceleration:` block (Constraint #15/#17); `guidance/` + `.agents/knowledge/architecture.md` doc updates (registry table + new module). - -### 6. Reuse (do not reinvent) -- Registry pattern: copy `rewards/registry.py` resolution + direct-path fallback. -- Hook seam: `RoutedComponentProxy.__getattr__` already exposes `enable_cache`/`cache_context`/`compile`. -- Rollout loop already isolated in `generate_samples()`/`sample_batch()` — single wrap point. -- diffusers native caches need no new dependency. - ---- - -## Verification - -1. **Correctness gating (unit):** assert `validator` raises (Constraint #26) for `lossy` accel on a - coupled trainer and for `lossy` on the `train` stage; passes for decoupled-rollout. -2. **Numerical (lossless):** with `torch.compile` only, a short GRPO run's per-step `log_prob` and loss - match the eager baseline within tolerance (confirms coupled-safety of the shared-module path). -3. **Rollout speed + quality (lossy):** run an NFT (or AWM/DGPO) image config on Qwen-Image with - caching vs baseline for N epochs; report rollout samples/s, peak VRAM, and **reward mean/std drift** - from the harness. Accept only if speedup ≥ target and reward drift within budget. -4. **End-to-end:** `ff-train examples//lora/qwen_image/default.yaml` with and without the - `acceleration:` block trains and checkpoints identically in shape; `black --check src/ && isort - --check src/` clean (Constraint #21); run `/ff-review` before commit. -5. **Video smoke (Phase 2 entry):** confirm the same plugin enables on Wan2/LTX2 (already use - `cache_context`) before building context parallelism. - ---- - -## Open Risks / Decisions -- **cache-dit vs diffusers-native** as the default lossy backend: cache-dit has richer policies - (DBCache/TaylorSeer) but adds a dep; diffusers-native (FBCache/FasterCache/PAB) is dep-free. Plan - ships **both** behind the registry; default = diffusers-native, cache-dit opt-in. -- **FLUX.1 cache_context wrap** is the only per-adapter code change in Phase 1 — keep it behavior-identical - when caching is off (empty context). -- **Reward-drift budget** for the lossy path must be set with the user (per reward model) — the harness - measures it but the accept/reject threshold is a product decision. -- Phase 2 context parallelism interacts with FSDP/`ModelBundle` sharding and the SDE scheduler's - per-token noise (`scheduler/flow_match_euler_discrete.py`) — needs its own design pass (separate plan). diff --git a/.gitignore b/.gitignore index e1973598..1f94df7b 100644 --- a/.gitignore +++ b/.gitignore @@ -139,3 +139,8 @@ dataset/**/*.mov *.zip *.gz *.tar + +# ==================== +# Cursor scratch plans (design/roadmap docs; keep .cursor/rules tracked) +# ==================== +.cursor/plans/ From 9ec997c5a75d29a73f86f37ebf3edc1435854dd5 Mon Sep 17 00:00:00 2001 From: Jayce-Ping <315229706@qq.com> Date: Tue, 30 Jun 2026 12:43:09 +0800 Subject: [PATCH 14/15] [acceleration,trainers,docs] fix: bound eval-loop memory under torch.compile Route BaseTrainer.evaluate() through _rollout_grad_context() instead of a bare torch.no_grad(). The compiled transformer's forward unconditionally forces enable_grad (for rollout/train graph consistency), so a no_grad eval still built an autograd graph that chained across the whole denoising loop (memory grew with num_inference_steps -> OOM). The shared context enables the per-step cast_latents detach so eval memory stays bounded. With no grad-forcing accelerator active it is exactly torch.no_grad() (no behavior change). evaluate()/generate_samples() are base-class-only, so this covers all trainers. Also fix the accelerator table in architecture.md: add a blank line so the trailing "Configured via the acceleration: block" paragraph renders below the table instead of being absorbed as a malformed single-column row. --- .agents/knowledge/architecture.md | 1 + src/flow_factory/trainers/abc.py | 21 ++++++++++++++++++--- 2 files changed, 19 insertions(+), 3 deletions(-) diff --git a/.agents/knowledge/architecture.md b/.agents/knowledge/architecture.md index 0600b6ca..d052a870 100644 --- a/.agents/knowledge/architecture.md +++ b/.agents/knowledge/architecture.md @@ -153,6 +153,7 @@ All four registries map string keys → lazy import paths. Resolution: registry | `attention_backend` | `AttentionBackendAccelerator` | lossless | both | Sets the diffusers attention backend on every transformer (requires a `backend` param). Listed as a `shared` entry (before `torch_compile`); this is the single code path for backend selection (the old `BaseAdapter._set_attention_backend` and the `model.attn_backend` knob were both removed — a config still setting `model.attn_backend` fails fast). Bagel forces flash_attention_2 at load and does not use it. | | `torch_compile` | `CompileAccelerator` | lossy | both | `torch.compile` of the shared transformer (regional/full); applied in-place after `post_init` so checkpoint keys / param identity stay stable. Marked `lossy` because it is applied symmetrically but is **not bit-exact across rollout vs training** (Inductor's grad/no-grad graph split → intermittent ~1e-5 on-policy residual, within `clip_range`); still allowed on coupled algos, validator warns. | | `diffusers_cache` | `DiffusersCacheAccelerator` | lossy | rollout | Diffusers `CacheMixin` feature caching (first_block/faster/pyramid/taylorseer/magcache). | + Configured via the `acceleration:` block (`hparams/acceleration_args.py`): two ordered lists of `{name, params}` entries — `shared` (lossless, both stages) and `rollout` (Stage-3 only). **List order is application order**: `shared` entries run their `setup()` in order (so `attention_backend` must precede `torch_compile`), and `rollout` entries nest their `rollout_context()` in order. The `acceleration/validator.py` enforces that a **lossy `rollout`** accelerator runs only on `decoupled`/`distillation` trainers (each trainer declares a `paradigm`), preserving train-inference consistency (constraints #7, #20a). A **lossy `shared` (stage=both)** accelerator (e.g. `torch_compile`, applied symmetrically but not bit-exact across stages) is allowed on any paradigm but the validator **warns** on coupled trainers that the on-policy ratio will be ≈1, not exactly 1. Off by default. --- diff --git a/src/flow_factory/trainers/abc.py b/src/flow_factory/trainers/abc.py index 16c5b606..2b1ec3ee 100644 --- a/src/flow_factory/trainers/abc.py +++ b/src/flow_factory/trainers/abc.py @@ -468,9 +468,9 @@ def _rollout_acceleration(self) -> Iterator[None]: @contextmanager def _rollout_grad_context(self) -> Iterator[None]: - """Grad context for the Stage-3 rollout loop. + """Grad context for an ``adapter.inference()`` loop (Stage-3 rollout and eval). - Default: ``torch.no_grad()`` (rollout needs no gradients — saves memory). + Default: ``torch.no_grad()`` (inference needs no gradients — saves memory). Exception: if any active **shared** accelerator declares ``requires_grad_rollout`` (only ``torch_compile`` does), rollout instead runs @@ -488,6 +488,15 @@ def _rollout_grad_context(self) -> Iterator[None]: (bit-exact on-policy ratio, with or without CFG — see ``guidance/acceleration.md``) while the autograd graph never chains across denoising steps (memory stays bounded). + + Reused by the eval loop (:meth:`evaluate`). Eval does not compute the PPO + ratio, so graph consistency is irrelevant there — but the compiled + transformer's forward forces ``enable_grad`` *unconditionally*, so a bare + ``torch.no_grad()`` eval would still build an autograd graph that chains across + the whole denoising loop (memory grows with ``num_inference_steps`` → OOM). + Routing eval through this context enables the same per-step ``cast_latents`` + detach, keeping eval memory bounded. With no grad-forcing accelerator active it + is exactly ``torch.no_grad()`` for both callers (no behavior change). """ needs_grad = any(acc.requires_grad_rollout for acc in self.shared_accelerators) if not needs_grad: @@ -960,7 +969,13 @@ def evaluate(self) -> None: self.adapter.eval() - with torch.no_grad(), self.autocast(), self.adapter.use_ema_parameters(): + # Use the rollout grad context (not a bare `torch.no_grad()`): when a + # grad-forcing shared accelerator (torch.compile) is active the compiled + # transformer forces `enable_grad`, so a plain `no_grad` here would let the + # autograd graph chain across the whole denoising loop (memory grows with + # num_inference_steps -> OOM). This context enables the per-step latent detach + # to keep eval memory bounded; with no such accelerator it is `torch.no_grad()`. + with self._rollout_grad_context(), self.autocast(), self.adapter.use_ema_parameters(): for dataset_name, dataloader in self.eval_dataloaders.items(): buffer = self.eval_dataset_reward_buffers.get(dataset_name) if buffer is None: From a58f32f8ffc58c99274c58dd4fefc4ea853a76c7 Mon Sep 17 00:00:00 2001 From: Jayce-Ping <315229706@qq.com> Date: Sun, 12 Jul 2026 08:04:40 +0800 Subject: [PATCH 15/15] [acceleration,docs] fix: align safety semantics and distributed logging Clarify torch.compile's non-bit-exact guarantees across code and examples, and keep accelerator lifecycle logs on the main process. --- .agents/knowledge/architecture.md | 2 +- examples/awm/lora/flux1/default.yaml | 2 +- .../awm/lora/flux2_klein_base/default.yaml | 2 +- examples/awm/lora/sd3_5/default.yaml | 2 +- examples/crd/lora/sd3_5/default.yaml | 2 +- examples/grpo/full/qwen_image/default.yaml | 6 +-- .../full/qwen_image_edit_plus/default.yaml | 6 +-- examples/grpo/lora/ltx2/t2av.yaml | 6 +-- examples/grpo/lora/ltx2/t2av_pickscore.yaml | 6 +-- examples/grpo/lora/qwen_image/default.yaml | 6 +-- .../lora/qwen_image_edit_plus/default.yaml | 6 +-- examples/grpo/lora/sd3_5/default.yaml | 6 +-- examples/grpo/lora/sd3_5/geneval.yaml | 6 +-- examples/grpo/lora/sd3_5/nocfg.yaml | 6 +-- examples/nft/lora/flux1/default.yaml | 2 +- .../lora/qwen_image/rational_rewards_t2i.yaml | 2 +- examples/nft/lora/sd3_5/default.yaml | 2 +- examples/nft/lora/wan21/i2v.yaml | 2 +- examples/nft/lora/wan21/t2v.yaml | 2 +- examples/nft/lora/wan22/t2v.yaml | 2 +- examples/template/sd3_5/async_reward.yaml | 6 +-- guidance/acceleration.md | 39 ++++++++++--------- multinode_examples/train.yaml | 6 +-- src/flow_factory/acceleration/abc.py | 15 +++---- .../acceleration/diffusers_cache.py | 3 +- src/flow_factory/acceleration/registry.py | 2 +- .../acceleration/torch_compile.py | 32 +++++++-------- src/flow_factory/acceleration/validator.py | 2 +- src/flow_factory/hparams/acceleration_args.py | 4 +- src/flow_factory/models/abc.py | 4 +- src/flow_factory/trainers/abc.py | 32 +++++++-------- 31 files changed, 113 insertions(+), 108 deletions(-) diff --git a/.agents/knowledge/architecture.md b/.agents/knowledge/architecture.md index d052a870..7dd103a6 100644 --- a/.agents/knowledge/architecture.md +++ b/.agents/knowledge/architecture.md @@ -154,7 +154,7 @@ All four registries map string keys → lazy import paths. Resolution: registry | `torch_compile` | `CompileAccelerator` | lossy | both | `torch.compile` of the shared transformer (regional/full); applied in-place after `post_init` so checkpoint keys / param identity stay stable. Marked `lossy` because it is applied symmetrically but is **not bit-exact across rollout vs training** (Inductor's grad/no-grad graph split → intermittent ~1e-5 on-policy residual, within `clip_range`); still allowed on coupled algos, validator warns. | | `diffusers_cache` | `DiffusersCacheAccelerator` | lossy | rollout | Diffusers `CacheMixin` feature caching (first_block/faster/pyramid/taylorseer/magcache). | -Configured via the `acceleration:` block (`hparams/acceleration_args.py`): two ordered lists of `{name, params}` entries — `shared` (lossless, both stages) and `rollout` (Stage-3 only). **List order is application order**: `shared` entries run their `setup()` in order (so `attention_backend` must precede `torch_compile`), and `rollout` entries nest their `rollout_context()` in order. The `acceleration/validator.py` enforces that a **lossy `rollout`** accelerator runs only on `decoupled`/`distillation` trainers (each trainer declares a `paradigm`), preserving train-inference consistency (constraints #7, #20a). A **lossy `shared` (stage=both)** accelerator (e.g. `torch_compile`, applied symmetrically but not bit-exact across stages) is allowed on any paradigm but the validator **warns** on coupled trainers that the on-policy ratio will be ≈1, not exactly 1. Off by default. +Configured via the `acceleration:` block (`hparams/acceleration_args.py`): two ordered lists of `{name, params}` entries — `shared` (persistent `stage='both'` accelerators applied to rollout and training) and `rollout` (Stage-3 only). **List order is application order**: `shared` entries run their `setup()` in order (so `attention_backend` must precede `torch_compile`), and `rollout` entries nest their `rollout_context()` in order. The `acceleration/validator.py` enforces that a **lossy `rollout`** accelerator runs only on `decoupled`/`distillation` trainers (each trainer declares a `paradigm`), preserving train-inference consistency (constraint #7). A **lossy `stage='both'` accelerator in the `shared` slot** (e.g. `torch_compile`, applied symmetrically but not bit-exact across stages) is allowed on any paradigm but the validator **warns** on coupled trainers that the on-policy ratio will be ≈1, not exactly 1. Off by default. --- diff --git a/examples/awm/lora/flux1/default.yaml b/examples/awm/lora/flux1/default.yaml index 00c2ee3a..5ad3d85d 100644 --- a/examples/awm/lora/flux1/default.yaml +++ b/examples/awm/lora/flux1/default.yaml @@ -34,7 +34,7 @@ model: # See guidance/acceleration.md. This is a decoupled trainer, so lossy rollout # feature-caching is allowed. # acceleration: -# shared: # lossless, applied to BOTH rollout and training (in order) +# shared: # persistent stage='both'; applied in list order # - name: attention_backend # set the attention backend first... # params: { backend: _flash_3_hub } # - name: torch_compile # ...so the compiled graph captures it diff --git a/examples/awm/lora/flux2_klein_base/default.yaml b/examples/awm/lora/flux2_klein_base/default.yaml index aa470d80..a7fd2142 100644 --- a/examples/awm/lora/flux2_klein_base/default.yaml +++ b/examples/awm/lora/flux2_klein_base/default.yaml @@ -34,7 +34,7 @@ model: # See guidance/acceleration.md. This is a decoupled trainer, so lossy rollout # feature-caching is allowed. # acceleration: -# shared: # lossless, applied to BOTH rollout and training (in order) +# shared: # persistent stage='both'; applied in list order # - name: attention_backend # set the attention backend first... # params: { backend: _flash_3_hub } # - name: torch_compile # ...so the compiled graph captures it diff --git a/examples/awm/lora/sd3_5/default.yaml b/examples/awm/lora/sd3_5/default.yaml index f7eaa4ed..37b06c75 100644 --- a/examples/awm/lora/sd3_5/default.yaml +++ b/examples/awm/lora/sd3_5/default.yaml @@ -35,7 +35,7 @@ model: # See guidance/acceleration.md. This is a decoupled trainer, so lossy rollout # feature-caching is allowed. # acceleration: -# shared: # lossless, applied to BOTH rollout and training (in order) +# shared: # persistent stage='both'; applied in list order # - name: attention_backend # set the attention backend first... # params: { backend: _flash_3_hub } # - name: torch_compile # ...so the compiled graph captures it diff --git a/examples/crd/lora/sd3_5/default.yaml b/examples/crd/lora/sd3_5/default.yaml index 7f490059..d242cb03 100644 --- a/examples/crd/lora/sd3_5/default.yaml +++ b/examples/crd/lora/sd3_5/default.yaml @@ -36,7 +36,7 @@ model: # See guidance/acceleration.md. This is a decoupled trainer, so lossy rollout # feature-caching is allowed. # acceleration: -# shared: # lossless, applied to BOTH rollout and training (in order) +# shared: # persistent stage='both'; applied in list order # - name: attention_backend # set the attention backend first... # params: { backend: _flash_3_hub } # - name: torch_compile # ...so the compiled graph captures it diff --git a/examples/grpo/full/qwen_image/default.yaml b/examples/grpo/full/qwen_image/default.yaml index dc982f6e..0f12b5cf 100644 --- a/examples/grpo/full/qwen_image/default.yaml +++ b/examples/grpo/full/qwen_image/default.yaml @@ -30,10 +30,10 @@ model: resume_path: null # Local path or HF repo id (e.g. 'owner/repo[/subdir][@rev]') for previous checkpoint/lora adapter resume_type: null # Options: lora, full, state. Null to auto-detect based on `finetune_type` # Optional acceleration plugins (off by default). Applied in list order. -# See guidance/acceleration.md. GRPO is a coupled trainer, so only lossless -# `shared` accelerators are allowed (lossy rollout caching would bias the PPO ratio). +# See guidance/acceleration.md. GRPO is coupled, so lossy rollout caching is unsafe; +# persistent `shared` accelerators run in both rollout and training. # acceleration: -# shared: # lossless, applied to BOTH rollout and training (in order) +# shared: # persistent stage='both'; applied in list order # - name: attention_backend # set the attention backend first... # params: { backend: _flash_3_varlen_hub } # - name: torch_compile # ...so the compiled graph captures it diff --git a/examples/grpo/full/qwen_image_edit_plus/default.yaml b/examples/grpo/full/qwen_image_edit_plus/default.yaml index 4e28bd9f..33ed422e 100644 --- a/examples/grpo/full/qwen_image_edit_plus/default.yaml +++ b/examples/grpo/full/qwen_image_edit_plus/default.yaml @@ -29,10 +29,10 @@ model: resume_path: null # Local path or HF repo id (e.g. 'owner/repo[/subdir][@rev]') for previous checkpoint/lora adapter resume_type: null # Options: lora, full, state. Null to auto-detect based on `finetune_type` # Optional acceleration plugins (off by default). Applied in list order. -# See guidance/acceleration.md. GRPO is a coupled trainer, so only lossless -# `shared` accelerators are allowed (lossy rollout caching would bias the PPO ratio). +# See guidance/acceleration.md. GRPO is coupled, so lossy rollout caching is unsafe; +# persistent `shared` accelerators run in both rollout and training. # acceleration: -# shared: # lossless, applied to BOTH rollout and training (in order) +# shared: # persistent stage='both'; applied in list order # - name: attention_backend # set the attention backend first... # params: { backend: _flash_3_varlen_hub } # - name: torch_compile # ...so the compiled graph captures it diff --git a/examples/grpo/lora/ltx2/t2av.yaml b/examples/grpo/lora/ltx2/t2av.yaml index 15ac742a..6611c1c4 100644 --- a/examples/grpo/lora/ltx2/t2av.yaml +++ b/examples/grpo/lora/ltx2/t2av.yaml @@ -34,10 +34,10 @@ model: resume_path: null # Local path or HF repo id (e.g. 'owner/repo[/subdir][@rev]') for previous checkpoint/lora adapter resume_type: null # Options: lora, full, state. Null to auto-detect based on `finetune_type` # Optional acceleration plugins (off by default). Applied in list order. -# See guidance/acceleration.md. GRPO is a coupled trainer, so only lossless -# `shared` accelerators are allowed (lossy rollout caching would bias the PPO ratio). +# See guidance/acceleration.md. GRPO is coupled, so lossy rollout caching is unsafe; +# persistent `shared` accelerators run in both rollout and training. # acceleration: -# shared: # lossless, applied to BOTH rollout and training (in order) +# shared: # persistent stage='both'; applied in list order # - name: attention_backend # set the attention backend first... # params: { backend: _flash_3_hub } # - name: torch_compile # ...so the compiled graph captures it diff --git a/examples/grpo/lora/ltx2/t2av_pickscore.yaml b/examples/grpo/lora/ltx2/t2av_pickscore.yaml index 21038e92..03d43865 100644 --- a/examples/grpo/lora/ltx2/t2av_pickscore.yaml +++ b/examples/grpo/lora/ltx2/t2av_pickscore.yaml @@ -32,10 +32,10 @@ model: resume_path: null # Local path or HF repo id (e.g. 'owner/repo[/subdir][@rev]') for previous checkpoint/lora adapter resume_type: null # Options: lora, full, state. Null to auto-detect based on `finetune_type` # Optional acceleration plugins (off by default). Applied in list order. -# See guidance/acceleration.md. GRPO is a coupled trainer, so only lossless -# `shared` accelerators are allowed (lossy rollout caching would bias the PPO ratio). +# See guidance/acceleration.md. GRPO is coupled, so lossy rollout caching is unsafe; +# persistent `shared` accelerators run in both rollout and training. # acceleration: -# shared: # lossless, applied to BOTH rollout and training (in order) +# shared: # persistent stage='both'; applied in list order # - name: attention_backend # set the attention backend first... # params: { backend: _flash_3_hub } # - name: torch_compile # ...so the compiled graph captures it diff --git a/examples/grpo/lora/qwen_image/default.yaml b/examples/grpo/lora/qwen_image/default.yaml index dd6747ea..d17f7df3 100644 --- a/examples/grpo/lora/qwen_image/default.yaml +++ b/examples/grpo/lora/qwen_image/default.yaml @@ -31,10 +31,10 @@ model: resume_path: null # Local path or HF repo id (e.g. 'owner/repo[/subdir][@rev]') for previous checkpoint/lora adapter resume_type: null # Options: lora, full, state. Null to auto-detect based on `finetune_type` # Optional acceleration plugins (off by default). Applied in list order. -# See guidance/acceleration.md. GRPO is a coupled trainer, so only lossless -# `shared` accelerators are allowed (lossy rollout caching would bias the PPO ratio). +# See guidance/acceleration.md. GRPO is coupled, so lossy rollout caching is unsafe; +# persistent `shared` accelerators run in both rollout and training. # acceleration: -# shared: # lossless, applied to BOTH rollout and training (in order) +# shared: # persistent stage='both'; applied in list order # - name: attention_backend # set the attention backend first... # params: { backend: _flash_3_varlen_hub } # - name: torch_compile # ...so the compiled graph captures it diff --git a/examples/grpo/lora/qwen_image_edit_plus/default.yaml b/examples/grpo/lora/qwen_image_edit_plus/default.yaml index bfdd6669..ddb7daa7 100644 --- a/examples/grpo/lora/qwen_image_edit_plus/default.yaml +++ b/examples/grpo/lora/qwen_image_edit_plus/default.yaml @@ -32,10 +32,10 @@ model: resume_path: null # Local path or HF repo id (e.g. 'owner/repo[/subdir][@rev]') for previous checkpoint/lora adapter resume_type: null # Options: lora, full, state. Null to auto-detect based on `finetune_type` # Optional acceleration plugins (off by default). Applied in list order. -# See guidance/acceleration.md. GRPO is a coupled trainer, so only lossless -# `shared` accelerators are allowed (lossy rollout caching would bias the PPO ratio). +# See guidance/acceleration.md. GRPO is coupled, so lossy rollout caching is unsafe; +# persistent `shared` accelerators run in both rollout and training. # acceleration: -# shared: # lossless, applied to BOTH rollout and training (in order) +# shared: # persistent stage='both'; applied in list order # - name: attention_backend # set the attention backend first... # params: { backend: _flash_3_varlen_hub } # - name: torch_compile # ...so the compiled graph captures it diff --git a/examples/grpo/lora/sd3_5/default.yaml b/examples/grpo/lora/sd3_5/default.yaml index 6399ce89..ad8b1cba 100644 --- a/examples/grpo/lora/sd3_5/default.yaml +++ b/examples/grpo/lora/sd3_5/default.yaml @@ -34,10 +34,10 @@ model: resume_path: null # Local path or HF repo id (e.g. 'owner/repo[/subdir][@rev]') for previous checkpoint/lora adapter resume_type: null # Options: lora, full, state. Null to auto-detect based on `finetune_type` # Optional acceleration plugins (off by default). Applied in list order. -# See guidance/acceleration.md. GRPO is a coupled trainer, so only lossless -# `shared` accelerators are allowed (lossy rollout caching would bias the PPO ratio). +# See guidance/acceleration.md. GRPO is coupled, so lossy rollout caching is unsafe; +# persistent `shared` accelerators run in both rollout and training. # acceleration: -# shared: # lossless, applied to BOTH rollout and training (in order) +# shared: # persistent stage='both'; applied in list order # - name: attention_backend # set the attention backend first... # params: { backend: _flash_3_hub } # - name: torch_compile # ...so the compiled graph captures it diff --git a/examples/grpo/lora/sd3_5/geneval.yaml b/examples/grpo/lora/sd3_5/geneval.yaml index 3116654b..0a60fd40 100644 --- a/examples/grpo/lora/sd3_5/geneval.yaml +++ b/examples/grpo/lora/sd3_5/geneval.yaml @@ -31,10 +31,10 @@ model: resume_path: null # Local path or HF repo id (e.g. 'owner/repo[/subdir][@rev]') for previous checkpoint/lora adapter resume_type: null # Options: lora, full, state. Null to auto-detect based on `finetune_type` # Optional acceleration plugins (off by default). Applied in list order. -# See guidance/acceleration.md. GRPO is a coupled trainer, so only lossless -# `shared` accelerators are allowed (lossy rollout caching would bias the PPO ratio). +# See guidance/acceleration.md. GRPO is coupled, so lossy rollout caching is unsafe; +# persistent `shared` accelerators run in both rollout and training. # acceleration: -# shared: # lossless, applied to BOTH rollout and training (in order) +# shared: # persistent stage='both'; applied in list order # - name: attention_backend # set the attention backend first... # params: { backend: _flash_3_hub } # - name: torch_compile # ...so the compiled graph captures it diff --git a/examples/grpo/lora/sd3_5/nocfg.yaml b/examples/grpo/lora/sd3_5/nocfg.yaml index dd7a8be5..79c356f2 100644 --- a/examples/grpo/lora/sd3_5/nocfg.yaml +++ b/examples/grpo/lora/sd3_5/nocfg.yaml @@ -34,10 +34,10 @@ model: resume_path: null # Local path or HF repo id (e.g. 'owner/repo[/subdir][@rev]') for previous checkpoint/lora adapter resume_type: null # Options: lora, full, state. Null to auto-detect based on `finetune_type` # Optional acceleration plugins (off by default). Applied in list order. -# See guidance/acceleration.md. GRPO is a coupled trainer, so only lossless -# `shared` accelerators are allowed (lossy rollout caching would bias the PPO ratio). +# See guidance/acceleration.md. GRPO is coupled, so lossy rollout caching is unsafe; +# persistent `shared` accelerators run in both rollout and training. # acceleration: -# shared: # lossless, applied to BOTH rollout and training (in order) +# shared: # persistent stage='both'; applied in list order # - name: attention_backend # set the attention backend first... # params: { backend: _flash_3_hub } # - name: torch_compile # ...so the compiled graph captures it diff --git a/examples/nft/lora/flux1/default.yaml b/examples/nft/lora/flux1/default.yaml index 6c53e19e..18d4ad46 100644 --- a/examples/nft/lora/flux1/default.yaml +++ b/examples/nft/lora/flux1/default.yaml @@ -34,7 +34,7 @@ model: # See guidance/acceleration.md. This is a decoupled trainer, so lossy rollout # feature-caching is allowed. # acceleration: -# shared: # lossless, applied to BOTH rollout and training (in order) +# shared: # persistent stage='both'; applied in list order # - name: attention_backend # set the attention backend first... # params: { backend: _flash_3_hub } # - name: torch_compile # ...so the compiled graph captures it diff --git a/examples/nft/lora/qwen_image/rational_rewards_t2i.yaml b/examples/nft/lora/qwen_image/rational_rewards_t2i.yaml index 875c264b..9f795280 100644 --- a/examples/nft/lora/qwen_image/rational_rewards_t2i.yaml +++ b/examples/nft/lora/qwen_image/rational_rewards_t2i.yaml @@ -86,7 +86,7 @@ scheduler: # See guidance/acceleration.md. NFT is a decoupled trainer, so lossy rollout # feature-caching is allowed. # acceleration: -# shared: # lossless, applied to BOTH rollout and training (in order) +# shared: # persistent stage='both'; applied in list order # - name: attention_backend # set the attention backend first... # params: { backend: _flash_3_varlen_hub } # - name: torch_compile # ...so the compiled graph captures it diff --git a/examples/nft/lora/sd3_5/default.yaml b/examples/nft/lora/sd3_5/default.yaml index ff4b7e95..926e6219 100644 --- a/examples/nft/lora/sd3_5/default.yaml +++ b/examples/nft/lora/sd3_5/default.yaml @@ -35,7 +35,7 @@ model: # See guidance/acceleration.md. This is a decoupled trainer, so lossy rollout # feature-caching is allowed. # acceleration: -# shared: # lossless, applied to BOTH rollout and training (in order) +# shared: # persistent stage='both'; applied in list order # - name: attention_backend # set the attention backend first... # params: { backend: _flash_3_hub } # - name: torch_compile # ...so the compiled graph captures it diff --git a/examples/nft/lora/wan21/i2v.yaml b/examples/nft/lora/wan21/i2v.yaml index 973f4e9b..c6209aaf 100644 --- a/examples/nft/lora/wan21/i2v.yaml +++ b/examples/nft/lora/wan21/i2v.yaml @@ -35,7 +35,7 @@ model: # See guidance/acceleration.md. This is a decoupled trainer, so lossy rollout # feature-caching is allowed. # acceleration: -# shared: # lossless, applied to BOTH rollout and training (in order) +# shared: # persistent stage='both'; applied in list order # - name: attention_backend # set the attention backend first... # params: { backend: _flash_3_hub } # - name: torch_compile # ...so the compiled graph captures it diff --git a/examples/nft/lora/wan21/t2v.yaml b/examples/nft/lora/wan21/t2v.yaml index 0026a9a8..04da0c14 100644 --- a/examples/nft/lora/wan21/t2v.yaml +++ b/examples/nft/lora/wan21/t2v.yaml @@ -35,7 +35,7 @@ model: # See guidance/acceleration.md. This is a decoupled trainer, so lossy rollout # feature-caching is allowed. # acceleration: -# shared: # lossless, applied to BOTH rollout and training (in order) +# shared: # persistent stage='both'; applied in list order # - name: attention_backend # set the attention backend first... # params: { backend: _flash_3_hub } # - name: torch_compile # ...so the compiled graph captures it diff --git a/examples/nft/lora/wan22/t2v.yaml b/examples/nft/lora/wan22/t2v.yaml index 8a12fe92..c04c744c 100644 --- a/examples/nft/lora/wan22/t2v.yaml +++ b/examples/nft/lora/wan22/t2v.yaml @@ -35,7 +35,7 @@ model: # See guidance/acceleration.md. This is a decoupled trainer, so lossy rollout # feature-caching is allowed. # acceleration: -# shared: # lossless, applied to BOTH rollout and training (in order) +# shared: # persistent stage='both'; applied in list order # - name: attention_backend # set the attention backend first... # params: { backend: _flash_3_hub } # - name: torch_compile # ...so the compiled graph captures it diff --git a/examples/template/sd3_5/async_reward.yaml b/examples/template/sd3_5/async_reward.yaml index 336b14d6..5ab71f2b 100644 --- a/examples/template/sd3_5/async_reward.yaml +++ b/examples/template/sd3_5/async_reward.yaml @@ -31,10 +31,10 @@ model: resume_path: null # Local path or HF repo id (e.g. 'owner/repo[/subdir][@rev]') for previous checkpoint/lora adapter resume_type: null # Options: lora, full, state. Null to auto-detect based on `finetune_type` # Optional acceleration plugins (off by default). Applied in list order. -# See guidance/acceleration.md. GRPO is a coupled trainer, so only lossless -# `shared` accelerators are allowed (lossy rollout caching would bias the PPO ratio). +# See guidance/acceleration.md. GRPO is coupled, so lossy rollout caching is unsafe; +# persistent `shared` accelerators run in both rollout and training. # acceleration: -# shared: # lossless, applied to BOTH rollout and training (in order) +# shared: # persistent stage='both'; applied in list order # - name: attention_backend # set the attention backend first... # params: { backend: _flash_3_hub } # - name: torch_compile # ...so the compiled graph captures it diff --git a/guidance/acceleration.md b/guidance/acceleration.md index 570dc15d..c2d86418 100644 --- a/guidance/acceleration.md +++ b/guidance/acceleration.md @@ -15,23 +15,23 @@ enforces them against the trainer's `paradigm` before training starts (fail-fast | marker | values | meaning | |--------|--------|---------| -| `stage` | `both` / `rollout` | `both`: persistent transform applied to the transformer shared by rollout `inference()` and training `forward()` → **consistent by construction, safe for any algorithm** (the `shared` slot). `rollout`: per-epoch context, torn down before training (the `rollout` slot). | -| `safety` | `lossless` / `lossy` | Only consulted for `stage='rollout'`. `lossy` = changes rollout outputs → diverges from training. `lossless` = bit-identical. | +| `stage` | `both` / `rollout` | `both`: persistent transform applied to both rollout `inference()` and training `forward()` (the `shared` slot). `rollout`: per-epoch context, torn down before training (the `rollout` slot). | +| `safety` | `lossless` / `lossy` | Numerical consistency class. `lossless` = bit-identical; `lossy` = not bit-identical. It gates rollout-only accelerators and warns for lossy `stage='both'` accelerators. | The single restriction: a **`lossy` rollout** accelerator is allowed **only on `decoupled` / `distillation`** trainers. Why: for **coupled** algorithms (GRPO, GRPO-Guard, DPPO) the rollout's per-step log-prob becomes the PPO "old log-prob"; changing the rollout while the training forward stays exact biases the importance ratio -and silently corrupts gradients (`.agents/knowledge/constraints.md` #7, #20a). For +and silently corrupts gradients (`.agents/knowledge/constraints.md` #7). For **decoupled** (NFT, AWM, DGPO, DPO, CRD) and **distillation** (diffusion-opd), the rollout log-prob does not enter the loss, so a lossy rollout only shifts the generated-sample distribution — acceptable and tunable. Monitor the reward mean/std when enabling it. -A `stage='both'` accelerator is **always safe** — even a numerically-approximate one. For -example, Sage int8 attention used as the attention backend runs in *both* rollout and -training, so the two stay consistent; there is no need to reject it or give it a special -"lossy" status. Numerical exactness only matters when a transform is applied to one stage -but not the other, which is exactly what `stage='rollout'` + `safety='lossy'` captures. +`stage='both'` means the transform persists on the module used by both rollout and +training; it does not imply numerical exactness. The separate `safety` marker records +measured cross-stage divergence. A lossy shared accelerator such as `torch_compile` is +allowed on coupled trainers when its residual stays within `clip_range`, but the validator +warns because the ratio is approximately 1 rather than bit-exact. ## Configuration @@ -41,7 +41,7 @@ application order**: ```yaml acceleration: - # Lossless, applied to BOTH rollout and the training forward, in list order. + # Persistent stage='both' accelerators, applied to rollout and training in list order. shared: - name: attention_backend # set the diffusers backend first... params: { backend: _flash_3_hub } @@ -95,18 +95,19 @@ that still sets `model.attn_backend` fails fast with a migration error. ### torch.compile train-inference consistency (coupled algorithms) -For coupled algorithms (GRPO / GRPO-Guard / DPPO) the on-policy PPO ratio on the first inner -step must be **1.0**. `torch.compile` (Inductor) threatens this because it compiles a +For coupled algorithms (GRPO / GRPO-Guard / DPPO), the first-inner-step PPO ratio must remain +**numerically on-policy around 1 and within `clip_range`**. `torch.compile` (Inductor) +threatens this because it compiles a **separate, numerically non-identical graph for grad vs no-grad mode** (Dynamo guards on `grad_mode`): rollout normally runs the transformer under `torch.no_grad()` and the training -forward under grad, so a naive compiled rollout would diverge from training (~1e-5) and bias -the ratio. - -`CompileAccelerator` handles this automatically (no user action needed): it declares -`requires_grad_rollout = True`, so the trainer runs the rollout transformer under -`torch.enable_grad()` (overriding the `@torch.no_grad()` on `inference()`) and detaches only -the per-step latent feedback (in `cast_latents`) to keep memory bounded. It also compiles the -**base** transformer under any PEFT/LoRA wrapper (not the wrapper itself). With this, the +forward under grad, so a naive compiled rollout would diverge from training and bias the ratio. + +`CompileAccelerator` handles this automatically (no user action needed): its wrapper forces +the compiled transformer to run under `torch.enable_grad()` (overriding the +`@torch.no_grad()` on `inference()`). `BaseTrainer._rollout_grad_context` sets +`adapter._rollout_detach`, and `BaseAdapter.cast_latents` detaches the per-step latent +feedback to keep memory bounded. The accelerator also compiles the **base** transformer +under any PEFT/LoRA wrapper (not the wrapper itself). With this, the on-policy ratio is driven to **≈1, well within `clip_range` (1e-4)** — but **not strictly bit-exact**: an intermittent **~1e-5** residual remains on a minority of samples/ranks (16×H20, ZeRO-2/FSDP2, SD3.5+Qwen-Image, regional/full, CFG and no-CFG). Forcing grad removes diff --git a/multinode_examples/train.yaml b/multinode_examples/train.yaml index d7cfe81a..4ae42c3b 100644 --- a/multinode_examples/train.yaml +++ b/multinode_examples/train.yaml @@ -32,10 +32,10 @@ model: resume_type: null # Options: lora, full, state. Null to auto-detect based on `finetune_type` # Optional acceleration plugins (off by default). Applied in list order. -# See guidance/acceleration.md. GRPO is a coupled trainer, so only lossless -# `shared` accelerators are allowed (lossy rollout caching would bias the PPO ratio). +# See guidance/acceleration.md. GRPO is coupled, so lossy rollout caching is unsafe; +# persistent `shared` accelerators run in both rollout and training. # acceleration: -# shared: # lossless, applied to BOTH rollout and training (in order) +# shared: # persistent stage='both'; applied in list order # - name: attention_backend # set the attention backend first... # params: { backend: _flash_3_hub } # - name: torch_compile # ...so the compiled graph captures it diff --git a/src/flow_factory/acceleration/abc.py b/src/flow_factory/acceleration/abc.py index bbbafe19..c5fcd308 100644 --- a/src/flow_factory/acceleration/abc.py +++ b/src/flow_factory/acceleration/abc.py @@ -23,9 +23,10 @@ * ``stage`` — *where/how* it is applied, and which config slot it belongs to: - ``"both"``: a persistent, one-time mutation via :meth:`setup` (e.g. - ``torch.compile``, attention backend). Because it transforms the module shared - by rollout ``inference()`` and training ``forward()``, the two stay CONSISTENT - by construction — safe for any algorithm. Belongs in the ``shared`` slot. + ``torch.compile``, attention backend) applied to the module used by both rollout + ``inference()`` and training ``forward()``. This marker describes application + scope, not numerical exactness; ``safety`` records measured cross-stage behavior. + Belongs in the ``shared`` slot. - ``"rollout"``: a per-epoch context via :meth:`rollout_context` (e.g. feature caching), torn down before the training forward. Belongs in the ``rollout`` slot. @@ -41,7 +42,7 @@ - ``stage='rollout'`` + lossy (e.g. feature caching): rollout diverges from the training forward, which cannot replicate it — only safe when the rollout log-prob never feeds the loss, i.e. **decoupled / distillation** algorithms - (validator *rejects* it on coupled; see ``constraints.md`` #7, #20a). + (validator *rejects* it on coupled; see ``constraints.md`` #7). - ``stage='both'`` + lossy (e.g. ``torch.compile``, whose grad/no-grad compiled-graph split leaves a ~1e-5 residual): applied symmetrically and stays within ``clip_range``, so it is *allowed* on any paradigm — but the validator @@ -80,11 +81,11 @@ class BaseAccelerator(ABC): # Whether this accelerator requires the Stage-3 rollout to run with autograd # ENABLED (instead of the default ``torch.no_grad()``) so the transformer - # executes the *same* graph in rollout and the training forward. Only + # uses the same grad-mode compiled path in rollout and training. Only # ``torch_compile`` needs this: Inductor compiles a separate, numerically # non-identical graph for grad vs no-grad mode (Dynamo guards on grad_mode), - # so a no-grad rollout would diverge from the grad training forward and break - # the on-policy PPO ratio==1 invariant for coupled algorithms. When set, + # so a no-grad rollout would diverge from the grad training forward and + # undermine coupled on-policy consistency. When set, # ``CompileAccelerator`` wraps the compiled transformer to force grad (returning # the grad-carrying output directly — an inner detach would let Inductor pick a # divergent inference kernel), and the trainer flags the rollout via diff --git a/src/flow_factory/acceleration/diffusers_cache.py b/src/flow_factory/acceleration/diffusers_cache.py index a1fe8f04..4f38e754 100644 --- a/src/flow_factory/acceleration/diffusers_cache.py +++ b/src/flow_factory/acceleration/diffusers_cache.py @@ -130,7 +130,8 @@ def rollout_context(self, adapter: "BaseAdapter") -> Iterator[None]: cache_hook = getattr(unwrapped, "_diffusers_hook", None) if cache_hook is not None: cache_hook._child_registries_cache = None - logger.info("DiffusersCacheAccelerator: cache enabled for '%s'.", name) + if adapter.accelerator.is_main_process: + logger.info("DiffusersCacheAccelerator: cache enabled for '%s'.", name) yield finally: for transformer in enabled: diff --git a/src/flow_factory/acceleration/registry.py b/src/flow_factory/acceleration/registry.py index 1e88fa18..d595a2e2 100644 --- a/src/flow_factory/acceleration/registry.py +++ b/src/flow_factory/acceleration/registry.py @@ -31,7 +31,7 @@ # Accelerator Registry Storage _ACCELERATOR_REGISTRY: Dict[str, str] = { - # Lossless (safe for any algorithm, applied to the shared transformer). + # Persistent stage='both' accelerators applied to rollout and training. # `attention_backend` is the single code path for attention-backend selection; # add it as a `shared` entry (before `torch_compile`) in the acceleration block. 'attention_backend': 'flow_factory.acceleration.attention_backend.AttentionBackendAccelerator', diff --git a/src/flow_factory/acceleration/torch_compile.py b/src/flow_factory/acceleration/torch_compile.py index a9cae4bc..0a7d9add 100644 --- a/src/flow_factory/acceleration/torch_compile.py +++ b/src/flow_factory/acceleration/torch_compile.py @@ -13,7 +13,7 @@ # limitations under the License. # src/flow_factory/acceleration/torch_compile.py -"""Lossless ``torch.compile`` accelerator for the shared transformer(s).""" +"""Apply ``torch.compile`` persistently to the rollout/training transformer(s).""" import functools from typing import TYPE_CHECKING, Any, Dict @@ -80,8 +80,8 @@ class CompileAccelerator(BaseAccelerator): stage = "both" # Inductor compiles a separate graph for grad vs no-grad mode whose fused # kernels are NOT bit-identical. To keep rollout (Stage 3) and the training - # forward (Stage 6) on the exact same graph — required for the coupled - # on-policy ratio==1 invariant — the rollout must run with grad enabled. + # forward (Stage 6) on the same grad-mode compiled path — preserving a + # numerically on-policy ratio within `clip_range` — rollout must run with grad enabled. requires_grad_rollout = True def setup(self, adapter: "BaseAdapter") -> None: @@ -113,7 +113,8 @@ def setup(self, adapter: "BaseAdapter") -> None: # `_repeated_blocks` so regional compile is unavailable under LoRA. # * the base transformer keeps its LoRA submodules inside the compiled # graph (they still train), exposes `_repeated_blocks`, and the - # grad-force wrap makes rollout/train bit-identical. + # grad-force wrap keeps rollout/training on the same grad-mode compiled + # path (ratio ≈ 1 within `clip_range`, but not bit-exact). inner = self._peel_peft(adapter._unwrap(module)) if mode == "regional": # `compile_repeated_blocks` exists on every diffusers ModelMixin but @@ -132,15 +133,14 @@ def setup(self, adapter: "BaseAdapter") -> None: inner.compile(**compile_kwargs) # Force every forward of the compiled transformer to run under # torch.enable_grad() (overriding the @torch.no_grad() on - # adapter.inference()), and detach the output during rollout. This keeps - # rollout (Stage 3) and the training forward (Stage 6) on the SAME - # compiled graph — Inductor emits numerically different kernels for the - # grad vs no-grad graph, so a no-grad rollout would break the coupled - # on-policy ratio==1 invariant. Detaching during rollout (gated by the - # trainer's `adapter._rollout_detach` flag) stops autograd from chaining - # across denoising steps, so rollout memory stays bounded. + # adapter.inference()). This keeps rollout (Stage 3) and the training + # forward (Stage 6) on the same grad-mode compiled path — Inductor emits + # numerically different kernels for the grad vs no-grad graph. During + # rollout, `BaseAdapter.cast_latents` detaches the per-step latent feedback + # (gated by `adapter._rollout_detach`) so memory stays bounded. self._wrap_forward_grad_consistent(adapter, inner) - logger.info("CompileAccelerator: compiled '%s' (mode=%s).", name, mode) + if adapter.accelerator.is_main_process: + logger.info("CompileAccelerator: compiled '%s' (mode=%s).", name, mode) @staticmethod def _peel_peft(module): @@ -165,10 +165,10 @@ def _wrap_forward_grad_consistent(adapter: "BaseAdapter", module) -> None: Wraps the module's call entry point(s) so every forward runs under ``torch.enable_grad()`` — overriding the ``@torch.no_grad()`` on ``adapter.inference()``. This pins rollout (Stage 3) and the training forward - (Stage 6) to the *same* Inductor graph: Inductor compiles a separate, + (Stage 6) to the same grad-mode compiled path: Inductor compiles a separate, numerically non-identical graph for grad vs no-grad mode (Dynamo guards on ``grad_mode``), so a no-grad rollout would diverge from the grad training - forward and break the coupled on-policy PPO ratio==1 invariant. + forward and undermine coupled on-policy consistency. Crucially the output is **NOT detached here**: detaching inside the wrapper lets Inductor see the result is unused-for-grad and pick a different @@ -176,8 +176,8 @@ def _wrap_forward_grad_consistent(adapter: "BaseAdapter", module) -> None: rollout's per-step latent feedback is detached in :meth:`BaseAdapter.cast_latents` (gated by ``adapter._rollout_detach``), which breaks the autograd graph chain across denoising steps so rollout memory stays - bounded while every transformer call still executes the identical grad-mode - graph (near-exact noise_pred → on-policy ratio ≈ 1). + bounded while every transformer call still uses the same grad-mode compiled + path (near-exact noise_pred → on-policy ratio ≈ 1 within ``clip_range``). Known residual (NOT strictly bit-exact): forcing grad removes the *dominant* divergence (the grad-vs-no-grad graph split), but it does not make the rollout diff --git a/src/flow_factory/acceleration/validator.py b/src/flow_factory/acceleration/validator.py index 56218aa0..6cc23003 100644 --- a/src/flow_factory/acceleration/validator.py +++ b/src/flow_factory/acceleration/validator.py @@ -15,7 +15,7 @@ # src/flow_factory/acceleration/validator.py """Paradigm-gated safety validation for accelerators (fail-fast, ``constraints.md`` #26). -The correctness contract (``constraints.md`` #7 + #20a) hinges on **symmetric +The correctness contract (``constraints.md`` #7) hinges on **symmetric application**, encoded by ``stage``, not on numerical bit-exactness: * ``stage='both'`` accelerators mutate the transformer persistently, so the same diff --git a/src/flow_factory/hparams/acceleration_args.py b/src/flow_factory/hparams/acceleration_args.py index 85031eff..2748c0b6 100644 --- a/src/flow_factory/hparams/acceleration_args.py +++ b/src/flow_factory/hparams/acceleration_args.py @@ -88,8 +88,8 @@ class AccelerationArguments(ArgABC): empty by default, so existing configs are unaffected). List order is the application order: - * ``shared`` — lossless accelerators applied to BOTH rollout and the training - forward, in order, via each accelerator's ``setup()`` (e.g. + * ``shared`` — persistent ``stage='both'`` accelerators applied to rollout and + the training forward, in order, via each accelerator's ``setup()`` (e.g. ``attention_backend`` then ``torch_compile`` — backend first so the compiled graph captures it). * ``rollout`` — accelerators applied only during Stage-3 rollout, nested in diff --git a/src/flow_factory/models/abc.py b/src/flow_factory/models/abc.py index 151102b3..ab0bf685 100644 --- a/src/flow_factory/models/abc.py +++ b/src/flow_factory/models/abc.py @@ -234,7 +234,9 @@ def cast_latents(self, latents: torch.Tensor, default_dtype: Optional[torch.dtyp per-step latent feedback would otherwise chain an autograd graph across the whole denoising loop. Detaching here — the single feedback chokepoint every adapter routes through — breaks that chain so rollout memory stays bounded while - each transformer call still executes the identical grad-mode graph. + each transformer call uses the same grad-mode compiled path. Distinct Inductor + invocations can still leave an intermittent ~1e-5 on-policy ratio residual within + ``clip_range``, so this path is not bit-exact. """ if self._rollout_detach and latents.requires_grad: latents = latents.detach() diff --git a/src/flow_factory/trainers/abc.py b/src/flow_factory/trainers/abc.py index ede88054..e6db0300 100644 --- a/src/flow_factory/trainers/abc.py +++ b/src/flow_factory/trainers/abc.py @@ -93,7 +93,7 @@ def __init__( self._initialization() self.adapter.post_init() - # Apply the shared (lossless) accelerator last: after prepare, state-resume, + # Apply persistent stage='both' accelerators last: after prepare, state-resume, # EMA, and reference-parameter setup, so e.g. torch.compile wraps the final # weights and keeps state_dict keys / parameter identity stable. self._apply_shared_acceleration() @@ -382,8 +382,8 @@ def _initialization(self): # Load inference modules, excluding all bundle members (already prepared). self._load_inference_components(bundle_names) - # Build + validate acceleration plugins. The shared (lossless) accelerator - # is *applied* later via _apply_shared_acceleration(), after post_init() + # Build + validate acceleration plugins. Persistent stage='both' accelerators + # are *applied* later via _apply_shared_acceleration(), after post_init() # finishes any state-resume / EMA / reference setup. self._init_acceleration() @@ -396,9 +396,9 @@ def _init_acceleration(self): Two independent slots, each an **ordered list** (both empty by default). List order is the application order: - * ``shared`` — lossless accelerators (e.g. ``attention_backend`` then - ``torch_compile``) that affect both rollout and the training forward. Only - built/validated here; they are *applied* later by + * ``shared`` — persistent ``stage='both'`` accelerators (e.g. + ``attention_backend`` then ``torch_compile``) applied to both rollout and + the training forward. Only built/validated here; they are *applied* later by :meth:`_apply_shared_acceleration` (after ``post_init`` finishes state-resume / EMA / reference setup), so they transform the final weights. * ``rollout`` — accelerators applied per-epoch in :meth:`generate_samples` @@ -435,7 +435,7 @@ def _init_acceleration(self): ) def _apply_shared_acceleration(self) -> None: - """Apply the shared (lossless) accelerators to the adapter, in config order. + """Apply persistent ``stage='both'`` accelerators in config order. Called from ``__init__`` AFTER ``adapter.post_init()`` so transforms wrap the final weights — i.e. after ``accelerator.prepare``, any ``state`` checkpoint @@ -480,16 +480,15 @@ def _rollout_grad_context(self) -> Iterator[None]: ``torch.compile`` (Inductor) emits a *separate, numerically non-identical* graph for grad vs no-grad mode (Dynamo guards on ``grad_mode``), so a no-grad rollout would diverge from the grad-mode training forward and break the - coupled on-policy PPO ratio==1 invariant. The compiled transformer is wrapped + coupled on-policy consistency. The compiled transformer is wrapped (see :class:`CompileAccelerator`) so its forward always executes under ``torch.enable_grad()`` — overriding the ``@torch.no_grad()`` on ``adapter.inference()``. The output is NOT detached inside that wrapper (an inner detach lets Inductor pick a divergent inference kernel); instead the per-step latent feedback is detached in ``BaseAdapter.cast_latents`` (gated by - this flag), so rollout and training share the identical grad-mode graph - (bit-exact on-policy ratio, with or without CFG — see - ``guidance/acceleration.md``) while the autograd graph never chains across - denoising steps (memory stays bounded). + this flag), so rollout and training use the same grad-mode compiled path and + remain numerically on-policy (ratio ≈ 1 within ``clip_range``, but not + bit-exact) while the autograd graph never chains across denoising steps. Reused by the eval loop (:meth:`evaluate`). Eval does not compute the PPO ratio, so graph consistency is irrelevant there — but the compiled @@ -505,9 +504,10 @@ def _rollout_grad_context(self) -> Iterator[None]: with torch.no_grad(): yield return - # Compile active: the compiled-transformer wrapper forces grad locally and - # detaches its output while this flag is set, so rollout matches the training - # graph bit-for-bit without chaining autograd across steps. + # Compile active: the compiled-transformer wrapper forces grad locally, while + # `BaseAdapter.cast_latents` detaches per-step latent feedback under this flag. + # Rollout and training use the same grad-mode compiled path without chaining + # autograd across denoising steps. prev = self.adapter._rollout_detach self.adapter._rollout_detach = True try: @@ -907,7 +907,7 @@ def generate_samples( # so its state never leaks into the Stage-6 training forward. # Grad context is normally no_grad, but flips to enable_grad when a # compile accelerator is active (see _rollout_grad_context) so rollout and - # training share the same compiled graph (bit-exact on-policy ratio). + # training use the same grad-mode compiled path and remain numerically on-policy. with self._rollout_acceleration(), self._rollout_grad_context(), self.autocast(): for _ in tqdm( range(self.training_args.num_batches_per_epoch),