[REQUEST] Dynamic/Adaptive Prefetching Window for ZeRO-3

Problem I have Faced
I am frustrated by the fact that __prefetch_bucket_sz in ZeRO-3 is a static, user-configured value. In large-scale training, network bandwidth can fluctuate and computation step times vary. A static bucket size creates two distinct bottlenecks:

Compute Starvation: If the bucket is too small, the network cannot fetch parameters fast enough, leading to CPU/GPU blocking when param.wait() is called.
Memory Pressure: If the bucket is too large, it aggressively materializes too many parameters in advance, causing memory spikes, eviction thrashing, or OOM errors.

Solution I thought:
I would like to introduce an Adaptive/Dynamic Prefetching Window to the PartitionedParameterCoordinator.

The proposed solution involves:

Stall Tracking: Using host-side time.time() around the .wait() loop to measure communication wait time, and tracking compute time as the duration between consecutive module fetches.
Dynamic Scaling: Calculating a moving average of the stall_ratio (wait time / compute time).
Adjustment Policy:
Scale Up: If stall_ratio > 0.05, we scale __prefetch_bucket_sz by a factor (e.g., 1.2x), bounding it to a maximum of __max_n_available_params * 0.5. This aggressively prefetches further ahead in the trace to hide latency.
Scale Down: If stall_ratio is near 0.0, we slowly decay the bucket size down to the baseline size to conserve memory pressure when the network is fast enough.

Alternatives I saw
The only current alternative is manually tuning prefetch_bucket_size for every specific hardware configuration and model, which does not account for real-time network fluctuations during training. We also considered tying prefetch scaling to CUDA events directly, but determined host-side timing avoids injecting unnecessary synchronization overhead.

Additional context Per the new feature contribution guidelines, I am planning the following experiments to demonstrate effectiveness before submitting the Pull Request:

Task: Fine-tuning BERT on SQuAD (NLP) and ResNet on CIFAR (CV) using the DeepSpeedExamples repo.
Metrics to validate:
Total throughput improvements (samples/sec) compared to the static baseline.
Peak GPU memory allocation (to ensure the dynamic scaling respects memory limits).
Validation loss over a subset of steps to verify convergence correctness is maintained.
I have already prototyped this logic locally and would appreciate feedback from the core team on this design before opening the Pull Requests!

Start with the PartitionedParameterCoordinator and inspect the __prefetch_bucket_sz, __max_n_available_params, and .wait() loop mentioned in the proposal. Evaluate the host-side stall and compute timing approach, then validate the design with BERT on SQuAD and ResNet on CIFAR in DeepSpeedExamples. Done means demonstrating throughput and memory effects while preserving validation loss and convergence correctness.