Can deepspeed's zero_optimization achieve model parallelism?

[ojipadeson]

My machine has 8 Tesla V100s. But when loading LLM (I loaded Qwen2-7B-Instruct), an OOM error will be reported when using a single card.

class Model(nn.Module):
    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        
        # Load model and tokenizer
        self.qwen2_config = AutoConfig.from_pretrained(
            "Qwen/Qwen2-7B",
            cache_dir=None,
        )
        self.qwen2_config.use_cache = False

        self.llm_model = AutoModelForCausalLM.from_pretrained(
            "Qwen/Qwen2-7B",
            config=self.qwen2_config,
            cache_dir=None,
            device_map=None,
            quantization_config=None
            # **model_load_kwargs,
        )
        self.tokenizer = AutoTokenizer.from_pretrained(
            "Qwen/Qwen2-7B",
            cache_dir=None,
            model_max_length=8192,
            padding_side="right",
            use_fast=False,
        )
        
        if self.tokenizer.eos_token:
            self.tokenizer.pad_token = self.tokenizer.eos_token
        else:
            pad_token = '[PAD]'
            self.tokenizer.add_special_tokens({'pad_token': pad_token})
            self.tokenizer.pad_token = pad_token

        for param in self.llm_model.parameters():
            param.requires_grad = False
            
    def forward(self, x_enc):
        prompt = "Hello"
        prompt = self.tokenizer(prompt, return_tensors="pt", padding=True, truncation=True, max_length=2048).input_ids
        prompt_embeddings = self.llm_model.get_input_embeddings()(prompt.to(x_enc.device))  # (batch, prompt_token, dim)
        dec_out = self.llm_model(inputs_embeds=prompt_embeddings).last_hidden_state
        return dec_out

I used deepspeed to divide the model parameters into 8 cards, but it was always unsuccessful (OOM error). I don't know if it can easily implement this function?

my config.json

{
  "bfloat16": {
      "enabled": false
  },
  "fp16": {
      "enabled": "auto",
      "loss_scale": 0,
      "loss_scale_window": 1000,
      "initial_scale_power": 16,
      "hysteresis": 2,
      "min_loss_scale": 1
  },
  "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": 1e9,
      "reduce_bucket_size": "auto",
      "stage3_prefetch_bucket_size": "auto",
      "stage3_param_persistence_threshold": "auto",
      "stage3_max_live_parameters": 1e9,
      "stage3_max_reuse_distance": 1e9,
      "stage3_gather_fp16_weights_on_model_save": true
  },
  "gradient_accumulation_steps": 1,
  "gradient_clipping": "auto",
  "steps_per_print": 1e5,
  "train_batch_size": 8,
  "train_micro_batch_size_per_gpu": 1,
  "wall_clock_breakdown": false
}