Fix lint and typing regressions in shared training refactor

Author: Kovbo

src/art/_backend_training.py

@@ -51,9 +51,7 @@ def build_rl_train_configs(
     }
 
     if allow_training_without_logprobs is not None:
-        dev_config["allow_training_without_logprobs"] = (
-            allow_training_without_logprobs
-        )
+        dev_config["allow_training_without_logprobs"] = allow_training_without_logprobs
     if plot_tensors is not None:
         dev_config["plot_tensors"] = plot_tensors
     if truncated_importance_sampling is not None:
@@ -63,9 +61,7 @@ def build_rl_train_configs(
             scale_learning_rate_by_reward_std_dev
         )
     if logprob_calculation_chunk_size is not None:
-        dev_config["logprob_calculation_chunk_size"] = (
-            logprob_calculation_chunk_size
-        )
+        dev_config["logprob_calculation_chunk_size"] = logprob_calculation_chunk_size
     if num_trajectories_learning_rate_multiplier_power is not None:
         dev_config["num_trajectories_learning_rate_multiplier_power"] = (
             num_trajectories_learning_rate_multiplier_power

src/art/megatron/service.py

@@ -2,6 +2,7 @@
 from dataclasses import asdict, dataclass
 import datetime
 from functools import cached_property
+import importlib
 import json
 import os
 from pathlib import Path
@@ -10,8 +11,6 @@
 from typing import Any, AsyncIterator
 
 from peft.tuners.lora.config import LoraConfig
-from safetensors import safe_open
-from safetensors.torch import load_file, save_file
 import torch
 from vllm import AsyncEngineArgs
 from vllm.lora.request import LoRARequest
@@ -31,6 +30,12 @@
     MegatronTrainingJob,
 )
 
+safetensors = importlib.import_module("safetensors")
+safetensors_torch = importlib.import_module("safetensors.torch")
+safe_open = safetensors.safe_open
+load_file = safetensors_torch.load_file
+save_file = safetensors_torch.save_file
+
 
 @dataclass
 class MegatronService:

src/art/megatron/shared.py

@@ -1,21 +1,25 @@
+from dataclasses import dataclass
 import gc
+import importlib
 import json
 import math
 import os
 import shutil
 import time
-from dataclasses import dataclass
 from typing import Any
 
 from megatron.core import parallel_state as ps
-from safetensors.torch import load_file, save_file
 import torch
 
 from ..loss import loss_fn, shift_tensor
 from ..preprocessing.pack import PackedTensors, packed_tensors_from_dir
 from .flex_attention import create_shared_prefix_attention_state
 from .jobs import MegatronSFTTrainingJob, MegatronTrainingJob
 
+safetensors_torch = importlib.import_module("safetensors.torch")
+load_file = safetensors_torch.load_file
+save_file = safetensors_torch.save_file
+
 
 @dataclass
 class MegatronTrainContext:
@@ -274,7 +278,9 @@ def run_megatron_sft_job(
             update_successful, grad_norm, num_zeros_in_grad = ctx.optimizer.step()
             ctx.optimizer.zero_grad()
 
-            torch.distributed.reduce(batch_loss, dst=0, op=torch.distributed.ReduceOp.SUM)
+            torch.distributed.reduce(
+                batch_loss, dst=0, op=torch.distributed.ReduceOp.SUM
+            )
             avg_loss = batch_loss / num_trainable_tokens
 
             batch_time = time.perf_counter() - batch_start_time
@@ -289,7 +295,9 @@ def run_megatron_sft_job(
                             "loss": avg_loss.item(),
                             "learning_rate": job.learning_rates[batch_idx],
                             "grad_norm": float(grad_norm),
-                            "num_trajectories": float(batch_metadata["num_trajectories"]),
+                            "num_trajectories": float(
+                                batch_metadata["num_trajectories"]
+                            ),
                             "num_trainable_tokens": float(num_trainable_tokens),
                             "tokens_per_second": tokens_per_second,
                         }

src/art/megatron/train.py

@@ -28,7 +28,9 @@
     torch.distributed.barrier()
     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")
+        job_name
+        for job_name in os.listdir(DEFAULT_JOBS_DIR)
+        if job_name.endswith(".json")
     )
     if not job_names:
         time.sleep(1)

