From aa90f544e3b82317a2073501653d9fa263172d49 Mon Sep 17 00:00:00 2001 From: Junjie Mao Date: Wed, 1 Oct 2025 11:42:51 +0800 Subject: [PATCH] DeepCompile: Fix IPG bucket clearing (#7610) PR #6993 replaces the flat IPG buffers with a dict maintaining type-indexed buckets. The member is also renamed from `_ipg_bucket_flat_buffer` to `ipg_buckets`. Update the bucket clearing logic in `init_z3` accordingly. Signed-off-by: Junjie Mao --- deepspeed/compile/init_z3.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/deepspeed/compile/init_z3.py b/deepspeed/compile/init_z3.py index dfacbd8a0..2b9d404b3 100644 --- a/deepspeed/compile/init_z3.py +++ b/deepspeed/compile/init_z3.py @@ -23,8 +23,8 @@ def init_z3(engine, backend, compile_config, compile_kwargs, schedule=None): optimizer = engine.optimizer use_opt = not isinstance(optimizer, DeepSpeedZeRoOffload) - if use_opt and hasattr(optimizer, '_DeepSpeedZeroOptimizer_Stage3__ipg_bucket_flat_buffer'): - optimizer._DeepSpeedZeroOptimizer_Stage3__ipg_bucket_flat_buffer = None + if use_opt and hasattr(optimizer, "ipg_buckets"): + optimizer.ipg_buckets.clear() get_accelerator().empty_cache() dc = get_deepcompile_handle()