Add RCCL and training warmup for HYBRID_SHARD stability #83
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change | ||||
|---|---|---|---|---|---|---|
| @@ -1,42 +1,133 @@ | ||||||
| #!/bin/bash | ||||||
| # ============================================================================= | ||||||
| # Global NCCL/RCCL environment variables for multi-node training | ||||||
| # Configured for MI350X cluster | ||||||
| # | ||||||
| # This file is the SINGLE SOURCE OF TRUTH for all NCCL/RCCL configuration. | ||||||
| # Edit variables here - local_launch.sh will automatically pick them up. | ||||||
| # | ||||||
| # NOTE: When adding a new environment variable, you MUST also add its name | ||||||
| # to the DOCKER_ENV_VARS array below, otherwise it won't be passed | ||||||
| # to the Docker container. | ||||||
| # ============================================================================= | ||||||
|
|
||||||
| # ----------------------------------------------------------------------------- | ||||||
| # NCCL Debug Settings | ||||||
| # ----------------------------------------------------------------------------- | ||||||
| export NCCL_DEBUG=WARN | ||||||
| export NCCL_DEBUG_SUBSYS= # Options: COLL,INIT,NET (empty = none) | ||||||
|
There was a problem hiding this comment. Setting NCCL_DEBUG_SUBSYS to an empty string may cause issues depending on how NCCL interprets empty environment variables. Some versions of NCCL may treat an empty value differently than an unset variable. Consider either unsetting the variable (unset NCCL_DEBUG_SUBSYS) when not needed, or omitting the export statement entirely when the value is empty, rather than exporting an empty string.
Suggested change
|
||||||
|
|
||||||
| # ----------------------------------------------------------------------------- | ||||||
| # RCCL-Specific Settings (ROCm) | ||||||
| # ----------------------------------------------------------------------------- | ||||||
| export RCCL_DIRECT_ALLGATHER_DISABLE=1 # Disable direct allgather | ||||||
| export RCCL_MSCCL_ENABLE=0 # Disable MSCCL | ||||||
| export RCCL_THREADS_PER_BLOCK=256 # Threads per block (override via --threads) | ||||||
|
|
||||||
| # ----------------------------------------------------------------------------- | ||||||
| # IB/RNIC Configuration for MI350X | ||||||
| # ----------------------------------------------------------------------------- | ||||||
| export NCCL_IB_HCA=bnxt_re0,bnxt_re1,bnxt_re2,bnxt_re3,bnxt_re4,bnxt_re5,bnxt_re6,bnxt_re7 | ||||||
| export NCCL_IB_GID_INDEX=3 | ||||||
| export NCCL_NCHANNELS_PER_NET_PEER=8 | ||||||
|
|
||||||
| # ----------------------------------------------------------------------------- | ||||||
| # HSA Settings for ROCm | ||||||
| # ----------------------------------------------------------------------------- | ||||||
| export HSA_ENABLE_IPC_MODE_LEGACY=1 | ||||||
| export HSA_ENABLE_SDMA=0 # Disable SDMA for stability | ||||||
|
|
||||||
| # ----------------------------------------------------------------------------- | ||||||
| # NCCL Protocol and Channels | ||||||
| # ----------------------------------------------------------------------------- | ||||||
| export NCCL_PROTO=Simple | ||||||
| #export NCCL_MIN_NCHANNELS=40 | ||||||
| export NCCL_MAX_NCHANNELS=56 # Override via --channels | ||||||
|
|
||||||
| # ----------------------------------------------------------------------------- | ||||||
| # Network Interface for MI350X cluster | ||||||
| # ----------------------------------------------------------------------------- | ||||||
| export NCCL_SOCKET_IFNAME=enp49s0f0np0,fenic0 | ||||||
|
|
||||||
| # ----------------------------------------------------------------------------- | ||||||
| # Timeout and Error Handling | ||||||
| # ----------------------------------------------------------------------------- | ||||||
| export NCCL_TIMEOUT_MS=12000 # 12 second timeout (legacy, not used by PyTorch) | ||||||
| export NCCL_TIMEOUT=100 # 100 second timeout - first backward can be slow due to JIT/init | ||||||
| export TORCH_DIST_INIT_TIMEOUT=150 # Match collective timeout for consistency | ||||||
|
Collaborator
There was a problem hiding this comment. timeout |
||||||
| export TORCH_NCCL_ASYNC_ERROR_HANDLING=1 | ||||||
| export TORCH_NCCL_TRACE_BUFFER_SIZE=10000 | ||||||
| export TORCH_NCCL_DUMP_ON_TIMEOUT=1 # Critical for hang debugging | ||||||
| #export AMD_LOG_LEVEL=5 | ||||||
| # AMD_LOG_LEVEL_FILE is set dynamically in local_launch.sh to point to experiment directory | ||||||
| # Default fallback (will be overridden): | ||||||
| #export AMD_LOG_LEVEL_FILE=trace_amd.log | ||||||
| # ----------------------------------------------------------------------------- | ||||||
| # PyTorch ROCm Profiler | ||||||
| # ----------------------------------------------------------------------------- | ||||||
| export PYTORCH_ROCM_PROFILER_ENABLE_TRACING=1 | ||||||
|
|
||||||
| # ----------------------------------------------------------------------------- | ||||||
| # List of environment variables to pass to Docker container | ||||||
| # Add/remove variables here to control what gets passed through | ||||||
| # ----------------------------------------------------------------------------- | ||||||
| DOCKER_ENV_VARS=( | ||||||
| # NCCL Debug | ||||||
| NCCL_DEBUG | ||||||
| NCCL_DEBUG_SUBSYS | ||||||
| # RCCL | ||||||
| RCCL_DIRECT_ALLGATHER_DISABLE | ||||||
| RCCL_MSCCL_ENABLE | ||||||
| RCCL_THREADS_PER_BLOCK | ||||||
| # IB/RNIC | ||||||
| NCCL_IB_HCA | ||||||
| NCCL_IB_GID_INDEX | ||||||
| NCCL_NCHANNELS_PER_NET_PEER | ||||||
| # HSA | ||||||
| HSA_ENABLE_IPC_MODE_LEGACY | ||||||
| HSA_ENABLE_SDMA | ||||||
| # Protocol/Channels | ||||||
| NCCL_PROTO | ||||||
| NCCL_MIN_NCHANNELS | ||||||
|
There was a problem hiding this comment. The DOCKER_ENV_VARS array includes NCCL_MIN_NCHANNELS (line 91), but the actual export statement for NCCL_MIN_NCHANNELS is commented out (line 44). This means the build_docker_env_flags function will pass an undefined or empty value for NCCL_MIN_NCHANNELS to the Docker container. Either uncomment the export on line 44 or remove NCCL_MIN_NCHANNELS from the DOCKER_ENV_VARS array.
Suggested change
|
||||||
| NCCL_MAX_NCHANNELS | ||||||
| # Network | ||||||
| NCCL_SOCKET_IFNAME | ||||||
| # Timeout/Error Handling | ||||||
| NCCL_TIMEOUT_MS | ||||||
| NCCL_TIMEOUT | ||||||
| TORCH_DIST_INIT_TIMEOUT | ||||||
| TORCH_NCCL_ASYNC_ERROR_HANDLING | ||||||
| TORCH_NCCL_TRACE_BUFFER_SIZE | ||||||
| TORCH_NCCL_DUMP_ON_TIMEOUT | ||||||
| # AMD Logging | ||||||
| AMD_LOG_LEVEL | ||||||
| AMD_LOG_LEVEL_FILE | ||||||
| # Profiler | ||||||
| PYTORCH_ROCM_PROFILER_ENABLE_TRACING | ||||||
| ) | ||||||
| export DOCKER_ENV_VARS | ||||||
|
|
||||||
| # ----------------------------------------------------------------------------- | ||||||
| # Helper function: Build docker -e flags from DOCKER_ENV_VARS | ||||||
| # Usage: DOCKER_ENV_FLAGS=$(build_docker_env_flags) | ||||||
| # ----------------------------------------------------------------------------- | ||||||
| build_docker_env_flags() { | ||||||
| local flags="" | ||||||
| for var in "${DOCKER_ENV_VARS[@]}"; do | ||||||
| local value="${!var}" | ||||||
| flags+=" -e ${var}=${value}" | ||||||
|
There was a problem hiding this comment. The build_docker_env_flags function does not quote the environment variable values when building the flags string. If any environment variable contains spaces or special characters (e.g., NCCL_SOCKET_IFNAME=enp49s0f0np0,fenic0), this could cause issues with shell word splitting when the flags are used in the docker exec command. Consider quoting the values: flags+=" -e ${var}="${value}""
Suggested change
|
||||||
| done | ||||||
| echo "$flags" | ||||||
| } | ||||||
| export -f build_docker_env_flags | ||||||
|
|
||||||
| # ============================================================================= | ||||||
| # Optional settings (uncomment to enable) | ||||||
| # ============================================================================= | ||||||
|
|
||||||
| # Force non-overlap for debugging (single HW queue) | ||||||
| # export GPU_MAX_HW_QUEUES=1 | ||||||
| # unset TORCH_NCCL_HIGH_PRIORITY | ||||||
|
|
||||||
| # Disable IB for Ethernet-only testing | ||||||
| # export NCCL_IB_DISABLE=1 | ||||||
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The expansion of
DOCKER_EXEChere is vulnerable to shell command injection because it is built from untrusted inputs (theDOCKER_CONTAINERCLI argument and environment-derived values inDOCKER_ENV_FLAGS) and then expanded unquoted in a command context. An attacker who can influenceDOCKER_CONTAINERor one of the environment variables (e.g., by including characters like;,&&or$(...)) can break out of the intendeddocker execinvocation and execute arbitrary additional commands on the host. To mitigate this, avoid constructing a full command string inDOCKER_EXECand instead passdocker execand its arguments as separate, properly quoted words (or an array), and ensureDOCKER_CONTAINERis strictly validated/whitelisted to be a simple container name.