[BUG] Training with Zero-2 slow down a lot from v0.18.4 to v0.18.5

[jue-jue-zi]

Describe the bug
Training with Zero-2 slow down a lot from v0.18.4 to v0.18.5 / v0.18.6.

ds_report output

--------------------------------------------------
DeepSpeed C++/CUDA extension op report
--------------------------------------------------
NOTE: Ops not installed will be just-in-time (JIT) compiled at
      runtime if needed. Op compatibility means that your system
      meet the required dependencies to JIT install the op.
--------------------------------------------------
JIT compiled ops requires ninja
ninja .................. [OKAY]
--------------------------------------------------
op name ................ installed .. compatible
--------------------------------------------------
 [WARNING]  async_io requires the dev libaio .so object and headers but these were not found.
 [WARNING]  async_io: please install the libaio-devel package with yum
 [WARNING]  If libaio is already installed (perhaps from source), try setting the CFLAGS and LDFLAGS environment variables to where it can be found.
async_io ............... [NO] ....... [NO]
fused_adam ............. [NO] ....... [OKAY]
cpu_adam ............... [NO] ....... [OKAY]
cpu_adagrad ............ [NO] ....... [OKAY]
cpu_lion ............... [NO] ....... [OKAY]
dc ..................... [NO] ....... [OKAY]
 [WARNING]  Please specify CUTLASS location directory as environment variable CUTLASS_PATH
 [WARNING]  Possible values are: a path, DS_IGNORE_CUTLASS_DETECTION and DS_USE_CUTLASS_PYTHON_BINDINGS
evoformer_attn ......... [NO] ....... [NO]
 [WARNING]  FP Quantizer is using an untested triton version (3.6.0), only 2.3.(0, 1) and 3.0.0 are known to be compatible with these kernels
fp_quantizer ........... [NO] ....... [NO]
fused_lamb ............. [NO] ....... [OKAY]
fused_lion ............. [NO] ....... [OKAY]
 [WARNING]  gds requires the dev libaio .so object and headers but these were not found.
 [WARNING]  gds: please install the libaio-devel package with yum
 [WARNING]  If libaio is already installed (perhaps from source), try setting the CFLAGS and LDFLAGS environment variables to where it can be found.
gds .................... [NO] ....... [NO]
transformer_inference .. [NO] ....... [OKAY]
inference_core_ops ..... [NO] ....... [OKAY]
cutlass_ops ............ [NO] ....... [OKAY]
quantizer .............. [NO] ....... [OKAY]
ragged_device_ops ...... [NO] ....... [OKAY]
ragged_ops ............. [NO] ....... [OKAY]
random_ltd ............. [NO] ....... [OKAY]
 [WARNING]  sparse_attn requires a torch version >= 1.5 and < 2.0 but detected 2.10
 [WARNING]  using untested triton version (3.6.0), only 1.0.0 is known to be compatible
sparse_attn ............ [NO] ....... [NO]
spatial_inference ...... [NO] ....... [OKAY]
transformer ............ [NO] ....... [OKAY]
stochastic_transformer . [NO] ....... [OKAY]
utils .................. [NO] ....... [OKAY]
--------------------------------------------------
DeepSpeed general environment info:
torch install path ............... ['/data/miniconda3/envs/runtime/lib/python3.12/site-packages/torch']
torch version .................... 2.10.0+cu126
deepspeed install path ........... ['/data/miniconda3/envs/runtime/lib/python3.12/site-packages/deepspeed']
deepspeed info ................... 0.18.6, unknown, unknown
torch cuda version ............... 12.6
torch hip version ................ None
nvcc version ..................... 12.6
deepspeed wheel compiled w. ...... torch 0.0, cuda 0.0
shared memory (/dev/shm) size .... 1005.66 GB

Screenshots

System info (please complete the following information):

• OS: CentOS
• GPU count and types: one machine with x8 H20s
• Interconnects (if applicable): NVLINK
• Python version: 3.12.12

Launcher context
accelerator + deepspeed zero2

compute_environment: LOCAL_MACHINE
debug: false
deepspeed_config:
  gradient_clipping: auto
  offload_optimizer_device: none
  offload_param_device: none
  zero3_init_flag: false
  zero_stage: 2
distributed_type: DEEPSPEED
downcast_bf16: 'no'
enable_cpu_affinity: false
machine_rank: 0
main_training_function: main
mixed_precision: 'no'
num_machines: 1
num_processes: 8
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false

[tohtana]

Hi @jue-jue-zi ,
Thank you for the report! Can you share the model and hyperparameters (batch size and sequence length)?

[tohtana]

Hi @jue-jue-zi,
I was able to reproduce the slow down and opened a PR #7886 (Fixed) to address it.
Can you try this fix?

[jue-jue-zi]

Hi, @tohtana
Thanks for the PRs. Is this meant to follow up on #7886 or #7882?

[tohtana]

@jue-jue-zi Sorry, I meant #7886 (fixed the comment)

[jue-jue-zi]

@tohtana I checkout v0.18.6 and cherry-pick e1b99f7^..afa45cd (#7886). It works!

[tohtana]

Thank you @jue-jue-zi! We will merge and release it asap.

[tohtana]

v0.18.7 was released and it should fix the issue. Let me close this. We appreciate you reporting this issue! @jue-jue-zi

[jue-jue-zi]

@tohtana Thanks for your contributions. After upgrading to v0.18.7, it seems that zero0 is slower than zero1 / zero2 for single machine training?

[tohtana]

@jue-jue-zi In theory, ZeRO-1 has more overhead than ZeRO-0, but in my experience, DeepSpeed hasn't focused much on the ZeRO-0 path. There might still be some underlying inefficiencies there.