code: 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
      if large:
          shape = ((1 * 2048, 4096), (4096, 256008))
      print("\nshapes are", shape)
  
      input, weight, target = get_data(shape, dtype=torch.float16)
      k = MEVO(weight, tile_factor=16, scale=lmcl_scale)
  
      o = k(input, target)
      o.backward()
      print("MEVO loss", o, o.shape)
      del o
  
      cur_mem = round(torch.cuda.memory_allocated() / 1024 / 1024)
      mem = round(torch.cuda.max_memory_allocated() / 1024 / 1024)
      print("cur and peak mem for tiled fwd+bwd =", cur_mem, mem)
  
      assert input.shape == input.grad.shape
      input_data = input.data.cpu()
      input_grad1 = input.grad.cpu()
      del input
  
      cur_mem = round(torch.cuda.memory_allocated() / 1024 / 1024)
      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("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, scale=lmcl_scale)
      o = refk(input, target)
      o.backward()
      print("Reference loss", o, o.shape)
      del o
      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 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
      print("input grad diff", diff.cuda().min(), diff.cuda().max())
  
file_path: fairscale/tests/experimental/nn/test_mevo.py
imports: import os
  import pytest
  import torch
  from fairscale.experimental.nn import MEVO
  from fairscale.experimental.nn.mevo import BaselineSoftmaxNllLoss, get_data
  from fairscale.fair_dev.testing.testing import skip_if_no_cuda
label_desc: Use image node_id 3 for calling a global function with example usage: test_mevo(lmcl_scale) without return types