Rename Megatron merge helper

Author: Kovbo

src/art/megatron/service.py

@@ -33,7 +33,7 @@
     DEFAULT_VLLM_WAKE_LOCK_PATH,
     MegatronTrainingJob,
 )
-from .shared import merge_sharded_lora_adapter
+from .shared import merge_lora_adapter
 
 safetensors = importlib.import_module("safetensors")
 safe_open = safetensors.safe_open
@@ -303,7 +303,7 @@ async def train(
                     for line in lines:
                         if line := line.strip():
                             if line == "all done":
-                                merge_sharded_lora_adapter(lora_path)
+                                merge_lora_adapter(lora_path)
                                 os.remove(DEFAULT_TRAINING_LOG_PATH)
                                 break
                             num_lines += 1

src/art/megatron/shared.py

@@ -270,7 +270,7 @@ def run_megatron_sft_job(
         torch.cuda.empty_cache()
 
 
-def merge_sharded_lora_adapter(lora_path: str) -> None:
+def merge_lora_adapter(lora_path: str) -> None:
     base_dir = Path(lora_path)
     shard_filenames = sorted(base_dir.glob("adapter_model-*-of-*.safetensors"))
     if not shard_filenames: