Repository: microsoft/DeepSpeed

Issue:
[BUG] pin_memory() raises an error when using pipeline parallelism
**Describe the bug**
when I use pipeline parallelism and ZeRo-1 offload at the same time, there will be an error:
```
│ /mnt/petrelfs/zhangshuo/miniconda3/envs/dl/lib/python3.9/site-packages/deepspeed/runtime/zero/st │
│ age_1_and_2.py:485 in __init__                                                                   │
│                                                                                                  │
│    482 │   │   self.dynamic_loss_scale = self.loss_scaler.dynamic                                │
│    483 │   │                                                                                     │
│    484 │   │   see_memory_usage("Before initializing optimizer states", force=True)              │
│ ❱  485 │   │   self.initialize_optimizer_states()                                                │
│    486 │   │   see_memory_usage("After initializing optimizer states", force=True)               │
│    487 │   │                                                                                     │
│    488 │   │   if dist.get_rank() == 0:                                                          │
│                                                                                                  │
│ /mnt/petrelfs/zhangshuo/miniconda3/envs/dl/lib/python3.9/site-packages/deepspeed/runtime/zero/st │
│ age_1_and_2.py:615 in initialize_optimizer_states                                                │
│                                                                                                  │
│    612 │   │   │   print(single_grad_partition.dtype)                                            │
│    613 │   │   │   print(single_grad_partition.device)                                           │
│    614 │   │   │   print(single_grad_partition.shape)                                            │
│ ❱  615 │   │   │   self.single_partition_of_fp32_groups[i].grad = get_accelerator().pin_memory(  │
│    616 │   │   │   │   single_grad_partition) if self.cpu_offload else single_grad_partition     │
│    617 │   │                                                                                     │
│    618 │   │   self.optimizer.step()                                                             │
│                                                                                                  │
│ /mnt/petrelfs/zhangshuo/miniconda3/envs/dl/lib/python3.9/site-packages/deepspeed/accelerator/cud │
│ a_accelerator.py:217 in pin_memory                                                               │
│                                                                                                  │
│   214 │   │   return torch.cuda.LongTensor                                                       │
│   215 │                                                                                          │
│   216 │   def pin_memory(self, tensor):                                                          │
│ ❱ 217 │   │   return tensor.pin_memory()                                                         │
│   218 │                                                                                          │
│   219 │   def on_accelerator(self, tensor):                                                      │
│   220 │   │   device_str = str(tensor.device)                                                    │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯
RuntimeError: CUDA error: OS call failed or operation not supported on this OS
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.

```

And this error only happens when my model is really large (maybe 7 billion parameters). 

**To Reproduce**

My code:

```python
# model2.py
import torch
import torch.nn as nn
from torch.utils.data import Dataset

import deepspeed
from deepspeed.pipe import LayerSpec
from deepspeed.pipe import PipelineModule
from deepspeed.runtime.pipe.topology import PipeModelDataParallelTopology

class Layer1(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.linear = nn.Linear(1, 1000000000)
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.linear(x)
    
class Layer2(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.linear = nn.Linear(1000000000, 1)
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.linear(x)

class DummyDataset(Dataset):
    def __init__(self):
        pass
        
    def __len__(self):
        return 100
    
    def __getitem__(self, idx):
        return torch.tensor([idx]).float(), torch.tensor([idx]).float()
    
deepspeed.init_distributed(dist_backend='nccl', init_method="env://")
model = PipelineModule(
    layers=[LayerSpec(Layer1), LayerSpec(Layer2)],
    num_stages=2,
    topology=PipeModelDataParallelTopology(num_pp=2, num_dp=1, num_mp=1),
    loss_fn=lambda x, y: (print(x), x.sum())[1]
)
dataset = DummyDataset()
engine, optimizer, training_dataloader, _ = deepspeed.initialize(
        model=model,
        model_parameters=[p for p in model.parameters() if p.requires_grad],
        training_data=dataset,
        config={
            "train_micro_batch_size_per_gpu": 1,
            "train_batch_size": 1,
            "gradient_accumulation_steps": 1,
            "optimizer": {
                "type": "Adam",
                "params": {
                "lr": 0.001,
                "betas": [
                    0.8,
                    0.999
                ],
                "eps": 1e-8,
                "weight_decay": 3e-7
                }
            },
            "zero_optimization": {
                "stage": 1,
                "offload_optimizer": {
                    "device": "cpu"
                },
                "contiguous_gradients": True,
                "overlap_comm": True,
                "sub_group_size": 1e12,
                "reduce_bucket_size": "auto"
            }
        }
    )
print(engine.train_batch())
```

