Add AnyFlow algorithm (any-step video diffusion via flow maps)#25
Add AnyFlow algorithm (any-step video diffusion via flow maps)#25Enderfga wants to merge 12 commits into
Conversation
AnyFlow is an any-step video diffusion method that trains a single model
u_theta(x_t, t, r) to predict the average velocity from t back to r, so
the same checkpoint supports arbitrary inference NFE.
Training has two stages, switched via config.loss_config.training_stage:
* pretrain — flow-map prediction with a central-difference target
target = (eps - x0) - (t - r) * dF/dt
with dF/dt estimated by central differences at (t ± delta).
Per-batch sampling assigns r=t to a `diffusion_ratio`
fraction (pure flow matching) and r=0 to a
`consistency_ratio` fraction (consistency to clean data).
* onpolicy — distribution-matching distillation with r=0 conditioning
on top of the pretrained flow-map weights. Inherits DMD2's
alternating fake_score / teacher / discriminator updates.
The backbone requirement (a secondary timestep r) is already satisfied by
the Wan transformer with r_timestep=True, which MeanFlow also exercises;
no Wan-side changes are needed.
New files:
fastgen/methods/distribution_matching/anyflow.py
fastgen/methods/distribution_matching/anyflow_scheduler.py
fastgen/configs/methods/config_anyflow.py
fastgen/configs/experiments/WanT2V/config_anyflow.py
tests/test_anyflowmodel.py
Modified:
fastgen/methods/__init__.py (+1 import)
fastgen/methods/distribution_matching/README.md (+1 algorithm entry)
The multi-step rollout-with-gradient training (matching
self_forcing.py's rollout_with_gradient) is intentionally left for a
follow-up PR — the on-policy stage here uses single-step student
generation.
Signed-off-by: Enderfga <qq2639135175@gmail.com>
|
Thanks a lot for the PR! Did you test the implementation and, if yes, do you have example videos or could you share the wandb run? |
AnyFlow's released HF checkpoints store the r-pathway as
``condition_embedder.delta_embedder.*`` inside the shared
``WanTwoTimeTextImageEmbedding`` module and use ONE shared ``time_proj``
for both t and (t, r). Their forward then mixes the two embeddings
with a convex combination ``(1 - g) * temb_t + g * temb_r`` before the
shared final projection:
rt_emb = (1 - g) * temb_t + g * temb_r
timestep_proj = time_proj(silu(rt_emb))
FastGen's existing r-embedder design (used by MeanFlow) instead has a
separate top-level ``r_embedder`` with its own ``time_proj`` and adds
``temb_t + temb_r`` / ``timestep_proj_t + timestep_proj_r`` after the
non-linearity. The two layouts are not functionally equivalent because
``silu`` is non-linear.
Two changes:
* ``Wan.__init__``: add ``r_embedder_fusion: str = "additive"`` (default
preserves MeanFlow's behaviour) and ``r_embedder_gate_value: float =
0.25``. When ``r_embedder_fusion="gated"``, ``classify_forward_prepare``
computes the convex-mix variant and uses ``r_embedder.time_proj``
(which ``init_embedder`` already deep-copies from
``condition_embedder.time_proj``) for the shared final projection.
* ``fastgen/methods/distribution_matching/anyflow.py``: add
``remap_anyflow_keys`` helper that rewrites AnyFlow's
``condition_embedder.delta_embedder.linear_{1,2}.*`` to FastGen's
``r_embedder.time_embedder.linear_{1,2}.*`` and duplicates
``condition_embedder.time_proj.*`` into ``r_embedder.time_proj.*``
so the two projections start identical. The function is a no-op when
no AnyFlow-format keys are present.
Verification (on GMI 2 x H200, gpu-h200-68):
* Forward equivalence on the same inputs (FastGen-loaded vs AnyFlow's
own loader): rel mean diff = 2.8% in bf16 (forward noise floor).
* Training-step loss equivalence (AnyFlow ``train_bidirection`` math
reproduced inline on both code paths, same seed): AnyFlow loss
0.381619 vs FastGen loss 0.397162, rel diff = 4.07%.
* 4-step Euler-flow inference end-to-end (text encoder + FastGen Wan +
VAE decode) produces a finite 81-frame 480x832 video matching the
AnyFlow paper's any-step inference pattern.
Signed-off-by: Enderfga <qq2639135175@gmail.com>
|
Thanks for the review! Verification is complete on both 1.3B and 14B — inference and training-step accuracy agree to bf16 noise on the published AnyFlow checkpoints. Inference correctnessLoaded On identical inputs the FastGen-loaded model agrees with AnyFlow's own loader to within bf16 forward noise (rel mean diff Training-step equivalenceInline replica of AnyFlow's
A stub-network compare of the central-difference target tensor (so the math is isolated from network weights) gives max abs diff Sample videosSame prompt +
1p3b_fastgen_nfe4.mp4
14b_fastgen_nfe4.mp4
14b_fastgen_nfe50.mp4What this PR changes
Both files are additive. Existing methods (MeanFlow, DMD2, CMs, …) keep their previous forward bit-identical. Re-pushed as commit 03ed6cd on top of the original ef13247 ("Add AnyFlow algorithm"). |
Enderfga
force-pushed
the
feature/anyflow-algorithm
branch
from
03ed6cd to
99c0415
Compare
Replace the single-step student forward in AnyFlow's on-policy stage with a multi-step Euler-flow rollout that enables gradients at one randomly-chosen step. This matches AnyFlow's ``WanAnyFlowPipeline.training_rollout`` (the published on-policy training mode in the reference repo) and gives the DMD generator update a usable gradient through a full denoising window instead of a single forward. Changes: * ``AnyFlowModel._rollout_with_gradient(batch_size, dtype, condition)``: start from pure noise at ``ns.max_t``, iterate ``student_sample_steps`` Euler-flow updates with ``r = t_next`` (mean-velocity, matching the reference default), and toggle ``torch.set_grad_enabled`` at the randomly-selected step. ``grad_step`` is broadcast from rank 0 in distributed runs so all ranks share the same gradient window. The step schedule honours ``sample_t_cfg.t_list`` when set, otherwise falls back to ``noise_scheduler.get_t_list``. * ``_onpolicy_student_update_step`` and ``_onpolicy_fake_score_discriminator_update_step``: source ``gen_data`` from the rollout instead of a single ``self.net(input_student, ...)`` forward. ``input_student`` / ``t_student`` from ``_generate_noise_and_time`` become unused for on-policy and are discarded explicitly. * ``_get_outputs``: when on-policy, always take the multi-step generator callable path (no longer special-cases ``student_sample_steps == 1`` for the validation hook, since the rollout output is always usable). * ``tests/test_anyflowmodel.py``: bump ``student_sample_steps`` to 2 in the on-policy fixtures and add ``test_onpolicy_rollout_propagates_gradient`` which asserts the rollout output keeps a usable autograd graph and that ``backward()`` reaches the student weights. All 13 unit tests pass (`make pytest tests/test_anyflowmodel.py`). Signed-off-by: Enderfga <qq2639135175@gmail.com>
|
Follow-up commit New
The rollout output replaces the single Unit tests bumped to 13. The new |
|
Hi @juliusberner — gentle ping. 🙏 The verification you asked for is in the follow-up comment (forward parity + training-step parity + sample videos on the published 1.3B and 14B checkpoints), and commit Happy to address any further feedback whenever you have a slot — thanks again for the early review! |
Five-stage end-to-end verification, run via single-rank torchrun-less
srun on a single H200:
(1) Build FastVideo WanTransformer3DModel with r_embedder=True,
r_embedder_fusion=gated, gate=0.25.
(2) Load nvidia/AnyFlow-Wan2.1-T2V-1.3B-Diffusers safetensors and
translate keys via WanVideoArchConfig.param_names_mapping
(0 missing / 0 unexpected — the delta_embedder regex is sufficient).
(3) Build AnyFlow's reference loader (FAR_Wan_Transformer3DModel).
(4) Forward parity on identical inputs — bf16 noise.
(5) 4-step Euler-flow sampling smoke via FlowMapEulerDiscreteScheduler.
(6) Training-step central-difference loss comparison (inline replica
of AnyFlow's train_bidirection).
Measured on Wan2.1-T2V-1.3B + nvidia/AnyFlow checkpoint:
forward rel mean diff : 2.55%
forward max abs diff : 7.81e-2
training loss diff : 1.33% (AnyFlow 0.381619 vs FastVideo 0.386694)
Both within bf16 kernel noise. Compare to the FastGen port at
NVlabs/FastGen#25 which reported 2.8% forward + 4.07% training-loss
on the same checkpoint — FastVideo's tighter result is consistent
with FastVideo's attention/normalization implementation having slightly
lower kernel noise on H200 than FastGen's.
|
Hi @Enderfga, Thanks a lot for all the evaluations and videos, this is in a great shape! We'll take a closer look soon, but I wanted to ask two questions first:
|
|
@Enderfga Thanks a lot for the PR and its follow-up! |
Addresses two pieces of PR NVlabs#25 reviewer feedback: (1) Code sharing with MeanFlow. The previous commit added AnyFlow's gated t/r mixing as an inline branch inside ``classify_forward_prepare``, which made it visually hard to tell which lines were MeanFlow's additive path and which were AnyFlow-specific. This commit factors both fusion modes into a single ``_fuse_r_embedding`` method bound on the transformer (parallel pattern to ``classify_forward_prepare`` and friends). Both paths still share ``r_embedder.time_embedder`` / ``time_proj`` / ``act_fn`` modules — the helper just makes that sharing explicit and shrinks the call site to three lines. Forward semantics are bit-identical to the previous commit for both additive (MeanFlow) and gated (AnyFlow) modes across all three ``encoder_depth`` cases. (2) Ship a paper-aligned on-policy stage config. Previously the only documented way to run Stage 3 was an inline tweak in the pretrain config docstring. New file ``fastgen/configs/experiments/WanT2V/config_anyflow_onpolicy.py`` inherits the pretrain config and flips the loss into "onpolicy" with the paper's Stage 3 hyperparameters (lr=2e-6, 1200 iter, GAN on at the DMD2-default 0.03, ``student_update_freq=5``). The docstring notes that the AnyFlow paper's rank-256 LoRA variant is not reproduced here because FastGen does not ship a PEFT/LoRA training path; this config is a full-rank fine-tune of a Stage 2 pretrain checkpoint. The AnyFlow method README is updated to (a) document the new ``r_embedder_fusion="gated"`` requirement when loading the released AnyFlow HF checkpoints, (b) replace the stale "multi-step rollout deferred to a follow-up" note (already landed in ab1174d) with an explicit acknowledgement that end-to-end convergence-scale validation on the paper's training corpus is deferred to a follow-up, and (c) cross-reference both pretrain and on-policy configs. Tests: all 13 AnyFlow + 3 MeanFlow unit tests pass. Signed-off-by: Enderfga <qq2639135175@gmail.com>
|
Thanks @juliusberner and @cxlcl — pushed commit (1) MeanFlow code sharing. Extracted (2) Convergence-scale validation. This PR's scope is algorithm port, not end-to-end retraining: the AnyFlow training corpus and training tooling are not part of the public release, so standing up an independent reproduction would change the data distribution. Correctness evidence is therefore algorithmic, not convergence-based:
The README now states this scope explicitly. Convergence-scale validation on the paper's training corpus is left as a follow-up. Please advise whether that's acceptable for merge or whether you'd prefer to block on end-to-end numbers. @cxlcl — re: config tuning. |
juliusberner
left a comment
juliusberner
left a comment
There was a problem hiding this comment.
Thanks again for the PR, I did a code review and added several comments.
There was a problem hiding this comment.
Do we need FlowMapDiscreteScheduler?
- Our FastGen noise scheduler supports sampling from a shifted distribution and we can just use the
sample_tmethod directly. - We can define the weightings in
anyflow.pydirectly as a function oftinstead of precomputing and nearest-neighbor search?
There was a problem hiding this comment.
Removed. (t, r) sampling goes through noise_scheduler.sample_t with the shifted distribution now (two draws + max/min — the shift map is monotonic so it commutes), and the weighting is evaluated directly as a function of t in mean_flow.py; the reference's normalization constant over the shifted grid is computed once and cached.
| import fastgen.utils.logging_utils as logger | ||
|
|
||
|
|
||
| def remap_anyflow_keys(state_dict: dict) -> dict: |
There was a problem hiding this comment.
This should live in networks/Wan/network.py, since it's Wan-specific.
There was a problem hiding this comment.
Moved into Wan/network.py and now applied inside Wan.load_state_dict, no-op for everything else. One caveat I hit and documented: loading the HF folder via diffusers' from_pretrained silently drops the delta_embedder weights, so the state dict has to go through load_state_dict.
| self, data: Dict[str, Any], iteration: int | ||
| ) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor | Callable]]: | ||
| if self.loss_config.training_stage == "pretrain": | ||
| return self._pretrain_single_train_step(data, iteration) |
There was a problem hiding this comment.
Can't we just add a little bit of functionality (e.g., loss weighting) to the mean_flow.py method that already exists in this repo and use that (instead of adding this functionality to this class)?
There was a problem hiding this comment.
Done — pretrain is now just MeanFlowModel with a config. I added the weighting (plus the guidance fusion and loss rebalancing the reference applies, see the top-level comment) as opt-in fields defaulting to off, so existing MeanFlow behavior is unchanged — the tests check the default path stays identical.
| ) | ||
| x_t_sg = self.net.noise_scheduler.forward_process(gen_data, eps, t) | ||
|
|
||
| from fastgen.methods.common_loss import ( |
There was a problem hiding this comment.
Can we move imports to the top?
| teacher_x0, fake_feat = self.teacher( | ||
| perturbed_data, | ||
| t, | ||
| r=r_zero_t, |
There was a problem hiding this comment.
Why does the teacher need an r argument?
There was a problem hiding this comment.
Good catch — this made me re-read the reference, and the r=0 I had was wrong: trainer_wan_anyflow_onpolicy.py queries both real and fake score with r_timestep=timesteps, i.e. r=t. Since a gated-fusion network always consumes the r pathway, I made r=t the network-level default when r isn't passed, which also means the DMD2 call sites stay untouched.
|
|
||
| return x | ||
|
|
||
| def _onpolicy_student_update_step( |
There was a problem hiding this comment.
Is it possible to just override the _generate_noise_and_time and gen_data_from_net methods instead of the _student_update_step method of DMD2?
There was a problem hiding this comment.
Mostly yes — _generate_noise_and_time and gen_data_from_net carry the rollout. I did end up keeping a thin single_train_step override because the reference co-trains the Stage-2 flow-map loss at every generator step (cotrain_forward_kl); it calls the stock DMD2 update and just adds that one term.
| if r_embedder_fusion not in ("additive", "gated"): | ||
| raise ValueError(f"r_embedder_fusion must be 'additive' or 'gated', got {r_embedder_fusion!r}") | ||
| self.transformer.r_embedder.fusion_mode = r_embedder_fusion | ||
| self.transformer.r_embedder.register_buffer( |
There was a problem hiding this comment.
Can we just store the float instead of adding it as a non-persistent buffer, since the latter would need to be added to the reset_parameters method for FSDP.
| """ | ||
| fusion = getattr(self.r_embedder, "fusion_mode", "additive") | ||
|
|
||
| if fusion == "gated": |
There was a problem hiding this comment.
Can we also respect the encoder_depth for the gated-fusion?
There was a problem hiding this comment.
Done — mirrors the additive branch: encoder blocks keep the t-only projection, the rest switch to the gated one.
Address review feedback on PR NVlabs#25: - Pretrain (Stage 2) now runs MeanFlowModel directly: add a fixed per-timestep loss weighting (loss_config.weight_type, evaluated as a function of t) and a consistency bucket (sample_t_cfg.consistency_ratio) to MeanFlow; both default off. Drop FlowMapDiscreteScheduler — (t, r) pair sampling uses noise_scheduler.sample_t with the shifted distribution, weights need no precomputed table. - On-policy (Stage 3) keeps only two DMD2 overrides: _generate_noise_and_time (start from pure noise at max_t) and gen_data_from_net (multi-step Euler-flow rollout with one gradient-enabled step). Teacher/fake_score are flow-map networks queried at the instantaneous velocity r=t — the reference passes r_timestep=timesteps to both (not r=0 as before) — implemented as the network-level default for dual-timestep nets when r is not passed. - Wan: store r_embedder.gate_value as a plain float (no buffer to re-materialize in reset_parameters for FSDP); gated fusion respects encoder_depth (mirrors additive); move remap_anyflow_keys here and apply it inside Wan.load_state_dict. - Configs: set r_embedder_fusion=gated + time_cond_type=abs on both stages, matching the published checkpoints (deltatime_type 'r', gate 0.25); imports at module top throughout.
Adversarial review against the reference surfaced several silent deviations in the training objective; all fixed: Pretrain (MeanFlow, all opt-in via config): - rebalance_to_diffusion: non-diffusion (flow-map / consistency) sample losses are rescaled by the detached factor mean(global diffusion losses) / (own loss + 1e-5), all-gathered across ranks, matching the reference's scale_weight. - guidance_fuse_scale: prediction-side guidance distillation — the conditional output learns the guided flow via (u_cond + (g-1) u_uncond) / g against the raw data velocity, with the unconditional branch queried at the SAME (t, r) slice, the finite-difference dF/dt divided by g, plain text dropout, and the probes extrapolated along the raw velocity. MeanFlow's target-side eq. 19 fusion (guidance_scale) is a different mechanism and stays untouched. - consistency bucket pins r = 0 (not min_t) and, together with the flow-matching head, uses the reference's deterministic global-batch partition by rank index instead of independent binomial draws (which biased the effective ratios at small per-GPU batches). - weight_type=uniform is exactly 1 (the reference applies no grid normalization to it). On-policy: - rollout NFE sampled per iteration from student_sample_steps_list ([2, 4, 8, 16, 50]) with rank-0 broadcast; schedule computed from the shifted grid per NFE. - rollout compressed to <= 3 flow-map forwards (jump t0->tg, fine step tg->tg+1, jump to 0) with gradient through ALL segments — the previous N-step loop with gradient on one step did not match training_rollout. - every student update co-trains the Stage-2 flow-map loss on the real batch (cotrain_pretrain_weight, reference cotrain_forward_kl). - no adversarial loss: the reference 'discriminator' is the fake score network; gan_loss_weight_gen=0. - student/fake-score updates alternate 1:1 (student_update_freq=2), teacher CFG strength corrected to the reference's cond + 3*(cond-uncond) (FastGen formula: guidance_scale=4), optimizer betas (0.0, 0.999), wd=0, grad clip 1.0, EMA 0.99. Configs also align the pretrain recipe (grad clip 1.0, 1000-step LR warmup, EMA 0.999, exact shifted 4-step eval schedule).
|
Thanks again for the careful review — all eight comments should be addressed now (replies in the threads below, changes in While reworking the code I also went back through the reference trainers line by line, and caught a few places where my original port deviated from what
Forward parity on the released checkpoints and the fp32 stub check of the central-difference target are unaffected by all this; the rest is verified by porting against the reference code plus the CPU unit tests (now 29). A few known deviations remain, noted in the config docstrings: full-rank fine-tuning instead of the rank-256 LoRA, shifted-uniform instead of shifted-logit-normal noising times for the fake score, constant EMA decay without the warmup, and the real/fake score init (your recipe starts from a separately fine-tuned flow-map teacher — users need to point the teacher path at one). And as discussed above, I still don't have spare GPUs for a convergence run, so that part stays out of scope for this PR. |
| else condition.reshape(-1, self.label_dim) | ||
| ) | ||
|
|
||
| # A dual-timestep network always consumes the r-pathway (the embedding |
There was a problem hiding this comment.
Why do we need this?
There was a problem hiding this comment.
The dual-timestep EDM nets size their embedding for the concat (cond_channels = noise_channels * (1 + r_timestep)), so a forward without r fails with a shape mismatch in map_layer0. DMD2's fake/real score call sites don't pass r, and the reference queries both score networks at r_timestep=timesteps (i.e. r=t) anyway — so this is the same network-level default as on the Wan side, needed here because the unit tests run the on-policy path on the tiny EDM backbone.
| """ | ||
| if getattr(self.loss_config, "rebalance_to_diffusion", False) and (~is_diffusion).any(): | ||
| with torch.no_grad(): | ||
| if torch.distributed.is_available() and torch.distributed.is_initialized(): |
There was a problem hiding this comment.
Can we use our internal world_size utils?
There was a problem hiding this comment.
Done — switched to fastgen.utils.distributed.world_size / get_rank, also for the rollout broadcast in anyflow.py.
| weights over its shifted discrete timestep grid; we compute the same | ||
| constant once and cache it. | ||
| """ | ||
| if self.loss_config.weight_type == "uniform": |
There was a problem hiding this comment.
Do we need this additional if (since this is handled by _timestep_weight_raw)?
There was a problem hiding this comment.
You're right — removed. Checking why I'd added it turned up the real issue: my normalization grid included the t=0 endpoint, which the reference's set_timesteps excludes. With the grid aligned (linspace(1, 0, 1001)[:-1]), uniform normalizes to exactly 1.0 and _timestep_weight_raw covers it.
| loss = torch.sum(loss, dim=list(range(1, loss.ndim))) | ||
| weight = self._compute_weight(loss) | ||
| loss = loss * weight | ||
| if self.loss_config.weight_type is not None: |
There was a problem hiding this comment.
Can we integrate that in _compute_weight (adding a time input to the method) and add optional time-weighting to all methods? We can introduce None as norm_method if we don't normalize.
There was a problem hiding this comment.
Done — _compute_weight(tensor, t) now multiplies the adaptive norm_method weight (None disables it) with the optional weight_type per-timestep weight, and both loss types go through it, so time-weighting is available everywhere. One subtlety: without the adaptive normalization the reduction scale matters, so norm_method=None reduces the l2 loss with a per-element mean (the reference's loss scale) instead of the sum; the AnyFlow configs set it explicitly and the default path is unchanged.
| zero_mask = torch.arange(batch_size, device=self.device) < flow_matching_size | ||
| r = torch.where(zero_mask, t, r) | ||
| # sample t and r, with per-batch buckets (flow matching / consistency) | ||
| t, r, is_diffusion = self._sample_t_r_buckets(batch_size) |
There was a problem hiding this comment.
is_diffusion might be confusing, perhaps r_eq_t_mask?
|
@coderabbitai review |
✅ Action performedReview finished.
|
WalkthroughThis PR introduces AnyFlow, a two-stage training pipeline that extends MeanFlow's consistency training for text-to-video generation with dual-timestep embeddings, deterministic bucket sampling, guidance fusion, and on-policy rollout distillation. The implementation includes config schema extensions, a new AnyFlowModel class, dual-timestep network architecture in Wan, checkpoint conversion utilities, and comprehensive test coverage. ChangesAnyFlow Two-Stage Training Pipeline
🎯 3 (Moderate) | ⏱️ ~25 minutes
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
fastgen/methods/distribution_matching/anyflow.py (1)
231-265:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winFix formatting to pass CI: add trailing newline.
The pipeline failure indicates
ruff format --checkwould reformat this file. The file is missing a trailing newline at line 265.Run the suggested command to fix:
python3 -m ruff format --exclude fastgen/third_party/ fastgen/methods/distribution_matching/anyflow.py🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@fastgen/methods/distribution_matching/anyflow.py` around lines 231 - 265, The file ends without a trailing newline which fails ruff format; open the function single_train_step in anyflow.py (and the file EOF) and add a newline at the end of the file (or run the suggested formatter command: python3 -m ruff format --exclude fastgen/third_party/ fastgen/methods/distribution_matching/anyflow.py) so the file ends with a single trailing newline and passes CI.Source: Pipeline failures
fastgen/networks/Wan/network.py (1)
1-1:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winFix formatting to pass CI.
The pipeline reports that
ruff format --checkfailed for this file. Run the formatter to fix:python3 -m ruff format fastgen/networks/Wan/network.py🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@fastgen/networks/Wan/network.py` at line 1, Run the code formatter on the module to fix ruff formatting failures: run `python3 -m ruff format fastgen/networks/Wan/network.py` (or apply equivalent formatting) so the SPDX header and entire file conform to ruff rules; ensure the top-of-file SPDX comment and any surrounding whitespace in network.py are corrected and committed.Source: Pipeline failures
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@fastgen/configs/methods/config_anyflow.py`:
- Around line 1-66: The cond_keys_no_dropout attribute in ModelConfig currently
uses a mutable default (empty list) which can be shared across instances; change
its declaration to use an attrs factory instead of a literal default by
replacing the current default value with attrs.field(factory=list) for the
cond_keys_no_dropout List[str] field in the ModelConfig class so each instance
gets its own list.
In `@tests/test_anyflowmodel.py`:
- Line 156: The test unpacks three values from model._sample_t_r_buckets(4) but
the first variable t is unused; change the unpack to use a throwaway name (e.g.,
_t or _) instead of t to silence the RUF059 lint warning and keep behavior
identical (locate the unpacking in tests/test_anyflowmodel.py where
model._sample_t_r_buckets is called).
- Around line 172-175: The test fails on CUDA because torch.zeros(n_consistency)
creates a CPU tensor while r[n_diffusion : n_diffusion + n_consistency] may be
on another device; update the assertion to create the zero tensor on the same
device and dtype as the slice (e.g. use r[n_diffusion : n_diffusion +
n_consistency].new_zeros(n_consistency) or torch.zeros(...,
device=that_slice.device, dtype=that_slice.dtype)) so torch.allclose compares
tensors on the same device.
---
Outside diff comments:
In `@fastgen/methods/distribution_matching/anyflow.py`:
- Around line 231-265: The file ends without a trailing newline which fails ruff
format; open the function single_train_step in anyflow.py (and the file EOF) and
add a newline at the end of the file (or run the suggested formatter command:
python3 -m ruff format --exclude fastgen/third_party/
fastgen/methods/distribution_matching/anyflow.py) so the file ends with a single
trailing newline and passes CI.
In `@fastgen/networks/Wan/network.py`:
- Line 1: Run the code formatter on the module to fix ruff formatting failures:
run `python3 -m ruff format fastgen/networks/Wan/network.py` (or apply
equivalent formatting) so the SPDX header and entire file conform to ruff rules;
ensure the top-of-file SPDX comment and any surrounding whitespace in network.py
are corrected and committed.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Enterprise
Run ID: 99311cde-946f-437d-b71e-879c8ac425e0
📒 Files selected for processing (11)
fastgen/configs/experiments/WanT2V/config_anyflow.pyfastgen/configs/experiments/WanT2V/config_anyflow_onpolicy.pyfastgen/configs/methods/config_anyflow.pyfastgen/configs/methods/config_mean_flow.pyfastgen/methods/__init__.pyfastgen/methods/consistency_model/mean_flow.pyfastgen/methods/distribution_matching/README.mdfastgen/methods/distribution_matching/anyflow.pyfastgen/networks/EDM/network.pyfastgen/networks/Wan/network.pytests/test_anyflowmodel.py
| # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | ||
| # SPDX-License-Identifier: Apache-2.0 | ||
|
|
||
| """Config schema for the AnyFlow on-policy method (paper Stage 3). | ||
|
|
||
| AnyFlow's on-policy stage is DMD2 with a flow-map student that generates via | ||
| a multi-step rollout-with-gradient, so the config is the DMD2 model config | ||
| unchanged. The flow-map pretrain stage (paper Stage 2) is MeanFlow with | ||
| AnyFlow's hyperparameters — see ``configs/experiments/WanT2V/config_anyflow.py``. | ||
| """ | ||
|
|
||
| from typing import List, Optional | ||
|
|
||
| import attrs | ||
| from omegaconf import DictConfig | ||
|
|
||
| from fastgen.configs.callbacks import ( | ||
| EMA_CALLBACK, | ||
| GPUStats_CALLBACK, | ||
| GradClip_CALLBACK, | ||
| ParamCount_CALLBACK, | ||
| TrainProfiler_CALLBACK, | ||
| WANDB_CALLBACK, | ||
| ) | ||
| from fastgen.configs.config import BaseConfig | ||
| from fastgen.configs.methods.config_dmd2 import ModelConfig as DMD2ModelConfig | ||
| from fastgen.configs.methods.config_mean_flow import ( | ||
| LossConfig as MeanFlowLossConfig, | ||
| SampleRConfig, | ||
| SampleTConfig as MeanFlowSampleTConfig, | ||
| ) | ||
| from fastgen.methods import AnyFlowModel | ||
| from fastgen.utils import LazyCall as L | ||
|
|
||
|
|
||
| @attrs.define(slots=False) | ||
| class ModelConfig(DMD2ModelConfig): | ||
| """AnyFlow on-policy model config — DMD2 plus the rollout / cotrain knobs. | ||
|
|
||
| The MeanFlow loss / sampling configs drive the co-trained Stage-2 | ||
| flow-map loss inside the student update (the reference's | ||
| ``cotrain_forward_kl``). | ||
| """ | ||
|
|
||
| # MeanFlow-style (t, r) sampling for the co-trained flow-map loss; the | ||
| # extra fields are ignored by the DMD2 noising-time sampling. | ||
| sample_t_cfg: MeanFlowSampleTConfig = attrs.field(factory=MeanFlowSampleTConfig) | ||
| sample_r_cfg: SampleRConfig = attrs.field(factory=SampleRConfig) | ||
| loss_config: MeanFlowLossConfig = attrs.field(factory=MeanFlowLossConfig) | ||
|
|
||
| # Weight of the co-trained Stage-2 flow-map loss in the student update. | ||
| # The reference runs it at weight 1 (cotrain_forward_kl: True); 0 disables. | ||
| cotrain_pretrain_weight: float = 1.0 | ||
|
|
||
| # Rollout NFE list, sampled uniformly per iteration with rank-0 broadcast | ||
| # (reference rollout_cfg.num_inference_steps_list). None falls back to | ||
| # the fixed student_sample_steps. | ||
| student_sample_steps_list: Optional[List[int]] = None | ||
|
|
||
| # Text dropout for the co-trained flow-map loss (reference drop_text_ratio). | ||
| cond_dropout_prob: Optional[float] = None | ||
| cond_keys_no_dropout: List[str] = [] | ||
|
|
||
| # Precision for autocast in the co-trained loss JVP (None = training precision). | ||
| precision_amp_jvp: str | None = None | ||
|
|
There was a problem hiding this comment.
Fix mutable default value for cond_keys_no_dropout.
Line 62 uses a mutable default value (empty list) for a dataclass attribute. This can lead to shared state across instances if the list is mutated.
🔧 Proposed fix
- cond_keys_no_dropout: List[str] = []
+ cond_keys_no_dropout: List[str] = attrs.field(factory=list)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | |
| # SPDX-License-Identifier: Apache-2.0 | |
| """Config schema for the AnyFlow on-policy method (paper Stage 3). | |
| AnyFlow's on-policy stage is DMD2 with a flow-map student that generates via | |
| a multi-step rollout-with-gradient, so the config is the DMD2 model config | |
| unchanged. The flow-map pretrain stage (paper Stage 2) is MeanFlow with | |
| AnyFlow's hyperparameters — see ``configs/experiments/WanT2V/config_anyflow.py``. | |
| """ | |
| from typing import List, Optional | |
| import attrs | |
| from omegaconf import DictConfig | |
| from fastgen.configs.callbacks import ( | |
| EMA_CALLBACK, | |
| GPUStats_CALLBACK, | |
| GradClip_CALLBACK, | |
| ParamCount_CALLBACK, | |
| TrainProfiler_CALLBACK, | |
| WANDB_CALLBACK, | |
| ) | |
| from fastgen.configs.config import BaseConfig | |
| from fastgen.configs.methods.config_dmd2 import ModelConfig as DMD2ModelConfig | |
| from fastgen.configs.methods.config_mean_flow import ( | |
| LossConfig as MeanFlowLossConfig, | |
| SampleRConfig, | |
| SampleTConfig as MeanFlowSampleTConfig, | |
| ) | |
| from fastgen.methods import AnyFlowModel | |
| from fastgen.utils import LazyCall as L | |
| @attrs.define(slots=False) | |
| class ModelConfig(DMD2ModelConfig): | |
| """AnyFlow on-policy model config — DMD2 plus the rollout / cotrain knobs. | |
| The MeanFlow loss / sampling configs drive the co-trained Stage-2 | |
| flow-map loss inside the student update (the reference's | |
| ``cotrain_forward_kl``). | |
| """ | |
| # MeanFlow-style (t, r) sampling for the co-trained flow-map loss; the | |
| # extra fields are ignored by the DMD2 noising-time sampling. | |
| sample_t_cfg: MeanFlowSampleTConfig = attrs.field(factory=MeanFlowSampleTConfig) | |
| sample_r_cfg: SampleRConfig = attrs.field(factory=SampleRConfig) | |
| loss_config: MeanFlowLossConfig = attrs.field(factory=MeanFlowLossConfig) | |
| # Weight of the co-trained Stage-2 flow-map loss in the student update. | |
| # The reference runs it at weight 1 (cotrain_forward_kl: True); 0 disables. | |
| cotrain_pretrain_weight: float = 1.0 | |
| # Rollout NFE list, sampled uniformly per iteration with rank-0 broadcast | |
| # (reference rollout_cfg.num_inference_steps_list). None falls back to | |
| # the fixed student_sample_steps. | |
| student_sample_steps_list: Optional[List[int]] = None | |
| # Text dropout for the co-trained flow-map loss (reference drop_text_ratio). | |
| cond_dropout_prob: Optional[float] = None | |
| cond_keys_no_dropout: List[str] = [] | |
| # Precision for autocast in the co-trained loss JVP (None = training precision). | |
| precision_amp_jvp: str | None = None | |
| # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | |
| # SPDX-License-Identifier: Apache-2.0 | |
| """Config schema for the AnyFlow on-policy method (paper Stage 3). | |
| AnyFlow's on-policy stage is DMD2 with a flow-map student that generates via | |
| a multi-step rollout-with-gradient, so the config is the DMD2 model config | |
| unchanged. The flow-map pretrain stage (paper Stage 2) is MeanFlow with | |
| AnyFlow's hyperparameters — see ``configs/experiments/WanT2V/config_anyflow.py``. | |
| """ | |
| from typing import List, Optional | |
| import attrs | |
| from omegaconf import DictConfig | |
| from fastgen.configs.callbacks import ( | |
| EMA_CALLBACK, | |
| GPUStats_CALLBACK, | |
| GradClip_CALLBACK, | |
| ParamCount_CALLBACK, | |
| TrainProfiler_CALLBACK, | |
| WANDB_CALLBACK, | |
| ) | |
| from fastgen.configs.config import BaseConfig | |
| from fastgen.configs.methods.config_dmd2 import ModelConfig as DMD2ModelConfig | |
| from fastgen.configs.methods.config_mean_flow import ( | |
| LossConfig as MeanFlowLossConfig, | |
| SampleRConfig, | |
| SampleTConfig as MeanFlowSampleTConfig, | |
| ) | |
| from fastgen.methods import AnyFlowModel | |
| from fastgen.utils import LazyCall as L | |
| `@attrs.define`(slots=False) | |
| class ModelConfig(DMD2ModelConfig): | |
| """AnyFlow on-policy model config — DMD2 plus the rollout / cotrain knobs. | |
| The MeanFlow loss / sampling configs drive the co-trained Stage-2 | |
| flow-map loss inside the student update (the reference's | |
| ``cotrain_forward_kl``). | |
| """ | |
| # MeanFlow-style (t, r) sampling for the co-trained flow-map loss; the | |
| # extra fields are ignored by the DMD2 noising-time sampling. | |
| sample_t_cfg: MeanFlowSampleTConfig = attrs.field(factory=MeanFlowSampleTConfig) | |
| sample_r_cfg: SampleRConfig = attrs.field(factory=SampleRConfig) | |
| loss_config: MeanFlowLossConfig = attrs.field(factory=MeanFlowLossConfig) | |
| # Weight of the co-trained Stage-2 flow-map loss in the student update. | |
| # The reference runs it at weight 1 (cotrain_forward_kl: True); 0 disables. | |
| cotrain_pretrain_weight: float = 1.0 | |
| # Rollout NFE list, sampled uniformly per iteration with rank-0 broadcast | |
| # (reference rollout_cfg.num_inference_steps_list). None falls back to | |
| # the fixed student_sample_steps. | |
| student_sample_steps_list: Optional[List[int]] = None | |
| # Text dropout for the co-trained flow-map loss (reference drop_text_ratio). | |
| cond_dropout_prob: Optional[float] = None | |
| cond_keys_no_dropout: List[str] = attrs.field(factory=list) | |
| # Precision for autocast in the co-trained loss JVP (None = training precision). | |
| precision_amp_jvp: str | None = None |
🧰 Tools
🪛 Ruff (0.15.15)
[warning] 62-62: Do not use mutable default values for dataclass attributes
(RUF008)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@fastgen/configs/methods/config_anyflow.py` around lines 1 - 66, The
cond_keys_no_dropout attribute in ModelConfig currently uses a mutable default
(empty list) which can be shared across instances; change its declaration to use
an attrs factory instead of a literal default by replacing the current default
value with attrs.field(factory=list) for the cond_keys_no_dropout List[str]
field in the ModelConfig class so each instance gets its own list.
| """With consistency_ratio=1.0 (and no flow-matching head), every sample's | ||
| r must be pinned to 0 (consistency to clean data, as in the reference).""" | ||
| model = _build_pretrain_model(consistency_ratio=1.0, r_sample_ratio=1.0) | ||
| t, r, is_diffusion = model._sample_t_r_buckets(4) |
There was a problem hiding this comment.
Remove unused unpacked variable to keep lint clean.
At Line 156, t is unpacked but never used (RUF059). Rename it to _t (or _) to avoid warning churn.
Proposed fix
- t, r, is_diffusion = model._sample_t_r_buckets(4)
+ _t, r, is_diffusion = model._sample_t_r_buckets(4)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| t, r, is_diffusion = model._sample_t_r_buckets(4) | |
| _t, r, is_diffusion = model._sample_t_r_buckets(4) |
🧰 Tools
🪛 Ruff (0.15.15)
[warning] 156-156: Unpacked variable t is never used
Prefix it with an underscore or any other dummy variable pattern
(RUF059)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/test_anyflowmodel.py` at line 156, The test unpacks three values from
model._sample_t_r_buckets(4) but the first variable t is unused; change the
unpack to use a throwaway name (e.g., _t or _) instead of t to silence the
RUF059 lint warning and keep behavior identical (locate the unpacking in
tests/test_anyflowmodel.py where model._sample_t_r_buckets is called).
Source: Linters/SAST tools
| assert torch.allclose( | ||
| r[n_diffusion : n_diffusion + n_consistency].float(), | ||
| torch.zeros(n_consistency), | ||
| ) |
There was a problem hiding this comment.
Cross-device tensor mismatch can fail this test on CUDA.
At Line 172-175, torch.zeros(n_consistency) is always CPU, but r[...] follows model.device and can be CUDA. This can raise a device mismatch error and make the test non-portable.
Proposed fix
assert torch.allclose(
r[n_diffusion : n_diffusion + n_consistency].float(),
- torch.zeros(n_consistency),
+ torch.zeros(
+ n_consistency,
+ device=r.device,
+ dtype=r.float().dtype,
+ ),
)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| assert torch.allclose( | |
| r[n_diffusion : n_diffusion + n_consistency].float(), | |
| torch.zeros(n_consistency), | |
| ) | |
| assert torch.allclose( | |
| r[n_diffusion : n_diffusion + n_consistency].float(), | |
| torch.zeros( | |
| n_consistency, | |
| device=r.device, | |
| dtype=r.float().dtype, | |
| ), | |
| ) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/test_anyflowmodel.py` around lines 172 - 175, The test fails on CUDA
because torch.zeros(n_consistency) creates a CPU tensor while r[n_diffusion :
n_diffusion + n_consistency] may be on another device; update the assertion to
create the zero tensor on the same device and dtype as the slice (e.g. use
r[n_diffusion : n_diffusion + n_consistency].new_zeros(n_consistency) or
torch.zeros(..., device=that_slice.device, dtype=that_slice.dtype)) so
torch.allclose compares tensors on the same device.
Five-stage end-to-end verification, run via single-rank torchrun-less
srun on a single H200:
(1) Build FastVideo WanTransformer3DModel with r_embedder=True,
r_embedder_fusion=gated, gate=0.25.
(2) Load nvidia/AnyFlow-Wan2.1-T2V-1.3B-Diffusers safetensors and
translate keys via WanVideoArchConfig.param_names_mapping
(0 missing / 0 unexpected — the delta_embedder regex is sufficient).
(3) Build AnyFlow's reference loader (FAR_Wan_Transformer3DModel).
(4) Forward parity on identical inputs — bf16 noise.
(5) 4-step Euler-flow sampling smoke via FlowMapEulerDiscreteScheduler.
(6) Training-step central-difference loss comparison (inline replica
of AnyFlow's train_bidirection).
Measured on Wan2.1-T2V-1.3B + nvidia/AnyFlow checkpoint:
forward rel mean diff : 2.55%
forward max abs diff : 7.81e-2
training loss diff : 1.33% (AnyFlow 0.381619 vs FastVideo 0.386694)
Both within bf16 kernel noise. Compare to the FastGen port at
NVlabs/FastGen#25 which reported 2.8% forward + 4.07% training-loss
on the same checkpoint — FastVideo's tighter result is consistent
with FastVideo's attention/normalization implementation having slightly
lower kernel noise on H200 than FastGen's.
- Use fastgen.utils.distributed world_size/get_rank instead of raw torch.distributed queries (mean_flow buckets/rebalancing, anyflow rollout broadcast). - Fold the fixed per-timestep weighting into _compute_weight(tensor, t): the adaptive norm_method weight (None disables it) multiplies the optional weight_type weight, for both l2 and opt_grad losses. With norm_method=None the l2 loss reduces with a per-element mean, matching the reference loss scale; the default path is unchanged. - Align the weight-normalization grid with the reference set_timesteps (1000 points, t=0 excluded), which makes the uniform special case redundant; drop it. - Rename the is_diffusion mask to r_eq_t_mask. - Use attrs.field(factory=list) for cond_keys_no_dropout (the plain [] default is shared across config instances). - Formatting fixes for ruff==0.6.9 (trailing newline, line join) and test cleanups (device-safe zeros, unused unpack). Signed-off-by: Enderfga <qq2639135175@gmail.com>
|
Sorry for the month of silence here — I completely missed the notifications for your June review and only caught up on it now. Not the turnaround this PR deserved after your careful comments, apologies. All five points are addressed in f93a980 (replies in the threads), along with the lint failure and CodeRabbit's comments. One extra fix that came out of re-checking the weighting question: my weight-normalization grid included the t=0 endpoint that the reference's |
Greptile SummaryThis PR adds AnyFlow support for any-step video diffusion. The main changes are:
Confidence Score: 5/5This looks safe to merge.
Important Files Changed
Reviews (5): Last reviewed commit: "anyflow: fix stale docs" | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| global_bsz = world_size() * batch_size | ||
| n_flow_matching = round((1.0 - self.sample_t_cfg.r_sample_ratio) * global_bsz) | ||
| n_consistency = round(self.sample_t_cfg.consistency_ratio * global_bsz) |
There was a problem hiding this comment.
When the Wan AnyFlow configs run with dataloader_train.batch_size = 1, this uses the per-forward local batch instead of the accumulated batch_size_global. With one rank, round(0.5 * 1) and round(0.25 * 1) both become zero, so the configured flow-matching and consistency buckets are never assigned and training silently uses only random (t, r) pairs.
There was a problem hiding this comment.
This is faithful to the reference — sample_timestep in trainer_wan_anyflow_pretrain.py partitions by rank index over world_size * per_rank_batch, and gradient accumulation doesn't enter the computation there either. The intended recipes run multi-GPU (global batch 32 at bs=1 per rank), where the buckets are assigned across ranks; a single-GPU bs=1 run degenerates identically in the reference.
| gathered_loss = [torch.zeros_like(mf_loss) for _ in range(world_size())] | ||
| gathered_mask = [torch.zeros_like(r_eq_t_mask) for _ in range(world_size())] | ||
| torch.distributed.all_gather(gathered_loss, mf_loss.contiguous()) | ||
| torch.distributed.all_gather(gathered_mask, r_eq_t_mask.contiguous()) |
There was a problem hiding this comment.
This collective gathers tensors shaped like each rank's local mf_loss. If the last distributed batch is uneven, or a sampler does not pad/drop to identical local counts, the ranks call all_gather with different tensor sizes and the AnyFlow rebalance path can fail during training.
There was a problem hiding this comment.
Same assumption as the reference, which calls dist.nn.all_gather(loss) on the local batch directly. FastGen's training loaders yield fixed-size per-rank batches, so the collective shapes always match.
| # dF/dt of the fused prediction: finite difference of the | ||
| # conditional output divided by g (the unconditional derivative is | ||
| # dropped, as in the reference's compute_central_difference). | ||
| u_theta_jvp = self._jvp(x_t, t, r, dxt_dt, condition=condition) / g |
There was a problem hiding this comment.
guidance_fuse_scale is accepted as any value, but this branch divides the JVP by g. Setting the config to 0.0 enters the guidance-fusion path and produces infinities or NaNs in the loss instead of disabling the feature or failing clearly.
There was a problem hiding this comment.
Added an assert in de4adde — a non-positive scale now fails fast with a message pointing to None for disabling the feature.
| assert x_t.dtype == real_data.dtype, f"x_t.dtype: {x_t.dtype}, real_data.dtype: {real_data.dtype}" | ||
| u_theta = self.net(x_t, t, r=r, condition=condition, fwd_pred_type="flow") | ||
| with torch.no_grad(): | ||
| u_uncond = self.net(x_t, t, r=r, condition=neg_condition, fwd_pred_type="flow") |
There was a problem hiding this comment.
When guidance_fuse_scale is enabled, _drop_condition tolerates neg_condition is None, but this unconditional forward still passes condition=None into the network. A dataset without neg_condition can crash here or train the unconditional branch with invalid conditioning.
There was a problem hiding this comment.
Added an assert in de4adde: enabling guidance_fuse_scale without neg_condition now fails with a clear message instead of crashing inside the network.
Fail fast with a clear message when guidance_fuse_scale is non-positive (the fused prediction divides by it) or when neg_condition is missing (the unconditional branch is queried at the same (t, r)). Signed-off-by: Enderfga <qq2639135175@gmail.com>
| assert torch.all(t >= r), "r cannot be larger than t" | ||
|
|
||
| if self.sample_t_cfg.consistency_ratio > 0: | ||
| global_bsz = world_size() * batch_size |
There was a problem hiding this comment.
dataloader_train.batch_size = 1 and trainer.batch_size_global = 32, a single-rank run computes global_bsz == 1, so both round(0.5 * 1) and round(0.25 * 1) become zero. The flow-matching and consistency buckets are never assigned, and training silently uses only random (t, r) pairs. This needs to include the accumulation batch size when deriving the deterministic bucket partition.
There was a problem hiding this comment.
The partition deliberately matches the reference, which spans ranks but not gradient-accumulation rounds — folding accumulation in would change the training math relative to the reference, and the accumulation round index isn't visible at this level anyway. What I did take from this: the degenerate case was silent, so 552f077 adds a one-time warning when the configured ratios produce empty buckets (e.g. a single rank at batch size 1).
There was a problem hiding this comment.
That's a fair point — the reference's bucket assignment is per-rank-collective step, not per-accumulation round, so folding in the accumulation factor would diverge from the reference training math. And 552f077 adds the degenerate-case warning, which directly addresses the silent failure mode. That resolves my concern.
Tip: You can customize Greptile's behavior for this repo with .greptile/rules.md and .greptile/config.json.
| torch.distributed.all_gather(gathered_loss, mf_loss.contiguous()) | ||
| torch.distributed.all_gather(gathered_mask, r_eq_t_mask.contiguous()) |
There was a problem hiding this comment.
zeros_like(mf_loss), so every rank must contribute the same local loss shape. If the final distributed batch is uneven, or the sampler does not pad or drop to identical per-rank counts, ranks enter all_gather with different tensor sizes and the AnyFlow rebalance step can fail or hang. The rebalance path needs to enforce equal local shapes before this collective or gather with explicit size handling.
There was a problem hiding this comment.
Reworked in 552f077 — the rebalance factor only needs the global flow-matching-loss mean, so it now all_reduces a sum and a count instead of gathering the per-sample losses. Same math as the reference's cat(all_gather(loss))[mask].mean(), cheaper, and independent of per-rank batch shapes.
The rebalance factor only needs the global flow-matching-loss mean, so all_reduce a sum and a count instead of all_gathering the per-sample losses — equivalent to the reference's cat(all_gather(loss)) math, cheaper, and independent of per-rank batch sizes. The deterministic (t, r) bucket partition spans ranks but not gradient-accumulation rounds (as in the AnyFlow reference); log a one-time warning when the configured ratios produce empty buckets so a degenerate setup (e.g. single rank at batch size 1) is visible. Signed-off-by: Enderfga <qq2639135175@gmail.com>
| if getattr(self.loss_config, "rebalance_to_diffusion", False) and (~r_eq_t_mask).any(): | ||
| with torch.no_grad(): | ||
| # The global flow-matching-loss mean only needs the global sum | ||
| # and count, so reduce two scalars instead of gathering the | ||
| # per-sample losses (equivalent to the reference's | ||
| # cat(all_gather(loss))[mask].mean(), and independent of the | ||
| # per-rank batch sizes). | ||
| fm_loss_sum = torch.where(r_eq_t_mask, mf_loss, torch.zeros_like(mf_loss)).sum() | ||
| fm_count = r_eq_t_mask.sum().to(mf_loss.dtype) | ||
| if world_size() > 1: | ||
| torch.distributed.all_reduce(fm_loss_sum) | ||
| torch.distributed.all_reduce(fm_count) | ||
| scale = torch.ones_like(mf_loss) | ||
| if fm_count > 0: | ||
| scale[~r_eq_t_mask] = (fm_loss_sum / fm_count) / (mf_loss[~r_eq_t_mask] + 1e-5) | ||
| mf_loss = mf_loss * scale |
There was a problem hiding this comment.
This rebalance path can still hang in distributed training because the all_reduce calls are behind a rank-local (~r_eq_t_mask).any() check. With AnyFlow buckets, one rank can hold only flow-matching samples while another rank has flow-map or consistency samples. For example, with two ranks, local batch size 4, and r_sample_ratio=0.5, rank 0 can skip this block while rank 1 waits in all_reduce. Every rank should enter the same collectives whenever rebalancing is enabled; only the local scaling needs to depend on whether that rank has non-flow samples.
| if getattr(self.loss_config, "rebalance_to_diffusion", False) and (~r_eq_t_mask).any(): | |
| with torch.no_grad(): | |
| # The global flow-matching-loss mean only needs the global sum | |
| # and count, so reduce two scalars instead of gathering the | |
| # per-sample losses (equivalent to the reference's | |
| # cat(all_gather(loss))[mask].mean(), and independent of the | |
| # per-rank batch sizes). | |
| fm_loss_sum = torch.where(r_eq_t_mask, mf_loss, torch.zeros_like(mf_loss)).sum() | |
| fm_count = r_eq_t_mask.sum().to(mf_loss.dtype) | |
| if world_size() > 1: | |
| torch.distributed.all_reduce(fm_loss_sum) | |
| torch.distributed.all_reduce(fm_count) | |
| scale = torch.ones_like(mf_loss) | |
| if fm_count > 0: | |
| scale[~r_eq_t_mask] = (fm_loss_sum / fm_count) / (mf_loss[~r_eq_t_mask] + 1e-5) | |
| mf_loss = mf_loss * scale | |
| if getattr(self.loss_config, "rebalance_to_diffusion", False): | |
| with torch.no_grad(): | |
| # The global flow-matching-loss mean only needs the global sum | |
| # and count, so reduce two scalars instead of gathering the | |
| # per-sample losses (equivalent to the reference's | |
| # cat(all_gather(loss))[mask].mean(), and independent of the | |
| # per-rank batch sizes). Every rank must participate in these | |
| # collectives, even if its local batch has no non-diffusion samples. | |
| fm_loss_sum = torch.where(r_eq_t_mask, mf_loss, torch.zeros_like(mf_loss)).sum() | |
| fm_count = r_eq_t_mask.sum().to(mf_loss.dtype) | |
| if world_size() > 1: | |
| torch.distributed.all_reduce(fm_loss_sum) | |
| torch.distributed.all_reduce(fm_count) | |
| scale = torch.ones_like(mf_loss) | |
| if fm_count > 0 and (~r_eq_t_mask).any(): | |
| scale[~r_eq_t_mask] = (fm_loss_sum / fm_count) / (mf_loss[~r_eq_t_mask] + 1e-5) | |
| mf_loss = mf_loss * scale |
There was a problem hiding this comment.
Good catch — that guard was mine, not the reference's (the reference runs its gather unconditionally every step), and with the deterministic partition some ranks hold only flow-matching samples by design, so this would have deadlocked the shipped 8-rank recipe. Fixed in 9aa47ff: the branch is gated on the config flag only — identical on every rank — and the empty-selection scale assignment is a no-op. Added a regression test for the all-flow-matching batch.
The all_reduce in _reduce_mf_loss was gated on the rank-local (~r_eq_t_mask).any(), so a rank holding only flow-matching samples (which the deterministic bucket partition produces by design) skipped the collective while other ranks entered it, deadlocking distributed training. Gate on the config flag only — identical on every rank — and let the empty-selection scale assignment be a no-op, as in the reference (which runs its gather unconditionally). Add a regression test for the all-flow-matching batch. Signed-off-by: Enderfga <qq2639135175@gmail.com>
The on-policy config disables the adversarial loss (gan_loss_weight_gen=0), so the README config line saying 'GAN on' contradicted both the config and the paragraph above it. The rollout gradient test docstring still described the old one-step-with-gradient scheme instead of the compressed all-segments rollout. Signed-off-by: Enderfga <qq2639135175@gmail.com>
3 participants
Summary
Adds AnyFlow as a new method under
fastgen/methods/distribution_matching/anyflow.py. AnyFlow trains a single flow-map modelu_θ(x_t, t, r)that predicts the average velocity fromtback tor, so the same checkpoint supports arbitrary inference NFE.Training has two stages:
Flow-map pretrain (paper Stage 2) is the MeanFlow objective with AnyFlow's hyperparameters, so it runs directly on
MeanFlowModelvia config — there is no AnyFlow-specific pretrain code. The AnyFlow pieces are opt-in extensions on MeanFlow (all defaulting to the original behavior): a fixed per-timestep loss weight (weight_type, normalized over the reference's shifted 1000-point grid), aconsistency_ratiobucket pinned tor = 0with the reference's deterministic rank-indexed partition, prediction-side guidance fusion (guidance_fuse_scale: the conditional output learns the guided flow directly), and global rebalancing of flow-map/consistency losses to the flow-matching-loss mean (rebalance_to_diffusion, implemented as two scalarall_reduces).On-policy (paper Stage 3) —
AnyFlowModel(DMD2Model), stock DMD2 with the reference's three deviations: the student generates via a flow-map rollout compressed into at most three network forwards (jumpt_0 → t_g, fine step, jump to 0) with gradient through all segments and the NFE sampled per iteration fromstudent_sample_steps_list(rank-0 broadcast); the student always starts from pure noise atmax_t; and every student update co-trains the Stage-2 flow-map loss (cotrain_pretrain_weight, the reference'scotrain_forward_kl). No adversarial loss — the reference's "discriminator" is the fake score network.Why the Wan backbone needs minimal changes
The Wan transformer already accepts a secondary timestep via its
r_embedder(r_timestep=True, exercised by MeanFlow). The additions: anr_embedder_fusionflag whose"gated"mode reproduces AnyFlow'sWanTwoTimeTextImageEmbedding.forward_timestep(rt_emb = (1−g)·temb + g·rembthrough the sharedtime_proj; default"additive"keeps MeanFlow/TCM/sCM bit-identical), and aremap_anyflow_keys()helper applied insideWan.load_state_dictthat rewrites the published-checkpoint layout (condition_embedder.delta_embedder.*→r_embedder.*, no-op for all other state dicts) so NVIDIA'sAnyFlow-Wan2.1-T2V-{1.3B,14B}-Diffusersreleases load as-is. Gated-fusion networks default tor = twhenrisn't passed (how the reference queries its score networks), so DMD2's update steps run unchanged.Files
New
fastgen/methods/distribution_matching/anyflow.py—AnyFlowModel(compressed rollout, co-trained flow-map loss)fastgen/configs/methods/config_anyflow.py— method config (DMD2's plus the rollout/cotrain knobs)fastgen/configs/experiments/WanT2V/config_anyflow.py/config_anyflow_onpolicy.py— Wan2.1-T2V-1.3B Stage 2 / Stage 3 reference experimentstests/test_anyflowmodel.py— 24 unit tests covering both stages, the rollout, the Wan fusion helpers, and the checkpoint remapModified (additive, defaults preserve existing behavior)
fastgen/methods/consistency_model/mean_flow.py— opt-in AnyFlow extensions listed abovefastgen/networks/Wan/network.py—_fuse_r_embeddinghelper + checkpoint remapfastgen/networks/EDM/network.py— dual-timestep nets default tor = t(they cannot run withr=None)fastgen/configs/methods/config_mean_flow.py,fastgen/methods/__init__.py,README.mdTest plan
pytest tests/test_anyflowmodel.py tests/test_meanflowmodel.py tests/test_dmd2model.py— 30/30 passing (no regression on MeanFlow/DMD2 defaults)ruff==0.6.9format + lint cleanOut of scope