Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Update NCCL_ASYNC_ERROR_HANDLING env variable #3267

Merged
merged 6 commits into from
May 13, 2024
Merged

Update NCCL_ASYNC_ERROR_HANDLING env variable #3267

Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
d6f823d
Update env variable
May 8, 2024
fc8dad9
Merge branch 'dev' into priba/env_nccl_update
May 8, 2024
d030fdf
Update cli variables according to version
May 9, 2024
a2f2dea
precommit formatting
May 9, 2024
5732aa2
Merge branch 'dev' into priba/env_nccl_update
priba May 10, 2024
8b59152
Merge branch 'dev' into priba/env_nccl_update
priba May 12, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 9 additions & 1 deletion composer/cli/launcher.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@

import psutil
import torch
from packaging import version

import composer
from composer.loggers.mosaicml_logger import (
Expand Down Expand Up @@ -317,6 +318,13 @@ def _launch_processes(
log.info('Starting distributed environment on local node for global_rank(%s-%s)', base_rank, base_rank + nproc - 1)
log.info('Distributed KV store: tcp://%s:%s', master_addr, master_port)

nccl_env_variable = {
(
'NCCL_ASYNC_ERROR_HANDLING' if version.parse(torch.__version__) < version.parse('2.2.0') else 'TORCH_NCCL_ASYNC_ERROR_HANDLING'
):
'1',
}

for local_rank in range(nproc):
global_rank = base_rank + local_rank
if command_mode and module_mode:
Expand All @@ -339,7 +347,7 @@ def _launch_processes(
MASTER_ADDR=master_addr,
MASTER_PORT=str(master_port),
PYTHONUNBUFFERED='1',
NCCL_ASYNC_ERROR_HANDLING='1',
**nccl_env_variable,
):
# Populate the distributed variables in all launcher args
for arg in training_script_args:
Expand Down
Loading