CUDA error: out of memory with zero3 offload
### System Info

WSL2

- `transformers` version: 4.30.2
- Platform: Linux-5.15.90.1-microsoft-standard-WSL2-x86_64-with-glibc2.35
- Python version: 3.11.4
- Huggingface_hub version: 0.15.1
- Safetensors version: 0.3.1
- PyTorch version (GPU?): 2.0.1+cu117 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: yes
- Using distributed or parallel set-up in script?: no


### Who can help?

@pacman100

### Information

- [X] The official example scripts
- [ ] My own modified scripts

### Tasks

- [X] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

1. the example notebook `https://huggingface.co/ctheodoris/Geneformer/blob/main/examples/cell_classification.ipynb`
2. modify here `training_args_init = TrainingArguments(**training_args, deepspeed ='ds_config_zero3.json')`
3. the example dataset `https://huggingface.co/datasets/ctheodoris/Genecorpus-30M/tree/main/example_input_files/cell_classification/cell_type_annotation/cell_type_train_data.dataset`

ds_config_zero3.json:
This json worked well with transformers deepspeed test `deepspeed examples/pytorch/translation/run_translation.py --deepspeed tests/deepspeed/ds_config_zero3.json --model_name_or_path t5-small --output_dir output_dir --do_eval --max_eval_samples 50 --warmup_steps 50  --max_source_length 128 --val_max_target_length 128 --overwrite_output_dir --per_device_eval_batch_size 4 --predict_with_generate --dataset_config "ro-en" --fp16 --source_lang en --target_lang ro --dataset_name wmt16 --source_prefix "translate English to Romanian: "`  
I also confirmed the CPU offload with `https://github.com/huggingface/transformers-bloom-inference` (the transfer from VRAM to CPU RAM)
```
{
    "fp16": {
        "enabled": "auto",
        "loss_scale": 0,
        "loss_scale_window": 1000,
        "initial_scale_power": 16,
        "hysteresis": 2,
        "min_loss_scale": 1
    },

    "bf16": {
        "enabled": "auto"
    },

    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": "auto",
            "betas": "auto",
            "eps": "auto",
            "weight_decay": "auto"
        }
    },

    "scheduler": {
        "type": "WarmupLR",
        "params": {
            "warmup_min_lr": "auto",
            "warmup_max_lr": "auto",
            "warmup_num_steps": "auto"
        }
    },

    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": true
        },
        "offload_param": {
            "device": "cpu",
            "pin_memory": true
        },
        "overlap_comm": true,
        "contiguous_gradients": true,
        "sub_group_size": 1e8,
        "reduce_bucket_size": "auto",
        "stage3_prefetch_bucket_size": "auto",
        "stage3_param_persistence_threshold": "auto",
        "stage3_max_live_parameters": 1e8,
        "stage3_max_reuse_distance": 1e8,
        "stage3_gather_16bit_weights_on_model_save": true
    },

    "gradient_accumulation_steps": "auto",
    "gradient_clipping": "auto",
    "steps_per_print": 2000,
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "wall_clock_breakdown": false
}
```

