how to avoid Half() and Float() type confilction of "loss" during backward pass?

I have read many related issues but I coudlnt find a clever way to solve my problem.

Initially my model was trained on FP32 so loss function also implemented considering float32 (using template
global void SigmoidFocalLossForward..

template
global void SigmoidFocalLossBackward...

)and to speed up, I'm trying to train it again on FP16.

I didn't call model.half(), of course.

-----------------1st Trial-----------------
[ train.py]
model.to('cuda')
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
optimizer = make_optimizer(cfg, model)
scheduler = make_lr_scheduler(cfg, optimizer)

use_mixed_precision = True#cfg.DTYPE == "float16"
amp_opt_level = 'O1' if use_mixed_precision else 'O0'
model, optimizer = amp.initialize(model, optimizer, opt_level=amp_opt_level)

model = torch.nn.parallel.DistributedDataParallel(
              model, device_ids=[local_rank], output_device=local_rank,
              broadcast_buffers=False,find_unused_parameters=True
)

then I got the first error:
RuntimeError: "SigmoidFocalLoss_forward" not implemented for 'Half'

-----------------2nd Trial-----------------
so I suspect the issue would be caused by loss function related to data type.

[loss.py]
// sigmoid focal loss for calculating classification loss!
// so "target"'s type must be "int" as target means the "class id of ground truth"
class _SigmoidFocalLoss(Function):

@staticmethod
def forward(ctx, logits, targets, gamma, alpha):
    **ctx.save_for_backward( logits.float() ,  targets.int() )**
    # as apex convert the dtype flows during 'model' in half() automatically, 
    # dtype of logits is half() implicitly. 
    # So, when I pass just 'logits.half()' it trigger 
    #                           -> "SigmoidFocalLoss_backward" not implemented for 'Half'
    # or ' logits.float()' -> _SigmoidFocalLossBackward returned an invalid gradient at index 0
    # - expected type torch.cuda.HalfTensor but got torch.cuda.FloatTensor
    # or 'logits.half()' -> "SigmoidFocalLoss_backward" not implemented for 'Half'
    # or 'logits.int()' ->  "SigmoidFocalLoss_backward" not implemented for 'Int'

    ...
    ...
    losses = _C.sigmoid_focalloss_forward(
        **logits.float(), targets.int(), num_classes, gamma, alpha**
    )
    return losses

@staticmethod
@once_differentiable
def backward(ctx, d_loss):
    logits, targets = ctx.saved_tensors
    ...
    ...
    d_loss = d_loss.contiguous()
    d_logits = _C.sigmoid_focalloss_backward(
        logits, targets, d_loss, num_classes, gamma, alpha
    )  
   ** return d_logits, None, None, None, None **

but it triggers the error as below;
...
...
losses.backward()
File "/usr/local/lib/python3.6/dist-packages/torch/tensor.py", line 107, in backward
torch.autograd.backward(self, gradient, retain_graph, create_graph)
File "/usr/local/lib/python3.6/dist-packages/torch/autograd/init.py", line 93, in backward
allow_unreachable=True) # allow_unreachable flag
RuntimeError: Function _SigmoidFocalLossBackward returned an invalid gradient at index 0 - expected type torch.cuda.HalfTensor but got torch.cuda.FloatTensor

-----------------3rd Trial-----------------
>> "d_logits" are float type and my model's parameters are half/float mixed type layer by layer by apex, so I thought I need to convert the dtype of d_logits to half() as below. even if I change d_logits to half type as below, it shows similar error again.

...
...
optimizer.step()
File "/usr/local/lib/python3.6/dist-packages/apex/amp/initialize.py", line 247, in new_step
output = old_step(*args, **kwargs)
File "/usr/local/lib/python3.6/dist-packages/torch/optim/sgd.py", line 93, in step
d_p.add(weight_decay, p.data)
File "/usr/local/lib/python3.6/dist-packages/apex/amp/wrap.py", line 101, in wrapper
return orig_fn(arg0, *args, **kwargs)
RuntimeError: expected backend CUDA and dtype Float but got backend CUDA and dtype Half

I can print the values of d_logits in backward(..) function of _SigmoidFocalLoss.
So, it is obvious that the error occurs after backward of the _SigmoidFocalLoss class.

-----------------4th Trial-----------------
"return d_logits.half(), None, None, None, None" is replaced wjth

    "d_logits=d_logits.half()
    return d_logits, None, None, None, None"

ERROR again!
...
...
...
File "/usr/local/lib/python3.6/dist-packages/torch/nn/parallel/distributed.py", line 390, in forward
self.reducer.prepare_for_backward(list(_find_tensors(output)))
RuntimeError: grad.type() == variable.type() ASSERT FAILED at /pytorch/torch/csrc/distributed/c10d/reducer.cpp:214, please report a bug to PyTorch. (mark_variable_ready at /pytorch/torch/csrc/distributed/c10d/reducer.cpp:214)

As far as I understand, model is run on mixed precision and only loss part run on fp32 during forward, and then in backward, it start from fp32 of loss results which is passed to model again, so I should change it to "half()" type again.

1.  [FP16] ---- [FP32/FP16]------[FP32]
   

Input Image tensor -> model(parameters) -> loss
__________________model(parameters) <- loss
during the Backpropagataion, weights are updated in FP16.

Uuuuuu after many trials to change the types of logits or target variables to float / int / half in both forward/backward fucntion, I can't find a nice solution...

Please anybody help me~!

Thank you!

Start with train.py and loss.py, then trace _SigmoidFocalLoss.forward and backward through _C.sigmoid_focalloss_forward and _C.sigmoid_focalloss_backward. Reproduce the mixed-precision run and inspect the tensor types passed between the custom loss, DistributedDataParallel, and the optimizer; done means the forward and backward passes complete without dtype or gradient assertions.