My command:

```bash
CUDA_VISIBLE_DEVICES=0,1 torchrun --standalone --nproc_per_node=2 model2.py
```

**Maybe a fix?**
I find that this error is caused by the function `pin_memory()` in `deepspeed/runtime/zero/stage_1_and_2.py:426` and `deepspeed/runtime/zero/stage_1_and_2.py:611`. I also find that if I keep executing this function, it has a chance of succeeding. So I changed the code to make it keep trying:

```python
# deepspeed/runtime/zero/stage_1_and_2.py:611
while True:
                try:
                    self.single_partition_of_fp32_groups[i].grad = get_accelerator().pin_memory(
                        single_grad_partition) if self.cpu_offload else single_grad_partition
                    break
                except RuntimeError as e:
                    continue
```

This did fix my problem, but I'm concerned that there may be some additional side effects.


Discussion:
1. I also encountered the same problem. Your solution is effective for opt-1.3B, But when training gpt-3.5B, stuck in the loop for a long time. The larger the model, the longer it will be stuck in the loop, the llama-7B has been stuck for an hour and has not successfully run. Do you have any other solutions now?
2. I think it may be related to the size of RAM. Can you tell me your GPU and RAM size? My device is 2*A6000 (48G) and 128G of RAM, I want to know the device GPU size and RAM size you used when training the 7B model
@tjruwase My devices is 8*A100 (80G) and 1024G of RAM. And I have fonud another solution. I found that the `pin_memory: false`  in `ds_config` didn't do anything. So I add a patch to make the ``get_accelerator().pin_memory`` useless. See [patch_deepspeed()](https://github.com/OpenLMLab/collie/blob/d4679aa483a613a773e4aac063ad1a85da1e62af/collie/utils/dist_utils.py#L180)
@00INDEX, apologies for the delayed response. The underlying problem is that cpu_offload was originally developed for zero-stage-2, and is better tested for that scenario. Did you try using zero-stage-2? 
@xiaotingyun, are you still having this problem?
@tjruwase I can run through the method of 00INDEX. However, But if I don't modify the source code, as long as offload is turned on, whether it is zero-2 or zero-3, there is a high probability that running a 1.3B model will report an error.
I found out that when the model is very large, the pinned tensor will be very large too (contains billions of elements). Seems that pinning this too-large tensor caused this error.
@tjruwase 
@x54-729, yes, memory pinning can incur high memory consumption and cause system instability. Additionally, as @00INDEX previously observed, there is a bug where the ds_config of `pin_memory: false` is ignored in DeepSpeed. We will fix this asap. 

Is your AMD? I find that AMD causes this problem.

Resolving patch:
diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py
--- a/deepspeed/runtime/engine.py
+++ b/deepspeed/runtime/engine.py
@@ -1452,7 +1452,7 @@ def _configure_zero_optimizer(self, optimizer):
                 expert_data_parallel_group=self.expert_data_parallel_group if self.has_moe_layers else None,
                 reduce_scatter=self.zero_reduce_scatter(),
                 overlap_comm=overlap_comm,
-                cpu_offload=self.zero_cpu_offload(),
+                offload_optimizer_config=self.zero_offload_optimizer(),
                 mpu=self.mpu,
                 postscale_gradients=self.postscale_gradients(),
                 gradient_predivide_factor=self.gradient_predivide_factor(),
