Add AnyFlow algorithm (any-step video diffusion via flow maps)

[Enderfga]

Summary

Adds AnyFlow as a new method under fastgen/methods/distribution_matching/anyflow.py. AnyFlow trains a single flow-map model u_θ(x_t, t, r) that predicts the average velocity from t back to r, 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 MeanFlowModel via 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), a consistency_ratio bucket pinned to r = 0 with 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 scalar all_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 (jump t_0 → t_g, fine step, jump to 0) with gradient through all segments and the NFE sampled per iteration from student_sample_steps_list (rank-0 broadcast); the student always starts from pure noise at max_t; and every student update co-trains the Stage-2 flow-map loss (cotrain_pretrain_weight, the reference's cotrain_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: an r_embedder_fusion flag whose "gated" mode reproduces AnyFlow's WanTwoTimeTextImageEmbedding.forward_timestep (rt_emb = (1−g)·temb + g·remb through the shared time_proj; default "additive" keeps MeanFlow/TCM/sCM bit-identical), and a remap_anyflow_keys() helper applied inside Wan.load_state_dict that rewrites the published-checkpoint layout (condition_embedder.delta_embedder.* → r_embedder.*, no-op for all other state dicts) so NVIDIA's AnyFlow-Wan2.1-T2V-{1.3B,14B}-Diffusers releases load as-is. Gated-fusion networks default to r = t when r isn'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 experiments
• tests/test_anyflowmodel.py — 24 unit tests covering both stages, the rollout, the Wan fusion helpers, and the checkpoint remap

Modified (additive, defaults preserve existing behavior)

• fastgen/methods/consistency_model/mean_flow.py — opt-in AnyFlow extensions listed above
• fastgen/networks/Wan/network.py — _fuse_r_embedding helper + checkpoint remap
• fastgen/networks/EDM/network.py — dual-timestep nets default to r = t (they cannot run with r=None)
• fastgen/configs/methods/config_mean_flow.py, fastgen/methods/__init__.py, README.md

Test 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.9 format + lint clean
• Forward parity and training-step parity against the published 1.3B/14B checkpoints + any-step sample videos: verification comment

Out of scope

• LoRA-only training. The reference's rank-256 LoRA mode needs a PEFT path across FastGen's model zoo — follow-up PR.

[juliusberner]

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?

[Enderfga]

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 correctness

Loaded nvidia/AnyFlow-Wan2.1-T2V-1.3B-Diffusers and nvidia/AnyFlow-Wan2.1-T2V-14B-Diffusers through FastGen's Wan wrapper using a small remap_anyflow_keys helper (rewrites condition_embedder.delta_embedder.* → r_embedder.time_embedder.* and copies condition_embedder.time_proj.* into r_embedder.time_proj.*).

On identical inputs the FastGen-loaded model agrees with AnyFlow's own loader to within bf16 forward noise (rel mean diff 2.8%, max abs diff 8.7e-2).

Training-step equivalence

Inline replica of AnyFlow's train_bidirection central-difference math on real weights, same seed, same mini-batch through both code paths:

variant	AnyFlow loss	FastGen loss	rel diff
1.3B	0.381619	0.397162	4.07%
14B	0.141866	0.146120	3.00%

A stub-network compare of the central-difference target tensor (so the math is isolated from network weights) gives max abs diff 1.9e-5 in fp32 — the algorithm is reproduced exactly.

Sample videos

Same prompt + seed=0 + 81 frames @ 480×832 + shift=5 + weight_type=beta08 + guidance_scale=1.0 (matching demo.py's default), AnyFlow's own pipeline vs FastGen-loaded:

• 1.3B NFE=4 — visually indistinguishable from AnyFlow/demo.py output

1p3b_fastgen_nfe4.mp4

• 14B NFE=4

14b_fastgen_nfe4.mp4

• 14B NFE=50

14b_fastgen_nfe50.mp4

What this PR changes

• fastgen/networks/Wan/network.py (+39 −5): adds r_embedder_fusion: str = "additive" (default unchanged, preserves MeanFlow / TCM / sCM forward bit-identical) / "gated" and r_embedder_gate_value: float = 0.25. When gated, classify_forward_prepare computes rt_emb = (1 − g)·temb_t + g·temb_r before SiLU and uses the shared r_embedder.time_proj (deep-copy of condition_embedder.time_proj per init_embedder) for the final projection — matching WanTwoTimeTextImageEmbedding.forward_timestep in the AnyFlow reference.
• fastgen/methods/distribution_matching/anyflow.py (+39): adds remap_anyflow_keys() helper. No-op for non-AnyFlow checkpoints, so safe to call unconditionally.

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]

Follow-up commit ab1174d replaces the on-policy student's single-step forward with a multi-step Euler-flow rollout, matching AnyFlow's WanAnyFlowPipeline.training_rollout (the published on-policy training mode).

New AnyFlowModel._rollout_with_gradient(batch_size, dtype, condition):

• Starts from pure noise at ns.max_t.
• Iterates student_sample_steps Euler-flow updates with r = t_next (mean-velocity sampling, matching AnyFlow's use_mean_velocity=True default).
• Toggles torch.set_grad_enabled so exactly one randomly-chosen step keeps an autograd record; the rest run under no_grad.
• Broadcasts grad_step from rank 0 in distributed runs so all ranks share the same gradient window (mirrors AnyFlow's broadcast(sample_step, src=0)).
• Honours sample_t_cfg.t_list when set (so configs can pin the AnyFlow paper's hand-tuned schedule, e.g. [0.999, 0.937, 0.833, 0.624, 0.0] for 4-step Wan); otherwise falls back to noise_scheduler.get_t_list.

The rollout output replaces the single self.net(input_student, ...) forward in both _onpolicy_student_update_step and _onpolicy_fake_score_discriminator_update_step, so the DMD generator update now receives the rollout's gradient through a full denoising window instead of a single forward.

Unit tests bumped to 13. The new test_onpolicy_rollout_propagates_gradient asserts gen_data.requires_grad and that backward() reaches the student weights through the chosen step. The forward-equivalence and training-step numbers reported above are unchanged (the rollout only changes the on-policy student-generation procedure; the central-difference target math and DMD2 distillation machinery are the same).

[Enderfga]

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 ab1174d adds the multi-step Euler-flow rollout to match AnyFlow's training_rollout (now 13/13 unit tests passing, no regression on DMD2/MeanFlow).

Happy to address any further feedback whenever you have a slot — thanks again for the early review!

[juliusberner]

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:

1. Do you think we could re-use more functionality from our MeanFlow implementation to not duplicate code?
2. Did you also try to train for a few hundred iterations (with a small batchsize) to check convergence?

[cxlcl]

@Enderfga Thanks a lot for the PR and its follow-up!
Are the config tuned for Anyflow, or is it only for demo and needs further tuning?

[Enderfga]

Thanks @juliusberner and @cxlcl — pushed commit 1671bb2 that addresses (1) and adds the on-policy config; (2) is scoped explicitly below.

(1) MeanFlow code sharing. Extracted _fuse_r_embedding on the Wan transformer (fastgen/networks/Wan/network.py) so the additive (MeanFlow) and gated (AnyFlow) fusion modes sit side-by-side in one helper instead of being an inline branch inside classify_forward_prepare. Both paths share the same r_embedder.time_embedder / time_proj / act_fn modules — the refactor makes that explicit and shrinks the call site to 3 lines. Forward semantics are bit-identical to the previous commit across all three encoder_depth cases; all 13 AnyFlow + 3 MeanFlow unit tests pass.

(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:

• forward parity within bf16 noise on the released 1.3B / 14B HF ckpts: 2.8% rel mean diff, 8.7e-2 max abs diff
• single-step training parity vs AnyFlow's train_bidirection central-difference math at 4.07% (1.3B) / 3.00% (14B) rel loss diff on real weights, same seed and mini-batch through both code paths
• stub-network central-difference target match to 1.9e-5 max abs in fp32 — the math is reproduced bit-for-bit; the bf16 numbers above are model-noise floor, not algorithm drift

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. fastgen/configs/experiments/WanT2V/config_anyflow.py is the paper's Stage 2 pretrain config 1-for-1 (shift=5, weight_type=beta08, ε=5e-3, lr=5e-5, 6k iter, batch_size_global=32, the 4-step Wan t_list, GAN off in pretrain). config_anyflow_onpolicy.py is added in this commit for Stage 3 (lr=2e-6, 1200 iter, GAN on at the DMD2-default 0.03). Caveat: the paper's Stage 3 uses a rank-256 LoRA adapter, but FastGen does not ship a PEFT/LoRA training path today, so the on-policy config does a full-rank fine-tune on top of a Stage 2 checkpoint — noted in the config docstring.

[juliusberner]

Thanks again for the PR, I did a code review and added several comments.

[juliusberner]

This should live in networks/Wan/network.py, since it's Wan-specific.

[Enderfga]

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.

[Enderfga]

Thanks again for the careful review — all eight comments should be addressed now (replies in the threads below, changes in 3c98dd1..2be625a).

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 trainer_wan_anyflow_pretrain.py / trainer_wan_anyflow_onpolicy.py actually do. Fixed in the same commits:

• The pretrain loss was missing the scale_weight rebalancing (non-diffusion sample losses rescaled to the global diffusion-loss mean), and I had wired the guidance through MeanFlow's eq.-19 target-side fusion, while the reference fuses on the prediction side ((u_cond + (g-1)·u_uncond) / g regressed against the raw velocity, with the uncond branch queried at the same (t, r)). Both are now opt-in flags on MeanFlow, so its defaults are untouched. The bucket assignment also now follows the reference's deterministic global partition (my per-rank binomial draws were noticeably biased at batch_size_per_gpu=1), with the consistency bucket at r=0.
• On the on-policy side I had misread training_rollout: it doesn't run N Euler steps with grad on one — it compresses the rollout into at most three flow-map forwards (jump to t_g, one fine step, jump to 0) with gradient through all of them, and samples the NFE per iteration from [2, 4, 8, 16, 50]. Also brought over cotrain_forward_kl (the Stage-2 loss co-trained at every generator step), dropped the GAN term (the reference "discriminator" is just the fake score — the 0.03 weight in my earlier config came from DMD2 defaults, not from your recipe), switched to 1:1 generator/fake-score updates, and fixed the CFG strength off-by-one (cond + 3·(cond - uncond) needs guidance_scale=4 in FastGen's convention). Optimizer betas, grad clip and EMA now match the yml.

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.

[juliusberner]

Why do we need this?

[Enderfga]

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.

[juliusberner]

@coderabbitai review

[coderabbitai]

✅ Action performed

Review finished.

Note: CodeRabbit is an incremental review system and does not re-review already reviewed commits. This command is applicable only when automatic reviews are paused.

[coderabbitai]

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 win

Fix formatting to pass CI: add trailing newline.

The pipeline failure indicates ruff format --check would 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 win

Fix formatting to pass CI.

The pipeline reports that ruff format --check failed 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

📥 Commits

Reviewing files that changed from the base of the PR and between 123e6a2 and 2be625a.

📒 Files selected for processing (11)

• fastgen/configs/experiments/WanT2V/config_anyflow.py
• fastgen/configs/experiments/WanT2V/config_anyflow_onpolicy.py
• fastgen/configs/methods/config_anyflow.py
• fastgen/configs/methods/config_mean_flow.py
• fastgen/methods/__init__.py
• fastgen/methods/consistency_model/mean_flow.py
• fastgen/methods/distribution_matching/README.md
• fastgen/methods/distribution_matching/anyflow.py
• fastgen/networks/EDM/network.py
• fastgen/networks/Wan/network.py
• tests/test_anyflowmodel.py

[coderabbitai]

⚠️ Potential issue | 🟡 Minor | ⚡ Quick win

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.

Suggested change

	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

[coderabbitai]

⚠️ Potential issue | 🟠 Major | ⚡ Quick win

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.

Suggested change

	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.

[Enderfga]

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 set_timesteps excludes — aligned now, which is also what made the uniform special case removable. Tests are 29/29 (AnyFlow + MeanFlow + DMD2), ruff==0.6.9 format/lint clean.

[greptile-apps]

Greptile Summary

This PR adds AnyFlow support for any-step video diffusion. The main changes are:

• New AnyFlow on-policy model and Wan experiment configs.
• MeanFlow extensions for AnyFlow timestep buckets, guidance fusion, and loss rebalancing.
• Wan gated r-embedding support and AnyFlow checkpoint key remapping.
• Tests for pretrain behavior, rollout behavior, and scheduler paths.

Confidence Score: 5/5

This looks safe to merge.

• No blocking issues found in the changed code.
• The latest distributed bucket and loss-reduction changes use matching collective paths across ranks.
• Guidance fusion now fails early for invalid scale or missing negative conditioning.

Important Files Changed

Filename	Overview
fastgen/methods/consistency_model/mean_flow.py	Adds AnyFlow bucket sampling, guidance fusion, timestep weighting, and loss rebalancing.
fastgen/methods/distribution_matching/anyflow.py	Adds the AnyFlow on-policy model with compressed flow-map rollout and cotrained flow-map loss.
fastgen/networks/Wan/network.py	Adds gated r-embedding fusion, r=t defaulting for gated flow-map networks, and AnyFlow checkpoint remapping.
fastgen/networks/EDM/network.py	Defaults r to t for dual-timestep EDM networks.
fastgen/configs/experiments/WanT2V/config_anyflow.py	Adds the Wan AnyFlow pretrain experiment config.
fastgen/configs/experiments/WanT2V/config_anyflow_onpolicy.py	Adds the Wan AnyFlow on-policy distillation experiment config.

_{Reviews (5): Last reviewed commit: "anyflow: fix stale docs" | Re-trigger Greptile}

[greptile-apps]

Accumulated Batch Buckets Disappear

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.

[Enderfga]

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.

[greptile-apps]

Uneven Rank Batches Break Gather

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.

[Enderfga]

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.

[greptile-apps]

Microbatch Buckets Remain This still sizes the AnyFlow buckets from the per-forward local batch, not the accumulated training batch. With the added Wan config using 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.

[Enderfga]

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).

[greptile-apps]

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.

[juliusberner]

@Enderfga thanks for the updated PR, I will take a look in the next few days!