OOM when launched with slurm

[mars747]

When I launch the trainer:

trainer = pl.Trainer(accelerator= 'gpu', devices = args.gpus, 
                     precision=32, callbacks=[logger],
                     default_root_dir=default_root_dir,
                     strategy = 'deepspeed')

if I launch it directly in the node's terminal, it works for the single node multi-gpu setup. However, if I use sbatch to launch it, it results in an OOM error for the same code. I have tried both deepspeed and deepspeed_stage_2 strategies. What can I do to resolve this?

The log printout:

╭───────────────────── Traceback (most recent call last) ──────────────────────╮
│ /mnt/storage/0/x/dev/ControlNet/tutorial_train.py:67 in <module>      │
│                                                                              │
│   64 │   │   │   │   │    strategy = 'deepspeed_stage_2')                    │
│   65                                                                         │
│   66 # Train!                                                                │
│ ❱ 67 trainer.fit(model, dataloader)                                          │
│   68                                                                         │
│                                                                              │
│ /mnt/home/x/.conda/envs/automatic1111_new/lib/python3.10/site-packages/pyt │
│ orch_lightning/trainer/trainer.py:520 in fit                                 │
│                                                                              │
│    517 │   │   """                                                           │
│    518 │   │   model = _maybe_unwrap_optimized(model)                        │
│    519 │   │   self.strategy._lightning_module = model                       │
│ ❱  520 │   │   call._call_and_handle_interrupt(                              │
│    521 │   │   │   self, self._fit_impl, model, train_dataloaders, val_datal │
│    522 │   │   )                                                             │
│    523                                                                       │
│                                                                              │
│ /mnt/home/x/.conda/envs/automatic1111_new/lib/python3.10/site-packages/pyt │
│ orch_lightning/trainer/call.py:42 in _call_and_handle_interrupt              │
│                                                                              │
│    39 │   """                                                                │
│    40 │   try:                                                               │
│    41 │   │   if trainer.strategy.launcher is not None:                      │
│ ❱  42 │   │   │   return trainer.strategy.launcher.launch(trainer_fn, *args, │
│    43 │   │   else:                                                          │
│    44 │   │   │   return trainer_fn(*args, **kwargs)                         │
│    45                                                                        │
│                                                                              │
│ /mnt/home/x/.conda/envs/automatic1111_new/lib/python3.10/site-packages/pyt │
│ orch_lightning/strategies/launchers/subprocess_script.py:92 in launch        │
│                                                                              │
│    89 │   │   """                                                            │
│    90 │   │   if not self.cluster_environment.creates_processes_externally:  │
│    91 │   │   │   self._call_children_scripts()                              │
│ ❱  92 │   │   return function(*args, **kwargs)                               │
│    93 │                                                                      │
│    94 │   def kill(self, signum: _SIGNUM) -> None:                           │
│    95 │   │   for proc in self.procs:                                        │
│                                                                              │
│ /mnt/home/x/.conda/envs/automatic1111_new/lib/python3.10/site-packages/pyt │
│ orch_lightning/trainer/trainer.py:559 in _fit_impl                           │
│                                                                              │
│    556 │   │   │   model_provided=True,                                      │
│    557 │   │   │   model_connected=self.lightning_module is not None,        │
│    558 │   │   )                                                             │
│ ❱  559 │   │   self._run(model, ckpt_path=ckpt_path)                         │
│    560 │   │                                                                 │
│    561 │   │   assert self.state.stopped                                     │
│    562 │   │   self.training = False                                         │
│                                                                              │
│ /mnt/home/x/.conda/envs/automatic1111_new/lib/python3.10/site-packages/pyt │
│ orch_lightning/trainer/trainer.py:935 in _run                                │
│                                                                              │
│    932 │   │   # ----------------------------                                │
│    933 │   │   # RUN THE TRAINER                                             │
│    934 │   │   # ----------------------------                                │
│ ❱  935 │   │   results = self._run_stage()                                   │
│    936 │   │                                                                 │
│    937 │   │   # ----------------------------                                │
│    938 │   │   # POST-Training CLEAN UP                                      │
│                                                                              │
│ /mnt/home/x/.conda/envs/automatic1111_new/lib/python3.10/site-packages/pyt │
│ orch_lightning/trainer/trainer.py:978 in _run_stage                          │
│                                                                              │
│    975 │   │   │   with isolate_rng():                                       │
│    976 │   │   │   │   self._run_sanity_check()                              │
│    977 │   │   │   with torch.autograd.set_detect_anomaly(self._detect_anoma │
│ ❱  978 │   │   │   │   self.fit_loop.run()                                   │
│    979 │   │   │   return None                                               │
│    980 │   │   raise RuntimeError(f"Unexpected state {self.state}")          │
│    981                                                                       │
│                                                                              │
│ /mnt/home/x/.conda/envs/automatic1111_new/lib/python3.10/site-packages/pyt │
│ orch_lightning/loops/fit_loop.py:201 in run                                  │
│                                                                              │
│   198 │   │   while not self.done:                                           │
│   199 │   │   │   try:                                                       │
│   200 │   │   │   │   self.on_advance_start()                                │
│ ❱ 201 │   │   │   │   self.advance()                                         │
│   202 │   │   │   │   self.on_advance_end()                                  │
│   203 │   │   │   │   self._restarting = False                               │
│   204 │   │   │   except StopIteration:                                      │
│                                                                              │
│ /mnt/home/x/.conda/envs/automatic1111_new/lib/python3.10/site-packages/pyt │
│ orch_lightning/loops/fit_loop.py:354 in advance                              │
│                                                                              │
│   351 │   │   assert self._data_fetcher is not None                          │
│   352 │   │   self._data_fetcher.setup(combined_loader)                      │
│   353 │   │   with self.trainer.profiler.profile("run_training_epoch"):      │
│ ❱ 354 │   │   │   self.epoch_loop.run(self._data_fetcher)                    │
│   355 │                                                                      │
│   356 │   def on_advance_end(self) -> None:                                  │
│   357 │   │   trainer = self.trainer                                         │
│                                                                              │
│ /mnt/home/x/.conda/envs/automatic1111_new/lib/python3.10/site-packages/pyt │
│ orch_lightning/loops/training_epoch_loop.py:133 in run                       │
│                                                                              │
│   130 │   │   self.on_run_start(data_fetcher)                                │
│   131 │   │   while not self.done:                                           │
│   132 │   │   │   try:                                                       │
│ ❱ 133 │   │   │   │   self.advance(data_fetcher)                             │
│   134 │   │   │   │   self.on_advance_end()                                  │
│   135 │   │   │   │   self._restarting = False                               │
│   136 │   │   │   except StopIteration:                                      │
│                                                                              │
│ /mnt/home/x/.conda/envs/automatic1111_new/lib/python3.10/site-packages/pyt │
│ orch_lightning/loops/training_epoch_loop.py:218 in advance                   │
│                                                                              │
│   215 │   │   │   with trainer.profiler.profile("run_training_batch"):       │
│   216 │   │   │   │   if trainer.lightning_module.automatic_optimization:    │
│   217 │   │   │   │   │   # in automatic optimization, there can only be one │
│ ❱ 218 │   │   │   │   │   batch_output = self.automatic_optimization.run(tra │
│   219 │   │   │   │   else:                                                  │
│   220 │   │   │   │   │   batch_output = self.manual_optimization.run(kwargs │
│   221                                                                        │
│                                                                              │
│ /mnt/home/x/.conda/envs/automatic1111_new/lib/python3.10/site-packages/pyt │
│ orch_lightning/loops/optimization/automatic.py:185 in run                    │
│                                                                              │
│   182 │   │   # ------------------------------                               │
│   183 │   │   # gradient update with accumulated gradients                   │
│   184 │   │   else:                                                          │
│ ❱ 185 │   │   │   self._optimizer_step(kwargs.get("batch_idx", 0), closure)  │
│   186 │   │                                                                  │
│   187 │   │   result = closure.consume_result()                              │
│   188 │   │   if result.loss is None:                                        │
│                                                                              │
│ /mnt/home/x/.conda/envs/automatic1111_new/lib/python3.10/site-packages/pyt │
│ orch_lightning/loops/optimization/automatic.py:261 in _optimizer_step        │
│                                                                              │
│   258 │   │   │   self.optim_progress.optimizer.step.increment_ready()       │
│   259 │   │                                                                  │
│   260 │   │   # model hook                                                   │
│ ❱ 261 │   │   call._call_lightning_module_hook(                              │
│   262 │   │   │   trainer,                                                   │
│   263 │   │   │   "optimizer_step",                                          │
│   264 │   │   │   trainer.current_epoch,                                     │
│                                                                              │
│ /mnt/home/x/.conda/envs/automatic1111_new/lib/python3.10/site-packages/pyt │
│ orch_lightning/trainer/call.py:142 in _call_lightning_module_hook            │
│                                                                              │
│   139 │   pl_module._current_fx_name = hook_name                             │
│   140 │                                                                      │
│   141 │   with trainer.profiler.profile(f"[LightningModule]{pl_module.__clas │
│ ❱ 142 │   │   output = fn(*args, **kwargs)                                   │
│   143 │                                                                      │
│   144 │   # restore current_fx when nested context                           │
│   145 │   pl_module._current_fx_name = prev_fx_name                          │
│                                                                              │
│ /mnt/home/x/.conda/envs/automatic1111_new/lib/python3.10/site-packages/pyt │
│ orch_lightning/core/module.py:1265 in optimizer_step                         │
│                                                                              │
│   1262 │   │   │   │   │   for pg in optimizer.param_groups:                 │
│   1263 │   │   │   │   │   │   pg["lr"] = lr_scale * self.learning_rate      │
│   1264 │   │   """                                                           │
│ ❱ 1265 │   │   optimizer.step(closure=optimizer_closure)                     │
│   1266 │                                                                     │
│   1267 │   def optimizer_zero_grad(self, epoch: int, batch_idx: int, optimiz │
│   1268 │   │   """Override this method to change the default behaviour of `` │
│                                                                              │
│ /mnt/home/x/.conda/envs/automatic1111_new/lib/python3.10/site-packages/pyt │
│ orch_lightning/core/optimizer.py:158 in step                                 │
│                                                                              │
│   155 │   │   │   raise MisconfigurationException("When `optimizer.step(clos │
│   156 │   │                                                                  │
│   157 │   │   assert self._strategy is not None                              │
│ ❱ 158 │   │   step_output = self._strategy.optimizer_step(self._optimizer, c │
│   159 │   │                                                                  │
│   160 │   │   self._on_after_step()                                          │
│   161                                                                        │
│                                                                              │
│ /mnt/home/x/.conda/envs/automatic1111_new/lib/python3.10/site-packages/pyt │
│ orch_lightning/strategies/ddp.py:259 in optimizer_step                       │
│                                                                              │
│   256 │   │   │   model: reference to the model, optionally defining optimiz │
│   257 │   │   │   **kwargs: Any extra arguments to ``optimizer.step``        │
│   258 │   │   """                                                            │
│ ❱ 259 │   │   optimizer_output = super().optimizer_step(optimizer, closure,  │
│   260 │   │                                                                  │
│   261 │   │   if self._model_averager is None:                               │
│   262 │   │   │   return optimizer_output                                    │
│                                                                              │
│ /mnt/home/x/.conda/envs/automatic1111_new/lib/python3.10/site-packages/pyt │
│ orch_lightning/strategies/strategy.py:224 in optimizer_step                  │
│                                                                              │
│   221 │   │   model = model or self.lightning_module                         │
│   222 │   │   # TODO(fabric): remove assertion once strategy's optimizer_ste │
│   223 │   │   assert isinstance(model, pl.LightningModule)                   │
│ ❱ 224 │   │   return self.precision_plugin.optimizer_step(optimizer, model=m │
│   225 │                                                                      │
│   226 │   def _setup_model_and_optimizers(self, model: Module, optimizers: L │
│   227 │   │   """Setup a model and multiple optimizers together.             │
│                                                                              │
│ /mnt/home/x/.conda/envs/automatic1111_new/lib/python3.10/site-packages/pyt │
│ orch_lightning/plugins/precision/deepspeed.py:92 in optimizer_step           │
│                                                                              │
│    89 │   ) -> Any:                                                          │
│    90 │   │   if isinstance(optimizer, LBFGS):                               │
│    91 │   │   │   raise MisconfigurationException("DeepSpeed and the LBFGS o │
│ ❱  92 │   │   closure_result = closure()                                     │
│    93 │   │   self._after_closure(model, optimizer)                          │
│    94 │   │   skipped_backward = closure_result is None                      │
│    95 │   │   # in manual optimization, the closure does not return a value  │
│                                                                              │
│ /mnt/home/x/.conda/envs/automatic1111_new/lib/python3.10/site-packages/pyt │
│ orch_lightning/loops/optimization/automatic.py:140 in __call__               │
│                                                                              │
│   137 │   │   return step_output                                             │
│   138 │                                                                      │
│   139 │   def __call__(self, *args: Any, **kwargs: Any) -> Optional[Tensor]: │
│ ❱ 140 │   │   self._result = self.closure(*args, **kwargs)                   │
│   141 │   │   return self._result.loss                                       │
│   142                                                                        │
│   143                                                                        │
│                                                                              │
│ /mnt/home/x/.conda/envs/automatic1111_new/lib/python3.10/site-packages/pyt │
│ orch_lightning/loops/optimization/automatic.py:135 in closure                │
│                                                                              │
│   132 │   │   │   self._zero_grad_fn()                                       │
│   133 │   │                                                                  │
│   134 │   │   if self._backward_fn is not None and step_output.closure_loss  │
│ ❱ 135 │   │   │   self._backward_fn(step_output.closure_loss)                │
│   136 │   │                                                                  │
│   137 │   │   return step_output                                             │
│   138                                                                        │
│                                                                              │
│ /mnt/home/x/.conda/envs/automatic1111_new/lib/python3.10/site-packages/pyt │
│ orch_lightning/loops/optimization/automatic.py:233 in backward_fn            │
│                                                                              │
│   230 │   │   │   return None                                                │
│   231 │   │                                                                  │
│   232 │   │   def backward_fn(loss: Tensor) -> None:                         │
│ ❱ 233 │   │   │   call._call_strategy_hook(self.trainer, "backward", loss, o │
│   234 │   │                                                                  │
│   235 │   │   return backward_fn                                             │
│   236                                                                        │
│                                                                              │
│ /mnt/home/x/.conda/envs/automatic1111_new/lib/python3.10/site-packages/pyt │
│ orch_lightning/trainer/call.py:288 in _call_strategy_hook                    │
│                                                                              │
│   285 │   │   return                                                         │
│   286 │                                                                      │
│   287 │   with trainer.profiler.profile(f"[Strategy]{trainer.strategy.__clas │
│ ❱ 288 │   │   output = fn(*args, **kwargs)                                   │
│   289 │                                                                      │
│   290 │   # restore current_fx when nested context                           │
│   291 │   pl_module._current_fx_name = prev_fx_name                          │
│                                                                              │
│ /mnt/home/x/.conda/envs/automatic1111_new/lib/python3.10/site-packages/pyt │
│ orch_lightning/strategies/strategy.py:199 in backward                        │
│                                                                              │
│   196 │   │   assert self.lightning_module is not None                       │
│   197 │   │   closure_loss = self.precision_plugin.pre_backward(closure_loss │
│   198 │   │                                                                  │
│ ❱ 199 │   │   self.precision_plugin.backward(closure_loss, self.lightning_mo │
│   200 │   │                                                                  │
│   201 │   │   closure_loss = self.precision_plugin.post_backward(closure_los │
│   202 │   │   self.post_backward(closure_loss)                               │
│                                                                              │
│ /mnt/home/x/.conda/envs/automatic1111_new/lib/python3.10/site-packages/pyt │
│ orch_lightning/plugins/precision/deepspeed.py:81 in backward                 │
│                                                                              │
│    78 │   │   │   │   " the backward logic internally."                      │
│    79 │   │   │   )                                                          │
│    80 │   │   deepspeed_engine: "deepspeed.DeepSpeedEngine" = model.trainer. │
│ ❱  81 │   │   deepspeed_engine.backward(tensor, *args, **kwargs)             │
│    82 │                                                                      │
│    83 │   def optimizer_step(  # type: ignore[override]                      │
│    84 │   │   self,                                                          │
│                                                                              │
│ /mnt/home/x/.conda/envs/automatic1111_new/lib/python3.10/site-packages/dee │
│ pspeed/utils/nvtx.py:15 in wrapped_fn                                        │
│                                                                              │
│   12 │                                                                       │
│   13 │   def wrapped_fn(*args, **kwargs):                                    │
│   14 │   │   get_accelerator().range_push(func.__qualname__)                 │
│ ❱ 15 │   │   ret_val = func(*args, **kwargs)                                 │
│   16 │   │   get_accelerator().range_pop()                                   │
│   17 │   │   return ret_val                                                  │
│   18                                                                         │
│                                                                              │
│ /mnt/home/x/.conda/envs/automatic1111_new/lib/python3.10/site-packages/dee │
│ pspeed/runtime/engine.py:1845 in backward                                    │
│                                                                              │
│   1842 │   │                                                                 │
│   1843 │   │   if self.zero_optimization():                                  │
│   1844 │   │   │   self.optimizer.is_gradient_accumulation_boundary = self.i │
│ ❱ 1845 │   │   │   self.optimizer.backward(loss, retain_graph=retain_graph)  │
│   1846 │   │   elif self.amp_enabled():                                      │
│   1847 │   │   │   # AMP requires delaying unscale when inside gradient accu │
│   1848 │   │   │   # https://nvidia.github.io/apex/advanced.html#gradient-ac │
│                                                                              │
│ /mnt/home/x/.conda/envs/automatic1111_new/lib/python3.10/site-packages/dee │
│ pspeed/runtime/zero/stage_1_and_2.py:1901 in backward                        │
│                                                                              │
│   1898 │   │   │   scaled_loss = self.external_loss_scale * loss             │
│   1899 │   │   │   scaled_loss.backward()                                    │
│   1900 │   │   else:                                                         │
│ ❱ 1901 │   │   │   self.loss_scaler.backward(loss.float(), retain_graph=reta │
│   1902 │                                                                     │
│   1903 │   def check_overflow(self, partition_gradients=True):               │
│   1904 │   │   self._check_overflow(partition_gradients)                     │
│                                                                              │
│ /mnt/home/x/.conda/envs/automatic1111_new/lib/python3.10/site-packages/dee │
│ pspeed/runtime/fp16/loss_scaler.py:62 in backward                            │
│                                                                              │
│    59 │                                                                      │
│    60 │   def backward(self, loss, retain_graph=False):                      │
│    61 │   │   scaled_loss = loss * self.loss_scale                           │
│ ❱  62 │   │   scaled_loss.backward(retain_graph=retain_graph)                │
│    63 │   │   # print(f'LossScalerBackward: {scaled_loss=}')                 │
│    64                                                                        │
│    65                                                                        │
│                                                                              │
│ /mnt/home/x/.conda/envs/automatic1111_new/lib/python3.10/site-packages/tor │
│ ch/_tensor.py:487 in backward                                                │
│                                                                              │
│    484 │   │   │   │   create_graph=create_graph,                            │
│    485 │   │   │   │   inputs=inputs,                                        │
│    486 │   │   │   )                                                         │
│ ❱  487 │   │   torch.autograd.backward(                                      │
│    488 │   │   │   self, gradient, retain_graph, create_graph, inputs=inputs │
│    489 │   │   )                                                             │
│    490                                                                       │
│                                                                              │
│ /mnt/home/x/.conda/envs/automatic1111_new/lib/python3.10/site-packages/tor │
│ ch/autograd/__init__.py:200 in backward                                      │
│                                                                              │
│   197 │   # The reason we repeat same the comment below is that              │
│   198 │   # some Python versions print out the first line of a multi-line fu │
│   199 │   # calls in the traceback and some print out the last line          │
│ ❱ 200 │   Variable._execution_engine.run_backward(  # Calls into the C++ eng │
│   201 │   │   tensors, grad_tensors_, retain_graph, create_graph, inputs,    │
│   202 │   │   allow_unreachable=True, accumulate_grad=True)  # Calls into th │
│   203                                                                        │
╰──────────────────────────────────────────────────────────────────────────────╯
OutOfMemoryError: CUDA out of memory. Tried to allocate 2.00 GiB (GPU 0; 31.75 
GiB total capacity; 28.59 GiB already allocated; 1.03 GiB free; 29.40 GiB 
reserved in total by PyTorch)