diff --git a/.agents/knowledge/architecture.md b/.agents/knowledge/architecture.md index 1507f835..7dd103a6 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,15 @@ 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 | +|-----|-------|--------|-------|-------| +| `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` (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. + --- ## Extension Points @@ -154,6 +163,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/.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/ 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/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..5ad3d85d 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: # 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 +# params: { mode: regional } # regional (compile_repeated_blocks) | full +# rollout: # rollout-only (Stage 3); lossy -> decoupled/distillation only +# - name: diffusers_cache # Options: diffusers_cache +# 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..a7fd2142 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: # 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 +# params: { mode: regional } # regional (compile_repeated_blocks) | full +# rollout: # rollout-only (Stage 3); lossy -> decoupled/distillation only +# - name: diffusers_cache # Options: diffusers_cache +# 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..37b06c75 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: # 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 +# params: { mode: regional } # regional (compile_repeated_blocks) | full +# rollout: # rollout-only (Stage 3); lossy -> decoupled/distillation only +# - name: diffusers_cache # Options: diffusers_cache +# 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..d242cb03 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: # 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 +# params: { mode: regional } # regional (compile_repeated_blocks) | full +# rollout: # rollout-only (Stage 3); lossy -> decoupled/distillation only +# - name: diffusers_cache # Options: diffusers_cache +# 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..0f12b5cf 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 coupled, so lossy rollout caching is unsafe; +# persistent `shared` accelerators run in both rollout and training. +# acceleration: +# 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 +# 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..33ed422e 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 coupled, so lossy rollout caching is unsafe; +# persistent `shared` accelerators run in both rollout and training. +# acceleration: +# 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 +# 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..6611c1c4 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 coupled, so lossy rollout caching is unsafe; +# persistent `shared` accelerators run in both rollout and training. +# acceleration: +# 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 +# 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..03d43865 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 coupled, so lossy rollout caching is unsafe; +# persistent `shared` accelerators run in both rollout and training. +# acceleration: +# 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 +# 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..d17f7df3 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 coupled, so lossy rollout caching is unsafe; +# persistent `shared` accelerators run in both rollout and training. +# acceleration: +# 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 +# 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..ddb7daa7 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 coupled, so lossy rollout caching is unsafe; +# persistent `shared` accelerators run in both rollout and training. +# acceleration: +# 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 +# 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..ad8b1cba 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 coupled, so lossy rollout caching is unsafe; +# persistent `shared` accelerators run in both rollout and training. +# acceleration: +# 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 +# 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..0a60fd40 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 coupled, so lossy rollout caching is unsafe; +# persistent `shared` accelerators run in both rollout and training. +# acceleration: +# 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 +# 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..79c356f2 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 coupled, so lossy rollout caching is unsafe; +# persistent `shared` accelerators run in both rollout and training. +# acceleration: +# 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 +# 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..18d4ad46 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: # 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 +# params: { mode: regional } # regional (compile_repeated_blocks) | full +# rollout: # rollout-only (Stage 3); lossy -> decoupled/distillation only +# - name: diffusers_cache # Options: diffusers_cache +# 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 bfffb2c1..9f795280 100644 --- a/examples/nft/lora/qwen_image/rational_rewards_t2i.yaml +++ b/examples/nft/lora/qwen_image/rational_rewards_t2i.yaml @@ -82,6 +82,19 @@ train: scheduler: dynamics_type: "ODE" +# 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: # 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 +# params: { mode: regional } # regional (compile_repeated_blocks) | full +# rollout: # rollout-only (Stage 3); lossy -> decoupled/distillation only +# - name: diffusers_cache # Options: diffusers_cache +# params: { policy: first_block, threshold: 0.08 } + eval: resolution: 512 condition_image_size: 512 diff --git a/examples/nft/lora/sd3_5/default.yaml b/examples/nft/lora/sd3_5/default.yaml index bd4fc8bd..926e6219 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: # 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 +# params: { mode: regional } # regional (compile_repeated_blocks) | full +# rollout: # rollout-only (Stage 3); lossy -> decoupled/distillation only +# - name: diffusers_cache # Options: diffusers_cache +# 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..c6209aaf 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: # 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 +# params: { mode: regional } # regional (compile_repeated_blocks) | full +# rollout: # rollout-only (Stage 3); lossy -> decoupled/distillation only +# - name: diffusers_cache # Options: diffusers_cache +# 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..04da0c14 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: # 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 +# params: { mode: regional } # regional (compile_repeated_blocks) | full +# rollout: # rollout-only (Stage 3); lossy -> decoupled/distillation only +# - name: diffusers_cache # Options: diffusers_cache +# 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..c04c744c 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: # 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 +# params: { mode: regional } # regional (compile_repeated_blocks) | full +# rollout: # rollout-only (Stage 3); lossy -> decoupled/distillation only +# - name: diffusers_cache # Options: diffusers_cache +# 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..5ab71f2b 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 coupled, so lossy rollout caching is unsafe; +# persistent `shared` accelerators run in both rollout and training. +# acceleration: +# 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 +# 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 new file mode 100644 index 00000000..c2d86418 --- /dev/null +++ b/guidance/acceleration.md @@ -0,0 +1,149 @@ +# 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) + +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): + +| marker | values | meaning | +|--------|--------|---------| +| `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). 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. + +`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 + +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: + # 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 } + - 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 | + params: { policy: first_block, threshold: 0.08 } +``` + +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. 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`). The single lossy rollout backend. | + +### Attention 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 +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`, 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). + +### torch.compile train-inference consistency (coupled algorithms) + +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 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 +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 + 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 +`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` is model-agnostic and applies 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). diff --git a/multinode_examples/train.yaml b/multinode_examples/train.yaml index 538a994d..4ae42c3b 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 coupled, so lossy rollout caching is unsafe; +# persistent `shared` accelerators run in both rollout and training. +# acceleration: +# 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 +# params: { mode: regional } # regional (compile_repeated_blocks) | full # Training Configuration train: diff --git a/pyproject.toml b/pyproject.toml index f4fc067f..e4e4c8b5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -64,6 +64,9 @@ quantization = [ "bitsandbytes>=0.45.3", ] +# 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 = [ "flash-attn>=2.5.8", diff --git a/src/flow_factory/acceleration/__init__.py b/src/flow_factory/acceleration/__init__.py new file mode 100644 index 00000000..62e22b00 --- /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 an accelerator +module (and any heavy import it needs) is only loaded 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..c5fcd308 --- /dev/null +++ b/src/flow_factory/acceleration/abc.py @@ -0,0 +1,152 @@ +# 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 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) 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. + +* ``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). + - ``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. +""" + +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"]] + + # Whether this accelerator requires the Stage-3 rollout to run with autograd + # ENABLED (instead of the default ``torch.no_grad()``) so the transformer + # 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 + # 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 + # ``_rollout_grad_context`` so the latent feedback is detached in ``cast_latents``. + # 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: + 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..6996b98b --- /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 +"""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 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 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 +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 the entry's ``params``): + backend: Backend name forwarded to ``transformer.set_attention_backend`` + (e.g. ``native`` / ``flash`` / ``_flash_3`` / ``_flash_3_hub`` / + ``sage`` / ``xformers``). Required. + + 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") + 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: + 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: + 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 new file mode 100644 index 00000000..4f38e754 --- /dev/null +++ b/src/flow_factory/acceleration/diffusers_cache.py @@ -0,0 +1,138 @@ +# 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 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 + 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) + # 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 + if adapter.accelerator.is_main_process: + 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 new file mode 100644 index 00000000..d595a2e2 --- /dev/null +++ b/src/flow_factory/acceleration/registry.py @@ -0,0 +1,125 @@ +# 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] = { + # 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', + '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', +} +_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. + + 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 _REMOVED_ACCELERATORS: + raise ValueError(_REMOVED_ACCELERATORS[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..0a7d9add --- /dev/null +++ b/src/flow_factory/acceleration/torch_compile.py @@ -0,0 +1,224 @@ +# 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 +"""Apply ``torch.compile`` persistently to the rollout/training transformer(s).""" + +import functools +from typing import TYPE_CHECKING, Any, Dict + +import torch + +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. + + 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 + ``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 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. + 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}``). + """ + + # `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 + # 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: + 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) + # 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 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 + # only works when the model declares `_repeated_blocks`; check the + # unwrapped module so the error is actionable. + 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." + ) + 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. + 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()). 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) + if adapter.accelerator.is_main_process: + 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): + return get_base() + 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 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 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 + (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 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 + 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: + * ``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/acceleration/validator.py b/src/flow_factory/acceleration/validator.py new file mode 100644 index 00000000..6cc23003 --- /dev/null +++ b/src/flow_factory/acceleration/validator.py @@ -0,0 +1,151 @@ +# 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) hinges on **symmetric +application**, encoded by ``stage``, not on numerical bit-exactness: + +* ``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 + 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 + +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 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.stage != "both": + 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 the `acceleration.rollout` list." + ) + # 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". + 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 in the " + "`acceleration.shared` list 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( + 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 " + "stage='both' 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..f7f29d2c 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, AccelerationSpec from .dataset_args import DatasetArguments, DatasetTrainSpec, DatasetEvalSpec from .log_args import LogArguments @@ -55,6 +56,8 @@ "get_training_args_class", "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 new file mode 100644 index 00000000..2748c0b6 --- /dev/null +++ b/src/flow_factory/hparams/acceleration_args.py @@ -0,0 +1,143 @@ +# 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, 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, each an **ordered list** of accelerator entries (both + empty by default, so existing configs are unaffected). List order is the + application order: + + * ``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 + 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: + - 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: 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 { + "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) + + 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/hparams/model_args.py b/src/flow_factory/hparams/model_args.py index 1f4977b7..4b40f86f 100644 --- a/src/flow_factory/hparams/model_args.py +++ b/src/flow_factory/hparams/model_args.py @@ -120,17 +120,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.trainable_parameters_dtype, str): self.trainable_parameters_dtype = dtype_map[self.trainable_parameters_dtype] if isinstance(self.frozen_parameters_dtype, str): diff --git a/src/flow_factory/models/abc.py b/src/flow_factory/models/abc.py index 8332fe78..ab0bf685 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) @@ -194,8 +198,11 @@ 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, 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: @@ -219,7 +226,20 @@ 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 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() target = self.latent_storage_dtype or default_dtype if target is None or latents.dtype == target: return latents @@ -848,25 +868,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 fe8602dc..e6db0300 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 ExitStack, contextmanager +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, @@ -84,6 +93,10 @@ def __init__( self._initialization() self.adapter.post_init() + # 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() self._init_logging_backend() self._patch_deepspeed_autocast(accelerator) @@ -369,9 +382,139 @@ def _initialization(self): # Load inference modules, excluding all bundle members (already prepared). self._load_inference_components(bundle_names) + # 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() + # Initialize reward model self._init_reward_model() + def _init_acceleration(self): + """Build and validate acceleration plugins from ``config.acceleration_args``. + + Two independent slots, each an **ordered list** (both empty by default). + List order is the application order: + + * ``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` + 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 = self.config.acceleration_args + self.shared_accelerators: List[BaseAccelerator] = [] + self.rollout_accelerators: List[BaseAccelerator] = [] + + trainer_name = type(self).__name__ + paradigm = type(self).paradigm + + 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_accelerators.append(accelerator) + + 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_accelerators.append(accelerator) + if self.accelerator.is_main_process: + logger.info( + "Acceleration: rollout accelerator '%s' (safety=%s) enabled.", + spec.name, + accelerator.safety, + ) + + def _apply_shared_acceleration(self) -> None: + """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 + 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. + """ + for accelerator in self.shared_accelerators: + accelerator.setup(self.adapter) + if self.accelerator.is_main_process: + logger.info( + "Acceleration: shared accelerator '%s' (safety=%s) applied to adapter.", + type(accelerator).__name__, + accelerator.safety, + ) + + @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 + + @contextmanager + def _rollout_grad_context(self) -> Iterator[None]: + """Grad context for an ``adapter.inference()`` loop (Stage-3 rollout and eval). + + 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 + 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 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 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 + 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: + with torch.no_grad(): + yield + return + # 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: + yield + finally: + self.adapter._rollout_detach = prev + def _synchronize_frozen_components(self): if self.accelerator.num_processes <= 1: return @@ -760,7 +903,12 @@ 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. + # 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 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), desc=f'Epoch {self.epoch} Sampling', @@ -823,7 +971,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: 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