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].
  
  Args:
      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-gather on output and make Y available 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: If True, do not add the bias term, instead return it to be added by the caller. This enables performance optimations where bias can be fused with other elementwise operations.
      skip_weight_param_allocation: If True, weight parameter is not allocated and must be passed as a keyword argument `weight` during the forward pass. Note that this does not affect bias, which will be allocated if bias is True. Defaults to False.
      is_expert: If True, the layer is treated as an MoE expert layer.
      config: ModelParallelConfig object
      tp_comm_buffer_name: Communication buffer name is not used in non-Transformer-Engine modules.
code: def __init__(
      self,
      input_size,
      output_size,
      *,
      config: ModelParallelConfig,
      init_method: Callable,
      bias=True,
      gather_output=False,
      stride=1,
      keep_master_weight_for_test=False,
      skip_bias_add=False,
      skip_weight_param_allocation: bool = False,
      is_expert: bool = False,
      tp_comm_buffer_name: str = None,  # Not used
  ):
      super(ColumnParallelLinear, self).__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
      self.is_expert = is_expert
      self.expert_parallel = config.expert_model_parallel_size > 1
      self.config = config
  
      # Parameters.
      # Note: torch.nn.functional.linear performs XA^T + b and as a result
      # we allocate the transpose.
      # Initialize weight.
      if not skip_weight_param_allocation:
          if config.use_cpu_initialization:
              self.weight = Parameter(
                  torch.empty(
                      self.output_size_per_partition,
                      self.input_size,
                      dtype=config.params_dtype,
                  )
              )
              if config.perform_initialization:
                  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,
                  )
          else:
              self.weight = Parameter(
                  torch.empty(
                      self.output_size_per_partition,
                      self.input_size,
                      device=torch.cuda.current_device(),
                      dtype=config.params_dtype,
                  )
              )
              if config.perform_initialization:
                  _initialize_affine_weight_gpu(
                      self.weight,
                      init_method,
                      partition_dim=0,
                      stride=stride,
                      expert_parallel=(
                          self.is_expert and self.expert_parallel
                      ),
                  )
  
          setattr(
              self.weight,
              "allreduce",
              not (self.is_expert and self.expert_parallel),
          )
      else:
          self.weight = None
  
      if bias:
          if config.use_cpu_initialization:
              self.bias = Parameter(
                  torch.empty(
                      self.output_size_per_partition,
                      dtype=config.params_dtype,
                  )
              )
          else:
              self.bias = Parameter(
                  torch.empty(
                      self.output_size_per_partition,
                      device=torch.cuda.current_device(),
                      dtype=config.params_dtype,
                  )
              )
          set_tensor_model_parallel_attributes(
              self.bias, True, 0, stride
          )
          if config.perform_initialization:
              # Always initialize bias to zero.
              with torch.no_grad():
                  self.bias.zero_()
          setattr(
              self.bias,
              "allreduce",
              not (self.is_expert and self.expert_parallel),
          )
      else:
          self.register_parameter("bias", None)
  
      self.async_tensor_model_parallel_allreduce = (
          config.async_tensor_model_parallel_allreduce
          and world_size > 1
      )
  
      self.sequence_parallel = config.sequence_parallel
      if self.sequence_parallel and world_size <= 1:
          warnings.warn(
              f"`sequence_parallel` is set to `True`, but tensor model parallel size is {world_size}. "
              f"Disabling sequence parallel."
          )
          self.sequence_parallel = False
  
      if (
          config.gradient_accumulation_fusion
          and not _grad_accum_fusion_available
      ):
          raise RuntimeError(
              "ColumnParallelLinear was called with gradient_accumulation_fusion set "
              "to True but the custom CUDA extension fused_weight_gradient_mlp_cuda "
              "module is not found. To use gradient_accumulation_fusion you must "
              "install APEX with --cpp_ext and --cuda_ext. For example: "
              'pip install --global-option="--cpp_ext" --global-option="--cuda_ext ." '
              "Note that the extension requires CUDA>=11. Otherwise, you must turn off "
              "gradient accumulation fusion."
          )
      self.gradient_accumulation_fusion = (
          config.gradient_accumulation_fusion
      )
  
      if (
          self.async_tensor_model_parallel_allreduce
          and self.sequence_parallel
      ):
          raise RuntimeError(
              "`async_tensor_model_parallel_allreduce` and `sequence_parallel` "
              "cannot be enabled at the same time."
          )
  
      self._forward_impl = (
          linear_with_grad_accumulation_and_async_allreduce
      )
      self.explicit_expert_comm = self.is_expert and (
          self.sequence_parallel or self.expert_parallel
      )
  
imports: import math
  import os
  import warnings
  from typing import Callable, Optional
  import torch
  import torch.nn.functional
  import torch.nn.init
  from torch.cuda.amp import custom_bwd, custom_fwd
  from torch.nn.parameter import Parameter
  from megatron.core.model_parallel_config import ModelParallelConfig
  from megatron.core.parallel_state import get_global_memory_buffer, get_tensor_model_parallel_group, get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size
  from .mappings import copy_to_tensor_model_parallel_region, gather_from_sequence_parallel_region, gather_from_tensor_model_parallel_region, reduce_from_tensor_model_parallel_region, reduce_scatter_to_sequence_parallel_region, scatter_to_tensor_model_parallel_region
  from .random import get_cuda_rng_tracker, get_expert_parallel_rng_tracker_name
  from .utils import VocabUtility, divide, split_tensor_along_last_dim
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)