Clean up prefetched parameters

[tohtana]

Parameters prefetched by ZeRO3 are sometimes not used. This occurs when the actual sub-module execution differs from previous tracing. As a result, the state of the allgather handle for such a parameter remains INFLIGHT, causing functions like empty_partition_cache to detect it and throw an error.
This PR resolves the issue by ensuring that communication finishes and the parameters are freed.

As this issue was mentioned in #6011, this includes the change of the branch. We need to merge #6011 first.

[tjruwase]

Please check if this PR fixes #5828.

[tohtana]

Please check if this PR fixes #5828.

@tjruwase Using this PR branch, the repro in #5828 shows the message below but exits without throwing an error. I think this is expected as the model has a conditional branch and the execution order of modules changes.

Invalidate trace cache @ step 3: expected module 2, but got module 4

[tohtana]

@tjruwase I added the cleaning of the inflight parameter registry in _invalidate_trace as you suggested. This allows us to free the gathered (but unused) parameters earlier. However, I also kept it in reset_step.
This is why we don't detect deviations from the trace when some modules at the end of the trace remain unvisited. The original assertion in reset_step will still be triggered in that case.