code: def test_dropout_layer_norm_training(
      hidden_size,
      input_dtype,
      residual_dtype,
      weight_dtype,
      dropout_p,
      has_residual,
      has_rowscale,
      has_colscale,
      is_rms_norm,
  ):
      if (
          weight_dtype == torch.float16
          and input_dtype == torch.bfloat16
      ):
          pytest.skip()  # Not supported
      if is_rms_norm and FusedRMSNorm is None:
          pytest.skip()  # We need Apex's FusedRMSNorm to test
      layer_norm_cls = (
          torch.nn.LayerNorm if not is_rms_norm else FusedRMSNorm
      )
      our_layer_norm_cls = (
          DropoutAddLayerNorm if not is_rms_norm else DropoutAddRMSNorm
      )
      our_layer_norm_func = (
          dropout_add_layer_norm
          if not is_rms_norm
          else dropout_add_rms_norm
      )
      device = "cuda"
      # rtol, atol = (1e-5, 1e-6) if input_dtype == torch.float32 else (1e-3, 1e-4)
      rtol, atol = (1e-3, 1e-4)
      # set seed
      torch.random.manual_seed(0)
      batch_size = 8
      seqlen = 512
      x0_pt = torch.randn(
          batch_size,
          seqlen,
          hidden_size,
          device=device,
          dtype=input_dtype,
          requires_grad=True,
      )
      x0 = x0_pt.detach().clone().requires_grad_()
      x0_ref = x0_pt.detach().clone().float().requires_grad_()
      if has_colscale:
          colscale = torch.randn(
              hidden_size,
              device=device,
              dtype=weight_dtype,
              requires_grad=True,
          )
          colscale_pt = colscale.detach().clone().requires_grad_()
          colscale_ref = (
              colscale.detach().clone().float().requires_grad_()
          )
      else:
          colscale = None
      if has_residual:
          res_pt = torch.randn_like(
              x0, dtype=residual_dtype, requires_grad=True
          )
          res = res_pt.detach().clone().requires_grad_()
          res_ref = res_pt.detach().clone().float().requires_grad_()
      else:
          res = None
      if has_rowscale:
          rowscale = torch.empty(
              batch_size, seqlen, device=device, dtype=input_dtype
          )
          survival_rate = 0.87
          rowscale = rowscale.bernoulli_(survival_rate) / survival_rate
          x0_scaled_pt = x0_pt * rearrange(rowscale, "... -> ... 1")
          x0_scaled_ref = x0_ref * rearrange(rowscale, "... -> ... 1")
      else:
          rowscale = None
          x0_scaled_pt = x0_pt
          x0_scaled_ref = x0_ref
      if has_colscale:
          x0_scaled_pt = x0_scaled_pt * colscale_pt
          x0_scaled_ref = x0_scaled_ref * colscale_ref
      model_pt = layer_norm_cls(hidden_size).to(
          device=device, dtype=weight_dtype
      )
      torch.nn.init.normal_(model_pt.weight)
      if not is_rms_norm:
          torch.nn.init.normal_(model_pt.bias)
      model_ref = layer_norm_cls(hidden_size).to(
          device=device, dtype=torch.float32
      )
      model = our_layer_norm_cls(
          hidden_size, p=dropout_p, device=device, dtype=weight_dtype
      )
      with torch.no_grad():
          model.weight.copy_(model_pt.weight)
          model_ref.weight.copy_(model_pt.weight)
          if not is_rms_norm:
              model.bias.copy_(model_pt.bias)
              model_ref.bias.copy_(model_pt.bias)
      residual_in_fp32 = (
          not has_residual
      ) and residual_dtype == torch.float32
      out, dmask = our_layer_norm_func(
          x0,
          res,
          model.weight,
          model.bias,
          model.p,
          model.eps,
          rowscale=rowscale,
          layerscale=colscale,
          residual_in_fp32=residual_in_fp32,
          return_dropout_mask=True,
      )
      assert out.dtype == input_dtype
      print(
          f"Actual dropout fraction: {1 - dmask.float().mean().item()}"
      )
      if has_residual:
          residual_pt = (
              (x0_scaled_pt.float() * dmask.float()) / (1 - dropout_p)
              + res_pt.float()
          ).to(dtype=residual_dtype)
          residual_ref = (x0_scaled_ref * dmask.float()) / (
              1 - dropout_p
          ) + res_ref
      else:
          residual_pt = (
              (x0_scaled_pt.float() * dmask.float()) / (1 - dropout_p)
          ).to(dtype=residual_dtype)
          residual_ref = (x0_scaled_ref * dmask.float()) / (
              1 - dropout_p
          )
      out_pt = model_pt(residual_pt.to(dtype=weight_dtype)).to(
          dtype=input_dtype
      )
      out_ref = model_ref(residual_ref)
      assert (out - out_ref).abs().max() <= 4 * (
          out_pt - out_ref
      ).abs().max() + 1e-4
  
      g = torch.randn_like(out) / batch_size
      out_pt.backward(g)
      out.backward(g)
      out_ref.backward(g)
      assert (x0.grad - x0_ref.grad).abs().max() <= 4 * (
          x0_pt.grad - x0_ref.grad
      ).abs().max() + 1e-4
      if has_residual:
          assert (res.grad - res_ref.grad).abs().max() <= 4 * (
              res_pt.grad - res_ref.grad
          ).abs().max() + 1e-4
      assert (
          model.weight.grad - model_ref.weight.grad
      ).abs().max() <= 3 * (
          model_pt.weight.grad - model_ref.weight.grad
      ).abs().max() + 3e-5
      if not is_rms_norm:
          assert (
              model.bias.grad - model_ref.bias.grad
          ).abs().max() <= 2 * (
              model_pt.bias.grad - model_ref.bias.grad
          ).abs().max() + 3e-5
      if has_colscale:
          assert (
              colscale.grad - colscale_ref.grad
          ).abs().max() <= 2 * (
              colscale_pt.grad - colscale_ref.grad
          ).abs().max() + 2e-4
  
file_path: flash-attention/tests/ops/test_dropout_layer_norm.py
imports: import math
  import pytest
  import torch
  import torch.nn.functional
  from einops import rearrange, repeat
  from flash_attn.ops.layer_norm import DropoutAddLayerNorm, dropout_add_layer_norm, dropout_add_layer_norm_parallel_residual, dropout_add_layer_norm_subset
  from flash_attn.ops.rms_norm import DropoutAddRMSNorm, dropout_add_rms_norm, dropout_add_rms_norm_parallel_residual, dropout_add_rms_norm_subset
label_desc: Use image node_id 1 for calling a global function with example usage: test_dropout_layer_norm_training(hidden_size, input_dtype, residual_dtype, weight_dtype, dropout_p, has_residual, has_rowscale, has_colscale, is_rms_norm) without return types