Spaces:
Running on Zero
Running on Zero
File size: 5,040 Bytes
fed6c68 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 | # Copyright (c) 2026 Bytedance Ltd. and/or its affiliate
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Sequence-parallel tensor ops.
Every op is a no-op when Ulysses is disabled (single-GPU) and otherwise
delegates to Open-VeOmni, imported lazily so single-GPU inference carries no
dependency on it.
"""
import torch
import torch.nn.functional as F
from .state import get_parallel_state
def gather_seq_scatter_heads(x, seq_dim, head_dim, unpadded_dim_size=0):
"""All-to-all: gather the sequence dim, scatter the head dim."""
if not get_parallel_state().ulysses_enabled:
return x
from veomni.distributed.sequence_parallel import gather_seq_scatter_heads as _f
return _f(x, seq_dim=seq_dim, head_dim=head_dim, unpadded_dim_size=unpadded_dim_size)
def gather_heads_scatter_seq(x, head_dim, seq_dim):
"""All-to-all: gather the head dim, scatter the sequence dim."""
if not get_parallel_state().ulysses_enabled:
return x
from veomni.distributed.sequence_parallel import gather_heads_scatter_seq as _f
return _f(x, head_dim=head_dim, seq_dim=seq_dim)
def slice_input_tensor(x, dim):
"""Keep only this rank's slice of `x` along `dim`."""
if not get_parallel_state().ulysses_enabled:
return x
from veomni.distributed.sequence_parallel import slice_input_tensor as _f
return _f(x, dim=dim)
def slice_input_tensor_scale_grad(x, dim):
"""`slice_input_tensor` variant used inside autograd-tracked code paths."""
if not get_parallel_state().ulysses_enabled:
return x
from veomni.distributed.sequence_parallel import slice_input_tensor_scale_grad as _f
return _f(x, dim=dim)
def gather_outputs(x, gather_dim, padding_dim=None, unpad_dim_size=None, group=None):
"""Gather a sequence-sharded tensor back to its full length."""
if not get_parallel_state().ulysses_enabled:
return x
from veomni.distributed.sequence_parallel import gather_outputs as _f
return _f(x, gather_dim=gather_dim, padding_dim=padding_dim, unpad_dim_size=unpad_dim_size)
def padding_tensor_for_seqeunce_parallel(x, dim):
"""Pad `x` along `dim` so its size is divisible by the Ulysses world size."""
if not get_parallel_state().ulysses_enabled:
return x
from veomni.distributed.sequence_parallel.utils import (
padding_tensor_for_seqeunce_parallel as _f,
)
return _f(x, dim=dim)
def pad_tensor(x, dim, padding_size, padding_value=0):
"""Append `padding_size` entries along `dim` (F.pad based, low peak memory)."""
pad_config = [0, 0] * x.ndim
pad_config[(x.ndim - 1 - dim) * 2 + 1] = padding_size
return F.pad(x, pad_config, value=padding_value)
def unpad_tensor(x, dim, padding_size):
"""Inverse of `pad_tensor`: drop the last `padding_size` entries along `dim`."""
slc = [slice(None)] * x.ndim
slc[dim] = slice(0, -padding_size)
return x[slc]
def gen_cu_seqlens_for_cross_attn(q_len, batch_seqlens_q, batch_seqlens_k, device="cpu"):
"""cu_seqlens / max_seqlens for cross-attention under Ulysses sequence parallel.
Each rank holds a contiguous ``q_len / sp_world`` slice of the query
sequence; this maps the per-sample query/key lengths onto that local slice.
"""
ps = get_parallel_state()
sp_world = ps.ulysses_size
rank = ps.ulysses_rank
rank_q_len = (q_len + ((sp_world - (q_len % sp_world)) % sp_world)) // sp_world
start = rank_q_len * rank
end = min(q_len, start + rank_q_len)
offset = 0
cu_seqlens_q = [start]
index = []
max_seqlen_q = -1
max_seqlen_k = -1
for i, length in enumerate(batch_seqlens_q):
offset = min(offset + length, end)
if offset <= start:
continue
cu_seqlens_q.append(offset)
index.append(i)
max_seqlen_q = max(max_seqlen_q, cu_seqlens_q[-1] - cu_seqlens_q[-2])
max_seqlen_k = max(max_seqlen_k, batch_seqlens_k[i])
if offset >= end:
break
cu_seqlens_q = torch.tensor(cu_seqlens_q, dtype=torch.int32, device=device)
max_seqlen_q = torch.tensor(max_seqlen_q, device=device)
cu_seqlens_q -= start
cu_seqlens_k = torch.zeros(len(batch_seqlens_k) + 1, dtype=torch.int32, device=device)
cu_seqlens_k[1:] = torch.tensor(batch_seqlens_k, dtype=torch.int32, device=device).cumsum(dim=0)
cu_seqlens_k = cu_seqlens_k[index[0] : index[-1] + 2]
max_seqlen_k = torch.tensor(max_seqlen_k, device=device)
return cu_seqlens_k, cu_seqlens_q, max_seqlen_k, max_seqlen_q, rank_q_len
|