Repository: facebookresearch/fairscale

Issue:
[feature] support normalization, scale and margin in MEVO
## 🚀 Feature
currently, vanilla FC + softmax + nllloss is supported in MEVO. In face recognition, cosface loss is commonly used which additionally normalize both the input feature and fc weights and add scale and margin to the logits.

## Motivation
face recognition models are typically trained with millions of IDs. MEVO will have broader impact in face recognition field if CosFace is supported.

## Additional context
CosFace paper: https://arxiv.org/abs/1801.09414


Discussion:
wow, this is very interesting. I haven't seen this before. Thanks @gongbudaizhe. Do you have some example code of cosface loss that's easily understandable? If you have time, I can help you with that. Otherwise, I will try to find time with this in the future depending on my other workloads.
Thanks for your interest. @min-xu-ai  Below is the sample code of cosface loss.

```python
# inputs shape and dtype
# x, shape (B, D), B batch size, D feature dim, float32 dtype
# w, shape (N, D), N num class, for example, 1 million, D feature dim, float32 dtype
# im_label, shape (B,), B batch size, long dtype. The labels of input x. 
 
x = backbone(img)  # output feature, (B, D), B batch size, D feature dim 
w = Tensor(N, D) # the weights of FC Layer, N num class, D feature dim.

# normalize feature and fc layer before multiplication
x = nn.functional.normalize(x, dim=1)
w = nn.functional.normalize(w, dim=1)
logits = torch.einsum('nc,kc->nk', [x, w])

# add margin and scale. They are hyperparameters.
margin, scale = 0.35, 64

# add margin
row_ind = torch.arange(x.shape[0], dtype=torch.long).to(x.device)
col_ind = im_label
logits[row_ind, col_ind] -= margin

# add scale
logits *= scale

loss = F.cross_entropy(logits, im_label)
```
Thanks @gongbudaizhe! Can you provide shape and dtype for x, w, im_label for big scale use cases for the code above?
@min-xu-ai updated in my last comment : )
Is the 8-64 a reasonable size for B and 128-1024 a good number for D?
> Is the 8-64 a reasonable size for B and 128-1024 a good number for D?

yes, these settings are good. Personally, I use B = 64, D = 256.  
I have one more question @gongbudaizhe.

Is this needed in your research or you are doing a product with this? The reason I am asking is that this has impact on the use cases of face rec. Will this enable massive surveillance of people without their consent? We all need to be mindful about social impact of the technologies that we work on. Thank you in advance for the discussion.
It is a good question. But I don't have a definitive answer.

Whether it is needed in my research or product doesn't really matter here. Once this feature is avaible, people will use it for various purposes.  Pens can kill, guns can protect.

CosFace loss is not only used in face recognition. It is also used in large scale image retrieval. 
>  It is also used in large scale image retrieval.

Thank you. Can you provide some references for the above?
CosFace used in landmark retrieval:
https://arxiv.org/pdf/1906.04087.pdf

ArcFace (a cross entropy variant which resembles to cosface), used in a more recent framework:
https://arxiv.org/pdf/2111.00775.pdf

