[Question]: bootstrap persistence & multi-module coexistence

[Enigmatisms]

Question

NVSHMEM 3.x | MoE training (DeepSeek DeepEP + custom modules in one process)

Background

We're building a training framework where multiple independent modules (e.g., DeepEP for MoE all-to-all, a custom overlap module) each call nvshmemx_init_attr(NVSHMEMX_INIT_WITH_UNIQUEID, ...) with potentially different rank/nranks. We ran into some issues and did a source audit — hoping to confirm our understanding and get advice.

Q1: Is it possible to change nranks via finalize + re-init?

From reading the source, it appears that nvshmem_finalize() resets the initialized flag but not the bootstrapped flag. So after finalize + re-init, the new rank/nranks/uid parameters are silently ignored — bootstrap is skipped entirely and the old boot_handle (pg_rank, pg_size, node topology, etc.) is reused.

Is this intended? Is there any supported way to fully reset bootstrap state within a process so that a subsequent init can join a different-sized communication world?

Q2: Multiple modules calling init with different nranks — is this safe?

For example:

• Module A inits with nranks=8 (intra-node only)
• Module B inits with nranks=32 (all expert-parallel ranks)

Since NVSHMEM is a process-global singleton, the second init just bumps the refcount and its parameters are discarded. Module B ends up in Module A's 32-PE world without knowing it. This seems fundamentally unsafe — wrong PE numbering, nvshmem_malloc blocking on a global barrier that Module B's subset can't satisfy, etc.

What's the recommended practice here? Our best guess is:

1. Coordinate a single init with the superset of all PEs
2. Use nvshmem_team_split_strided() for per-module sub-groups
3. Use team-based collectives + nvshmem_team_translate_pe() for RMA

Is this the right approach? Any guidance on coordinating NVSHMEM across independently-developed modules would be very helpful.

Context

This came up while integrating DeepEP (which has its own init/finalize lifecycle) into PaddlePaddle alongside other NVSHMEM-based modules. Thanks in advance!

[benhg]

For Q1:

nvshmem_finalize() clears nvshmemi_is_nvshmem_initialized but only clears nvshmemi_is_nvshmem_bootstrapped when using UID bootstrap; for MPI, PMI, and SHMEM bootstrap it leaves bootstrapped true. On re-init, if bootstrapped is still true, bootstrap is skipped and nvshmemi_boot_handle (with the old pg_rank and pg_size) is reused, so new rank/nranks arguments are ignored. This is intentional: with launcher-based bootstraps (MPI/PMI/SHMEM) the job world is fixed for the process, so init/finalize loops are meant for tearing down and reinitializing NVSHMEM state within the same world, not for changing world size. To change nranks without exiting the process you must use UID bootstrap: in that mode finalize fully tears down bootstrap, so a subsequent init with new nvshmemx_uniqueid_args_t (including new myrank and nranks) performs a fresh bootstrap. For MPI/PMI/SHMEM there is no supported way to change the world in-process; you must relaunch the job (or otherwise exit the process and re-start the process).

Note that you must use nvshmemx_init_attr with NVSHMEMX_INIT_WITH_UNIQUEID and a new attr. Plain nvshmem_init() does not pass UID args and will not work for UID bootstrap.

For Q2:

Your understanding is correct. There's no way to init multiple times in the same process with different nranks, and your proposed solution is also the one I'd recommend. Create a single init (e.g. TEAM_WORLD) with all the PEs you're going to use, and then use team features (you don't have to use team_split_strided specifically - we have arbitrary team management so you can create a team out of any subset of PEs). Then you can use team_translate for RMA and collectives with a team argument.

[Enigmatisms]

Thanks for the answer! Clearly I've missed the critical point on Q1 when reading the source code, I thought the bootstrap was just fixed when initialized.

For Q2 the problem actually gets messier: modules like DeepEP (and also, our comm module) assume the world is the comm group being passed to the module, which is usually the subset of all the PEs (for example, DeepEP uses EP group). In order not to have conflicting NVSHMEM comm singleton status, we tried statically link nvshmem against two modules, so that we have two NVSHMEM comm singletons. Conflicts resolved! However, ibrc.so and ibgda.so are dyn-libs, which must be shared along with their internal static variables. Seems like this is causing some strange racing effects (like hang, or crash during finalizing).

So do you have any suggestion on incorporating several independent (meaning that using one single NVSHMEM takes too much effort to coordinate between modules) modules (which depends on NVSHMEM) safely?

[benhg]

Thanks for explaining your use case. Because NVSHMEM uses a singleton, it's not well adapted to that kind of use case at the moment. This suggestion feels a bit hacky, but perhaps you could try it: Could you have multiple copies of the .so objects and have each process open its own version?

[Enigmatisms]

Could you have multiple copies of the .so objects and have each process open its own version?

The best guess would be: I need to rename different .so file, so that dyn-lib loader won't consider those as being duplicated. It really is hacky. I guess the safer route should be building the entire NVSHMEM as dynamic lib and link all the dependent packages against it; then write a single NVSHMEM context manager to create a set of all PEs so that different modules can then extract subsets of PEs using team.

Anyway, thanks for your help! I am still working on the hacky solution, since this is seemingly simpler.

[benhg]

I guess the safer route should be building the entire NVSHMEM as dynamic lib and link all the dependent packages against it; then write a single NVSHMEM context manager to create a set of all PEs so that different modules can then extract subsets of PEs using team.

Yes, this is what we would recommend. Let me know either way. I am interested to hear about use cases that push our boundaries like this.

[Enigmatisms]

Here is the update. I tried using independent NVSHMEM dyn-libs, and it worked. Turns out we should:

