Multiple GPU training

[lcoandrade]

Do I need to do something extra to be able to train with multiple GPUs or just setting a Lighting Trainer with strategy, accelerator and devices like the code below is enough:

trainer = pl.Trainer(
        accelerator='gpu',
        devices='auto',
        max_epochs=EPOCHS,
        callbacks=[checkpoint_callback, ],
        logger=logger,
        strategy='ddp_notebook'
    )

I'm mentioning strategy='ddp_notebook' because I'm working on Kaggle.

I'm asking this because when I set my notebook to use GPU T4x2, I get the following error:

INFO: [rank: 0] Received SIGTERM: 15
---------------------------------------------------------------------------
ProcessRaisedException                    Traceback (most recent call last)
Cell In[24], line 5
      3 if os.path.isfile(checkpoint_file):
      4     print('Resuming training from previous checkpoint...')
----> 5     trainer.fit(
      6         model=task, 
      7         datamodule=datamodule,
      8         ckpt_path=checkpoint_file
      9     )
     10 else:
     11     print('Starting training from scratch...')

File /opt/conda/lib/python3.10/site-packages/lightning/pytorch/trainer/trainer.py:529, in Trainer.fit(self, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path)
    527 model = _maybe_unwrap_optimized(model)
    528 self.strategy._lightning_module = model
--> 529 call._call_and_handle_interrupt(
    530     self, self._fit_impl, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path
    531 )

File /opt/conda/lib/python3.10/site-packages/lightning/pytorch/trainer/call.py:41, in _call_and_handle_interrupt(trainer, trainer_fn, *args, **kwargs)
     39 try:
     40     if trainer.strategy.launcher is not None:
---> 41         return trainer.strategy.launcher.launch(trainer_fn, *args, trainer=trainer, **kwargs)
     42     return trainer_fn(*args, **kwargs)
     44 except _TunerExitException:

File /opt/conda/lib/python3.10/site-packages/lightning/pytorch/strategies/launchers/multiprocessing.py:124, in _MultiProcessingLauncher.launch(self, function, trainer, *args, **kwargs)
    116 process_context = mp.start_processes(
    117     self._wrapping_function,
    118     args=process_args,
   (...)
    121     join=False,  # we will join ourselves to get the process references
    122 )
    123 self.procs = process_context.processes
--> 124 while not process_context.join():
    125     pass
    127 worker_output = return_queue.get()

File /opt/conda/lib/python3.10/site-packages/torch/multiprocessing/spawn.py:160, in ProcessContext.join(self, timeout)
    158 msg = "\n\n-- Process %d terminated with the following error:\n" % error_index
    159 msg += original_trace
--> 160 raise ProcessRaisedException(msg, error_index, failed_process.pid)

ProcessRaisedException: 

-- Process 1 terminated with the following error:
Traceback (most recent call last):
  File "/opt/conda/lib/python3.10/site-packages/torch/multiprocessing/spawn.py", line 69, in _wrap
    fn(i, *args)
  File "/opt/conda/lib/python3.10/site-packages/lightning/pytorch/strategies/launchers/multiprocessing.py", line 147, in _wrapping_function
    results = function(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/lightning/pytorch/trainer/trainer.py", line 568, in _fit_impl
    self._run(model, ckpt_path=ckpt_path)
  File "/opt/conda/lib/python3.10/site-packages/lightning/pytorch/trainer/trainer.py", line 973, in _run
    results = self._run_stage()
  File "/opt/conda/lib/python3.10/site-packages/lightning/pytorch/trainer/trainer.py", line 1014, in _run_stage
    self._run_sanity_check()
  File "/opt/conda/lib/python3.10/site-packages/lightning/pytorch/trainer/trainer.py", line 1043, in _run_sanity_check
    val_loop.run()
  File "/opt/conda/lib/python3.10/site-packages/lightning/pytorch/loops/utilities.py", line 177, in _decorator
    return loop_run(self, *args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/lightning/pytorch/loops/evaluation_loop.py", line 115, in run
    self._evaluation_step(batch, batch_idx, dataloader_idx)
  File "/opt/conda/lib/python3.10/site-packages/lightning/pytorch/loops/evaluation_loop.py", line 375, in _evaluation_step
    output = call._call_strategy_hook(trainer, hook_name, *step_kwargs.values())
  File "/opt/conda/lib/python3.10/site-packages/lightning/pytorch/trainer/call.py", line 291, in _call_strategy_hook
    output = fn(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/lightning/pytorch/strategies/ddp.py", line 336, in validation_step
    return self.model(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/torch/nn/parallel/distributed.py", line 1156, in forward
    output = self._run_ddp_forward(*inputs, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/torch/nn/parallel/distributed.py", line 1103, in _run_ddp_forward
    inputs, kwargs = _to_kwargs(
  File "/opt/conda/lib/python3.10/site-packages/torch/distributed/utils.py", line 100, in _to_kwargs
    _recursive_to(inputs, device_id, use_side_stream_for_tensor_copies)
  File "/opt/conda/lib/python3.10/site-packages/torch/distributed/utils.py", line 92, in _recursive_to
    res = to_map(inputs)
  File "/opt/conda/lib/python3.10/site-packages/torch/distributed/utils.py", line 83, in to_map
    return list(zip(*map(to_map, obj)))
  File "/opt/conda/lib/python3.10/site-packages/torch/distributed/utils.py", line 87, in to_map
    return [type(obj)(i) for i in zip(*map(to_map, obj.items()))]
  File "/opt/conda/lib/python3.10/site-packages/torch/distributed/utils.py", line 87, in <listcomp>
    return [type(obj)(i) for i in zip(*map(to_map, obj.items()))]
TypeError: first argument must be callable or None

Thanks!

[adamjstewart]

There's no reason on the TorchGeo side why this wouldn't work, we use multiple GPUs all the time. I'm not familiar with all of the Lightning Trainer settings. Might be worth opening a discussion in the Lightning repo to see if there's anything obvious.

[lcoandrade]

I’ll do that. But what I consider interesting is that I use the same trainer configuration (gpu, auto and ddp_notebbok) in other pure lightning notebook and it works. That’s why I asked here.