From d143141afb1b4a9bb1b38b9de73e94a8a8c5dea6 Mon Sep 17 00:00:00 2001 From: Logan Adams Date: Mon, 30 Dec 2024 09:17:00 -0800 Subject: [PATCH] Formatting fix --- deepspeed/checkpoint/ds_to_universal.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/deepspeed/checkpoint/ds_to_universal.py b/deepspeed/checkpoint/ds_to_universal.py index cbbbef6bcbb4..2b5dc81b2f66 100755 --- a/deepspeed/checkpoint/ds_to_universal.py +++ b/deepspeed/checkpoint/ds_to_universal.py @@ -111,7 +111,10 @@ def _save_checkpoint(file_path, chkpt_sd): def extract_zero_shards(dir, ds_checkpoint, indices_3D): pp_index, tp_index, dp_index = indices_3D - sd = ds_checkpoint.get_zero_checkpoint_state(pp_index=pp_index, tp_index=tp_index, dp_index=dp_index, strip_tensor_paddings=False) + sd = ds_checkpoint.get_zero_checkpoint_state(pp_index=pp_index, + tp_index=tp_index, + dp_index=dp_index, + strip_tensor_paddings=False) # pprint(f"Processing {dp_index=} {pp_index=}, {tp_index=}")