(fix): handle is_last_step in DeepSpeedTrainStrategy.save_checkpoint

Previously, is_last_step was passed via **kwargs and transparently
forwarded to DeepSpeedEngine.save_checkpoint(), which does not accept
this argument, causing a TypeError at checkpoint time.

Fix by explicitly declaring is_last_step=None in the signature (consistent
with megatron_strategy and fsdp2_strategy), and applying the same
async_upload guard logic as the other strategies.

Signed-off-by: Xuchun Shang <xuchun.shang@linux.alibaba.com>

Author: XucSh

roll/distributed/strategy/deepspeed_strategy.py

@@ -503,7 +503,7 @@ def train_step(
                     self.offload_states(include=[OffloadStateType.optimizer_states], non_blocking=True)
         return metrics
 
-    def save_checkpoint(self, save_dir, global_step, ckpt_id, tag="checkpoint", local_state_path=None, **kwargs):
+    def save_checkpoint(self, save_dir, global_step, ckpt_id, tag="checkpoint", local_state_path=None, is_last_step=None, **kwargs):
         """
         save ckpt/hf model/tokenizer to local dir
         save_dir/actor_train/{hf files}
@@ -540,7 +540,7 @@ def save_checkpoint(self, save_dir, global_step, ckpt_id, tag="checkpoint", loca
             # save tokenizer
         self.model.save_checkpoint(save_dir, tag=tag, **kwargs)
 
-        if self.worker_config.checkpoint_config.get("async_upload", True):
+        if self.worker_config.checkpoint_config.get("async_upload", True) and not is_last_step:
             self.thread_executor.submit(self.checkpoint_manager.upload, ckpt_id=ckpt_id, local_state_path=local_state_path)
         else:
             self.checkpoint_manager.upload(ckpt_id=ckpt_id, local_state_path=local_state_path)