← Về thư mục
verl/models/llama/megatron/modeling_llama_megatron.py
- _make_causal_mask · function · L45-L54 — def _make_causal_mask(input_ids_shape: torch.Size, dtype: torch.dtype, device: torch.device)
- _expand_mask · function · L58-L69 — def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None)
- ParallelLlamaModel · class · L72-L152 — class ParallelLlamaModel(nn.Module)
- init · method · L80-L94 — def init(self, config: LlamaConfig, megatron_config: ModelParallelConfig)
- _prepare_decoder_attention_mask · method · L97-L115 — def _prepare_decoder_attention_mask(self, attention_mask, input_shape, inputs_embeds): # create causal mask # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
- forward · method · L117-L152 — def forward( self, input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, ) -> Union[Tuple, BaseModelOutputWithPast]
- ParallelLlamaForCausalLM · class · L155-L209 — class ParallelLlamaForCausalLM(nn.Module)
- init · method · L157-L172 — def init(self, config: LlamaConfig, megatron_config: ModelParallelConfig)
- forward · method · L174-L209 — def forward( self, input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, ) -> Union[Tuple, CausalLMOutputWithPast]
- ParallelLlamaModelRmPad · class · L215-L276 — class ParallelLlamaModelRmPad(nn.Module)
- init · method · L223-L238 — def init(self, config: LlamaConfig, megatron_config: ModelParallelConfig)
- forward · method · L240-L276 — def forward(self, input_ids: torch.Tensor, position_ids: Optional[torch.LongTensor] = None, sequence_length: int = None, indices: torch.Tensor = None, cu_seqlens: int = None, max_seqlen_in_batch: int = None) -> Union[Tuple, BaseModelOutputWithPast]
- ParallelLlamaForCausalLMRmPad · class · L279-L363 — class ParallelLlamaForCausalLMRmPad(nn.Module)
- init · method · L281-L287 — def init(self, config: LlamaConfig, megatron_config: ModelParallelConfig)
- _init_head · method · L289-L299 — def _init_head(self)
- _forward_head · method · L301-L306 — def _forward_head(self, hidden_states): # all_gather from sequence parallel region is performed inside lm_head
- forward · method · L308-L363 — def forward( self, input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, ) -> Union[Tuple, CausalLMOutputWithPast]
- ParallelLlamaForValueRmPad · class · L366-L392 — class ParallelLlamaForValueRmPad(ParallelLlamaForCausalLMRmPad)
- _init_head · method · L368-L375 — def _init_head(self)
- _forward_head · method · L377-L382 — def _forward_head(self, hidden_states)
- forward · method · L384-L392 — def forward( self, input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, ) -> Union[Tuple, CausalLMOutputWithPast]
- ParallelLlamaModelRmPadPP · class · L400-L511 — class ParallelLlamaModelRmPadPP(nn.Module)
- init · method · L410-L455 — def init(self, config: LlamaConfig, megatron_config: ModelParallelConfig, pre_process, post_process)
- set_input_tensor · method · L457-L465 — def set_input_tensor(self, input_tensor)
- forward · method · L467-L511 — def forward(self, input_ids: torch.Tensor, position_ids: Optional[torch.LongTensor] = None, sequence_length: int = None, indices: torch.Tensor = None, cu_seqlens: int = None, max_seqlen_in_batch: int = None) -> Union[Tuple, BaseModelOutputWithPast]
- ParallelLlamaForCausalLMRmPadPP · class · L514-L623 — class ParallelLlamaForCausalLMRmPadPP(nn.Module)
- init · method · L516-L529 — def init(self, config: LlamaConfig, megatron_config: ModelParallelConfig, pre_process, post_process)
- set_input_tensor · method · L531-L540 — def set_input_tensor(self, input_tensor)
- _init_head · method · L542-L552 — def _init_head(self)
- _forward_head · method · L554-L560 — def _forward_head(self, hidden_states): # all_gather from sequence parallel region is performed inside lm_head # logits shape before forward_head hidden_states.shape: [4, 32, 4096]
- forward · method · L562-L623 — def forward( self, # original input *, input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, ) -> Union[Tuple, CausalLMOutputWithPast]
- ParallelLlamaForValueRmPadPP · class · L626-L656 — class ParallelLlamaForValueRmPadPP(ParallelLlamaForCausalLMRmPadPP)
- _init_head · method · L628-L635 — def _init_head(self)
- _forward_head · method · L637-L642 — def _forward_head(self, hidden_states)
- forward · method · L644-L656 — def forward( self, *, input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, ) -> Union[Tuple, CausalLMOutputWithPast]