Add option to normalize loss per target

[Muennighoff]

No description provided.

[thomasw21]

There's a fun hack you can do, view have the same storage space as the initial model. so you can probably write something like:

def fast_normalize(loss_mask: torch.Tensor):
    """
    Turn loss_mask from [0,0,0,1,1,0,0,1,0,0,1,1,1] > [0,0,0,0.5,0.5,0,0,1,0,0,0.3,0.3,0.3]
    """
    flatten_view = loss_mask.view(-1)
    _, inverse_indices, counts = torch.unique_consecutive(loss_mask, return_inverse=True, return_counts=True)
    counts = torch.gather(dim=0, index=inverse_indices, input=counts)
    flatten_view.div_(counts)
    return loss_mask

[thomasw21]

you could also clone before doing this operation so that you actually don't make fast_normalize a in-place operation.

[Muennighoff]

Why is

def fast_normalize(loss_mask: torch.Tensor):
    """
    Turn loss_mask from [0,0,0,1,1,0,0,1,0,0,1,1,1] > [0,0,0,0.5,0.5,0,0,1,0,0,0.3,0.3,0.3]
    """
    flatten_view = loss_mask.view(-1)
    _, inverse_indices, counts = torch.unique_consecutive(loss_mask, return_inverse=True, return_counts=True)
    counts = torch.gather(dim=0, index=inverse_indices, input=counts)
    flatten_view.div_(counts)
    return loss_mask

better than

def fast_normalize(loss_mask: torch.Tensor):
    """
    Turn loss_mask from [0,0,0,1,1,0,0,1,0,0,1,1,1] > [0,0,0,0.5,0.5,0,0,1,0,0,0.3,0.3,0.3]
    """
    _, inverse_indices, counts = torch.unique_consecutive(loss_mask, return_inverse=True, return_counts=True)
    counts = torch.gather(dim=0, index=inverse_indices, input=counts)
    return loss_mask / counts

?

[thomasw21]

Does the latter work if loss_mask is not 1D?