Fix Wrong SP4 Images From NCCL Graph Registration

docs/environment_variables.rst

@@ -165,6 +165,12 @@ Communication
        the conductor sets ``lo`` for its workers and NCCL's bootstrap
        sockets listen on loopback only. A value you set, or a set
        ``NCCL_SOCKET_FAMILY``, is kept as is.
+   * - ``NCCL_GRAPH_REGISTER``
+     - ``0``
+     - NCCL's own variable. On, NCCL registers the buffers a CUDA graph
+       captures, which succeeds on ``expandable_segments`` memory, and
+       Cosmos3 SP4 with batched captures then gave wrong images (#897).
+       Workers on CUDA set ``0`` unless you set it.
    * - ``GLOO_SOCKET_IFNAME``
      - ``lo``
      - PyTorch's variable for the Gloo groups next to each TP/SP group.

mstar/platforms/interface.py

@@ -80,6 +80,10 @@ def arch(self, device=None) -> str:
     def cuda_capability(self, device=None) -> tuple[int, int]:
         return torch.cuda.get_device_capability(device)
 
+    def process_env_defaults(self) -> dict[str, str]:
+        # NCCL registering the buffers captured in CUDA graphs gave wrong Cosmos3 SP4 images
+        return {"NCCL_GRAPH_REGISTER": "0"}
+
 
 class RocmPlatform(Platform):
     name = "rocm"