(
self,
input: torch.Tensor,
*,
all_reduce_params: Optional[AllReduceParams] = None,
mapping_lm_head_tp: Optional[Mapping] = None,
is_spec_decoding_head: bool = False,
)
| 88 | return self.out_features |
| 89 | |
| 90 | def forward( |
| 91 | self, |
| 92 | input: torch.Tensor, |
| 93 | *, |
| 94 | all_reduce_params: Optional[AllReduceParams] = None, |
| 95 | mapping_lm_head_tp: Optional[Mapping] = None, |
| 96 | is_spec_decoding_head: bool = False, |
| 97 | ) -> torch.Tensor: |
| 98 | if is_spec_decoding_head and self.enable_lm_head_tp_in_adp: |
| 99 | # For LM head TP in ADP, we need to slice the weight for the LM head |
| 100 | tp_rank = mapping_lm_head_tp.tp_rank |
| 101 | tp_size = mapping_lm_head_tp.tp_size |
| 102 | slice_width = ceil_div(self.out_features, tp_size) |
| 103 | slice_start = tp_rank * slice_width |
| 104 | slice_end = min((tp_rank + 1) * slice_width, self.out_features) |
| 105 | output = F.linear(input, self.weight[slice_start:slice_end, :], |
| 106 | None) |
| 107 | else: |
| 108 | output = super().forward(input, all_reduce_params=all_reduce_params) |
| 109 | if (self.tp_mode == TensorParallelMode.COLUMN and self.gather_output |
| 110 | and self.padding_size > 0): |
| 111 | output = output[..., :-self.padding_size] |
| 112 | |
| 113 | return output |
| 114 | |
| 115 | def skip_forward( |
| 116 | self, |
nothing calls this directly
no test coverage detected