[TRTLLM-6826][feat] Allow sending more than 2GiB through MPI by using mpi4py.util.pkl5 (NVIDIA#6522)

Signed-off-by: Amit Zuker <203509407+amitz-nv@users.noreply.github.com>

Author: amitz-nv

tensorrt_llm/_utils.py

@@ -32,6 +32,7 @@
 import numpy as np
 import nvtx
 from mpi4py import MPI
+from mpi4py.util import pkl5
 from packaging import version
 
 # isort: off
@@ -455,7 +456,7 @@ def dim_resolve_negative(dim, ndim):
 # mpi4py only exports MPI_COMM_TYPE_SHARED, so we define OMPI_COMM_TYPE_HOST here
 OMPI_COMM_TYPE_HOST = 9
 
-comm = MPI.COMM_WORLD
+comm = pkl5.Intracomm(MPI.COMM_WORLD)
 
 
 def set_mpi_comm(new_comm):