Fix: deepspeed_utils.py

The function for handling data imbalance has been renamed to "deepspeed_join," and a bug where this function was not called has been fixed.

Author: teamtee

src/slam_llm/utils/deepspeed_utils.py

@@ -107,7 +107,7 @@ def decorated_main(cfg_passthrough: Optional[DictConfig] = None) -> Any:
 
     return main_decorator
 
-def slam_join(group_join):
+def deepspeed_join(group_join):
     """
     Copy from wenet:https://github.com/wenet-e2e/wenet/blob/main/wenet/utils/executor.py#L64
     """
@@ -204,6 +204,8 @@ def train(
             else:
                 pbar = tqdm(colour="blue", desc=f"Training Epoch: {epoch+1}", dynamic_ncols=True)
             for step, batch in enumerate(train_dataloader):
+                if train_config.batching_strategy == "dynamic" and deepspeed_join(group_join):
+                    break
                 for key in batch.keys():
                     batch[key] = (
                         batch[key].to(local_rank).half()