message and error:
```
DESKTOP-6FHRRIO:1110179:1110179 [0] NCCL INFO Bootstrap : Using eth0:172.31.110.212<0>
DESKTOP-6FHRRIO:1110179:1110179 [0] NCCL INFO NET/Plugin : No plugin found (libnccl-net.so), using internal implementation

DESKTOP-6FHRRIO:1110179:1110179 [0] misc/cudawrap.cc:90 NCCL WARN Failed to find CUDA library in (null) (NCCL_CUDA_PATH=(null))
NCCL version 2.14.3+cuda11.7
DESKTOP-6FHRRIO:1110179:1110525 [0] NCCL INFO Failed to open libibverbs.so[.1]
DESKTOP-6FHRRIO:1110179:1110525 [0] NCCL INFO NET/Socket : Using [0]eth0:172.31.110.212<0>
DESKTOP-6FHRRIO:1110179:1110525 [0] NCCL INFO Using network Socket
DESKTOP-6FHRRIO:1110179:1110525 [0] NCCL INFO Channel 00/32 :    0
DESKTOP-6FHRRIO:1110179:1110525 [0] NCCL INFO Channel 01/32 :    0
DESKTOP-6FHRRIO:1110179:1110525 [0] NCCL INFO Channel 02/32 :    0
DESKTOP-6FHRRIO:1110179:1110525 [0] NCCL INFO Channel 03/32 :    0
DESKTOP-6FHRRIO:1110179:1110525 [0] NCCL INFO Channel 04/32 :    0
DESKTOP-6FHRRIO:1110179:1110525 [0] NCCL INFO Channel 05/32 :    0
DESKTOP-6FHRRIO:1110179:1110525 [0] NCCL INFO Channel 06/32 :    0
DESKTOP-6FHRRIO:1110179:1110525 [0] NCCL INFO Channel 07/32 :    0
DESKTOP-6FHRRIO:1110179:1110525 [0] NCCL INFO Channel 08/32 :    0
DESKTOP-6FHRRIO:1110179:1110525 [0] NCCL INFO Channel 09/32 :    0
DESKTOP-6FHRRIO:1110179:1110525 [0] NCCL INFO Channel 10/32 :    0
DESKTOP-6FHRRIO:1110179:1110525 [0] NCCL INFO Channel 11/32 :    0
DESKTOP-6FHRRIO:1110179:1110525 [0] NCCL INFO Channel 12/32 :    0
DESKTOP-6FHRRIO:1110179:1110525 [0] NCCL INFO Channel 13/32 :    0
DESKTOP-6FHRRIO:1110179:1110525 [0] NCCL INFO Channel 14/32 :    0
DESKTOP-6FHRRIO:1110179:1110525 [0] NCCL INFO Channel 15/32 :    0
DESKTOP-6FHRRIO:1110179:1110525 [0] NCCL INFO Channel 16/32 :    0
DESKTOP-6FHRRIO:1110179:1110525 [0] NCCL INFO Channel 17/32 :    0
DESKTOP-6FHRRIO:1110179:1110525 [0] NCCL INFO Channel 18/32 :    0
DESKTOP-6FHRRIO:1110179:1110525 [0] NCCL INFO Channel 19/32 :    0
DESKTOP-6FHRRIO:1110179:1110525 [0] NCCL INFO Channel 20/32 :    0
DESKTOP-6FHRRIO:1110179:1110525 [0] NCCL INFO Channel 21/32 :    0
DESKTOP-6FHRRIO:1110179:1110525 [0] NCCL INFO Channel 22/32 :    0
DESKTOP-6FHRRIO:1110179:1110525 [0] NCCL INFO Channel 23/32 :    0
DESKTOP-6FHRRIO:1110179:1110525 [0] NCCL INFO Channel 24/32 :    0
DESKTOP-6FHRRIO:1110179:1110525 [0] NCCL INFO Channel 25/32 :    0
DESKTOP-6FHRRIO:1110179:1110525 [0] NCCL INFO Channel 26/32 :    0
DESKTOP-6FHRRIO:1110179:1110525 [0] NCCL INFO Channel 27/32 :    0
DESKTOP-6FHRRIO:1110179:1110525 [0] NCCL INFO Channel 28/32 :    0
DESKTOP-6FHRRIO:1110179:1110525 [0] NCCL INFO Channel 29/32 :    0
DESKTOP-6FHRRIO:1110179:1110525 [0] NCCL INFO Channel 30/32 :    0
DESKTOP-6FHRRIO:1110179:1110525 [0] NCCL INFO Channel 31/32 :    0
DESKTOP-6FHRRIO:1110179:1110525 [0] NCCL INFO Trees [0] -1/-1/-1->0->-1 [1] -1/-1/-1->0->-1 [2] -1/-1/-1->0->-1 [3] -1/-1/-1->0->-1 [4] -1/-1/-1->0->-1 [5] -1/-1/-1->0->-1 [6] -1/-1/-1->0->-1 [7] -1/-1/-1->0->-1 [8] -1/-1/-1->0->-1 [9] -1/-1/-1->0->-1 [10] -1/-1/-1->0->-1 [11] -1/-1/-1->0->-1 [12] -1/-1/-1->0->-1 [13] -1/-1/-1->0->-1 [14] -1/-1/-1->0->-1 [15] -1/-1/-1->0->-1 [16] -1/-1/-1->0->-1 [17] -1/-1/-1->0->-1 [18] -1/-1/-1->0->-1 [19] -1/-1/-1->0->-1 [20] -1/-1/-1->0->-1 [21] -1/-1/-1->0->-1 [22] -1/-1/-1->0->-1 [23] -1/-1/-1->0->-1 [24] -1/-1/-1->0->-1 [25] -1/-1/-1->0->-1 [26] -1/-1/-1->0->-1 [27] -1/-1/-1->0->-1 [28] -1/-1/-1->0->-1 [29] -1/-1/-1->0->-1 [30] -1/-1/-1->0->-1 [31] -1/-1/-1->0->-1
DESKTOP-6FHRRIO:1110179:1110525 [0] NCCL INFO Connected all rings
DESKTOP-6FHRRIO:1110179:1110525 [0] NCCL INFO Connected all trees
DESKTOP-6FHRRIO:1110179:1110525 [0] NCCL INFO 32 coll channels, 32 p2p channels, 32 p2p channels per peer
DESKTOP-6FHRRIO:1110179:1110525 [0] NCCL INFO comm 0x560d4b0d62b0 rank 0 nranks 1 cudaDev 0 busId 3000 - Init COMPLETE
#################come on tensor([0., 0., 0.,  ..., 0., 0., 0.])
#################come on tensor([0., 0., 0.,  ..., 0., 0., 0.])
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[9], line 1
----> 1 trainer.train()

File /home/pc/miniconda3/envs/Transformers/lib/python3.11/site-packages/transformers/trainer.py:1645, in Trainer.train(self, resume_from_checkpoint, trial, ignore_keys_for_eval, **kwargs)
   1640     self.model_wrapped = self.model
   1642 inner_training_loop = find_executable_batch_size(
   1643     self._inner_training_loop, self._train_batch_size, args.auto_find_batch_size
   1644 )
-> 1645 return inner_training_loop(
   1646     args=args,
   1647     resume_from_checkpoint=resume_from_checkpoint,
   1648     trial=trial,
   1649     ignore_keys_for_eval=ignore_keys_for_eval,
   1650 )

File /home/pc/miniconda3/envs/Transformers/lib/python3.11/site-packages/transformers/trainer.py:1759, in Trainer._inner_training_loop(self, batch_size, args, resume_from_checkpoint, trial, ignore_keys_for_eval)
   1756             model, self.optimizer = self.accelerator.prepare(self.model, self.optimizer)
   1757     else:
   1758         # to handle cases wherein we pass "DummyScheduler" such as when it is specified in DeepSpeed config.
-> 1759         model, self.optimizer, self.lr_scheduler = self.accelerator.prepare(
   1760             self.model, self.optimizer, self.lr_scheduler
   1761         )
   1763 if self.is_fsdp_enabled:
   1764     self.model = model

File /home/pc/miniconda3/envs/Transformers/lib/python3.11/site-packages/accelerate/accelerator.py:1178, in Accelerator.prepare(self, device_placement, *args)
   1176         args = self._prepare_ipex(*args)
   1177 if self.distributed_type == DistributedType.DEEPSPEED:
-> 1178     result = self._prepare_deepspeed(*args)
   1179 elif self.distributed_type == DistributedType.MEGATRON_LM:
   1180     result = self._prepare_megatron_lm(*args)

File /home/pc/miniconda3/envs/Transformers/lib/python3.11/site-packages/accelerate/accelerator.py:1505, in Accelerator._prepare_deepspeed(self, *args)
   1502             if type(scheduler).__name__ in deepspeed.runtime.lr_schedules.VALID_LR_SCHEDULES:
   1503                 kwargs["lr_scheduler"] = scheduler
-> 1505 engine, optimizer, _, lr_scheduler = deepspeed.initialize(**kwargs)
   1506 if optimizer is not None:
   1507     optimizer = DeepSpeedOptimizerWrapper(optimizer)

File /home/pc/miniconda3/envs/Transformers/lib/python3.11/site-packages/deepspeed/__init__.py:165, in initialize(args, model, optimizer, model_parameters, training_data, lr_scheduler, mpu, dist_init_required, collate_fn, config, config_params)
    153         engine = DeepSpeedHybridEngine(args=args,
    154                                        model=model,
    155                                        optimizer=optimizer,
   (...)
    162                                        config=config,
    163                                        config_class=config_class)
    164     else:
--> 165         engine = DeepSpeedEngine(args=args,
    166                                  model=model,
    167                                  optimizer=optimizer,
    168                                  model_parameters=model_parameters,
    169                                  training_data=training_data,
    170                                  lr_scheduler=lr_scheduler,
    171                                  mpu=mpu,
    172                                  dist_init_required=dist_init_required,
    173                                  collate_fn=collate_fn,
    174                                  config=config,
    175                                  config_class=config_class)
    176 else:
    177     assert mpu is None, "mpu must be None with pipeline parallelism"

File /home/pc/miniconda3/envs/Transformers/lib/python3.11/site-packages/deepspeed/runtime/engine.py:309, in DeepSpeedEngine.__init__(self, args, model, optimizer, model_parameters, training_data, lr_scheduler, mpu, dist_init_required, collate_fn, config, config_class, dont_change_device)
    306     model_parameters = list(model_parameters)
    308 if has_optimizer:
--> 309     self._configure_optimizer(optimizer, model_parameters)
    310     self._configure_lr_scheduler(lr_scheduler)
    311     self._report_progress(0)

File /home/pc/miniconda3/envs/Transformers/lib/python3.11/site-packages/deepspeed/runtime/engine.py:1184, in DeepSpeedEngine._configure_optimizer(self, client_optimizer, model_parameters)
   1181 optimizer_wrapper = self._do_optimizer_sanity_check(basic_optimizer)
   1183 if optimizer_wrapper == ZERO_OPTIMIZATION:
-> 1184     self.optimizer = self._configure_zero_optimizer(basic_optimizer)
   1185 elif optimizer_wrapper == AMP:
   1186     amp_params = self.amp_params()

File /home/pc/miniconda3/envs/Transformers/lib/python3.11/site-packages/deepspeed/runtime/engine.py:1474, in DeepSpeedEngine._configure_zero_optimizer(self, optimizer)
   1472         log_dist(f'Creating {model_dtype} ZeRO stage {zero_stage} optimizer', ranks=[0])
   1473         from deepspeed.runtime.zero.stage3 import DeepSpeedZeroOptimizer_Stage3
-> 1474         optimizer = DeepSpeedZeroOptimizer_Stage3(
   1475             self.module,
   1476             optimizer,
   1477             timers=timers,
   1478             ds_config=self.config,
   1479             static_loss_scale=self.loss_scale(),
   1480             dynamic_loss_scale=self.dynamic_loss_scale(),
   1481             dynamic_loss_args=self.dynamic_loss_scale_args(),
   1482             clip_grad=self.gradient_clipping(),
   1483             contiguous_gradients=self.zero_contiguous_gradients(),
   1484             reduce_bucket_size=self.zero_reduce_bucket_size(),
   1485             prefetch_bucket_size=self.zero_prefetch_bucket_size(),
   1486             max_reuse_distance=self.zero_max_reuse_distance(),
   1487             max_live_parameters=self.zero_max_live_parameters(),
   1488             param_persistence_threshold=self.zero_param_persistence_threshold(),
   1489             model_persistence_threshold=self.zero_model_persistence_threshold(),
   1490             dp_process_group=self.data_parallel_group,
   1491             reduce_scatter=self.zero_reduce_scatter(),
   1492             overlap_comm=self.zero_overlap_comm(),
   1493             offload_optimizer_config=self.zero_offload_optimizer(),
   1494             offload_param_config=self.zero_offload_param(),
   1495             sub_group_size=self.zero_sub_group_size(),
   1496             mpu=self.mpu,
   1497             postscale_gradients=self.postscale_gradients(),
   1498             gradient_predivide_factor=self.gradient_predivide_factor(),
   1499             gradient_accumulation_steps=self.gradient_accumulation_steps(),
   1500             aio_config=self.aio_config(),
   1501             communication_data_type=self.communication_data_type)
   1503 else:
   1504     raise NotImplementedError("ZeRO stage {} not implemented".format(zero_stage))

File /home/pc/miniconda3/envs/Transformers/lib/python3.11/site-packages/deepspeed/runtime/zero/stage3.py:149, in DeepSpeedZeroOptimizer_Stage3.__init__(self, module, init_optimizer, timers, ds_config, static_loss_scale, dynamic_loss_scale, dynamic_loss_args, verbose, contiguous_gradients, reduce_bucket_size, prefetch_bucket_size, max_reuse_distance, max_live_parameters, param_persistence_threshold, model_persistence_threshold, dp_process_group, reduce_scatter, overlap_comm, offload_optimizer_config, offload_param_config, sub_group_size, mpu, clip_grad, communication_data_type, postscale_gradients, gradient_predivide_factor, gradient_accumulation_steps, elastic_checkpoint, aio_config)
    146 self.params_in_nvme_and_cpu = False
    147 self.max_params_in_cpu = 0
--> 149 self.parameter_offload = self.initialize_ds_offload(module=module,
    150                                                     timers=timers,
    151                                                     ds_config=ds_config,
    152                                                     overlap_comm=overlap_comm,
    153                                                     prefetch_bucket_size=prefetch_bucket_size,
    154                                                     max_reuse_distance=max_reuse_distance,
    155                                                     max_live_parameters=max_live_parameters,
    156                                                     param_persistence_threshold=param_persistence_threshold,
    157                                                     model_persistence_threshold=model_persistence_threshold,
    158                                                     offload_param_config=offload_param_config,
    159                                                     mpu=mpu)
    161 self.persistent_parameters = self.parameter_offload.persistent_parameters
    162 self._configure_offloading(offload_optimizer_config, offload_param_config)

File /home/pc/miniconda3/envs/Transformers/lib/python3.11/site-packages/deepspeed/runtime/zero/stage3.py:352, in DeepSpeedZeroOptimizer_Stage3.initialize_ds_offload(self, module, timers, ds_config, overlap_comm, prefetch_bucket_size, max_reuse_distance, max_live_parameters, param_persistence_threshold, model_persistence_threshold, offload_param_config, mpu)
    338 def initialize_ds_offload(
    339     self,
    340     module,
   (...)
    350     mpu,
    351 ):
--> 352     return DeepSpeedZeRoOffload(module=module,
    353                                 timers=timers,
    354                                 ds_config=ds_config,
    355                                 overlap_comm=overlap_comm,
    356                                 prefetch_bucket_size=prefetch_bucket_size,
    357                                 max_reuse_distance=max_reuse_distance,
    358                                 max_live_parameters=max_live_parameters,
    359                                 param_persistence_threshold=param_persistence_threshold,
    360                                 model_persistence_threshold=model_persistence_threshold,
    361                                 offload_param_config=offload_param_config,
    362                                 mpu=mpu)

File /home/pc/miniconda3/envs/Transformers/lib/python3.11/site-packages/deepspeed/runtime/zero/parameter_offload.py:229, in DeepSpeedZeRoOffload.__init__(self, module, timers, ds_config, overlap_comm, prefetch_bucket_size, max_reuse_distance, max_live_parameters, param_persistence_threshold, model_persistence_threshold, offload_param_config, mpu)
    226     self.offload_device = offload_param_config.device
    227     self.offload_param_pin_memory = offload_param_config.pin_memory
--> 229 self._convert_to_zero_parameters(ds_config, module, mpu)
    231 for m in module.modules():
    232     _init_external_params(m)

File /home/pc/miniconda3/envs/Transformers/lib/python3.11/site-packages/deepspeed/runtime/zero/parameter_offload.py:297, in DeepSpeedZeRoOffload._convert_to_zero_parameters(self, ds_config, module, mpu)
    294 if mpu:
    295     group = mpu.get_data_parallel_group()
--> 297 Init(module=module,
    298      data_parallel_group=group,
    299      dtype=self.dtype,
    300      config_dict_or_path=ds_config,
    301      remote_device=self.offload_device,
    302      pin_memory=self.offload_param_pin_memory,
    303      mpu=mpu)

File /home/pc/miniconda3/envs/Transformers/lib/python3.11/site-packages/deepspeed/runtime/zero/partition_parameters.py:782, in Init.__init__(self, module, data_parallel_group, mem_efficient_linear, remote_device, pin_memory, config_dict_or_path, config, enabled, dtype, mpu)
    780 if module is not None:
    781     assert isinstance(module, torch.nn.Module)
--> 782     self._convert_to_zero_parameters(module.parameters(recurse=True))
    784 self.use_all_gather_into_tensor = dist.has_all_gather_into_tensor()
    785 if not self.use_all_gather_into_tensor:

File /home/pc/miniconda3/envs/Transformers/lib/python3.11/site-packages/deepspeed/runtime/zero/partition_parameters.py:798, in Init._convert_to_zero_parameters(self, param_list)
    796     continue
    797 self._convert_to_deepspeed_param(param)
--> 798 param.partition()

File /home/pc/miniconda3/envs/Transformers/lib/python3.11/site-packages/deepspeed/runtime/zero/partition_parameters.py:966, in Init._convert_to_deepspeed_param.<locals>.partition(param_list, hierarchy, has_been_updated)
    964 if param_list is None:
    965     param_list = [cls]
--> 966 self._partition(param_list, has_been_updated=has_been_updated)

File /home/pc/miniconda3/envs/Transformers/lib/python3.11/site-packages/deepspeed/runtime/zero/partition_parameters.py:1104, in Init._partition(self, param_list, force, has_been_updated)
   1100 def _partition(self, param_list, force=False, has_been_updated=False):
   1101     for param in param_list:
   1102         #print_rank_0(f"Before Partitioning Param {param.ds_id}")
   1103         # self._param_status(param)
-> 1104         self._partition_param(param, has_been_updated=has_been_updated)
   1105         param.ds_status = ZeroParamStatus.NOT_AVAILABLE

File /home/pc/miniconda3/envs/Transformers/lib/python3.11/site-packages/deepspeed/utils/nvtx.py:15, in instrument_w_nvtx.<locals>.wrapped_fn(*args, **kwargs)
     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

File /home/pc/miniconda3/envs/Transformers/lib/python3.11/site-packages/deepspeed/runtime/zero/partition_parameters.py:1186, in Init._partition_param(self, param, buffer, has_been_updated)
   1183 if start < param.ds_numel and end <= param.ds_numel:
   1184     src_tensor = one_dim_param.narrow(0, start, partition_size)
-> 1186     param.ds_tensor.copy_(src_tensor)
   1187     #partitioned_tensor = src_tensor.clone().detach().to(self.remote_device)
   1188
   1189 else:
   1190     # partitioned_tensor = torch.zeros(partition_size,
   1191     #                                  dtype=param.dtype,
   1192     #                                  device=self.remote_device )
   1194     if start < param.ds_numel:

RuntimeError: CUDA error: out of memory
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.
```

