[Bug] Wan MeanFlow zero initialization causes degenerate gradients in r_embedder

[YunkangTao]

Description

In fastgen/configs/experiments/WanT2V/config_mf.py, the Wan MeanFlow model uses:

config.model.net.r_embedder_init = "zero"

This triggers the following code in fastgen/networks/Wan/network.py:

elif embedder_init == "zero":
    for param in embedder.parameters():
        param.data.zero_()

This sets every parameter of the multi-layer r_embedder to zero. As a result, several weight tensors receive zero gradients because of the zero symmetry (only some bias parameters may update). The r-conditioning pathway therefore remains effectively inactive and is difficult to learn.

Reproduction

Run the Wan MeanFlow configuration and inspect the gradients of model.net.transformer.r_embedder after the first backward pass.

The weights of layers such as time_embedder.linear_1 and time_embedder.linear_2.weight remain zero or receive zero gradients.

Expected behavior

The r pathway should initially produce zero output without initializing all of its trainable layers to zero.

Suggested fix

The existing "random" branch appears to implement the intended behavior: it randomly initializes the hidden layers and zero-initializes only the output projections.

Could the "zero" branch be changed to use this initialization strategy, or could the configuration be changed to:

config.model.net.r_embedder_init = "random"

?