add int64 to int32

Signed-off-by: noemotiovon <757486878@qq.com>

Author: noemotiovon

roll/distributed/scheduler/protocol.py

@@ -19,6 +19,9 @@
 
 from roll.utils.functionals import union_two_dict, divide_by_chunk_size
 from roll.platforms import current_platform
+from roll.utils.logging import get_logger
+
+logger = get_logger()
 
 try:
     tensordict.set_lazy_legacy(False).set()
@@ -248,6 +251,9 @@ def from_single_dict(cls, data: Dict[str, Union[torch.Tensor, np.ndarray]], meta
 
         for key, val in data.items():
             if isinstance(val, torch.Tensor):
+                if current_platform.is_npu and val.dtype == torch.int64:
+                    logger.debug(f"[NPU] Converting Tensor {key} from int64 -> int32, shape={val.shape}")
+                    val = val.to(torch.int32)
                 tensors[key] = val
             elif isinstance(val, np.ndarray):
                 non_tensors[key] = val

roll/pipeline/distill/distill_worker.py

@@ -20,6 +20,7 @@
 from roll.utils.cuda_ipc_utils import MultiprocessingSerializer
 from roll.utils.offload_states import OffloadStateType
 from roll.pipeline.distill.various_divergence import VariousDivergence, GPTLMLoss
+from roll.platforms import current_platform
 
 
 
@@ -71,7 +72,7 @@ def train_step(self, data: DataProto):
                 is_offload_states=is_offload_states,
                 load_kwargs={"include": None},
         ):
-            data = data.to("cuda")
+            data = data.to(current_platform.device_type)
             data = self.strategy.get_data_input(data)
             self.logger.info(f"global_step: {data.meta_info.get('global_step',0)}")
             per_device_train_batch_size = self.worker_config.training_args.per_device_train_batch_size
@@ -176,7 +177,7 @@ def forward(self, data: DataProto):
                 is_offload_states=is_offload_states,
                 load_kwargs={"include": None},
         ):
-            data = data.to("cuda")
+            data = data.to(current_platform.device_type)
             data.meta_info["micro_batch_size"] = self.pipeline_config.student.training_args.per_device_train_batch_size
             data.meta_info["output_on_all_tp_ranks"] = True
             self.logger.info(f"global_step: {data.meta_info.get('global_step', 0)}")

roll/pipeline/dpo/actor_worker.py

@@ -6,6 +6,7 @@
 from roll.distributed.scheduler.decorator import Dispatch, register
 from roll.distributed.scheduler.protocol import DataProto
 from roll.pipeline.base_worker import ActorWorker as BaseActorWorker
+from roll.platforms import current_platform
 from roll.utils.context_managers import state_offload_manger
 from roll.utils.functionals import append_to_dict
 from roll.utils.offload_states import OffloadStateType
@@ -75,7 +76,7 @@ def train_step(self, data: DataProto):
             metric_infix=f"{self.cluster_name}/train_step",
             is_offload_states=is_offload_states,
         ):
-            data = data.to("cuda")
+            data = data.to(current_platform.device_type)
             data = self.strategy.get_data_input(data)
             per_device_train_batch_size = self.worker_config.training_args.per_device_train_batch_size
             backward_batch_size = (
@@ -124,7 +125,7 @@ def compute_log_probs(self, data: DataProto):
             metric_infix=f"{self.cluster_name}/compute_log_probs",
             is_offload_states=is_offload_states,
         ):
-            data = data.to("cuda")
+            data = data.to(current_platform.device_type)
             data.meta_info["micro_batch_size"] = self.worker_config.infer_batch_size
 
             with torch.no_grad():

roll/platforms/npu.py

@@ -13,9 +13,13 @@ class NpuPlatform(Platform):
     ray_experimental_noset: str = "RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES"
     communication_backend: str = "hccl"
 
+    @classmethod
+    def is_npu(cls) -> bool:
+        return True
+
     @classmethod
     def clear_cublas_workspaces(cls) -> None:
-        pass
+        return
 
     @classmethod
     def get_vllm_worker_class(clas):
@@ -38,7 +42,7 @@ def get_vllm_worker_class(clas):
 
     @classmethod
     def set_allocator_settings(cls) -> None:
-        pass
+        return
 
     @classmethod
     def get_custom_env_vars(cls) -> dict:

roll/platforms/platform.py

@@ -87,6 +87,10 @@ def __getattr__(self, key: str):
         else:
             logger.warning("Current platform %s does not have '%s'" " attribute.", self.device_type, key)
             return None
+    
+    @classmethod
+    def is_npu(cls) -> bool:
+        return False
 
     @classmethod
     def clear_cublas_workspaces(cls) -> None: