Ensure NVLS and NVLSTree chunksizes are matched

This would prevent failures when NVLS size is smaller than the NVLSTree chunksize or silent fallback to a lower value, if NCCL were to change defaults. Signed-off-by: Raghu Raja <[email protected]>
aws · Apr 4, 2024 · 3defe9e · 3defe9e
1 parent c5e9b22
commit 3defe9e
Showing 1 changed file with 10 additions and 0 deletions.
diff --git a/src/platform-aws.c b/src/platform-aws.c
@@ -447,6 +447,9 @@ int platform_init(const char **provider_filter)
 	 * Setting this unconditionally without relying on ncclGetVersion symbol
 	 * being available, since the parameter did not exist in versions prior
 	 * to v2.20.
+	 *
+	 * The NVLSTree chunk size can not be larger than the NVLS chunk size,
+	 * so we ensure both are set to 512KiB.
 	 */
 	NCCL_OFI_INFO(NCCL_INIT | NCCL_NET, "Setting NCCL_NVLSTREE_MAX_CHUNKSIZE to 512KiB");
 	ret = setenv("NCCL_NVLSTREE_MAX_CHUNKSIZE", "524288", 0);
@@ -456,6 +459,13 @@ int platform_init(const char **provider_filter)
 		goto exit;
 	}
 
+	NCCL_OFI_INFO(NCCL_INIT | NCCL_NET, "Setting NCCL_NVLS_CHUNKSIZE to 512KiB");
+	ret = setenv("NCCL_NVLS_CHUNKSIZE", "524288", 0);
+	if (ret != 0) {
+		NCCL_OFI_WARN("Unable to set NCCL_NVLS_CHUNKSIZE");
+		ret = -errno;
+		goto exit;
+	}
 #endif
 
 	/*