feat: add NPU device_memory_used and vllm support

Author: UsernameFull

roll/distributed/strategy/strategy.py

@@ -23,6 +23,7 @@ def __init__(self, worker: "Worker"):
         self.worker = worker
         self.model = None
         self.tokenizer = None
+        self.running = False
 
         self.worker_config = self.worker.worker_config
         self.thread_executor: futures.ThreadPoolExecutor = futures.ThreadPoolExecutor(max_workers=5)

roll/pipeline/rlvr/actor_pg_worker.py

@@ -275,6 +275,10 @@ def _compute_topr_loss(self, ratio: torch.Tensor, log_probs: torch.Tensor, old_l
             "topr_negative_total_clipfrac": negative_total_clipped.mean().detach().item(),
             "topr_scores_mean": scores.mean().detach().item(),
             "topr_scores_std": scores.std().detach().item(),
+            "topr_positive_loss": positive_loss,
+            "topr_negative_loss": negative_loss,
+            "topr_weighted_positive_loss": weighted_positive_loss,
+            "topr_weighted_negative_loss": weighted_negative_loss,
         })
 
         return topr_loss

roll/platforms/npu.py

@@ -1,6 +1,8 @@
 from .platform import Platform
 from ..utils.logging import get_logger
 
+import torch
+
 logger = get_logger()
 
 
@@ -74,3 +76,7 @@ def get_vllm_run_time_env_vars(cls, gpu_rank: str) -> dict:
     @classmethod
     def apply_ulysses_patch(cls) -> None:
         return
+
+    @classmethod
+    def device_memory_used(cls) -> None:
+        return torch.npu.mem_get_info()[0]