[BUG] `reduce_bucket_size` influences training convergence of Zero2

Describe the bug
I launch deepspeed training for a 600M parameter diffusion model, and only vary reduce_bucket_size.
I tried the following values:

• reduce_bucket_size: 500_000_000 — converges poorly
• reduce_bucket_size: 1_000_000_000 — converges sllightly better in the beginning, but then still worse than Zero Stage 1.
• reduce_bucket_size: 10_000_000 — almost does not converge at all, the losses are several times higher.
• reduce_bucket_size: 1_000_000 — I start getting NaNs for my loss values almost immediately during training.

I use deepspeed 0.14.2 (upgrading to 0.14.5 didn't help).
The rest of my config looks like this:

deepspeed:
  train_micro_batch_size_per_gpu: 1
  steps_per_print: 1
  bf16:
    enabled: True
  checkpoint:
    load_universal: false

  gradient_clipping: 1.0
  communication_data_type: bf16
  data_types: {grad_accum_dtype: fp32}

  zero_optimization:
    stage: 2
    allgather_partitions: True
    reduce_scatter: True
    allgather_bucket_size: 500_000_000
    reduce_bucket_size: 500_000_000
    overlap_comm: True
    contiguous_gradients: True

  wall_clock_breakdown: True
  csv_monitor:
    enabled: true
    output_path: "data/deepspeed_logs/"
    job_name: "${logging.run_name}"

  flops_profiler:
    enabled: true
    profile_step: 3
    module_depth: -1
    top_modules: 1
    detailed: true
    output_file: "results/profiling/${logging.run_name}.flops_profiler.txt"

  # Configuration for activation checkpointing
  activation_checkpointing:
    partition_activations: false
    cpu_checkpointing: false
    contiguous_memory_optimization: false
    number_checkpoints: null
    synchronize_checkpoint_boundary: false
    profile: false

To Reproduce
As per now, I cannot provide a simple reproducible example since it's deep in the internal codebase.
I want to ask you where can I look at to isolate/locate an issue.

Expected behavior
reduce_bucket_size shouldn't influence training at all.

Screenshots

Here is the screenshot with my training losses:

• Greenish Blue: Zero Stage 2, reduce_bucket_size: 500_000_000 (default)
• Blue: Zero Stage 2, reduce_bucket_size: 1_000_000_000
• Purple: Zero Stage 2, reduce_bucket_size: 10_000_000
• Dark Green (NaN losse): Zero Stage 2, reduce_bucket_size: 1_000_000
• Green: Zero Stage 1

System info (please complete the following information):

• Ubuntu 20.04.6 LTS
• 8x A100s 80 gb
• 4x EFA
• Python 3.11.9

Launcher context
I launch via torch.distributed.run

Docker context
It's an internal one, cannot share it.

Additional context
I use activations checkpointing with use_reentrant=True (use_reentrant=False makes the model converge worse), mixed precision training.
For debugging purposes, I tried to do backward with allreduce_gradients=False and then doing optimizer.reduce_gradients(pipeline_parallel=False) manually (to follow Zero Stage 1), but it didn't help.

Start at the Zero Stage 2 gradient-reduction path, comparing the allreduce_gradients=False flow with optimizer.reduce_gradients(pipeline_parallel=False) and the torch.distributed.run launcher context. Build a minimal controlled training case that varies only reduce_bucket_size and mixed-precision settings. Done means training convergence and loss behavior no longer depend materially on reduce_bucket_size.