docstr: A replacement for ``torch.nn.Linear`` that works with ZeRO-3 to reduce
  memory requirements via tiling.
  
  TiledLinear breaks the input and output dimensions of a linear layer
  into tiles that are processed in sequence. This class enables huge
  linear layers when combined with ZeRO-3 because inactive tiles can be
  partitioned and offloaded.
  
  .. note::
      We recommend using as few tiles as necessary. Tiling
      significantly reduces memory usage, but can reduce throughput
      for inexpensive layers. This due to the smaller kernels having
      less parallelism and lower arithmetic intensity, while
      introducing more frequent synchronization and communication.
  
  Args:
      in_features (int): See ``torch.nn.Linear``
      out_features (int): See ``torch.nn.Linear``
      bias (bool, optional): See ``torch.nn.Linear``
      in_splits (int, optional): The number of tiles along the input dimension. Defaults to 1.
      out_splits (int, optional): The number of tiles along the output dimension. Defaults to 1.
      input_is_already_split (bool, optional): If set to ``True``, assume that the ``input_`` in
          to ``forward()`` is already split into ``in_splits`` chunks. Defaults to ``False``.
      combine_out_splits (bool, optional): If set to ``False``, do not combine the ``out_splits`` outputs
          into a single tensor. Defaults to ``True``.
      linear_cls (class, optional): The underlying class to build individual tiles.
          Defaults to ``torch.nn.Linear``.
      init_linear (``torch.nn.Linear``, optional): If set, copy the parameters of
          ``init_linear``. Useful for debugging. Defaults to ``None``.
      kwargs (dict, optional): additional keyword arguments to provide to ``linear_cls()``.
  
  Raises:
      RuntimeError: ``in_splits`` must be within the range [1, in_features).
      RuntimeError: ``out_splits`` must be within the range of [1, out_features).