• Compile two copies of NVSHMEM independently, so that we can patch one of the copy. The patch applied is used to change the load path of the dynlib, like, instead of using fixed plugin_name = nvshmemi_options.BOOTSTRAP_UID_PLUGIN; or "nvshmem_transport_ibrc.so.%d", we will be loading the load path from environment variable, at least for one of the module. Otherwise the load path collides and libs are not loaded separately.
• Let's say there are two modules A and B. A uses the no-patching NVSHME, while we apply patches on B.
• For B, we need brand new nvshmem_transport_ibrc.so, nvshmem_transport_ibga.so and nvshmem_bootstrap_uid.so. Rename them, and DO NOT forget to patch the ELF sector:
  • SONAME in the ELF sector is actually responsible for glibc de-duplication. Try: readelf -d renamed_nvshmem_bootstrap_uid.so.3 | grep SONAME and you will find that SONAME is still nvshmem_bootstrap_uid.so. If this is not handled correctly, the dynlib will still be considered duplicated. HANG!
• Using patchelf --set-soname your_lib.so new_soname to patch your dynlib. Then everything works perfectly, no hang, no crash --- A and B coexistence!

However, this only proves that some of my guesses were correct: we need to do static lib symbol visibility guarding (LOCAL export), and we need to patch another copy of NVSHMEM to make a real dupilcation of NVSHMEM dynlibs. With all these done: no more static variable status conflicts, no shared network settings, but it feels unsafe, at least for me.

This is way too hacky and bug-prone. I wonder if NVSHMEM has plans for building copy-safe communication management instances, maybe, like what NCCL did?

[benhg]

Hi @Enigmatisms , thank you very much for the update. I don't have an update about our roadmap on this topic right now, but I will definitely discuss your feedback with the rest of the team.

In the mean time, as you observed, NCCL device APIs provide devComms and Windows which allow you to do this all within one process.

[Enigmatisms]

Yes, and some recommended using NCCL GIN instead of NVSHMEM to avoid this problem. idk, but I suppose the will be more and more cases that need modules like DeepEP being used along with other fine-grained communication modules in LLM infrastructures, so this is kinda important tbh.

[Enigmatisms]

So now that the question is clear and there are several viable solutions, consider this issue some sort of feature request. Thank you for your reply @benhg ! You can close this issue if you see fit.

[benhg]

Hi @Enigmatisms , thanks. I've created an internal tracking entry and I'll leave this open as an RFE for now.

[Enigmatisms]

Hi. While investigating how to create a top-level NVSHMEM management module, I ran into a few questions that need clarification:

• nvshmem_malloc appears to require participation from all PEs in the communicator (i.e., a global barrier). I checked the older codebase (3.2.5) and didn’t find anything like a team_malloc. This suggests that the management module may need a symmetric memory allocator, since letting dependent submodules (e.g., DeepEP) allocate buffers independently could risk hangs. On the other hand, this might not be... so necessary? Most distributed modules that use NVSHMEM already call nvshmem_malloc across all PEs, so we shouldn’t encounter cases where only a subset of PEs performs the allocation. I’m not fully sure whether introducing an allocator is worth it, given the added complexity.

• libnvshmem_device.a may not be usable as a dynamic library. There seem to be extern variables such as __constant__ nvshmemi_device_state_d that require definitions during nvcc -dlink. This could prevent fully relying on the NVSHMEM dynamic library. We might still need to account for static library symbol interposition, though I’m not certain.

• I’m also unsure whether the following is best practice: using all PEs, instead of sub-communicators, to construct the NVSHMEM world. This might consume more IBGDA QP resources. My understanding is that this effectively creates a complete graph communication pattern, leading to denser connectivity and higher QP usage, which could become a bottleneck at scale.

I’d appreciate any thoughts or suggestions on these points. Thanks in advance.

[benhg]

nvshmem_malloc appears to require participation from all PEs in the communicator (i.e., a global barrier). I checked the older codebase (3.2.5) and didn’t find anything like a team_malloc. This suggests that the management module may need a symmetric memory allocator, since letting dependent submodules (e.g., DeepEP) allocate buffers independently could risk hangs. On the other hand, this might not be... so necessary? Most distributed modules that use NVSHMEM already call nvshmem_malloc across all PEs, so we shouldn’t encounter cases where only a subset of PEs performs the allocation. I’m not fully sure whether introducing an allocator is worth it, given the added complexity.

Yes, NVSHMEM only has a collective malloc. There is no team-specific malloc implementation. This is again related to the fact that NVSHMEM's runtime state is a global singleton. We have had requests for a team-specific malloc before and I will note the additional request here.

libnvshmem_device.a may not be usable as a dynamic library. There seem to be extern variables such as constant nvshmemi_device_state_d that require definitions during nvcc -dlink. This could prevent fully relying on the NVSHMEM dynamic library. We might still need to account for static library symbol interposition, though I’m not certain.

Correct. libnvshmem_device is only meant to be used as a static library. The recommended way to link to NVSHMEM is to dynamically link to nvshmem_host.so, and statically link to libnvshmem_device.a

I’m also unsure whether the following is best practice: using all PEs, instead of sub-communicators, to construct the NVSHMEM world. This might consume more IBGDA QP resources. My understanding is that this effectively creates a complete graph communication pattern, leading to denser connectivity and higher QP usage, which could become a bottleneck at scale.

The NVSHMEM world is a global singleton, so communication from each PE to each other PE is required for NVSHMEM to initialize. If you're interested in a communicator-focused method, I'd suggest looking into NCCL device APIs.

[Enigmatisms]

Thanks for your replies. These issues seem... indeed, a bit tricky. Therefore I’m currently working on NVSHMEM and also taking a look at NCCL GIN.