verl/utils/ulysses.py

• set_ulysses_sequence_parallel_group · function · L29-L34 — def set_ulysses_sequence_parallel_group(group: dist.ProcessGroup)
• get_ulysses_sequence_parallel_group · function · L37-L42 — def get_ulysses_sequence_parallel_group() -> Optional[dist.ProcessGroup]
• get_ulysses_sequence_parallel_world_size · function · L45-L50 — def get_ulysses_sequence_parallel_world_size(group: ProcessGroup = None) -> int
• get_ulysses_sequence_parallel_rank · function · L53-L58 — def get_ulysses_sequence_parallel_rank(group: ProcessGroup = None) -> int
• gather_seq_scatter_heads · function · L61-L82 — def gather_seq_scatter_heads( x: Tensor, seq_dim: int, head_dim: int, unpadded_dim_size: int = 0, group: ProcessGroup = None, ) -> Tensor
• gather_heads_scatter_seq · function · L85-L100 — def gather_heads_scatter_seq(x: Tensor, head_dim: int, seq_dim: int, group: ProcessGroup = None) -> Tensor
• _pad_tensor · function · L103-L107 — def _pad_tensor(x: Tensor, dim: int, padding_size: int) -> Tensor
• _unpad_tensor · function · L110-L113 — def _unpad_tensor(x: Tensor, dim: int, padding_size: int) -> Tensor
• slice_input_tensor · function · L116-L129 — def slice_input_tensor(x: Tensor, dim: int, padding: bool = True, group: ProcessGroup = None) -> Tensor
• all_to_all_tensor · function · L132-L151 — def all_to_all_tensor( local_input: Tensor, scatter_dim: int, gather_dim: int, group: Optional[dist.ProcessGroup] = None, async_op: bool = False, )
• wait · function · L146-L148 — def wait()
• all_gather_tensor · function · L154-L161 — def all_gather_tensor(local_tensor: Tensor, group: Optional[dist.ProcessGroup] = None, async_op: bool = False)
• SeqAllToAll · class · L164-L194 — class SeqAllToAll(torch.autograd.Function)
• forward · method · L167-L179 — def forward( ctx: Any, group: dist.ProcessGroup, local_input: Tensor, scatter_dim: int, gather_dim: int, async_op: bool = False, ) -> Tensor
• backward · method · L182-L194 — def backward(ctx: Any, *grad_output: Tensor) -> Tuple[None, Tensor, None, None]
• Gather · class · L197-L230 — class Gather(torch.autograd.Function)
• forward · method · L200-L223 — def forward(ctx: Any, group: dist.ProcessGroup, local_tensor: Tensor, gather_dim: int, grad_scaler: bool = True, async_op=False) -> Tensor
• backward · method · L226-L230 — def backward(ctx: Any, grad_output: Tensor) -> Any
• gather_outpus_and_unpad · function · L233-L249 — def gather_outpus_and_unpad(x: Tensor, gather_dim: int, unpad_dim: int = None, padding_size: int = 0, grad_scaler: bool = True, group: Optional[dist.ProcessGroup] = None)
• ulysses_pad_and_slice_inputs · function · L252-L288 — def ulysses_pad_and_slice_inputs(input_ids_rmpad: torch.Tensor, position_ids_rmpad: Optional[torch.Tensor] = None, sp_size: int = 1)