2 tests fail: assert isinstance(model.llm[0], DistributedAttention)

[yurivict]

============================================================ FAILURES =============================================================
___________________________ TestAutoWrapModelForSP.test_llm_layers_replaced_with_distributed_attention ____________________________

self = <deepspeed.sequence.test_autosp.TestAutoWrapModelForSP object at 0x1e9a5777e00>

    def test_llm_layers_replaced_with_distributed_attention(self):
        """LLM attention layers must be wrapped with DistributedAttention."""
        pg = _make_mock_process_group(world_size=2, rank=0)
        model = _FakeLLMOnlyModel(num_layers=3)
        auto_wrap_model_for_sp(model, pg)
        for layer in model.layers:
>           assert isinstance(layer, DistributedAttention)
E           assert False
E            +  where False = isinstance(_FakeLLMAttn(), DistributedAttention)

deepspeed/sequence/test_autosp.py:253: AssertionError
-------------------------------------------------------- Captured log call --------------------------------------------------------
WARNING  deepspeed.sequence.auto_sp:auto_sp.py:105 AutoSP: LLM attention 'layers.0' (class _FakeLLMAttn) uses a HuggingFace hidden_states interface that is incompatible with DistributedAttention's Q/K/V interface. Skipping auto-wrap. Configure sequence parallelism for this layer manually.
WARNING  deepspeed.sequence.auto_sp:auto_sp.py:105 AutoSP: LLM attention 'layers.1' (class _FakeLLMAttn) uses a HuggingFace hidden_states interface that is incompatible with DistributedAttention's Q/K/V interface. Skipping auto-wrap. Configure sequence parallelism for this layer manually.
WARNING  deepspeed.sequence.auto_sp:auto_sp.py:105 AutoSP: LLM attention 'layers.2' (class _FakeLLMAttn) uses a HuggingFace hidden_states interface that is incompatible with DistributedAttention's Q/K/V interface. Skipping auto-wrap. Configure sequence parallelism for this layer manually.
________________________________ TestAutoWrapModelForSP.test_multimodal_model_wraps_both_branches _________________________________

self = <deepspeed.sequence.test_autosp.TestAutoWrapModelForSP object at 0x1e9a57755e0>

    def test_multimodal_model_wraps_both_branches(self):
        """Both ViT and LLM attention layers must be replaced in a combined model."""
        pg = _make_mock_process_group(world_size=2, rank=0)
        model = _FakeMultimodalModel()
        returned = auto_wrap_model_for_sp(model, pg)
        # auto_wrap_model_for_sp must return the same object (in-place)
        assert returned is model
        assert isinstance(model.vision_encoder[0], UlyssesSPViTAttention)
>       assert isinstance(model.llm[0], DistributedAttention)
E       assert False
E        +  where False = isinstance(_FakeLLMAttn(), DistributedAttention)

deepspeed/sequence/test_autosp.py:263: AssertionError
-------------------------------------------------------- Captured log call --------------------------------------------------------

Version: 0.19.6
Python-3.12
FreeBSD 15.1