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.

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

Add cuda to multigpu (xpu) bench #8386

Merged
merged 14 commits into from
Nov 16, 2023
Prev Previous commit
Next Next commit
adding cuda to bench
  • Loading branch information
puririshi98 committed Nov 15, 2023
commit e57818735fb845c21ee3721553155e0e94eaa127
7 changes: 6 additions & 1 deletion benchmark/multi_gpu/training/training_benchmark.py
Original file line number Diff line number Diff line change
@@ -88,6 +88,11 @@ def run(rank: int, world_size: int, args: argparse.ArgumentParser):
if args.device == 'xpu':
import intel_extension_for_pytorch as ipex
import oneccl_bindings_for_pytorch # noqa
else:
# CUDA
os.environ['MASTER_ADDR'] = 'localhost'
os.environ['MASTER_PORT'] = '12355'
dist.init_process_group('nccl', rank=rank, world_size=world_size)

if not device_conditions[args.device]():
raise RuntimeError(f'{args.device.upper()} is not available')
@@ -308,7 +313,7 @@ def get_dist_params() -> Tuple[int, int, str]:
world_size=world_size, rank=rank)
run(rank, world_size, args)
else:
# use mp spawn
import torch.multiprocessing as mp
max_world_size = torch.cuda.device_count()
chosen_world_size = args.n_gpus
if chosen_world_size <= max_world_size: