Local Parallel State Registry and Scoping#

VeOmni stores distributed topology in ParallelState. A process can register multiple named states and temporarily select the state used by APIs that read the ambient parallel context. This supports composite models whose modules may use different sequence-parallel groups while preserving the simple get_parallel_state() interface for ordinary trainers.

Public API#

API	Purpose
`init_parallel_state(..., name="base")`	Build a topology, register it under `name`, and establish the first state as the ambient default.
`get_parallel_state_by_name(name)`	Retrieve a registered state without changing the ambient state.
`use_parallel_state(name_or_state)`	Temporarily make a registered name or `ParallelState` object ambient, then restore the previous state on exit.
`get_parallel_state()`	Return the current ambient state; before initialization it returns a single-process state and logs a warning.
`clear_parallel_state()`	Clear the ambient state, topology cache, and named registry after distributed teardown.

init_parallel_state maintains two independent mappings:

The named registry maps logical module names to states. Registering an existing name logs a warning and returns the existing state.
The topology cache reuses a state when every topology-defining argument matches. Different names can therefore refer to the same state without creating duplicate process groups.

Basic usage#

from veomni.distributed.parallel_state import (
    get_parallel_state_by_name,
    init_parallel_state,
    use_parallel_state,
)

init_parallel_state(
    dp_size=4,
    dp_shard_size=4,
    ulysses_size=2,
    name="base",
)

base_state = get_parallel_state_by_name("base")

with use_parallel_state("base"):
    output = model(input_ids)

The context manager accepts either the registered name or the state object. It always restores the previous ambient state, including when the scoped operation raises.

Multiple module topologies#

Register each logical module on every rank, in the same order, before its first scoped operation. For an eight-rank process group:

init_parallel_state(dp_size=4, dp_shard_size=4, ulysses_size=2, name="thinker")
init_parallel_state(dp_size=8, dp_shard_size=8, ulysses_size=1, name="talker")

with use_parallel_state("thinker"):
    thinker_output = thinker(batch)

with use_parallel_state("talker"):
    talker_output = talker(thinker_output)

Sequence-parallel communication helpers resolve their process groups from the current ambient state. A sequence-parallel state must therefore be created by init_parallel_state; constructing a meshless ParallelState with SP enabled raises an error.

Trainer lifecycle#

Current built-in trainers register the main topology as "base" during BaseTrainer._setup(). Model, dataloader, optimizer, and scheduler construction run inside one use_parallel_state("base") build scope. At run time, only operations that depend on ambient groups are scoped: model forward, post-forward loss handling, backward, and gradient clipping. Callbacks retain their state explicitly rather than depending on an ambient scope.

When an API accepts an explicit process group, prefer passing the group from get_parallel_state_by_name("base") instead of opening a broader context.

Teardown#

Destroy distributed process groups before clearing the state registry:

import torch.distributed as dist

from veomni.distributed.parallel_state import clear_parallel_state

dist.destroy_process_group()
clear_parallel_state()

Clearing prevents a later distributed session from reusing cached device meshes or process groups that have already been destroyed.

Local Parallel State Registry and Scoping

Contents