[distributed] NotImplementedError: Operator aten._scaled_dot_product_fused_attention_overrideable.default does not have a sharding strategy registered.

[PenghuiCheng]

🚀 The feature, motivation and pitch

Error:
NotImplementedError: Operator aten._scaled_dot_product_fused_attention_overrideable.default does not have a sharding strategy registered.

cases:
test/distributed/tensor/parallel/test_tp_examples.py::DistTensorParallelExampleTest::test_transformer_req_grad_seq_parallel_float32_thaw_norm__output

tp_examples.log

Alternatives

Operator aten._scaled_dot_product_fused_attention_overrideable.default does not have a sharding strategy registered.

Additional context

No response

[zhangxiaoli73]

@ashokei Could you please assign someone to take a look?

[daisyden]

@githubsgi need further investigation to see why the kernel is called.

[githubsgi]

Looks like it is happening because the UT is setting requires_grad_(False) for attention.wk/q/v.weight .

[PenghuiCheng]

@githubsgi , UT is PyTorch code. Does XPU behavior align with CUDA?

[githubsgi]

Looks similar to Cuda, but we will need a fused kernel that supports both forward and backward to completely verify.