@mrluin, to optimize parameter partitioning, we trace the sequence of submodule/parameter fetches during a training/inference iteration. We store this fetch sequence in a trace cache. The trace cache allows us to prefetch parameters ahead of use and overlap fetch with computation for future iterations with the same fetch sequence. The message indicates that we detected a change in the fetch sequence, and so the fetching optimizations will be disabled.

I don't there is anything to be fixed since the message is indicating that optimization is being disabled to preserve correctness. Are you observing any incorrectness in your model results?

I am not saying there is an error in your training. All this message is indicating is that your module sequence is changing across training iterations, so you can decide whether this change ins intentional or not.

Here are two ideas for checking zero-3 correctness:

1. Compare the loss results with zero-2 run (i.e., without parameter partitioning). In this case, you may need to reduce model or batch size to fit with zero-2.
2. Disable the prefetching optimizations of zero-3, which will disable the trace cache, setting the following in your ds_config:

"stage3_prefetch_bucket_size": 0,
"stage3_max_live_parameters": 0,
"stage3_max_reuse_distance": 0,

@yES, MoE could cause this issue since different iterations could execute different expert modules. I don't believe we have evaluated MoE with zero-3, so I am not sure if that works correctly. However, our zero-2 works well with MoE.

@mrluin Hello. I encountered the similar warning when implementing deepspeed in torch lightning. For my project, this error was because the torchmetrics override Comparison Operators like == and !=, and then produce a CompositionalMetric rather than boolean result. Consequently, the error occurs at line 169 in partitioned_param_coordinator.py.

Hope my experience helps.