class_docstr: Model parallelization processing engine
  
  Args:
      num_gpus (int): number of gpus
      backend (str): distributed backend (default=nccl)
      custom_policies (Union[Policy, List[Policy]]): user customized policy objects
  
  Notes:
      The parallelization process is performed through the following process.
  
      1) slice parallelizable tensors and replace original tensors on CPU
      2) upload parallelizable (replaced) tensors to multiple GPUs simultaneously
      3) upload non-parallelizable tensors to multiple GPUs (e.g. embedding, lm_head, ...)
class_docstr_tok: ["Model","parallelization","processing","engine","Args",":","num_gpus","(","int",")",":","number","of","gpus","backend","(","str",")",":","distributed","backend","(","default=nccl",")","custom_policies","(","Union","[","Policy",",","List","[","Policy","]","]",")",":","user","customized","policy","objects","Notes",":","The","parallelization","process","is","performed","through","the","following","process",".","1",")","slice","parallelizable","tensors","and","replace","original","tensors","on","CPU","2",")","upload","parallelizable","(","replaced",")","tensors","to","multiple","GPUs","simultaneously","3",")","upload","non-parallelizable","tensors","to","multiple","GPUs","(","e.g",".","embedding",",","lm_head",",","...",")"]
docstr: Parallelize model to multiple GPUs
  
  Args:
      model (nn.Module): Huggingface pre-trained transformer model.
      fp16: (bool): whether use FP16 or not.
  
  Returns:
      nn.Module: parallelized model
docstr_tok: ["Parallelize","model","to","multiple","GPUs","Args",":","model","(","nn.Module",")",":","Huggingface","pre-trained","transformer","model",".","fp16",":","(","bool",")",":","whether","use","FP16","or","not",".","Returns",":","nn.Module",":","parallelized","model"]
code: def parallelize(self, model: nn.Module, fp16: bool) -> nn.Module:
      """
      Parallelize model to multiple GPUs
  
      Args:
          model (nn.Module): Huggingface pre-trained transformer model.
          fp16: (bool): whether use FP16 or not.
  
      Returns:
          nn.Module: parallelized model
      """
  
      super().__init__()
      replacer = TensorReplacer(
          model=model,
          fp16=fp16,
          mp_group=self.mp_group,
          num_gpus=self.num_gpus,
          custom_policies=self.custom_policies,
      )
  
      # Replace original layer to tensor sliced (megatron) layer.
      replacer.replace_modules()
  
      # Lazy GPU memory allocation (Only cpu tensors are loaded onto all gpus)
      # It's more memory-efficient than original implementation of DeepSpeed.
      for k, v in dict(model.state_dict()).items():
          if not v.is_cuda:
              if torch.is_tensor(v):
                  rsetattr(
                      model,
                      k + ".data",
                      v.to(torch.cuda.current_device()),
                  )
  
      return model
  
file_path: parallelformers/parallelformers/parallel/engine.py
imports: import os
  from typing import List, Union
  import torch
  import torch.distributed
  import torch.nn
  from parallelformers.parallel.replacing import TensorReplacer
  from parallelformers.policies.base import Policy
  from parallelformers.utils import rsetattr
label_desc: Use image node_id 2 for calling the ParallelEngine obj's underlying member method code with example usage: obj.parallelize(model, fp16) and returns: model