docstr_tok: ["A","replacement","for","``","torch.nn.Linear","``","that","works","with","ZeRO-3","to","reduce","memory","requirements","via","tiling",".","TiledLinear","breaks","the","input","and","output","dimensions","of","a","linear","layer","into","tiles","that","are","processed","in","sequence",".","This","class","enables","huge","linear","layers","when","combined","with","ZeRO-3","because","inactive","tiles","can","be","partitioned","and","offloaded",".","..","note",":",":","We","recommend","using","as","few","tiles","as","necessary",".","Tiling","significantly","reduces","memory","usage",",","but","can","reduce","throughput","for","inexpensive","layers",".","This","due","to","the","smaller","kernels","having","less","parallelism","and","lower","arithmetic","intensity",",","while","introducing","more","frequent","synchronization","and","communication",".","Args",":","in_features","(","int",")",":","See","``","torch.nn.Linear","``","out_features","(","int",")",":","See","``","torch.nn.Linear","``","bias","(","bool",",","optional",")",":","See","``","torch.nn.Linear","``","in_splits","(","int",",","optional",")",":","The","number","of","tiles","along","the","input","dimension",".","Defaults","to","1.","out_splits","(","int",",","optional",")",":","The","number","of","tiles","along","the","output","dimension",".","Defaults","to","1.","input_is_already_split","(","bool",",","optional",")",":","If","set","to","``","True","``",",","assume","that","the","``","input_","``","in","to","``","forward","(",")","``","is","already","split","into","``","in_splits","``","chunks",".","Defaults","to","``","False","``",".","combine_out_splits","(","bool",",","optional",")",":","If","set","to","``","False","``",",","do","not","combine","the","``","out_splits","``","outputs","into","a","single","tensor",".","Defaults","to","``","True","``",".","linear_cls","(","class",",","optional",")",":","The","underlying","class","to","build","individual","tiles",".","Defaults","to","``","torch.nn.Linear","``",".","init_linear","(","``","torch.nn.Linear","``",",","optional",")",":","If","set",",","copy","the","parameters","of","``","init_linear","``",".","Useful","for","debugging",".","Defaults","to","``","None","``",".","kwargs","(","dict",",","optional",")",":","additional","keyword","arguments","to","provide","to","``","linear_cls","(",")","``",".","Raises",":","RuntimeError",":","``","in_splits","``","must","be","within","the","range","[","1",",","in_features",")",".","RuntimeError",":","``","out_splits","``","must","be","within","the","range","of","[","1",",","out_features",")","."]
code: def __init__(
      self,
      in_features,
      out_features,
      bias=True,
      in_splits=1,
      out_splits=1,
      input_is_already_split=False,
      combine_out_splits=True,
      linear_cls=torch.nn.Linear,
      init_linear=None,
      **kwargs,
  ):
      """A replacement for ``torch.nn.Linear`` that works with ZeRO-3 to reduce
      memory requirements via tiling.
  
      TiledLinear breaks the input and output dimensions of a linear layer
      into tiles that are processed in sequence. This class enables huge
      linear layers when combined with ZeRO-3 because inactive tiles can be
      partitioned and offloaded.
  
      .. note::
          We recommend using as few tiles as necessary. Tiling
          significantly reduces memory usage, but can reduce throughput
          for inexpensive layers. This due to the smaller kernels having
          less parallelism and lower arithmetic intensity, while
          introducing more frequent synchronization and communication.
  
      Args:
          in_features (int): See ``torch.nn.Linear``
          out_features (int): See ``torch.nn.Linear``
          bias (bool, optional): See ``torch.nn.Linear``
          in_splits (int, optional): The number of tiles along the input dimension. Defaults to 1.
          out_splits (int, optional): The number of tiles along the output dimension. Defaults to 1.
          input_is_already_split (bool, optional): If set to ``True``, assume that the ``input_`` in
              to ``forward()`` is already split into ``in_splits`` chunks. Defaults to ``False``.
          combine_out_splits (bool, optional): If set to ``False``, do not combine the ``out_splits`` outputs
              into a single tensor. Defaults to ``True``.
          linear_cls (class, optional): The underlying class to build individual tiles.
              Defaults to ``torch.nn.Linear``.
          init_linear (``torch.nn.Linear``, optional): If set, copy the parameters of
              ``init_linear``. Useful for debugging. Defaults to ``None``.
          kwargs (dict, optional): additional keyword arguments to provide to ``linear_cls()``.
  
      Raises:
          RuntimeError: ``in_splits`` must be within the range [1, in_features).
          RuntimeError: ``out_splits`` must be within the range of [1, out_features).
      """
  
      super().__init__()
  
      if (in_splits < 1) or (in_splits > in_features):
          raise RuntimeError(
              "in splits must be in range [1, in_features]."
          )
      if (out_splits < 1) or (out_splits > out_features):
          raise RuntimeError(
              "out splits must be in range [1, out_features]."
          )
  
      # global, not necessarily local
      self.in_features = in_features
      self.out_features = out_features
      self.use_bias = bias
  
      self.out_splits = out_splits
      self.in_splits = in_splits
      self.input_is_already_split = input_is_already_split
      self.combine_out_splits = combine_out_splits
  
      # Build partition-lists. These are CSR-style splits [0, part0, part1, ..., features]
      # For example, row_parts[p] gives the start of partition p and row_parts[p+1]
      # is the exclusive end.
      self.in_parts = partition(
          num_items=in_features, num_parts=in_splits
      )
      self.out_parts = partition(
          num_items=out_features, num_parts=out_splits
      )
  
      assert len(self.out_parts) == out_splits + 1
      assert len(self.in_parts) == in_splits + 1
      assert self.out_parts[0] == 0
      assert self.out_parts[out_splits] == out_features
      assert self.in_parts[in_splits] == in_features
  
      self.linears = torch.nn.ModuleList()
      for out_id in range(out_splits):
          self.linears.append(torch.nn.ModuleList())
  
          local_out_dim = (
              self.out_parts[out_id + 1] - self.out_parts[out_id]
          )
  
          for in_id in range(in_splits):
              # if input_size is split, we only need one bias
              local_bias = bias if in_id == (in_splits - 1) else False
  
              local_in_dim = (
                  self.in_parts[in_id + 1] - self.in_parts[in_id]
              )
              local = linear_cls(
                  local_in_dim, local_out_dim, bias=local_bias, **kwargs
              )
              self.linears[out_id].append(local)
  
      # Optionally initialize with a known tensor
      if init_linear is not None:
          self.copy_params_from(init_linear)
  
file_path: DeepSpeed/deepspeed/runtime/zero/tiling.py
imports: import torch
  import deepspeed
  from deepspeed.runtime.utils import partition_uniform
label_desc: Use image node_id 1 to create a new TiledLinear object from inherited base classes: torch.nn with example: obj = TiledLinear(in_features, out_features, bias, in_splits, out_splits, input_is_already_split, combine_out_splits, linear_cls, init_linear)