npu grpo fix

Author: addsubmuldiv

src/twinkle/dataset/base.py

@@ -4,7 +4,7 @@
 from dataclasses import dataclass
 from typing import Callable, Type, Union, Dict, Any
 
-from datasets import interleave_datasets, concatenate_datasets, load_dataset, IterableDataset
+from datasets import interleave_datasets, concatenate_datasets, load_dataset, IterableDataset, DatasetDict
 from torch.utils.data import Dataset as TorchDataset
 
 import twinkle
@@ -132,6 +132,19 @@ def _load_dataset(dataset_meta: DatasetMeta, **kwargs):
                 dataset = load_dataset(file_type, data_files=dataset_id, **kwargs)
             else:
                 dataset = HubOperation.load_dataset(dataset_id, subset_name, split, **kwargs)
+
+        if isinstance(dataset, DatasetDict):
+            if split in dataset:
+                dataset = dataset[split]
+            elif 'train' in dataset:
+                dataset = dataset['train']
+            else:
+                available_splits = list(dataset.keys())
+                raise KeyError(
+                    f"Split '{split}' not found for dataset '{dataset_id}'. "
+                    f'Available splits: {available_splits}'
+                )
+
         if isinstance(dataset_meta.data_slice, Iterable) and hasattr(dataset, '__len__'):
 
             iter_list = []

src/twinkle/model/base.py

@@ -138,11 +138,18 @@ def upload_to_hub(self, checkpoint_dir: str, hub_model_id: str, hub_token: Optio
             )
 
     def _try_init_process_group(self):
+        import os
         import torch
         import torch.distributed as dist
         if not dist.is_initialized() and Platform.get_world_size() > 1:
             torch_util.set_device()
             backend = Platform.device_backend()