I

### Expected behavior

I expect the CPU offload to work so I can use a larger batch number (current 2 working without `deepspeed ='ds_config_zero3.json'`). However, with DeepSpeed, even batch 1 did not work with the same error.
Hello, are you running the notebook as is? or are you running it as a script with distributed launcher such as `deepspeed`/`torchrun`/`accelerate launch`?

You can't run DeepSpeed in a notebook. You need to convert the notebook to a script and run the script via a distributed launcher similar to the translation example that you are running<|||||>Worked!
Thank you @pacman100<|||||>Hi @pacman100

I encountered another issue. The .py ran fine alone but with deepspeed it encountered `server socket has failed to bind to [::]:29500 (errno: 98 - Address already in use).` It could not be resolved by specifying a port `os.environ["MASTER_PORT"] = "9994"`. Was it because ray assigns multiple ports at the same time and I only have 1 GPU?  Thank you.

```

from ray.air import session

def train(config):
    # ...
    session.report({"metric": metric}, checkpoint=checkpoint)

For more information please see https://docs.ray.io/en/latest/tune/api/trainable.html

  warnings.warn(
== Status ==
Current time: 2023-07-05 16:13:18 (running for 00:00:00.64)
Using FIFO scheduling algorithm.
Logical resource usage: 0/48 CPUs, 0/1 GPUs
Result logdir: /root/ray_results/_objective_2023-07-05_16-13-18
Number of trials: 1/100 (1 PENDING)
+---------------------+----------+-------+-----------------+---------------------+--------------------+------------------------+---------+----------------+----------------+
| Trial name          | status   | loc   |   learning_rate | lr_scheduler_type   |   num_train_epochs |   per_device_train_bat |    seed |   warmup_steps |   weight_decay |
|                     |          |       |                 |                     |                    |                ch_size |         |                |                |
|---------------------+----------+-------+-----------------+---------------------+--------------------+------------------------+---------+----------------+----------------|
| _objective_fd2e0c55 | PENDING  |       |     4.36572e-06 | polynomial          |                  1 |
    12 | 59.5864 |        1832.57 |      0.0619684 |
+---------------------+----------+-------+-----------------+---------------------+--------------------+------------------------+---------+----------------+----------------+


(pid=1172302) [2023-07-05 16:13:25,088] [INFO] [real_accelerator.py:110:get_accelerator] Setting ds_accelerator to cuda (auto detect)
(_objective pid=1172302) /home/pc/miniconda3/envs/Transformers/lib/python3.11/site-packages/loompy/bus_file.py:67: NumbaDeprecationWarning: The 'nopython' keyword argument was not supplied to the 'numba.jit' decorator. The implicit default value for this argument is currently False, but it will be changed to True in Numba 0.59.0. See https://numba.readthedocs.io/en/stable/reference/deprecation.html#deprecation-of-object-mode-fall-back-behaviour-when-using-jit for details.
(_objective pid=1172302)   @jit
(_objective pid=1172302) /home/pc/miniconda3/envs/Transformers/lib/python3.11/site-packages/loompy/bus_file.py:84: NumbaDeprecationWarning: The 'nopython' keyword argument was not supplied to the 'numba.jit' decorator. The implicit default value for this argument is currently False, but it will be changed to True in Numba 0.59.0. See https://numba.readthedocs.io/en/stable/reference/deprecation.html#deprecation-of-object-mode-fall-back-behaviour-when-using-jit for details.
(_objective pid=1172302)   @jit
(_objective pid=1172302) /home/pc/miniconda3/envs/Transformers/lib/python3.11/site-packages/loompy/bus_file.py:101: NumbaDeprecationWarning: The 'nopython' keyword argument was not supplied to the 'numba.jit' decorator. The implicit default value for this argument is currently False, but it will be changed to True in Numba 0.59.0. See https://numba.readthedocs.io/en/stable/reference/deprecation.html#deprecation-of-object-mode-fall-back-behaviour-when-using-jit for details.
(_objective pid=1172302)   @jit
2023-07-05 16:13:34,406 ERROR tune_controller.py:873 -- Trial task failed for trial _objective_fd2e0c55
Traceback (most recent call last):
  File "/home/pc/miniconda3/envs/Transformers/lib/python3.11/site-packages/ray/air/execution/_internal/event_manager.py", line 110, in resolve_future
    result = ray.get(future)
             ^^^^^^^^^^^^^^^
  File "/home/pc/miniconda3/envs/Transformers/lib/python3.11/site-packages/ray/_private/auto_init_hook.py", line 18, in auto_init_wrapper
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/home/pc/miniconda3/envs/Transformers/lib/python3.11/site-packages/ray/_private/client_mode_hook.py", line 103, in wrapper
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/pc/miniconda3/envs/Transformers/lib/python3.11/site-packages/ray/_private/worker.py", line 2540, in get
    raise value.as_instanceof_cause()
ray.exceptions.RayTaskError(RuntimeError): ray::ImplicitFunc.train() (pid=1172302, ip=172.31.110.212, actor_id=ffa19b5f202ac72158b2946001000000, repr=_objective)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/pc/miniconda3/envs/Transformers/lib/python3.11/site-packages/ray/tune/trainable/trainable.py", line 389, in train
    raise skipped from exception_cause(skipped)
  File "/home/pc/miniconda3/envs/Transformers/lib/python3.11/site-packages/ray/tune/trainable/function_trainable.py", line 336, in entrypoint
    return self._trainable_func(
           ^^^^^^^^^^^^^^^^^^^^^
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/pc/miniconda3/envs/Transformers/lib/python3.11/site-packages/ray/tune/trainable/function_trainable.py", line 653, in _trainable_func
    output = fn()
             ^^^^
  File "/home/pc/miniconda3/envs/Transformers/lib/python3.11/site-packages/transformers/integrations.py", line 357, in dynamic_modules_import_trainable
    return trainable(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/pc/miniconda3/envs/Transformers/lib/python3.11/site-packages/ray/tune/trainable/util.py", line 324, in inner
    return trainable(config, **fn_kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/pc/miniconda3/envs/Transformers/lib/python3.11/site-packages/transformers/integrations.py", line 258, in _objective
    local_trainer.train(resume_from_checkpoint=checkpoint, trial=trial)
  File "/home/pc/miniconda3/envs/Transformers/lib/python3.11/site-packages/transformers/trainer.py", line 1614, in train    self._hp_search_setup(trial)
  File "/home/pc/miniconda3/envs/Transformers/lib/python3.11/site-packages/transformers/trainer.py", line 1330, in _hp_search_setup
    self.create_accelerator_and_postprocess()
  File "/home/pc/miniconda3/envs/Transformers/lib/python3.11/site-packages/transformers/trainer.py", line 3968, in create_accelerator_and_postprocess
    self.accelerator = Accelerator(
                       ^^^^^^^^^^^^
  File "/home/pc/miniconda3/envs/Transformers/lib/python3.11/site-packages/accelerate/accelerator.py", line 345, in __init__
    self.state = AcceleratorState(
                 ^^^^^^^^^^^^^^^^^
  File "/home/pc/miniconda3/envs/Transformers/lib/python3.11/site-packages/accelerate/state.py", line 680, in __init__
    PartialState(cpu, **kwargs)
  File "/home/pc/miniconda3/envs/Transformers/lib/python3.11/site-packages/accelerate/state.py", line 191, in __init__
    torch.distributed.init_process_group(backend=self.backend, **kwargs)
  File "/home/pc/miniconda3/envs/Transformers/lib/python3.11/site-packages/torch/distributed/distributed_c10d.py", line 900, in init_process_group
    store, rank, world_size = next(rendezvous_iterator)
                              ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/pc/miniconda3/envs/Transformers/lib/python3.11/site-packages/torch/distributed/rendezvous.py", line 245, in _env_rendezvous_handler
    store = _create_c10d_store(master_addr, master_port, rank, world_size, timeout)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/pc/miniconda3/envs/Transformers/lib/python3.11/site-packages/torch/distributed/rendezvous.py", line 176, in _create_c10d_store
    return TCPStore(
           ^^^^^^^^^
RuntimeError: The server socket has failed to listen on any local network address. The server socket has failed to bind to [::]:29500 (errno: 98 - Address already in use). The server socket has failed to bind to DESKTOP-6FHRRIO:29500 (errno: 98 - Address already in use).
Result for _objective_fd2e0c55:
  date: 2023-07-05_16-13-25
  hostname: DESKTOP-6FHRRIO
  node_ip: 172.31.110.212
  pid: 1172302
  timestamp: 1688544805
  trial_id: fd2e0c55

(_objective pid=1172302) [W socket.cpp:426] [c10d] The server socket has failed to bind to [::]:29500 (errno: 98 - Address already in use).
(_objective pid=1172302) [W socket.cpp:426] [c10d] The server socket has failed to bind to DESKTOP-6FHRRIO:29500 (errno: 98 - Address already in use).
(_objective pid=1172302) [E socket.cpp:462] [c10d] The server socket has failed to listen on any local network address.
(pid=1172435) [2023-07-05 16:13:41,310] [INFO] [real_accelerator.py:110:get_accelerator] Setting ds_accelerator to cuda (auto detect)
(_objective pid=1172435) /home/pc/miniconda3/envs/Transformers/lib/python3.11/site-packages/loompy/bus_file.py:67: NumbaDeprecationWarning: The 'nopython' keyword argument was not supplied to the 'numba.jit' decorator. The implicit default value for this argument is currently False, but it will be changed to True in Numba 0.59.0. See https://numba.readthedocs.io/en/stable/reference/deprecation.html#deprecation-of-object-mode-fall-back-behaviour-when-using-jit for details.
(_objective pid=1172435)   @jit
(_objective pid=1172435) /home/pc/miniconda3/envs/Transformers/lib/python3.11/site-packages/loompy/bus_file.py:84: NumbaDeprecationWarning: The 'nopython' keyword argument was not supplied to the 'numba.jit' decorator. The implicit default value for this argument is currently False, but it will be changed to True in Numba 0.59.0. See https://numba.readthedocs.io/en/stable/reference/deprecation.html#deprecation-of-object-mode-fall-back-behaviour-when-using-jit for details.
(_objective pid=1172435)   @jit
(_objective pid=1172435) /home/pc/miniconda3/envs/Transformers/lib/python3.11/site-packages/loompy/bus_file.py:101: NumbaDeprecationWarning: The 'nopython' keyword argument was not supplied to the 'numba.jit' decorator. The implicit default value for this argument is currently False, but it will be changed to True in Numba 0.59.0. See https://numba.readthedocs.io/en/stable/reference/deprecation.html#deprecation-of-object-mode-fall-back-behaviour-when-using-jit for details.
(_objective pid=1172435)   @jit
^C[2023-07-05 16:13:47,161] [INFO] [launch.py:315:sigkill_handler] Killing subprocess 1168420
*
```<|||||>@pchiang5 Perhaps the process has not been fully shut down, please use `ps -aux` to find the remaining process, then use `kill -9 process_pid` to kill them totally.<|||||>@2033329616 Thank you for your feedback. Yes, it shall be due to an open process not shut down and could be resolved by randomly assigning a new port before the previous call. 

