class_docstr: Shard the original huggingface model according to the policy
  
  Args:
      policy (:class:`Policy`): The policy to shard the model
      model (:class:`torch.Module`): The model to shard
      shard_config: The setting of distributed model
class_docstr_tok: ["Shard","the","original","huggingface","model","according","to","the","policy","Args",":","policy","(",":","class",":","`","Policy","`",")",":","The","policy","to","shard","the","model","model","(",":","class",":","`","torch.Module","`",")",":","The","model","to","shard","shard_config",":","The","setting","of","distributed","model"]
docstr: Shard one layer according to the policy, the layer should be the same class as the key in policy's argument_policy return dict
  
  Args:
      org_layer (torch.nn.Module): The origin layer object to shard
      sub_module_replacement (List[SubModuleReplacementDescription]): The sub module replacement description list
      include (Set[nn.Module], optional): The set of modules to keep on current device when pipeline parallel is enabled. Defaults to None
docstr_tok: ["Shard","one","layer","according","to","the","policy",",","the","layer","should","be","the","same","class","as","the","key","in","policy","'s","argument_policy","return","dict","Args",":","org_layer","(","torch.nn.Module",")",":","The","origin","layer","object","to","shard","sub_module_replacement","(","List","[","SubModuleReplacementDescription","]",")",":","The","sub","module","replacement","description","list","include","(","Set","[","nn.Module","]",",","optional",")",":","The","set","of","modules","to","keep","on","current","device","when","pipeline","parallel","is","enabled",".","Defaults","to","None"]
code: def _replace_sub_module(
      self,
      org_layer: nn.Module,
      sub_module_replacement: List[SubModuleReplacementDescription],
      include: Optional[Set[nn.Module]] = None,
  ) -> None:
      r"""
      Shard one layer according to the policy, the layer should be the same class as the key in policy's argument_policy return dict
  
      Args:
          org_layer (torch.nn.Module): The origin layer object to shard
          sub_module_replacement (List[SubModuleReplacementDescription]): The sub module replacement description list
          include (Set[nn.Module], optional): The set of modules to keep on current device when pipeline parallel is enabled. Defaults to None
      """
      for description in sub_module_replacement:
          suffix = description.suffix
          target_module = description.target_module
          kwargs = (
              {} if description.kwargs is None else description.kwargs
          )
  
          assert (
              target_module is not None
          ), "target_module should not be None"
  
          native_sub_module = getattr_(org_layer, suffix, ignore=True)
          # Skip replacement if submodule is not kept by current device when pipeline parallel is enabled.
          if (
              (include is not None)
              and (native_sub_module is not None)
              and (native_sub_module not in include)
          ):
              continue
  
          assert not isinstance(
              native_sub_module, target_module
          ), f"The module with suffix {suffix} has been replaced, please check the policy"
  
          # if it is None and we are allowed to ignore this module
          # just skip
          if (
              description.ignore_if_not_exist
              and native_sub_module is None
          ):
              continue
  
          try:
              replace_layer = target_module.from_native_module(
                  native_sub_module,
                  process_group=self.shard_config.tensor_parallel_process_group,
                  **kwargs,
              )
          except Exception as e:
              raise RuntimeError(
                  f"Failed to replace {suffix} of type {native_sub_module.__class__.__qualname__}"
                  f" with {target_module.__qualname__} with the exception: {e}. "
                  "Please check your model configuration or sharding policy, you can set up an issue for us to help you as well."
              )
  
          setattr_(org_layer, suffix, replace_layer)
  
imports: from types import MethodType
  from typing import Any, Callable, Dict, List, Optional, Set, Union
  import torch.nn
  from torch import Tensor
  from colossalai.lazy import LazyInitContext
  from .._utils import getattr_, setattr_
  from ..policies.auto_policy import get_autopolicy
  from ..policies.base_policy import Policy, SubModuleReplacementDescription
  from .shard_config import ShardConfig
  from .utils import set_tensors_to_none
label_desc: Use image node_id 10 for calling the ModelSharder obj's underlying member method code with example usage: obj._replace_sub_module(org_layer, sub_module_replacement, include) without return types