Fix Megatron job finalization ordering

Author: Kovbo

src/art/megatron/service.py

@@ -29,7 +29,6 @@
 from ..vllm import get_llm, openai_server_task, run_on_workers
 from .jobs import (
     DEFAULT_JOBS_DIR,
-    DEFAULT_TRAINING_LOG_PATH,
     DEFAULT_VLLM_WAKE_LOCK_PATH,
     MegatronTrainingJob,
 )
@@ -277,6 +276,8 @@ async def train(
                 "moe_routing_replay_bundle is only supported for in-process/runtime APIs; "
                 "MegatronService subprocess jobs must use moe_routing_replay_path."
             )
+        log_dir = "/tmp/megatron_training_logs"
+        os.makedirs(log_dir, exist_ok=True)
         job = MegatronTrainingJob(
             lora_path=lora_path,
             optimizer_state_path=self._optimizer_state_path,
@@ -285,6 +286,9 @@ async def train(
             experimental_config=_config,
             moe_routing_replay_path=_config.get("moe_routing_replay_path"),
             moe_routing_replay_strict=_config.get("moe_routing_replay_strict", True),
+            log_path=os.path.join(
+                log_dir, f"{datetime.datetime.now().isoformat()}.jsonl"
+            ),
         )
         job_path = os.path.join(
             DEFAULT_JOBS_DIR,
@@ -297,14 +301,14 @@ async def train(
         while True:
             await asyncio.sleep(0.1)
             try:
-                with open(DEFAULT_TRAINING_LOG_PATH, "a+") as log_file:
+                with open(job.log_path, "a+") as log_file:
                     log_file.seek(0)
                     lines = log_file.readlines()[num_lines:]
                     for line in lines:
                         if line := line.strip():
                             if line == "all done":
                                 merge_lora_adapter(lora_path)
-                                os.remove(DEFAULT_TRAINING_LOG_PATH)
+                                os.remove(job.log_path)
                                 break
                             num_lines += 1
                             yield json.loads(line)

src/art/megatron/shared.py

@@ -51,8 +51,6 @@ def create_megatron_train_context(
 def run_megatron_rl_job(
     ctx: MegatronTrainContext,
     job: MegatronTrainingJob,
-    *,
-    job_path: str | None = None,
 ) -> None:
     packed_tensors = None
     adapter_model = None
@@ -125,12 +123,6 @@ def run_megatron_rl_job(
             lora_path=job.lora_path,
             optimizer_state_path=job.optimizer_state_path,
         )
-        _complete_job(
-            ctx,
-            job_path=job_path,
-            log_path=job.log_path,
-            cleanup_path=job.disk_packed_tensors["dir"],
-        )
     finally:
         if packed_tensors is not None:
             del packed_tensors
@@ -149,8 +141,6 @@ def run_megatron_rl_job(
 def run_megatron_sft_job(
     ctx: MegatronTrainContext,
     job: MegatronSFTTrainingJob,
-    *,
-    job_path: str | None = None,
 ) -> None:
     adapter_model = None
 
@@ -257,12 +247,6 @@ def run_megatron_sft_job(
             lora_path=job.lora_path,
             optimizer_state_path=job.optimizer_state_path,
         )
-        _complete_job(
-            ctx,
-            job_path=job_path,
-            log_path=job.log_path,
-            cleanup_path=job.sft_data_dir,
-        )
     finally:
         if adapter_model is not None:
             del adapter_model
@@ -381,19 +365,11 @@ def _load_lora_and_optimizer(
     optimizer_state_path: str,
 ) -> dict[str, torch.Tensor]:
     adapter_model_path = os.path.join(lora_path, "adapter_model.safetensors")
-    if os.path.exists(adapter_model_path):
-        print0(ctx.rank, "Loading adapter model from", adapter_model_path)
-        adapter_model = load_file(adapter_model_path)
-        load_adapter_into_model(ctx.model, adapter_model, ctx.optimizer)
-    else:
-        print0(ctx.rank, "No adapter model found at", adapter_model_path)
-        adapter_model = {}
-        with torch.no_grad():
-            for chunk in ctx.model:
-                for module in chunk.modules():
-                    if hasattr(module, "reset_lora_parameters"):
-                        module.reset_lora_parameters()  # type: ignore[attr-defined]
-        ctx.optimizer.reload_model_params()
+    if not os.path.exists(adapter_model_path):
+        raise FileNotFoundError(f"No adapter model found at {adapter_model_path}")
+    print0(ctx.rank, "Loading adapter model from", adapter_model_path)
+    adapter_model = load_file(adapter_model_path)
+    load_adapter_into_model(ctx.model, adapter_model, ctx.optimizer)
 
     optimizer_shard_path = os.path.join(
         optimizer_state_path,
@@ -449,7 +425,7 @@ def _save_lora_and_optimizer(
     torch.save(ctx.optimizer.state_dict(), optimizer_shard_path)
 
 
-def _complete_job(
+def finalize_megatron_job(
     ctx: MegatronTrainContext,
     *,
     job_path: str | None,
@@ -462,9 +438,10 @@ def _complete_job(
 
     if job_path is not None and os.path.exists(job_path):
         os.remove(job_path)
+    if os.path.exists(cleanup_path):
+        shutil.rmtree(cleanup_path)
     with open(log_path, "a+", encoding="utf-8") as log_file:
         log_file.write("all done\n")
-    shutil.rmtree(cleanup_path)
 
 
 def _placeholder_attention_mask(device: torch.device) -> torch.Tensor:

src/art/megatron/train.py

@@ -26,7 +26,6 @@
 from art.megatron.flex_attention import create_shared_prefix_attention_state
 from art.megatron.jobs import (
     DEFAULT_JOBS_DIR,
-    DEFAULT_TRAINING_LOG_PATH,
     DEFAULT_VLLM_WAKE_LOCK_PATH,
     MegatronTrainingJob,
 )
@@ -565,7 +564,7 @@ def _run_service_loop(runtime: TrainingRuntime) -> None:
     offload_to_cpu(runtime.model, runtime.optimizer, runtime.rank, offload_state)
 
     while True:
-        from .shared import run_megatron_rl_job
+        from .shared import finalize_megatron_job, run_megatron_rl_job
 
         torch.distributed.barrier()  # ty: ignore[possibly-missing-attribute]
         os.makedirs(DEFAULT_JOBS_DIR, exist_ok=True)
@@ -590,8 +589,18 @@ def _run_service_loop(runtime: TrainingRuntime) -> None:
 
         print0(runtime.rank, "Loaded job from", job_path)
         print0(runtime.rank, "Job:", job)
-        run_megatron_rl_job(runtime, job, job_path=job_path)
-        offload_to_cpu(runtime.model, runtime.optimizer, runtime.rank, offload_state)
+        try:
+            run_megatron_rl_job(runtime, job)
+        finally:
+            offload_to_cpu(
+                runtime.model, runtime.optimizer, runtime.rank, offload_state
+            )
+        finalize_megatron_job(
+            runtime,
+            job_path=job_path,
+            log_path=job.log_path,
+            cleanup_path=job.disk_packed_tensors["dir"],
+        )
 
 
 def main() -> None: