[BUG] ZenFlow: loss values becomes NaN after 'update_interval' number of steps.

[karthik2603-theBrogrammer]

I was training a Llama model using the script released in the Deepspeed-Examples repository using ZenFlow. Interestingly, I noticed the loss values becoming NaN after the preconfigured (in the deepspeed config) update_interval number of steps. Ex: the update interval is set to 4 in the deepspeed config, loss becomes nan from the 5th step. The code used is available here: https://github.com/deepspeedai/DeepSpeedExamples/tree/master/training/DeepSpeed-ZenFlow/finetuning

My software configurations and versions:

torch==2.5.0+cu118
Transformers==4.57.3
Cloned the latest version of deepspeed from github
Datasets==4.4.1

The jobs are run with DGX H200 GPUs and AMD EPYC 7742 processor.

I wish to understand why the loss turns NaN and more importantly why it happens only after the update_interval step. Can I please get some help with respect to this ?

[Antlera]

Hi @karthik2603-theBrogrammer Thanks for reporting this issue. I was able to get access to DGX H200 resources to reproduce the problem today.

I noticed that there have been some API changes in the DeepSpeed engine and optimizer, and it’s likely that there is a mismatch between the current ZenFlow optimizer code and the updated engine/optimizer APIs. This may explain why the NaNs appear right after the update_interval boundary.

I also try to reproduce the NaN issue on other hardware, so this does not seem to be H200-specific. I’ll need a bit more time to track down exactly where the mismatch happens and fix it.

I’ll follow up once I have more concrete findings.

[Antlera]

It appears that recent DeepSpeed changes unified the backward path to loss.backward(), bypassing optimizer.backward(). This means ZenFlow’s backward logic is not being triggered, which could explain the NaNs observed after the update_interval step.

[Antlera]

Hi @karthik2603-theBrogrammer, I’ve opened a follow-up PR that fixes the ZenFlow backward mismatch.
If you get a chance, could you help test whether it works on H200? That would be very helpful.