diff --git a/deepspeed/runtime/zero/stage_1_and_2.py b/deepspeed/runtime/zero/stage_1_and_2.py
--- a/deepspeed/runtime/zero/stage_1_and_2.py
+++ b/deepspeed/runtime/zero/stage_1_and_2.py
@@ -120,7 +120,7 @@ def __init__(self,
                  expert_data_parallel_group=None,
                  reduce_scatter=True,
                  overlap_comm=False,
-                 cpu_offload=False,
+                 offload_optimizer_config=None,
                  mpu=None,
                  clip_grad=0.0,
                  gradient_accumulation_dtype=torch.float32,
@@ -135,10 +135,17 @@ def __init__(self,
                  fp16_master_weights_and_gradients=False,
                  elastic_checkpoint=False):
 
+        if offload_optimizer_config is not None and offload_optimizer_config.device != OffloadDeviceEnum.none:
+            self.cpu_offload = True
+            self.cpu_offload_pin_memory = offload_optimizer_config.pin_memory
+        else:
+            self.cpu_offload = False
+            self.cpu_offload_pin_memory = False
+
         if dist.get_rank() == 0:
             logger.info(f"Reduce bucket size {reduce_bucket_size}")
             logger.info(f"Allgather bucket size {allgather_bucket_size}")
-            logger.info(f"CPU Offload: {cpu_offload}")
+            logger.info(f"CPU Offload: {self.cpu_offload}")
             logger.info(f'Round robin gradient partitioning: {round_robin_gradients}')
         # The fused optimizer does all the work. We need this layer for two reason:
         # 1. maintain same user API from apex.fp16_utils
@@ -153,7 +160,7 @@ def __init__(self,
         # - flat by groups, not keeping state. TODO: remove state explicitly?
         # - master grad and unflat master weight never exist. TODO: a way to save out unflat master?
         if not get_accelerator().is_available():
-            raise SystemError("Cannot use fp16 without accelerator.")
+            raise SystemError("Accelerator is not detected, cannot perform low precision training (e.g., fp16, bf16).")
         self.optimizer = init_optimizer
 
         # Use torch (un)flatten ops
@@ -170,9 +177,7 @@ def __init__(self,
 
         self.overlap_comm = overlap_comm
 
-        self.cpu_offload = cpu_offload
-
-        self.deepspeed_adam_offload = cpu_offload
+        self.deepspeed_adam_offload = self.cpu_offload
 
         self.device = get_accelerator().current_device_name() if not self.cpu_offload else 'cpu'
 
@@ -195,7 +200,7 @@ def __init__(self,
         self.is_gradient_accumulation_boundary = True
 
         # CPU-Offload requires contiguous gradients
-        self.contiguous_gradients = contiguous_gradients or cpu_offload
+        self.contiguous_gradients = contiguous_gradients or self.cpu_offload
 
         self.has_moe_layers = has_moe_layers
         if self.has_moe_layers:
@@ -440,8 +445,12 @@ def __init__(self,
             self.norm_for_param_grads = {}
             self.local_overflow = False
             self.grad_position = {}
-            self.temp_grad_buffer_for_cpu_offload = get_accelerator().pin_memory(
-                torch.zeros(largest_param_numel, device=self.device, dtype=self.dtype))
+            self.temp_grad_buffer_for_cpu_offload = torch.zeros(largest_param_numel,
+                                                                device=self.device,
+                                                                dtype=self.dtype)
+            if self.cpu_offload_pin_memory:
+                self.temp_grad_buffer_for_cpu_offload = get_accelerator().pin_memory(
+                    self.temp_grad_buffer_for_cpu_offload)
             self.temp_grad_buffer_for_gpu_offload = torch.zeros(largest_param_numel,
                                                                 device=get_accelerator().current_device_name(),
                                                                 dtype=self.dtype)
@@ -631,7 +640,7 @@ def initialize_optimizer_states(self):
                                                 dtype=self.single_partition_of_fp32_groups[i].dtype,
                                                 device=self.device)
             self.single_partition_of_fp32_groups[i].grad = get_accelerator().pin_memory(
-                single_grad_partition) if self.cpu_offload else single_grad_partition
+                single_grad_partition) if self.cpu_offload_pin_memory else single_grad_partition
 
         # Initialize the optimizer states with the flattened fp32 partition.
         # State initialization for the Adagrad optimizer occurs at construction as opposed to other optimizers
@@ -1101,7 +1110,8 @@ def async_accumulate_grad_in_cpu_via_gpu(self, param):
         #buffer for storing gradients for this parameter in CPU
         def buffer_to_accumulate_to_in_cpu():
             if not self.fp16_master_weights_and_gradients:
-                return get_accelerator().pin_memory(torch.zeros(param.numel(), dtype=param.dtype, device=self.device))
+                buffer = torch.zeros(param.numel(), dtype=param.dtype, device=self.device)
+                return get_accelerator().pin_memory(buffer) if self.cpu_offload_pin_memory else buffer
             else:
                 return self.single_partition_of_fp32_groups[i].grad.view(-1).narrow(0, dest_offset, num_elements)
 


Test patch: