From 7c937262492b2ac3e1e6387cb2927a6169747c25 Mon Sep 17 00:00:00 2001 From: Stas Bekman Date: Thu, 1 Sep 2022 14:25:05 -0700 Subject: [PATCH] [ds-inference] fix progress bar when loading the non-sharded checkpoint update the progress bar (fix by @RezaYazdaniAminabadi) - I've just tested it to work. --- deepspeed/module_inject/replace_module.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepspeed/module_inject/replace_module.py b/deepspeed/module_inject/replace_module.py index ccda3c8132b1..f42a973d8b21 100755 --- a/deepspeed/module_inject/replace_module.py +++ b/deepspeed/module_inject/replace_module.py @@ -857,7 +857,6 @@ def replace_fn(child, _policy, layer_id=0): desc=f"Loading {len(checkpoint)} checkpoint shards") for i in range(len(checkpoint)): - sd = [ torch.load(os.path.join(base_dir1, checkpoint[i]), @@ -870,6 +869,7 @@ def replace_fn(child, _policy, layer_id=0): ckpt_type, quantizer, ) + pbar.update(1) else: import gc num_checkpoints = len(ckpt_list) // ckpt_mp_size