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.
docstr: Forward of ColumnParallelLinear
  
  Args:
      input_: 3D tensor whose order of dimension is [sequence, batch, hidden]
  
      weight (optional): weight tensor to use, compulsory when
          skip_weight_param_allocation is True.
  
  Returns:
      - output
      - bias
docstr_tok: ["Forward","of","ColumnParallelLinear","Args",":","input_",":","3D","tensor","whose","order","of","dimension","is","[","sequence",",","batch",",","hidden","]","weight","(","optional",")",":","weight","tensor","to","use",",","compulsory","when","skip_weight_param_allocation","is","True",".","Returns",":","-","output","-","bias"]
code: def forward(
      self, input_: torch.Tensor, weight: Optional[torch.Tensor] = None
  ):
      """Forward of ColumnParallelLinear
  
      Args:
          input_: 3D tensor whose order of dimension is [sequence, batch, hidden]
  
          weight (optional): weight tensor to use, compulsory when
              skip_weight_param_allocation is True.
  
      Returns:
          - output
          - bias
  
      """
      if weight is None:
          if self.weight is None:
              raise RuntimeError(
                  "weight was not supplied to ColumnParallelLinear forward pass "
                  "and skip_weight_param_allocation is True."
              )
          weight = self.weight
      else:
          # Check the weight passed in is the correct shape
          expected_shape = (
              self.output_size_per_partition,
              self.input_size,
          )
          if weight.shape != expected_shape:
              raise RuntimeError(
                  f"supplied weight's shape is {tuple(weight.shape)}, "
                  f"not {expected_shape} as expected"
              )
  
      bias = self.bias if not self.skip_bias_add else None
  
      if (
          self.async_tensor_model_parallel_allreduce
          or self.sequence_parallel
          or self.explicit_expert_comm
      ):
          input_parallel = input_
      else:
          input_parallel = copy_to_tensor_model_parallel_region(input_)
  
      # Matrix multiply.
      if not weight.requires_grad:
          self._forward_impl = linear_with_frozen_weight
      else:
          self._forward_impl = (
              linear_with_grad_accumulation_and_async_allreduce
          )
      output_parallel = self._forward_impl(
          input=input_parallel,
          weight=weight,
          bias=bias,
          gradient_accumulation_fusion=self.gradient_accumulation_fusion,
          async_grad_allreduce=False
          if self.explicit_expert_comm
          else self.async_tensor_model_parallel_allreduce,
          sequence_parallel=False
          if self.explicit_expert_comm
          else self.sequence_parallel,
      )
      if self.gather_output:
          # All-gather across the partitions.
          assert not self.sequence_parallel
          output = gather_from_tensor_model_parallel_region(
              output_parallel
          )
      else:
          output = output_parallel
      output_bias = self.bias if self.skip_bias_add else None
      return output, output_bias
  
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 2 for calling the ColumnParallelLinear obj's underlying member method code with example usage: obj.forward(input_, weight) and returns: output,  output_bias