Conversation
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
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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_transformersrejected every string model pathdeepspeed/runtime/sequence_parallel/ulysses_sp.pycomparedcore_attn_implementationagainst the config's_attn_implementationeven when the caller passed a model path string:In transformers,
_attn_implementationis a property resolved at model load time —from_pretrainedresolves which implementation the runtime can actually use and writes it back to the config. A bareAutoConfignever goes through that resolution, so the property returns its fallback value for "unresolved":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, andTestUlyssesSPHFAttnImplMismatch(which passes a model object) still raises as intended.2.
UlyssesSPDataLoaderAdapterused a 0-dim tensor inall_gatherc10d.all_gatherrequires 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
TestUlyssesSPHFFlexAttentionneeds the CUDA flex kernel -> skip on non-CUDA accelerators (same guard idiom astests/unit/util.py, the nvme/triton CUDA-only tests).TestUlyssesSPHFPEFT: the mock attached a bareAutoConfig, whose_attn_implementationis 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 siblingTestUlyssesSPHFAttnImplMismatchin the same file already sets it the same way).TestUlyssesSPHFDisableInEvalpinned its tensors/models tocuda:<rank>although it runs sdpa and nothing in it is CUDA-specific. Useget_accelerator().current_device_name()like the rest of the file (on CUDA that is exactlycuda:<rank>, since the test harness sets the device per rank intests/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.TestUlyssesSPHFDisableInEvalnow actually executes on CPU rather than being skipped.TestUlyssesSPHFAttnImplMismatchstill raises the expectedValueError(the mismatch check is preserved where a resolved config exists).(
TestUlyssesSPHFHubKernelis not exercised by the version pin above — it needs thelazy_import_flash_attentionAPI 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.