scale parameter for torch.nn.functional.scaled_dot_product_attention

[atiorh]

apple/ml-stable-diffusion#345

[atiorh]

Tested for apple/ml-stable-diffusion#345 with:

python -m python_coreml_stable_diffusion.torch2coreml --convert-text-encoder --model-version runwayml/stable-diffusion-v1-5 -o /tmp  --check-output-correctness

which yields:

...
Converting PyTorch Frontend ==> MIL Ops: 100%██████████▍| 447/449 [00:00<00:00, 8624.97 ops/s]
Running MIL frontend_pytorch pipeline: 100%██████████| 5/5 [00:00<00:00, 267.80 passes/s]
Running MIL default pipeline: 100%|██████████▍ 79/79 [00:02<00:00, 36.21 passes/s]
Running MIL backend_mlprogram pipeline: 100%██████████| 12/12 [00:00<00:00, 310.13 passes/s]
INFO:__main__:Saved text_encoder into /tmp/Stable_Diffusion_version_runwayml_stable-diffusion-v1-5_text_encoder.mlpackage
INFO:__main__:text_encoder baseline PyTorch to reference CoreML: PSNR changed by -173.8 dB (230.4 -> 56.6)
INFO:__main__:56.6 dB > 35 dB (minimum allowed) parity check passed
...

[junpeiz]

@atiorh Great! Thank you for the PR!
Could you add a corresponding test case in TestScaledDotProductAttention? Then I will kickoff a CI run for your PR. Thanks!

[atiorh]

@junpeiz Just pushed, is this aligned with your test coverage expectations?

[junpeiz]

@junpeiz Just pushed, is this aligned with your test coverage expectations?

Looks good! Thank you so much for adding the test. I kicked off a CI run: https://gitlab.com/coremltools1/coremltools/-/pipelines/1404611382

[junpeiz]

@junpeiz Just pushed, is this aligned with your test coverage expectations?

Looks good! Thank you so much for adding the test. I kicked off a CI run: https://gitlab.com/coremltools1/coremltools/-/pipelines/1404611382

@atiorh Looks like flake8 failed. Could you pip install flake8 and then run flake8 locally to fix that? Thanks!

[atiorh]

@junpeiz Updated, thanks!

[junpeiz]

Thanks! Kicked off CI: https://gitlab.com/coremltools1/coremltools/-/pipelines/1404713912

[atiorh]

Thanks! Kicked off CI: https://gitlab.com/coremltools1/coremltools/-/pipelines/1404713912

Looks like python-3.10 PyTorch tests failed. I do not see much info about the failure in the logs.

[junpeiz]

Thanks! Kicked off CI: https://gitlab.com/coremltools1/coremltools/-/pipelines/1404713912

Looks like python-3.10 PyTorch tests failed. I do not see much info about the failure in the logs.

If you search TestScaledDotProductAttention.test_scale_argument in the log, you will find this error message:

args = (<function assert_allclose.<locals>.compare at 0x17f578670>, array([[[[0.670439  , 0.6824666 , 0.46654922, 0.5271193 ,...     [0.5374129 , 0.5065092 , 0.5429782 , 0.4860599 , 0.40717188,
          0.59352744, 0.6726432 ]]]], dtype=float32))
kwds = {'equal_nan': True, 'err_msg': '', 'header': 'Not equal to tolerance rtol=1e-05, atol=0.0001', 'verbose': True}
    @wraps(func)
    def inner(*args, **kwds):
        with self._recreate_cm():
>           return func(*args, **kwds)
E           AssertionError: 
E           Not equal to tolerance rtol=1e-05, atol=0.0001
E           
E           Mismatched elements: 210 / 210 (100%)
E           Max absolute difference: 0.110039
E           Max relative difference: 0.35954633
E            x: array([[[[0.670439, 0.682467, 0.466549, 0.527119, 0.476547, 0.548151,
E                     0.633629],
E                    [0.686291, 0.695975, 0.489841, 0.538309, 0.476038, 0.5612  ,...
E            y: array([[[[0.705857, 0.713525, 0.517099, 0.522854, 0.469099, 0.58429 ,
E                     0.670839],
E                    [0.687978, 0.711702, 0.482552, 0.566067, 0.483409, 0.567483,...
../../envs/coremltools-py3.10/lib/python3.10/contextlib.py:79: AssertionError

[YifanShenSZ]

Will be fixed in 8.0 release