frozenleaves/DeepSpeed
1
0
Fork 0
mirror of https://github.com/deepspeedai/DeepSpeed.git synced 2025-10-20 15:33:51 +08:00
Code Issues Packages Projects Releases Wiki Activity

DeepCompile: Fix IPG bucket clearing (#7610)

Browse Source 
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 <junjie.mao@linux.alibaba.com>
This commit is contained in:
Junjie Mao
2025-10-01 11:42:51 +08:00
committed by GitHub
parent e32e817306
commit aa90f544e3
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
deepspeed/compile/init_z3.py
View File

@ -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()