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

[delock]

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).

[FU-max-boop]

I have a stacked implementation ready on top of #8241 at FU-max-boop/DeepSpeed@05878c4.

The implementation removes the process-global dependency from DistributedAttention by threading two per-call values through _SeqAllToAll and its autograd context:

• total_heads: the tensor's total Q/K/V head count;
• partition_heads: the KV-group count that determines uneven rank partitions.

That distinction matters for GQA: Q=6/KV=3 over SP=2 must split Q as [4, 2], matching KV groups, rather than [3, 3]. Q/K/V/O carry (Q,KV)/(KV,KV)/(KV,KV)/(Q,KV), including the backward pass. The low-level global remains only as a compatibility fallback for callers such as FPDT that do not supply metadata; evenly divisible legacy Q=8 then K=4 does not set or read the fallback.

Local evidence:

• deterministic two-DistributedAttention forward/backward regression: red on the [AutoTP] Replace tp_shard process-wide globals with per-model AutoTPMeta #8241 head, green with the patch;
• real 2-rank Gloo GQA round trip for Q=6/KV=3: [4,2] shards, exact values, and all-one gradients on both ranks;
• real 2-rank Gloo legacy uneven-head round trips for both batch layouts;
• focused tests and all changed-file pre-commit hooks pass;
• independent diff review found no P0-P2 issues.

I have not opened an upstream PR yet because doing so against master would include the entire unmerged #8241 dependency. Once #8241 lands, I can rebase this single commit and open a clean PR immediately. If you prefer reviewing the stacked commit before then, I am happy to keep it rebased on the exact #8241 head.

[alanhuangyoo]

Opened #8316 for this. It threads the KV head count through DistributedAttention instead of reading the tp_shard slot, and touches only deepspeed/sequence/layer.py, so it does not depend on #8241.

Two things that came out of building it:

AutoTP on its own is enough to trigger this — no second Ulysses model needed. replace_module.py:306 and engine.py:770 both write the slot, and single_all_to_all reads "slot is set" as "this process does uneven-head SP". A 12-head model on sp=4 then gets [4, 4, 4, 0] instead of [3, 3, 3, 3], rank 3 attends over an empty tensor, and the forward hangs in the output all-to-all.

Threading a value through is also not sufficient by itself. get_shard_size_list falls back to the same module-level global internally, so it has to be passed num_kv_heads explicitly or the leak comes back through the helper.

@FU-max-boop your GQA point caught a real bug in my first revision — I was reading the count off the query tensor, so Q=6 / KV=3 over sp=2 came out [3, 3] and tore the KV group apart. Now it reads the key tensor and there is a regression test for that case. If you would rather land your stacked version, say so and I will close mine; with #8241 as the base yours would be the cleaner place for the total_heads / partition_heads split anyway.

[FU-max-boop]

Thanks for the careful follow-up and for crediting the GQA observation. A dependency-free #8316 is the preferable landing path if we close two regressions I reproduced on its exact current head: the legacy direct odd-head _SeqAllToAll backward loses the inferred count and desynchronizes the collective, and Q=2 / KV=1 / SP=2 produces a zero-head rank (ZeroDivisionError on one rank, peer hang) instead of synchronized fail-fast.

I left the exact base-vs-head evidence, public harness, and minimal fix directions on #8316. I will not open a competing PR while it is being addressed; I will keep 05878c4 as a fallback/reference and am happy to rerun the reproducer and review the next head.