amp does not work with model parallel

Background:
In my problem, I need to compute the cost over the whole global batch, i.e., the batch-correlated cost instead of the average of sample-independent losses (like L1), for example, the negative linear correlation loss in regression problems.

I am not familiar with distributed learning, and I do not know how to implement my code to compute the batch-correlated cost instead of computing them separately on each device.

Problem:
So I turned to model parallel. Then I got the following error:

RuntimeError: Expected tensor for argument #1 'input' to have the same device as tensor for argument #2 'weight'; but device 15 does not equal 14 (while checking arguments for cudnn_convolution)

So is there any way for using both amp and model parallel?
Or is there any simple way to compute the batch-correlated cost in the distributed training?

Thanks!

Start with the PyTorch model parallel tutorial linked in the issue and reproduce the reported cudnn_convolution device-mismatch error when using AMP. A useful resolution would establish whether AMP supports model parallel here and clarify a supported approach for computing the global batch-correlated cost, or document the limitation.