+            if backend == "hccl":
+                # Keep training-side HCCL sockets on a per-job port layout to
+                # avoid collisions with other jobs on the same host.
+                from twinkle.utils.network import _ensure_hccl_socket_env
+                master_port = int(os.environ.get("MASTER_PORT", "29500"))
+                _ensure_hccl_socket_env(master_port)
             init_kwargs = {
                 "backend": backend,
                 "init_method": "env://",

src/twinkle/sampler/vllm_sampler/vllm_engine.py

@@ -7,6 +7,7 @@
 from twinkle import get_logger
 from twinkle.sampler.base_engine import BaseSamplerEngine
 from twinkle.data_format.sampling import StopReason, SamplingParams, SampleResponse, SampledSequence
+from twinkle.utils.platform import get_vllm_device_uuid
 
 import inspect
 logger = get_logger()
@@ -569,8 +570,9 @@ async def _sync_iter():
         use_gpu_ipc = first_tensor.is_cuda
         use_shm = not use_gpu_ipc
 
-        # Get device UUID for ZMQ handle
-        device_uuid = current_platform.get_device_uuid(0)
+        # Get device UUID for ZMQ handle.
+        # For NPU, this is resolved from `npu-smi info` Bus-Id when needed.
+        device_uuid = get_vllm_device_uuid(0)
         zmq_handle = f"ipc:///tmp/twinkle-ipc-{device_uuid}.sock"
 
         bucket_size = bucket_size_mb << 20

src/twinkle/sampler/vllm_sampler/vllm_sampler.py

@@ -138,7 +138,9 @@ def __init__(
         self.engine: VLLMEngine = self._run_in_loop(
             self._create_engine_async(VLLMEngine, model_id, engine_kwargs)
         )
-        self._run_in_loop(self.engine.engine.collective_rpc("monkey_patch_model"))
+        # NPU platform may trigger triton errors with monkey_patch_model
+        if not Platform.get_platform().device_prefix().upper() == 'NPU':
+            self._run_in_loop(self.engine.engine.collective_rpc("monkey_patch_model"))
 
         VLLMLoraWeights()(self)
 

src/twinkle/sampler/vllm_sampler/vllm_worker_extension.py

@@ -21,6 +21,7 @@
 from twinkle import get_logger
 import torch
 from twinkle.utils.framework import Torch
+from twinkle.utils.platform import get_vllm_device_uuid
 
 logger = get_logger()
 
@@ -64,12 +65,6 @@ def _rebuild_shared_memory(name: str, size: int):
     return tensor, shm
 
 
-def _get_device_uuid(device_id: int) -> str:
-    """Get unique device identifier."""
-    from vllm.platforms import current_platform
-    return current_platform.get_device_uuid(device_id)
-
-
 class TwinkleWorkerExtension:
     """Extension class for vLLM workers to support weight synchronization.
 
@@ -122,7 +117,7 @@ def update_weights_from_ipc(
         import torch.distributed as dist
 
         if self.device is None:
-            self.device = torch.device(Torch.get_device())
+            self.device = torch.device(Torch.get_device(getattr(self, "local_rank", None)))
 
         if peft_config and base_sync_done:
             self.remove_lora(VLLM_LORA_INT_ID)
@@ -257,7 +252,7 @@ def load_synced_weights(
             base_sync_done: If True with peft_config, load as LoRA adapter.
         """
         if self.device is None:
-            self.device = torch.device(Torch.get_device())
+            self.device = torch.device(Torch.get_device(getattr(self, "local_rank", None)))
 
         weight_list = list(weights.items())
         self._load_weights(weight_list, peft_config=peft_config, base_sync_done=base_sync_done)
@@ -374,5 +369,5 @@ def _load_weights(
     def _get_zmq_handle(self) -> str:
         """Get ZMQ handle for IPC communication."""
         if not hasattr(self, '_device_uuid') or not self._device_uuid:
-            self._device_uuid = _get_device_uuid(self.device.index)
+            self._device_uuid = get_vllm_device_uuid(self.device.index)
         return f"ipc:///tmp/twinkle-ipc-{self._device_uuid}.sock"

src/twinkle/utils/network.py

@@ -1,11 +1,37 @@
 # Copyright (c) ModelScope Contributors. All rights reserved.
+import os
 import socket
 from datetime import timedelta
 from typing import Optional
 
 import torch
 
 
+def _ensure_hccl_socket_env(master_port: int, environ: Optional[dict] = None) -> None:
+    """Set deterministic HCCL socket env defaults to avoid port collisions.
+
+    In multi-job environments, HCCL's default base port (60000) can collide
+    across concurrent jobs and lead to:
+    `ra_hdc_socket_listen_start ... ret(-98)`.
+
+    We derive a per-job port layout from `master_port` so all ranks use the
+    same values while reducing cross-job conflicts. Explicit user settings are
+    preserved and never overwritten.
+    """
+    env = os.environ if environ is None else environ
+    if "HCCL_IF_BASE_PORT" not in env:
+        # 20000-39999, with an offset to avoid colliding with TCPStore port.
+        env["HCCL_IF_BASE_PORT"] = str(20000 + ((master_port + 997) % 20000))
+    if "HCCL_HOST_SOCKET_PORT_RANGE" not in env:
+        # 40000-40511 ... 47999-48510
+        start = 40000 + (master_port % 8000)
+        env["HCCL_HOST_SOCKET_PORT_RANGE"] = f"{start}-{start + 511}"
+    if "HCCL_NPU_SOCKET_PORT_RANGE" not in env:
+        # 50000-50511 ... 57999-58510
+        start = 50000 + (master_port % 8000)
+        env["HCCL_NPU_SOCKET_PORT_RANGE"] = f"{start}-{start + 511}"
+
+
 def is_valid_ipv6_address(ip: str) -> bool:
     """Check if the given string is a valid IPv6 address."""
     try:
@@ -87,6 +113,7 @@ def stateless_init_process_group(
     from vllm.distributed.utils import StatelessProcessGroup
 
     if backend == "hccl":
+        _ensure_hccl_socket_env(master_port)
         from vllm_ascend.distributed.device_communicators.pyhccl import (
             PyHcclCommunicator as Communicator,
         )

src/twinkle/utils/platform.py

@@ -1,7 +1,10 @@
 # Copyright (c) ModelScope Contributors. All rights reserved.
 import os
 import platform
+import hashlib
+import re
 import shutil
+import socket
 import subprocess
 from abc import ABC
 from dataclasses import dataclass, field
@@ -641,5 +644,67 @@ def is_last_rank():
         return True
     return dist.get_rank() == dist.get_world_size() - 1
 
+
+def _resolve_ascend_physical_device_id(device_id: int) -> int:
+    """Map local NPU device index to physical device id via visible devices."""
+    visible = os.environ.get("ASCEND_RT_VISIBLE_DEVICES", "").strip()
+    if not visible:
+        return device_id
+    parts = [p.strip() for p in visible.split(",") if p.strip()]
+    if device_id < 0 or device_id >= len(parts):
+        return device_id
+    return int(parts[device_id])
+
+
+def _get_npu_bus_id_from_npu_smi(device_id: int) -> Optional[str]:
+    """Get NPU Bus-Id from `npu-smi info` output."""
+    try:
+        physical_id = _resolve_ascend_physical_device_id(device_id)
+    except Exception:
+        physical_id = device_id
+
+    try:
+        output = subprocess.check_output(
+            ["npu-smi", "info"],
+            text=True,
+            stderr=subprocess.STDOUT,
+            timeout=5,
+        )
+    except Exception:
+        return None
+
+    # Typical line:
+    # | 0     0                   | 0000:9D:00.0  | ...
+    pattern = re.compile(
+        r"^\|\s*\d+\s+(\d+)\s*\|\s*"
+        r"([0-9A-Fa-f]{4}:[0-9A-Fa-f]{2}:[0-9A-Fa-f]{2}\.[0-9A-Fa-f])\s*\|",
+        re.MULTILINE,
+    )
+    for match in pattern.finditer(output):
+        phy_id = int(match.group(1))
+        if phy_id == physical_id:
+            return match.group(2).lower()
+    return None
+
+
+def get_vllm_device_uuid(device_id: int = 0) -> str:
+    """Get vLLM device uuid with NPU Bus-Id special handling."""
+    from vllm.platforms import current_platform
+
+    try:
+        return current_platform.get_device_uuid(device_id)
+    except NotImplementedError:
+        # NPU special case: prefer stable PCIe Bus-Id from npu-smi.
+        bus_id = _get_npu_bus_id_from_npu_smi(device_id)
+        if bus_id:
+            return bus_id
+        # Generic deterministic fallback to keep sender/receiver socket names aligned.
+        visible = os.environ.get("ASCEND_RT_VISIBLE_DEVICES") or os.environ.get(
+            "CUDA_VISIBLE_DEVICES", ""
+        )
+        raw = f"{socket.gethostname()}:{visible}:{device_id}"
+        return hashlib.sha1(raw.encode("utf-8")).hexdigest()[:16]
+
+
 def is_master():
     return Platform.is_master()