[Code scan] Use selected GPU count for DDP world size and spawning

[njzjz]

This issue is a result of a Codex global repository scan.

Summary

The trainer derives WORLD_SIZE and mp.spawn nprocs from torch.cuda.device_count(), even when use_gpu selects a subset such as "0,1". CUDA_VISIBLE_DEVICES is also mutated after querying CUDA, making selected-device and spawned-rank counts diverge.

Code references

unimol_tools/unimol_tools/tasks/trainer.py

Lines 71 to 94 in 4596596

	if torch.cuda.is_available() and self.cuda:
	if self.amp:
	self.scaler = torch.amp.GradScaler("cuda")
	else:
	self.scaler = None
	self.device = torch.device("cuda")
	world_size = torch.cuda.device_count()
	logger.info(f"Number of GPUs available: {world_size}")
	if self.gpu is not None:
	if self.gpu == "all":
	gpu = ",".join(str(i) for i in range(world_size))
	else:
	gpu = self.gpu
	else:
	gpu = "0"
	gpu_count = len(str(gpu).split(","))
	if world_size > 1 and self.ddp and gpu_count > 1:
	gpu = str(gpu).replace(" ", "")
	os.environ['MASTER_ADDR'] = os.getenv('MASTER_ADDR', 'localhost')
	os.environ['MASTER_PORT'] = os.getenv('MASTER_PORT', '19198')
	os.environ['CUDA_VISIBLE_DEVICES'] = gpu
	os.environ['WORLD_SIZE'] = str(world_size)
	logger.info(f"Using DistributedDataParallel for multi-GPU. GPUs: {gpu}")

unimol_tools/unimol_tools/tasks/trainer.py

Lines 185 to 203 in 4596596

	if torch.cuda.device_count() and self.ddp:
	with mp.Manager() as manager:
	shared_queue = manager.Queue()
	mp.spawn(
	self.fit_predict_with_ddp,
	args=(
	shared_queue,
	model,
	train_dataset,
	valid_dataset,
	loss_func,
	activation_fn,
	dump_dir,
	fold,
	target_scaler,
	feature_name,
	),
	nprocs=torch.cuda.device_count(),
	)

Impact

On hosts with more GPUs than requested, DDP can spawn too many processes or set a WORLD_SIZE that does not match the selected visible devices. This can fail initialization or leave ranks bound to unintended GPUs.

Suggested fix

Parse use_gpu before DDP setup and set WORLD_SIZE and nprocs to the selected GPU count. Alternatively require users to set CUDA_VISIBLE_DEVICES or launch with torchrun, and avoid mutating CUDA_VISIBLE_DEVICES after CUDA has already been queried.