However, I found the incompatibility of ray tune + hyperopt with deepspeed launcher is the main issue: Without ray tune + hyperopt, it ran with successful CPU offload. With deepspeed and ray+hyperopt as below, the zero3 offload did not work because the amount of VRAM consumption was identical to that without deepspeed. 

```

# create the trainer
trainer = Trainer(
    model_init=model_init,
    args=training_args_init,
    data_collator=DataCollatorForCellClassification(),
    train_dataset=organ_trainset,
    eval_dataset=organ_evalset,
    compute_metrics=compute_metrics,
    callbacks = [EarlyStoppingCallback(early_stopping_patience=3)]
)

# specify raytune hyperparameter search space
ray_config = {
    "num_train_epochs": tune.choice([epochs]),
    "learning_rate": tune.loguniform(1e-6, 1e-3),
    "weight_decay": tune.uniform(0.0, 0.3),
    "lr_scheduler_type": tune.choice(["linear","cosine","polynomial"]),
    "warmup_steps": tune.uniform(100, 2000),
    "seed": tune.uniform(0,100),
    "per_device_train_batch_size": tune.choice([geneformer_batch_size])
}

hyperopt_search = HyperOptSearch(
    metric="eval_macro_f1", mode="max")

early_stop = {
    "training_iteration": 10
}

# optimize hyperparameters
trainer.hyperparameter_search(
    direction="maximize",
    backend="ray",
    resources_per_trial={"cpu":18,"gpu":1},
    hp_space=lambda _: ray_config,
    stop=early_stop,
    search_alg=hyperopt_search,
    n_trials=100, # number of trials
    progress_reporter=tune.CLIReporter(max_report_frequency=600,
                                                   sort_by_metric=True,
                                                   max_progress_rows=100,
                                                   mode="max",
                                                   metric="eval_macro_f1",
                                                   metric_columns=["loss", "eval_loss", "eval_accuracy", "eval_macro_f1"])
```<|||||>Hi @pacman100,

> You can't run DeepSpeed in a Jupyter notebook. You need to convert the notebook to a script and run the script via a distributed launcher similar to the translation example that you are running.

I am also having a CUDA OOM error running DeepSpeed in a notebook on a single node with a single GPU (training [Segformer](https://huggingface.co/docs/transformers/model_doc/segformer) on a GPU with 8 GB RAM). I expected it to work, given [the deployment excerpt](https://huggingface.co/docs/transformers/main_classes/deepspeed#deployment-in-notebooks) from the docs. Why would you say so?

I injected the env variables as reported in the docs using stage 3 with CPU offloading, but still the error remains.