Resolving patch:
diff --git a/fairscale/experimental/nn/mevo.py b/fairscale/experimental/nn/mevo.py
--- a/fairscale/experimental/nn/mevo.py
+++ b/fairscale/experimental/nn/mevo.py
@@ -68,6 +68,10 @@ def get_data(
 class BaselineSoftmax(nn.Module):
     """Baseline softmax that does an output linear projection and a softmax.
 
+
+        We also support LMCL (Large Margin Cosine Loss) from the CosFace paper. See
+        more detailed comment in the MEVO class below.
+
         This is intended to be used with an embedding layer with shared weights.
 
     Args:
@@ -77,10 +81,23 @@ class BaselineSoftmax(nn.Module):
             Unused. It is here to make kernel init easier with MEVO.
         log_softmax (bool):
             If True, use log_softmax instead of softmax.
+        margin (float):
+            Used in LMCL (when scale != None). See MEVO comments for
+            more details.
+        scale (Optional[float]):
+            Used in LMCL. If scale is None, LMCL is turned off. See
+            MEVO comments for more details.
 
     """
 
-    def __init__(self, proj_weight: nn.Parameter, tile_factor: int = 0, log_softmax: bool = True):
+    def __init__(
+        self,
+        proj_weight: nn.Parameter,
+        tile_factor: int = 0,
+        log_softmax: bool = True,
+        margin: float = 0.35,
+        scale: Optional[float] = None,
+    ):
         super().__init__()
         out_dim, in_dim = proj_weight.shape
         assert "cuda" in str(proj_weight.device), "weight should be on GPU"
@@ -92,6 +109,27 @@ def __init__(self, proj_weight: nn.Parameter, tile_factor: int = 0, log_softmax:
         assert self.fc.weight.dtype in [torch.float16, torch.float32], self.fc.weight.dtype
         self.fp16 = self.fc.weight.dtype == torch.float16
         self.log_softmax = log_softmax
+        self.margin = margin
+        self.scale = scale
+
+    def lmcl_pre_softmax(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+        # normalize feature and fc layer before multiplication
+        #   n: number of features (tokens)
+        #   k: number of classes (vocab size)
+        #   c: hidden dimension (d_model)
+        x = F.normalize(input, dim=1)
+        w = F.normalize(self.fc.weight, dim=1)
+        logits = torch.einsum("nc,kc->nk", x, w)
+
+        # add margin
+        row_ind = torch.arange(x.shape[0], dtype=torch.long).to(x.device)
+        col_ind = target
+        logits[row_ind, col_ind] -= self.margin
+
+        # add scale
+        logits *= self.scale
+
+        return logits
 
     def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:  # type: ignore
         """Forward function that computes softmax output with the input and target."""
@@ -100,7 +138,10 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:  #
         input, target = _reshape_inputs(input, target)
         if self.fp16:
             assert input.dtype == torch.float16
-        x = self.fc(input)
+        if self.scale is not None:
+            x = self.lmcl_pre_softmax(input, target)
+        else:
+            x = self.fc(input)
         # Note that we do softmax in FP32, which is important for numerical stability.
         if self.log_softmax:
             x = F.log_softmax(x, dim=-1, dtype=torch.float32)
@@ -119,8 +160,15 @@ class BaselineSoftmaxNllLoss(BaselineSoftmax):
     This class is used for testing and benchmarking.
     """
 
-    def __init__(self, proj_weight: nn.Parameter, tile_factor: int = 0, log_softmax: bool = True):
-        super().__init__(proj_weight, tile_factor, log_softmax)
+    def __init__(
+        self,
+        proj_weight: nn.Parameter,
+        tile_factor: int = 0,
+        log_softmax: bool = True,
+        margin: float = 0.35,
+        scale: Optional[float] = None,
+    ):
+        super().__init__(proj_weight, tile_factor, log_softmax, margin, scale)
 
     def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:  # type: ignore
         """Forward that directly compute the loss."""
@@ -131,18 +179,49 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:  #
         return F.nll_loss(x, target, reduction="sum")
 
 
+def lmcl_matmul(
+    i: torch.Tensor, w: torch.Tensor, tgt: torch.Tensor, w_idx: int, margin: float, scale: Optional[float]
+) -> torch.Tensor:
+    """LMCL variation of matmul with normalization, margin and scale."""
+    # normalize and matmul
+    logits = torch.matmul(F.normalize(i, dim=1), F.normalize(w, dim=1).T)
+
+    # add margin using a mask since tgt might be out of the the weight split's range.
+    mask = torch.arange(w_idx * w.shape[0], (w_idx + 1) * w.shape[0], dtype=torch.long, device=i.device).expand(
+        i.shape[0], -1
+    )
+    logits[mask == tgt.reshape(-1, 1)] -= margin
+
+    # add scale
+    logits *= scale
+
+    return logits
+
+
 class GetMaxFunction(torch.autograd.Function):
     """Custom checkpointed function to get max-per-token from an input and a weight"""
 
     @staticmethod
-    def get_max(i: torch.Tensor, w: torch.Tensor, full_precision: bool) -> torch.Tensor:
+    def get_max(
+        i: torch.Tensor,
+        w: torch.Tensor,
+        tgt: torch.Tensor,
+        w_idx: int,
+        full_precision: bool,
+        margin: float,
+        scale: Optional[float],
+    ) -> torch.Tensor:
         """
         Throughout this code:
 
           i: input data with shape = (split-of-tokens, d_model)
           w: weight data with shape = (split-of-vocabs, d_model)
+          tgt: target prediction data with shape = (split-of-tokens,)
         """
-        _m = torch.matmul(i, w.T)
+        if scale is not None:
+            _m = lmcl_matmul(i, w, tgt, w_idx, margin, scale)
+        else:
+            _m = torch.matmul(i, w.T)
         if full_precision:
             _m = _m.float()
         _m = _m.max(dim=1)[0]
@@ -153,6 +232,7 @@ def forward(  # type: ignore
         ctx: Any,
         i: torch.Tensor,
         w: torch.Tensor,
+        tgt: torch.Tensor,
         kernel_obj: "MemoryEfficientVocabOutput",
         w_idx: int,
         w_split_size: int,
@@ -161,7 +241,7 @@ def forward(  # type: ignore
         """Forward function that computes the max, without saving activations."""
         if DEBUG and dist.is_initialized() and dist.get_rank() == 0:
             print("DEBUG max fwd")
-        ctx.save_for_backward(i, w)
+        ctx.save_for_backward(i, w, tgt)
         ctx.kernel_obj = kernel_obj
         ctx.w_idx = w_idx
         ctx.w_split_size = w_split_size
@@ -171,7 +251,7 @@ def forward(  # type: ignore
         # The activations will be recomputed in backward below and freed
         # immediately after use. This saves the overall GPU peak memory of this layer.
         with torch.no_grad():
-            return GetMaxFunction.get_max(i, w, kernel_obj.fp_max)
+            return GetMaxFunction.get_max(i, w, tgt, w_idx, kernel_obj.fp_max, kernel_obj.margin, kernel_obj.scale)
 
     @staticmethod
     def backward(ctx: Any, *args: Any) -> Any:
@@ -186,7 +266,7 @@ def backward(ctx: Any, *args: Any) -> Any:
         assert ctx.kernel_obj.proj_weight.grad is not None
 
         # Get saved i and w.
-        i, w = ctx.saved_tensors
+        i, w, tgt = ctx.saved_tensors
         assert i.requires_grad
         assert w.requires_grad
         # We use ``detach()'' to ensure the backward call below does not
@@ -199,7 +279,9 @@ def backward(ctx: Any, *args: Any) -> Any:
         # Forward + backward again.
         with torch.enable_grad():
             # This saves the activations.
-            maxs = GetMaxFunction.get_max(i, w, ctx.kernel_obj.fp_max)
+            maxs = GetMaxFunction.get_max(
+                i, w, tgt, ctx.w_idx, ctx.kernel_obj.fp_max, ctx.kernel_obj.margin, ctx.kernel_obj.scale
+            )
         # This will use the activations and free them immediately.
         torch.autograd.backward(maxs, *args)
 
@@ -208,15 +290,27 @@ def backward(ctx: Any, *args: Any) -> Any:
         with torch.no_grad():
             grads = torch.split(ctx.kernel_obj.proj_weight.grad, ctx.w_split_size)
             grads[ctx.w_idx].add_(w.grad)
-        return i.grad, None, None, None, None, None
+        return i.grad, None, None, None, None, None, None
 
 
 class GetSumFunction(torch.autograd.Function):
     """Custom checkpointed function to get sum-per-token from an input and a weight."""
 
     @staticmethod
-    def get_sum(i: torch.Tensor, w: torch.Tensor, maxs: torch.Tensor, full_precision: bool) -> torch.Tensor:
-        _s = torch.matmul(i, w.T)
+    def get_sum(
+        i: torch.Tensor,
+        w: torch.Tensor,
+        tgt: torch.Tensor,
+        maxs: torch.Tensor,
+        w_idx: int,
+        full_precision: bool,
+        margin: float,
+        scale: Optional[float],
+    ) -> torch.Tensor:
+        if scale is not None:
+            _s = lmcl_matmul(i, w, tgt, w_idx, margin, scale)
+        else:
+            _s = torch.matmul(i, w.T)
         if full_precision:
             _s = _s.float()
         _s = (_s - maxs.reshape(-1, 1)).exp().sum(dim=1)
@@ -227,6 +321,7 @@ def forward(  # type: ignore
         ctx: Any,
         i: torch.Tensor,
         w: torch.Tensor,
+        tgt: torch.Tensor,
         maxs: torch.Tensor,
         kernel_obj: "MemoryEfficientVocabOutput",
         w_idx: int,
@@ -236,13 +331,15 @@ def forward(  # type: ignore
         """Forward function that computes the sum, without saving activations."""
         if DEBUG and dist.is_initialized() and dist.get_rank() == 0:
             print("DEBUG sum fwd")
-        ctx.save_for_backward(i, w, maxs)
+        ctx.save_for_backward(i, w, tgt, maxs)
         ctx.kernel_obj = kernel_obj
         ctx.w_idx = w_idx
         ctx.w_split_size = w_split_size
         assert split_dim == 0
         with torch.no_grad():
-            return GetSumFunction.get_sum(i, w, maxs, kernel_obj.fp_sum)
+            return GetSumFunction.get_sum(
+                i, w, tgt, maxs, w_idx, kernel_obj.fp_sum, kernel_obj.margin, kernel_obj.scale
+            )
 
     @staticmethod
     def backward(ctx: Any, *args: Any) -> Any:
@@ -257,7 +354,7 @@ def backward(ctx: Any, *args: Any) -> Any:
         assert ctx.kernel_obj.proj_weight.grad is not None
 
         # Get saved i, w, and maxs.
-        i, w, maxs = ctx.saved_tensors
+        i, w, tgt, maxs = ctx.saved_tensors
         assert i.requires_grad
         assert w.requires_grad
         assert maxs.requires_grad
@@ -267,7 +364,9 @@ def backward(ctx: Any, *args: Any) -> Any:
 
         # Forward + backward again.
         with torch.enable_grad():
-            sums = GetSumFunction.get_sum(i, w, maxs, ctx.kernel_obj.fp_sum)
+            sums = GetSumFunction.get_sum(
+                i, w, tgt, maxs, ctx.w_idx, ctx.kernel_obj.fp_sum, ctx.kernel_obj.margin, ctx.kernel_obj.scale
+            )
         torch.autograd.backward(sums, *args)
 
         # Accumulate the grads.
@@ -275,22 +374,35 @@ def backward(ctx: Any, *args: Any) -> Any:
         with torch.no_grad():
             grads = torch.split(ctx.kernel_obj.proj_weight.grad, ctx.w_split_size)
             grads[ctx.w_idx].add_(w.grad)
-        return i.grad, None, maxs.grad, None, None, None, None
+        return i.grad, None, None, maxs.grad, None, None, None, None
 
 
 class TargetScoreFunction(torch.autograd.Function):
     """Custom checkpointed function to compute the target score."""
 
     @staticmethod
-    def get_target_score(i: torch.Tensor, w: torch.Tensor, target: torch.Tensor, full_precision: bool) -> torch.Tensor:
+    def get_target_score(
+        i: torch.Tensor,
+        w: torch.Tensor,
+        target: torch.Tensor,
+        full_precision: bool,
+        margin: float,
+        scale: Optional[float],
+    ) -> torch.Tensor:
         tokens, d_model = i.shape
         assert d_model == w.shape[1]
         tw = w.gather(dim=0, index=target.reshape(target.shape[0], 1).expand(target.shape[0], d_model))
         assert tw.shape == (tokens, d_model)
-        target_score = i * tw
+        if scale is not None:
+            target_score = F.normalize(i, dim=1) * F.normalize(tw, dim=1)
+        else:
+            target_score = i * tw
         if full_precision:
             target_score = target_score.float()
         target_score = target_score.sum(dim=1)  # sum into target scores with shape (tokens,)
+        if scale is not None:
+            target_score -= margin
+            target_score *= scale
         return target_score
 
     @staticmethod
@@ -303,7 +415,9 @@ def forward(  # type: ignore
         ctx.save_for_backward(i, w, target)
         ctx.kernel_obj = kernel_obj
         with torch.no_grad():
-            x = TargetScoreFunction.get_target_score(i, w, target, kernel_obj.fp_target)
+            x = TargetScoreFunction.get_target_score(
+                i, w, target, kernel_obj.fp_target, kernel_obj.margin, kernel_obj.scale
+            )
         return x
 
     @staticmethod
@@ -319,7 +433,9 @@ def backward(ctx: Any, *args: Any) -> Any:
         i = i.detach().requires_grad_(True)
         w = w.detach().requires_grad_(True)
         with torch.enable_grad():
-            scores = TargetScoreFunction.get_target_score(i, w, target, ctx.kernel_obj.fp_target)
+            scores = TargetScoreFunction.get_target_score(
+                i, w, target, ctx.kernel_obj.fp_target, ctx.kernel_obj.margin, ctx.kernel_obj.scale
+            )
         torch.autograd.backward(scores, *args)
         if ctx.kernel_obj.proj_weight.grad is not None:
             # This means we accumulate full grad between iters. Not memory efficient.
@@ -388,18 +504,55 @@ def forward(self) -> torch.Tensor:  # type: ignore
 class MemoryEfficientVocabOutput(nn.Module):  # AKA. MEVO
     """Fused fc + softmax + nll_loss in a tiled fashion.
 
-        This uses much less memory but is quite a bit slower.
+        MEVO uses much less memory but is quite a bit slower.
+
+        MEVO also implements the LMCL (Large Margin Cosine Loss) function introduced by
+        highly cited
+        `CosFace: Large Margin Cosine Loss for Deep Face Recognition [Wang et al.]`_.
+
+        .. _`CosFace: Large Margin Cosine Loss for Deep Face Recognition [Wang et al.]`: https://arxiv.org/abs/1801.09414
+
+        LMCL can be turned on using the ``margin`` and ``scale`` parameters below. These
+        hyperparameters most likely require tuning, depending on the number of classes etc.
+
+        MEVO LMCL can be suitable for face recognition and image retrieval tasks, esp. when
+        the number prediction target classes is large. MEVO is slower but can use much
+        less GPU memory in that case, which enables training with larger batches. We
+        hope this is helpful but we strongly recommend users (AI researchers
+        and engineers) to carefully consider their applications of this technology. This
+        types of technology should not be used by small group of people exclusively to
+        potentially harm the general public.
 
     Args:
         proj_weight (nn.Parameter):
             Sharing this weight with an embedding layer.
         tile_factor (int):
             Number of splits to use on the input sequence and vocab dimensions.
+            Default: 16
         reduction (str):
             Reduction OP (sum or mean).
+            Default: sum
+        margin (float):
+            Hyperparameter of the separation margin between classes. See the
+            appendix of the CosFace paper for a formula on how to compute its
+            value properly. The default value is unlikely to be suitable in all
+            cases.
+            Default: 0.35
+        scale (Optional[float]):
+            Hyperparameter of the feature-vector-scaling for LMCL. When not
+            supplied, LMCL is turned off. See the appendix of the CosFace paper for
+            a formula on how to compute its value properly.
+            Default: None
     """
 
-    def __init__(self, proj_weight: nn.Parameter, tile_factor: int = 16, reduction: str = "sum"):
+    def __init__(
+        self,
+        proj_weight: nn.Parameter,
+        tile_factor: int = 16,
+        reduction: str = "sum",
+        margin: float = 0.35,
+        scale: Optional[float] = None,
+    ):
         super().__init__()
         self.proj_weight = proj_weight
         # TODO (Min): these two factors doesn't have to be the same. More tuning can be done.
@@ -410,12 +563,14 @@ def __init__(self, proj_weight: nn.Parameter, tile_factor: int = 16, reduction:
         self.log_softmax = True
         self.reduction = reduction
         assert self.reduction in ["sum", "mean"]
+        self.margin = margin
+        self.scale = scale
         self.trigger = BackwardTrigger(self.proj_weight)
         if DEBUG and dist.is_initialized() and dist.get_rank() == 0:
             print(
                 f"DEBUG cfg tf_in={self.tf_in} tf_w={self.tf_w} fp_max={self.fp_max} "
                 f"fp_sum={self.fp_sum} fp_target={self.fp_target} log_softmax={self.log_softmax} "
-                f"reduction={self.reduction}"
+                f"reduction={self.reduction} margin={self.margin} scale={self.scale}"
             )
 
     def get_target_nlprob(
@@ -432,6 +587,8 @@ def get_target_nlprob(
 
     def eval_forward(self, input: torch.Tensor) -> torch.Tensor:
         """Eval time forward that doesn't fuse the softmax and NLL Loss kernels."""
+        # Margin, scaling and normalization of LMCL does not apply to eval time as far as
+        # I can tell. Therefore, we just do a matmul like the standard output layer.
         return torch.matmul(input, self.proj_weight.T)
 
     def forward(self, input: torch.Tensor, target: Optional[torch.Tensor]) -> torch.Tensor:  # type: ignore
@@ -449,8 +606,10 @@ def forward(self, input: torch.Tensor, target: Optional[torch.Tensor]) -> torch.
         input, target = _reshape_inputs(input, target)
 
         tokens, d_model = input.shape
+        (t2,) = target.shape
         vocab, d2 = self.proj_weight.shape
-        assert d_model == d2
+        assert d_model == d2, f"incorrect shape {d_model} vs {d2}"
+        assert tokens == t2, f"incorrect shape {tokens} vs {t2}"
         split_dim = 0
         input_split_size = _next_power_of_2_or_max(tokens // self.tf_in, tokens)
         weight_split_size = _next_power_of_2_or_max(vocab // self.tf_w, vocab)
@@ -458,12 +617,16 @@ def forward(self, input: torch.Tensor, target: Optional[torch.Tensor]) -> torch.
         weight = self.trigger()
         weights = torch.split(weight, weight_split_size, split_dim)
 
+        targets = tuple([torch.Tensor()] * len(inputs))
+        if self.scale is not None:
+            targets = torch.split(target, input_split_size, split_dim)
+
         # Get maxs
         maxs = []
-        for i in inputs:
+        for i, tgt in zip(inputs, targets):
             m = None  # max with (tokens_tile,) shape
             for w_idx, w in enumerate(weights):
-                _m = GetMaxFunction.apply(i, w, self, w_idx, weight_split_size, split_dim)
+                _m = GetMaxFunction.apply(i, w, tgt, self, w_idx, weight_split_size, split_dim)
                 if m is None:
                     m = _m
                 else:
@@ -475,10 +638,10 @@ def forward(self, input: torch.Tensor, target: Optional[torch.Tensor]) -> torch.
 
         # Get sums.
         sums = []
-        for idx, i in enumerate(inputs):
+        for i, tgt, debase_max in zip(inputs, targets, maxs):
             s = None  # sum with (tokens_tile,) shape
             for w_idx, w in enumerate(weights):
-                _s = GetSumFunction.apply(i, w, maxs[idx], self, w_idx, weight_split_size, split_dim)
+                _s = GetSumFunction.apply(i, w, tgt, debase_max, self, w_idx, weight_split_size, split_dim)
                 if s is None:
                     s = _s
                 else:


Test patch:
diff --git a/tests/experimental/nn/test_mevo.py b/tests/experimental/nn/test_mevo.py
--- a/tests/experimental/nn/test_mevo.py
+++ b/tests/experimental/nn/test_mevo.py
@@ -38,10 +38,15 @@ def test_mevo_eval():
     assert out.shape == (1, 5, 4)
 
 
+# Note for the lmcl_scale, overly large value, like 64 for small shape input
+# will cause inf/nan in mevo. Larger scale value is only needed for large shape inputs.
 @skip_if_no_cuda
-def test_mevo():
-    """Test the MEVO kernel by itself."""
+@pytest.mark.parametrize("lmcl_scale", [None, 8])
+def test_mevo(lmcl_scale):
+    """Test the MEVO kernel in a single process (no DDP/FSDP)."""
+    # Set seed and reset peak mem so that peak measure below is correct.
     torch.random.manual_seed(os.getpid())
+    torch.cuda.reset_peak_memory_stats()
     shape = ((5, 3), (3, 7))
     # Turn on large data for local testing.
     large = False
@@ -50,11 +55,11 @@ def test_mevo():
     print("\nshapes are", shape)
 
     input, weight, target = get_data(shape, dtype=torch.float16)
-    k = MEVO(weight, tile_factor=16)
+    k = MEVO(weight, tile_factor=16, scale=lmcl_scale)
 
     o = k(input, target)
     o.backward()
-    print(o, o.shape)
+    print("MEVO loss", o, o.shape)
     del o
 
     cur_mem = round(torch.cuda.memory_allocated() / 1024 / 1024)
@@ -70,22 +75,22 @@ def test_mevo():
     mem = round(torch.cuda.max_memory_allocated() / 1024 / 1024)
     print("after moving input and its grad, cur and peak mem for tiled fwd+bwd =", cur_mem, mem)
 
-    print(weight.grad.norm(), weight.grad)
+    print("MEVO grad norm and grad", weight.grad.norm(), weight.grad)
     g1 = weight.grad.clone()
     weight.grad = None
 
     input = input_data.cuda().requires_grad_(True)
-    refk = BaselineSoftmaxNllLoss(weight)
+    refk = BaselineSoftmaxNllLoss(weight, scale=lmcl_scale)
     o = refk(input, target)
     o.backward()
-    print(o, o.shape)
+    print("Reference loss", o, o.shape)
     del o
-    print(weight.grad.norm(), weight.grad)
+    print("Reference grad norm and grad", weight.grad.norm(), weight.grad)
     g2 = weight.grad.clone()
     input_grad2 = input.grad.cpu()
 
-    # Print the diff. We use .cuda() since in 1.7 and 1.8, min() and max() are not
-    # implemented for cpu float16.
+    # Print the diff. We use .cuda() since in torch 1.7 and 1.8, min() and max() are not
+    # implemented for cpu float16. The diff should in general be below 0.01 in magnitude.
     diff = g1 - g2
     print("weight grad diff", diff.cuda().min(), diff.cuda().max())
     diff = input_grad1 - input_grad2