src/art/unsloth/service.py

@@ -8,7 +8,7 @@
 import os
 import subprocess
 import sys
-from typing import Any, AsyncIterator
+from typing import Any, AsyncIterator, cast
 
 from trl import GRPOTrainer
 from vllm import AsyncEngineArgs
@@ -18,6 +18,7 @@
 from .. import dev, types
 from ..dev.validate import is_dedicated_mode
 from ..local.checkpoints import get_last_checkpoint_dir
+from ..preprocessing.inputs import TrainInputs
 from ..preprocessing.pack import DiskPackedTensors
 from ..preprocessing.tokenize import SFTBatch
 from ..utils.convert_moe_lora import convert_checkpoint_if_needed
@@ -34,6 +35,7 @@
 
 logger = logging.getLogger(__name__)
 
+
 def save_checkpoint(
     trainer: GRPOTrainer,
     output_dir: str,
@@ -558,7 +560,7 @@ async def train_sft(
 
     @cached_property
     def _state(self) -> UnslothTrainContext:
-        init_args = dict(self.config.get("init_args", {}))
+        init_args = dict(cast(dict[str, Any], self.config.get("init_args") or {}))
         checkpoint_dir = get_last_checkpoint_dir(self.output_dir)
         if checkpoint_dir:
             init_args["model_name"] = checkpoint_dir
@@ -567,8 +569,11 @@ def _state(self) -> UnslothTrainContext:
 
         return create_unsloth_train_context(
             init_args=init_args,
-            peft_args=dict(self.config.get("peft_args", {})),
-            trainer_args=dict(self.config.get("trainer_args", {})),
+            peft_args=cast(dict[str, Any], self.config.get("peft_args") or {}),
+            trainer_args=cast(
+                dict[str, Any],
+                self.config.get("trainer_args") or {},
+            ),
         )
 
     @cached_property

src/art/unsloth/shared.py

@@ -73,7 +73,10 @@ def offload_to_cpu(self) -> None:
         if optimizer is not None and hasattr(optimizer, "state"):
             for param_id, state in optimizer.state.items():
                 for key, value in state.items():
-                    if not isinstance(value, torch.Tensor) or value.device.type != "cuda":
+                    if (
+                        not isinstance(value, torch.Tensor)
+                        or value.device.type != "cuda"
+                    ):
                         continue
                     buffer_key = f"opt_{id(param_id)}_{key}"
                     if (
@@ -108,9 +111,14 @@ def reload_to_gpu(self, device: str = "cuda:0") -> None:
         if optimizer is not None and hasattr(optimizer, "state"):
             for state in optimizer.state.values():
                 for key, value in state.items():
-                    if not isinstance(value, torch.Tensor) or value.device.type != "cpu":
+                    if (
+                        not isinstance(value, torch.Tensor)
+                        or value.device.type != "cpu"
+                    ):
                         continue
-                    gpu_tensor = torch.empty(value.shape, dtype=value.dtype, device=device)
+                    gpu_tensor = torch.empty(
+                        value.shape, dtype=value.dtype, device=device
+                    )
                     gpu_tensor.copy_(value, non_blocking=True)
                     state[key] = gpu_tensor
 
@@ -224,7 +232,10 @@ def create_unsloth_train_context(
         loader_cls.from_pretrained(**init_args),
     )
 
-    if hasattr(model, "peft_config") and getattr(model, "peft_config", None) is not None:
+    if (
+        hasattr(model, "peft_config")
+        and getattr(model, "peft_config", None) is not None
+    ):
         peft_model = cast(peft.peft_model.PeftModelForCausalLM, model)
     else:
         peft_model = cast(
@@ -301,7 +312,9 @@ def _precalculate_new_logprobs(
                         if isinstance(value, torch.Tensor)
                     },
                     pixel_values=packed_tensors["pixel_values"][offset : offset + 1],
-                    image_grid_thw=packed_tensors["image_grid_thw"][offset : offset + 1],
+                    image_grid_thw=packed_tensors["image_grid_thw"][
+                        offset : offset + 1
+                    ],
                     config=config,
                     _config=_config,
                     return_new_logprobs=True,