Spaces:
Running on Zero
Running on Zero
File size: 3,913 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 | # 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.
"""Minimal parallel state: data parallel + Ulysses sequence parallel.
Single-GPU inference uses the default state (Ulysses disabled) and pulls in no
extra dependency. Multi-GPU sequence parallel delegates the actual collective
communication to Open-VeOmni, which is imported lazily only when enabled.
"""
import os
import torch.distributed as dist
class ParallelState:
"""Process-group layout for a single inference run.
Rank layout: consecutive blocks of ``ulysses_size`` ranks form one Ulysses
sequence-parallel group; ranks sharing the same in-block offset form one
data-parallel group.
"""
def __init__(self, ulysses_size: int = 1):
self.world_size = int(os.environ.get("WORLD_SIZE", "1"))
self.rank = int(os.environ.get("RANK", "0"))
if self.world_size % ulysses_size != 0:
raise ValueError(
f"world_size ({self.world_size}) must be divisible by "
f"ulysses_size ({ulysses_size})"
)
self.ulysses_size = ulysses_size
self.dp_size = self.world_size // ulysses_size
self.ulysses_rank = self.rank % ulysses_size
self.dp_rank = self.rank // ulysses_size
self.ulysses_group = None
self.dp_group = None
@property
def ulysses_enabled(self) -> bool:
return self.ulysses_size > 1
@property
def sp_enabled(self) -> bool:
return self.ulysses_enabled
@property
def sp_size(self) -> int:
return self.ulysses_size
@property
def sp_group(self):
return self.ulysses_group
_PARALLEL_STATE = ParallelState(ulysses_size=1)
def get_parallel_state() -> ParallelState:
return _PARALLEL_STATE
def init_parallel_state(ulysses_size: int = 1) -> ParallelState:
"""Build the data-parallel / Ulysses process groups.
``torch.distributed`` must already be initialized when ``ulysses_size > 1``.
"""
global _PARALLEL_STATE
ps = ParallelState(ulysses_size)
if ps.ulysses_enabled:
# Ulysses groups: consecutive blocks of `ulysses_size` ranks.
for i in range(ps.dp_size):
ranks = list(range(i * ulysses_size, (i + 1) * ulysses_size))
group = dist.new_group(ranks)
if ps.rank in ranks:
ps.ulysses_group = group
# Data-parallel groups: one rank from each Ulysses group.
for i in range(ulysses_size):
ranks = list(range(i, ps.world_size, ulysses_size))
group = dist.new_group(ranks)
if ps.rank in ranks:
ps.dp_group = group
# Register the Ulysses group with Open-VeOmni so its native sequence-
# parallel primitives use the exact same communication semantics as
# veomni_editing Gradio inference.
from veomni.distributed.parallel_state import init_parallel_state as init_veomni_parallel_state
from veomni.distributed.sequence_parallel.comm import (
set_ulysses_sequence_parallel_group,
set_unified_sequence_parallel_group,
)
init_veomni_parallel_state(dp_size=ps.dp_size, ulysses_size=ulysses_size)
set_ulysses_sequence_parallel_group(ps.ulysses_group)
set_unified_sequence_parallel_group(ps.ulysses_group)
_PARALLEL_STATE = ps
return ps
|