load_format = None
# convert megabytes to bytes
import torch_npu
# print(f"333333333 self.inference_engine.result_pid {self.inference_engine.result_pid}")
if self.device_mesh["infer_tp"].get_local_rank() != 0:
tp_size = self.device_mesh["infer_tp"].size()
result_pid = torch.zeros(tp_size, dtype=torch.int64)
result_pid = result_pid.to("npu")
else:
result_pid = self.inference_engine.result_pid
result_pid = torch.tensor(result_pid, dtype=torch.int64, device="npu")
dist.broadcast(result_pid, src=self.device_mesh["infer_tp"].mesh.tolist()[0], group=self.device_mesh["infer_tp"].get_group(),)
for pid in result_pid:
torch_npu._C._add_ipc_pid(pid)