class_docstr: Linear layer with column parallelism.
  
  The linear layer is defined as Y = XA + b. A is parallelized along
  its second dimension as A = [A_1, ..., A_p].
  
  .. note::
      Input is supposed to be three dimensional and each dimension
      is expected to be sequence, batch, and hidden feature, respectively.
  
  Arguments:
      input_size: first dimension of matrix A.
      output_size: second dimension of matrix A.
      bias: If true, add bias
      gather_output: If true, call all-gether on output and make Y avaiable
                     to all GPUs, otherwise, every GPU will have its output
                     which is Y_i = XA_i
      init_method: method to initialize weights. Note that bias is always set
                   to zero.
      stride: For the strided linear layers.
      keep_master_weight_for_test: This was added for testing and should be
                                   set to False. It returns the master weights
                                   used for initialization.
      skip_bias_add: This was added to enable performance optimations where bias
                     can be fused with other elementwise operations. we skip
                     adding bias but instead return it.
  
  Keyword Arguments:
      no_async_tensor_model_parallel_allreduce:
      params_dtype:
      use_cpu_initialization:
      gradient_accumulation_fusion:
      sequence_parallel_enabled:
      accumulation_in_fp16: Deprecated
code: def __init__(
      self,
      input_size,
      output_size,
      bias=True,
      gather_output=True,
      init_method=init.xavier_normal_,
      stride=1,
      keep_master_weight_for_test=False,
      skip_bias_add=False,
      *,
      no_async_tensor_model_parallel_allreduce=False,
      params_dtype=torch.float32,
      use_cpu_initialization=False,
      gradient_accumulation_fusion=False,
      sequence_parallel_enabled: bool = False,
      accumulation_in_fp16: Optional[bool] = None,
  ):
      super().__init__()
  
      # Keep input parameters
      self.input_size = input_size
      self.output_size = output_size
      self.gather_output = gather_output
      # Divide the weight matrix along the last dimension.
      world_size = get_tensor_model_parallel_world_size()
      self.output_size_per_partition = divide(output_size, world_size)
      self.skip_bias_add = skip_bias_add
  
      if accumulation_in_fp16 is not None:
          warnings.warn(
              f"Deprecated option `accumulation_in_fp16` is set to {accumulation_in_fp16}"
          )
  
      # Parameters.
      # Note: torch.nn.functional.linear performs XA^T + b and as a result
      # we allocate the transpose.
      # Initialize weight.
      if use_cpu_initialization:
          self.weight = Parameter(
              torch.empty(
                  self.output_size_per_partition,
                  self.input_size,
                  dtype=params_dtype,
              )
          )
          self.master_weight = _initialize_affine_weight_cpu(
              self.weight,
              self.output_size,
              self.input_size,
              self.output_size_per_partition,
              0,
              init_method,
              stride=stride,
              return_master_weight=keep_master_weight_for_test,
              params_dtype=params_dtype,
          )
      else:
          self.weight = Parameter(
              torch.empty(
                  self.output_size_per_partition,
                  self.input_size,
                  device=torch.cuda.current_device(),
                  dtype=params_dtype,
              )
          )
          _initialize_affine_weight_gpu(
              self.weight, init_method, partition_dim=0, stride=stride
          )
  
      if bias:
          if use_cpu_initialization:
              self.bias = Parameter(
                  torch.empty(
                      self.output_size_per_partition, dtype=params_dtype
                  )
              )
          else:
              self.bias = Parameter(
                  torch.empty(
                      self.output_size_per_partition,
                      device=torch.cuda.current_device(),
                      dtype=params_dtype,
                  )
              )
          set_tensor_model_parallel_attributes(
              self.bias, True, 0, stride
          )
          # Always initialize bias to zero.
          with torch.no_grad():
              self.bias.zero_()
      else:
          self.register_parameter("bias", None)
  
      self.async_tensor_model_parallel_allreduce = (
          not no_async_tensor_model_parallel_allreduce
          and world_size > 1
      )
      if sequence_parallel_enabled:
          if world_size <= 1:
              warnings.warn(
                  f"`sequence_parallel_enabled` is set to `True`, but got world_size of {world_size}"
              )
              # sequence_parallel_enabled = False
      self.sequence_parallel_enabled = sequence_parallel_enabled
      if gradient_accumulation_fusion:
          if not _grad_accum_fusion_available:
              # Basically, apex.transformer module users are expected to install APEX's
              # `--cpp_ext` and `--cuda_ext`. The example installation command is as follows:
              # `pip install --global-option="--cpp_ext" --global-option="--cuda_ext ."
              # at the root of APEX repository.
              warnings.warn(
                  "`gradient_accumulation_fusion` is set to `True` but "
                  "the custom CUDA extension of `fused_weight_gradient_mlp_cuda` module not "
                  "found. Thus `gradient_accumulation_fusion` set to `False`. "
                  "Note that the extension requires CUDA>=11."
              )
              gradient_accumulation_fusion = False
      self.gradient_accumulation_fusion = gradient_accumulation_fusion
  
      if (
          self.async_tensor_model_parallel_allreduce
          and self.sequence_parallel_enabled
      ):
          raise RuntimeError(
              "`async_tensor_model_parallel_allreduce` and `sequence_parallel_enabled` cannot be enabled at the same time."
          )
  
      self._forward_impl = (
          linear_with_grad_accumulation_and_async_allreduce
      )
  
imports: from typing import Optional, Dict, Tuple, List
  import warnings
  import torch
  import torch.nn.functional
  import torch.nn.init
  from torch.nn.parameter import Parameter
  from apex._autocast_utils import _cast_if_autocast_enabled
  from apex.transformer.parallel_state import get_tensor_model_parallel_group
  from apex.transformer.parallel_state import get_tensor_model_parallel_rank
  from apex.transformer.parallel_state import get_tensor_model_parallel_world_size
  from apex.transformer.utils import divide
  from apex.transformer.tensor_parallel.mappings import copy_to_tensor_model_parallel_region
  from apex.transformer.tensor_parallel.mappings import gather_from_tensor_model_parallel_region
  from apex.transformer.tensor_parallel.mappings import reduce_from_tensor_model_parallel_region
  from apex.transformer.tensor_parallel.mappings import scatter_to_tensor_model_parallel_region
  from apex.transformer.tensor_parallel.mappings import reduce_scatter_to_sequence_parallel_region
  from apex.transformer.tensor_parallel.random import get_cuda_rng_tracker
  from apex.transformer.tensor_parallel.utils import VocabUtility
  from apex.transformer.log_util import get_transformer_logger
label_desc: Use image node_id 1 to create a new ColumnParallelLinear object from inherited base classes: torch.nn with example: obj = ColumnParallelLinear(input_size, output_size, bias, gather_output, init_method, stride, keep_master_weight_for_test, skip_bias_add)