Skip to content

Fix Ulysses SP registration for string model paths and gloo collectives - #8702

Open
delock wants to merge 3 commits into
deepspeedai:masterfrom
delock:pr-p-ulysses-sp-string-path-and-gloo
Open

delock wants to merge 3 commits into
deepspeedai:masterfrom
delock:pr-p-ulysses-sp-string-path-and-gloo

Conversation

@delock

@delock delock commented Sep 29, 2026

Copy link
Copy Markdown
Collaborator

Description

Two latent bugs in Ulysses SP, both backend-independent (neither is CPU-specific — the CPU/gloo multi-rank run in #8381 is simply what surfaced them), plus test adaptations so the file can also run on non-CUDA accelerators.

1. register_with_transformers rejected every string model path

deepspeed/runtime/sequence_parallel/ulysses_sp.py compared core_attn_implementation against the config's _attn_implementation even when the caller passed a model path string:

hf_model_config = AutoConfig.from_pretrained(model_name_or_path)   # string branch
model_attn_implementation = getattr(hf_model_config, "_attn_implementation", None)
if model_attn_implementation is not None and model_attn_implementation != core_attn_implementation:
    raise ValueError(...)

In transformers, _attn_implementation is a property resolved at model load time — from_pretrained resolves which implementation the runtime can actually use and writes it back to the config. A bare AutoConfig never goes through that resolution, so the property returns its fallback value for "unresolved":

AutoConfig.from_pretrained(path)._attn_implementation      -> eager   (fallback, not a choice)
AutoModelForCausalLM.from_pretrained(path).config._attn... -> sdpa    (resolved)

So every string-path caller asking for sdpa/flex_attention/FA2 was rejected with a spurious mismatch against 'eager'. The check is only meaningful for a config that was resolved, i.e. when the caller handed us a constructed model (or a PEFT wrapper exposing .config); scope it to that case. The comparison expression itself is unchanged, and TestUlyssesSPHFAttnImplMismatch (which passes a model object) still raises as intended.

2. UlyssesSPDataLoaderAdapter used a 0-dim tensor in all_gather

seqlen = torch.tensor(n, ...)                    # 0-dim
seqlens = [torch.zeros(1, ...) for _ in range(sp_world_size)]   # shape (1,)
dist.all_gather(seqlens, seqlen, group=self.sp_group)

c10d.all_gather requires each receive tensor to have the same size as the send tensor. NCCL moves flattened bytes and tolerates the mismatch; gloo validates shapes and rejects it (ProcessGroupGloo::allgather: invalid tensor size at index 0 (expected (), got (1))). Send a 1-element tensor to match.

3. Tests

  • TestUlyssesSPHFFlexAttention needs the CUDA flex kernel -> skip on non-CUDA accelerators (same guard idiom as tests/unit/util.py, the nvme/triton CUDA-only tests).
  • TestUlyssesSPHFPEFT: the mock attached a bare AutoConfig, whose _attn_implementation is the unresolved 'eager' — so this test failed on every backend, not just CPU. Give the mock the load-resolved value a real PEFT wrapper carries (the sibling TestUlyssesSPHFAttnImplMismatch in the same file already sets it the same way).
  • TestUlyssesSPHFDisableInEval pinned its tensors/models to cuda:<rank> although it runs sdpa and nothing in it is CUDA-specific. Use get_accelerator().current_device_name() like the rest of the file (on CUDA that is exactly cuda:<rank>, since the test harness sets the device per rank in tests/unit/common.py), so the test runs on any backend instead of erroring.

Validation

Hardware: 20-core x86_64 CPU without AVX512-FP16, gloo backend, 2 ranks (LOCAL_SIZE=2), transformers 4.51.3, torch 2.13.

  • tests/unit/ulysses_alst/test_ulysses_sp_hf.py: 5 passed, 2 skipped (flex skipped by the new guard). The sdpa classes pass end to end including their numerical-parity assertions; before the fix they died at registration.
  • TestUlyssesSPHFDisableInEval now actually executes on CPU rather than being skipped.
  • TestUlyssesSPHFAttnImplMismatch still raises the expected ValueError (the mismatch check is preserved where a resolved config exists).
  • Full multi-rank CPU CI in [DON'T MERGE] Run multi-rank CPU unit tests in CI via LOCAL_SIZE #8381: this file went from 5 failures to 0, with the overall job green.

(TestUlyssesSPHFHubKernel is not exercised by the version pin above — it needs the lazy_import_flash_attention API from a newer transformers; it passes in CI, which installs transformers from git main.)

Sibling PRs from the same series: #8397, #8398, #8399, #8407, #8409, #8559, #8648, #8684.

Two latent bugs surfaced by the multi-rank CPU run (gloo):

- register_with_transformers validated core_attn_implementation against
  a bare AutoConfig's _attn_implementation when given a string model
  path. A config that has not gone through model loading still holds
  the unresolved 'eager' default, so every string-path caller asking
  for sdpa/flex/FA2 was rejected with a spurious mismatch. Only compare
  when the caller passed an actual model, whose config carries the
  implementation resolved at load time.

- UlyssesSPDataLoaderAdapter exchanged the local sequence length as a
  0-dim tensor but allocated 1-element receive buffers. Backends that
  move raw bytes (nccl) tolerate that; gloo validates gather shapes
  strictly and rejects it. Send a 1-element tensor to match.

With both fixed the sdpa Ulysses SP tests pass end to end on cpu/gloo
(verified 2/2 locally, incl. the numerical-parity assertions).

Signed-off-by: Guokai Ma <guokai.ma@intel.com>
With the engine's string-path registration fixed, the sdpa classes run
on cpu/gloo. Guard the CUDA-only remainder (the disable-in-eval test
hardcodes cuda:<rank> tensors, flex_attention needs CUDA) and make the
mock PEFT model carry the load-resolved attention implementation a real
PEFT wrapper would have, instead of the bare-config 'eager' default
that trips the consistency check.

Signed-off-by: Guokai Ma <guokai.ma@intel.com>
The test pinned its tensors and models to cuda:<rank> although nothing in
it is CUDA-specific (it runs sdpa); use the accelerator device like the
rest of the file, so it runs on any backend instead of being skipped.

Signed-off-by: Guokai Ma <guokai.ma@intel.com>

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant