[Ulysses SP] process-wide _ulysses_num_kv_heads global breaks a second model with a different head count in the same process #8291

Description

deepspeed/sequence/layer.py keeps a process-wide global for Ulysses sequence parallelism:

# deepspeed/sequence/layer.py
_ulysses_num_kv_heads = None

def set_ulysses_num_kv_heads(num):
    global _ulysses_num_kv_heads
    _ulysses_num_kv_heads = num

def get_ulysses_num_kv_heads():
    return _ulysses_num_kv_heads

It is memoized on the first uneven-head all-to-all:

def single_all_to_all(input, scatter_idx, gather_idx, batch_dim_idx, group, ...):
    ...
    if get_ulysses_num_kv_heads() is not None or (num_heads % seq_world_size != 0 and not scatter_idx < 2):
        if get_ulysses_num_kv_heads() is None:
            ...
            set_ulysses_num_kv_heads(num_heads)   # first model locks the value for the whole process
        return uneven_heads_all2all(input, scatter_idx, gather_idx, batch_dim_idx, group)

and every later call reads it (get_shard_size_list(get_ulysses_num_kv_heads(), ...) in uneven_heads_all2all, the heads_scale_coeff, chunk math, etc.).

So the first Ulysses model to take the uneven-heads path locks the total head count for the whole process; any second model in the same process with a different head count is then split against the first model's value, silently corrupting its all-to-all layout. This is the same family of bug as #8231 (AutoTP's process-wide kv-head globals, fixed per-model in #8241 by threading a per-model AutoTPMeta), and blocks the same use cases: teacher + student (OPD/distillation) and RL actor + reference policy with sequence parallelism enabled.

Note: #8241 deliberately left this global in place (its own TODO comment in layer.py documents the issue); @sfc-gh-truwase asked there to remove it as well ("yes, lets remove ulysses global variable").

Why it is a global today (the hard part)

The total head count can be recovered from the tensor shape on the scatter direction (num_heads = input.shape[2]), which is why the first call can memoize it. But:

• the gather direction (scatter_idx < 2), and
• _SeqAllToAll.backward (which re-applies the op with scatter/gather swapped)

cannot recover the total head count from the local tensor shape alone. That is why the value had to live somewhere reachable from both directions — today, a global.

Suggested approach (mirror the AutoTP per-model fix in #8241)

1. DistributedAttention (the per-model entry point that wraps each model's attention) records the total head count per instance:
   • explicit constructor arg preferred (callers have the config, same source AutoTP's tp_meta uses), or
   • probe local_attn attributes via the shared attribute lists (_KV_HEAD_ATTRS / _ATTN_HEAD_ATTRS from tp_shard), or
   • lazy capture on first scatter call — captured onto self, not a global.
2. Thread the head count through _SeqAllToAll.apply(...) → ctx → single_all_to_all(..., total_heads=...) and uneven_heads_all2all. The backward path must restore it from ctx when re-applying with scatter/gather swapped — this is the piece that cannot be dropped.
3. Keep a deprecated shim for set_ulysses_num_kv_heads/get_ulysses_num_kv_heads (external repos such as Megatron-DeepSpeed import them) that falls back to the lazy-capture behavior when no explicit value is threaded.
4. Regression test mirroring TestAutoTPMultipleModels: two DistributedAttention instances with different head counts (uneven, e.g. 3 vs 2 over sp=2) run forward+backward in the same process; the first instance's split must not change after the second is exercised.

Related: #8241 (AutoTP per-model meta), #8231 (original multi-model bug).