verl/models/llama/megatron/layers/parallel_linear.py

• QKVParallelLinear · class · L21-L49 — class QKVParallelLinear(tensor_parallel.ColumnParallelLinear)
• init · method · L23-L49 — def init(self, input_size, num_heads, num_key_value_heads, head_dim, , bias=True, gather_output=True, skip_bias_add=False, *kwargs): # Keep input parameters, and already restrict the head numbers
• MergedColumnParallelLinear · class · L52-L74 — class MergedColumnParallelLinear(tensor_parallel.ColumnParallelLinear)
• init · method · L54-L74 — def init(self, input_size, gate_ouput_size, up_output_size, , bias=True, gather_output=True, skip_bias_add=False, *kwargs): # Keep input parameters, and already restrict the head numbers