code: def forward(self, input):
      input_shard_size = get_shard_size(
          input.shape[-1], self.world_size, "lm_head"
      )
      input_shard_offset = sum(
          get_shard_size_list(
              input.shape[-1], self.world_size, "lm_head"
          )[0 : self.rank]
      )
      output = torch.matmul(
          input[
              :,
              :,
              input_shard_offset : input_shard_offset
              + input_shard_size,
          ],
          self.weight.transpose(-1, -2),
      )
      if self.mp_group is not None:
          dist.inference_all_reduce(output, group=self.mp_group)
      if self.bias is not None:
          output += self.bias
      return output
  
imports: import torch
  from deepspeed import comm
  from torch import nn
  from torch.nn import functional
  from torch.nn.parameter import Parameter
  from deepspeed.accelerator import get_accelerator
  from deepspeed.module_inject.tp_shard import get_shard_size, get_shard_size_list
label_desc: Use image node_id 2 for calling the LmHeadLinearAllreduce obj's underlying member method code with example usage: obj.forward(input) and returns: output