Share Megatron worker loop

Author: Kovbo

src/art/megatron/shared.py

@@ -6,7 +6,7 @@
 from pathlib import Path
 import shutil
 import time
-from typing import Any
+from typing import Any, Callable
 
 from megatron.core import parallel_state as ps
 import torch
@@ -15,7 +15,7 @@
 from ..preprocessing.pack import PackedTensors, packed_tensors_from_dir
 from .finalize_grads import finalize_model_grads_extended
 from .flex_attention import create_shared_prefix_attention_state
-from .jobs import MegatronSFTTrainingJob, MegatronTrainingJob
+from .jobs import DEFAULT_JOBS_DIR, MegatronSFTTrainingJob, MegatronTrainingJob
 from .offload import clear_optimizer_state
 from .train import (
     DEFAULT_MODEL_IDENTIFIER,
@@ -40,6 +40,7 @@
 save_file = safetensors_torch.save_file
 
 MegatronTrainContext = TrainingRuntime
+MegatronJob = MegatronTrainingJob | MegatronSFTTrainingJob
 
 
 def create_megatron_train_context(
@@ -48,6 +49,50 @@ def create_megatron_train_context(
     return build_training_runtime(model_identifier=model_identifier)
 
 
+def run_megatron_worker_loop(
+    ctx: MegatronTrainContext,
+    *,
+    supports_sft: bool,
+    wait_until_ready: Callable[[], None] | None = None,
+    before_job: Callable[[], None] | None = None,
+    after_job: Callable[[], None] | None = None,
+) -> None:
+    while True:
+        torch.distributed.barrier()  # type: ignore[possibly-missing-attribute]
+        os.makedirs(DEFAULT_JOBS_DIR, exist_ok=True)
+        job_names = sorted(
+            job_name
+            for job_name in os.listdir(DEFAULT_JOBS_DIR)
+            if job_name.endswith(".json")
+        )
+        if not job_names:
+            time.sleep(1)
+            continue
+
+        if wait_until_ready is not None:
+            wait_until_ready()
+        if before_job is not None:
+            before_job()
+
+        job_path = os.path.join(DEFAULT_JOBS_DIR, job_names[0])
+        job = _load_megatron_job(job_path, supports_sft=supports_sft)
+        print0(ctx.rank, "Loaded job from", job_path)
+        print0(ctx.rank, "Job:", job)
+
+        try:
+            _run_megatron_job(ctx, job)
+        finally:
+            if after_job is not None:
+                after_job()
+
+        finalize_megatron_job(
+            ctx,
+            job_path=job_path,
+            log_path=job.log_path,
+            cleanup_path=_job_cleanup_path(job),
+        )
+
+
 def run_megatron_rl_job(
     ctx: MegatronTrainContext,
     job: MegatronTrainingJob,
@@ -254,6 +299,29 @@ def run_megatron_sft_job(
         torch.cuda.empty_cache()
 
 
+def _load_megatron_job(job_path: str, *, supports_sft: bool) -> MegatronJob:
+    with open(job_path, "rb") as handle:
+        job_data = json.loads(handle.read())
+    if job_data.get("job_type") == "sft":
+        if not supports_sft:
+            raise NotImplementedError("SFT jobs are not supported in this worker loop")
+        return MegatronSFTTrainingJob.model_validate(job_data)
+    return MegatronTrainingJob.model_validate(job_data)
+
+
+def _run_megatron_job(ctx: MegatronTrainContext, job: MegatronJob) -> None:
+    if isinstance(job, MegatronSFTTrainingJob):
+        run_megatron_sft_job(ctx, job)
+        return
+    run_megatron_rl_job(ctx, job)
+
+
+def _job_cleanup_path(job: MegatronJob) -> str:
+    if isinstance(job, MegatronSFTTrainingJob):
+        return job.sft_data_dir
+    return job.disk_packed_tensors["dir"]
+
+
 def merge_lora_adapter(lora_path: str) -> None:
     base_dir = Path(lora_path)
     shard_filenames = sorted(base_dir.glob("adapter_model-*-of-*.safetensors"))

src/art/megatron/train.py

@@ -25,9 +25,7 @@
 from art.megatron.finalize_grads import finalize_model_grads_extended
 from art.megatron.flex_attention import create_shared_prefix_attention_state
 from art.megatron.jobs import (
-    DEFAULT_JOBS_DIR,
     DEFAULT_VLLM_WAKE_LOCK_PATH,
-    MegatronTrainingJob,
 )
 from art.megatron.lora import apply_lora_adapters
 from art.megatron.offload import (
@@ -562,45 +560,23 @@ def run_training_step(
 def _run_service_loop(runtime: TrainingRuntime) -> None:
     offload_state = OffloadState()
     offload_to_cpu(runtime.model, runtime.optimizer, runtime.rank, offload_state)
+    from .shared import run_megatron_worker_loop
 
-    while True:
-        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)
-        job_names = sorted(
-            job_name
-            for job_name in os.listdir(DEFAULT_JOBS_DIR)
-            if job_name.endswith(".json")
-        )
-        if not job_names:
-            time.sleep(1)
-            continue
-
+    def wait_until_ready() -> None:
         while os.path.exists(DEFAULT_VLLM_WAKE_LOCK_PATH):
             time.sleep(0.2)
 
-        reload_to_gpu(runtime.model, runtime.optimizer, runtime.rank, offload_state)
-
-        job_name = job_names[0]
-        job_path = os.path.join(DEFAULT_JOBS_DIR, job_name)
-        with open(job_path, "rb") as handle:
-            job = MegatronTrainingJob.model_validate_json(handle.read())
-
-        print0(runtime.rank, "Loaded job from", job_path)
-        print0(runtime.rank, "Job:", job)
-        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"],
-        )
+    run_megatron_worker_loop(
+        runtime,
+        supports_sft=False,
+        wait_until_ready=wait_until_ready,
+        before_job=lambda: reload_to_gpu(
+            runtime.model, runtime.optimizer, runtime.rank, offload_state
+        ),
+        after_job=lambda: offload_to_cpu(
+            runtime.model, runtime.optimizer, runtime.rank, offload_state
+        ),
+    )
 
 
 def main() -> None: