loss NAN when load original zimage weights

[pyx-ustc]

While reproducing TwinFlow training on Z-Image, we found that when the transformer loads the original weights , the loss becomes NaN. After debugging, we traced the issue to the forward function in transformer_z_image.py where the NaNs are introduced by the variables at lines 95–96: t_emb = t_emb + \ t_emb_2 * delta_t_abs.unsqueeze(1)

[QAQdev]

Did you use BF16 training? We suggest using mixed precision training, otherwise this could introduce NAN when t=tt:

TwinFlow/src/methodes/twinflow/twinflow.py

Line 257 in 4250938

cof_r = torch.where(mask, 1 / (t - tt), torch.ones_like(t) * 100)

[pyx-ustc]

1. I am using BF16 mixed-precision training：accelerate launch --mixed_precision='bf16' , and the training loop is

2. When the transformer loads the pretrained weights you provided, the loss remains normal and no NaNs are observed.

[QAQdev]

This problem might come from initialization of temb_2, since we add this new embedder

[pyx-ustc]

1. Yes, that was the cause. After initializing t_embedder_2, the loss is also normal when the transformer loads the original pretrained weights.

2. Below is the loss curve for the first ~200 training steps when loading the original weights. It appears to oscillate—does this behavior match your expectations?

3. Additionally, if possible, would you mind sharing the loss curve you observed during your reproduction experiment on zimage? It would be very helpful as a reference.

[QAQdev]

2. Yes, the loss doesn't decrease much.
3. I did not plot the loss curve, However, based on the curve you provided, could you try reducing the learning rate a bit and see what happens? We used 1e-5 during training.