code: def test_dropout_layer_norm_parallel_residual_prenorm_training(
      hidden_size,
      input_dtype,
      residual_dtype,
      weight_dtype,
      dropout_p,
      has_x1,
      has_residual,
      tied_norm,
      is_rms_norm,
  ):
      if (
          weight_dtype == torch.float16
          and input_dtype == torch.bfloat16
      ):
          pytest.skip()  # Not supported
      if is_rms_norm and fused_rms_norm_affine is None:
          pytest.skip()  # We need Apex's FusedRMSNorm to test
      our_layer_norm_func = (
          dropout_add_layer_norm_parallel_residual
          if not is_rms_norm
          else dropout_add_rms_norm_parallel_residual
      )
      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_x1:
          x1_pt = torch.randn(
              batch_size,
              seqlen,
              hidden_size,
              device=device,
              dtype=input_dtype,
              requires_grad=True,
          )
          x1 = x1_pt.detach().clone().requires_grad_()
          x1_ref = x1_pt.detach().clone().float().requires_grad_()
      else:
          x1 = 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
      weight0 = torch.randn(
          hidden_size,
          device=device,
          dtype=weight_dtype,
          requires_grad=True,
      )
      bias0 = (
          torch.randn(
              hidden_size,
              device=device,
              dtype=weight_dtype,
              requires_grad=True,
          )
          if not is_rms_norm
          else None
      )
      weight0_pt = weight0.detach().clone().requires_grad_()
      weight0_ref = weight0.detach().clone().float().requires_grad_()
      bias0_pt = (
          bias0.detach().clone().requires_grad_()
          if bias0 is not None
          else None
      )
      bias0_ref = (
          bias0.detach().clone().float().requires_grad_()
          if bias0 is not None
          else None
      )
      if not tied_norm:
          weight1 = torch.randn(
              hidden_size,
              device=device,
              dtype=weight_dtype,
              requires_grad=True,
          )
          bias1 = (
              torch.randn(
                  hidden_size,
                  device=device,
                  dtype=weight_dtype,
                  requires_grad=True,
              )
              if not is_rms_norm
              else None
          )
          weight1_pt = weight1.detach().clone().requires_grad_()
          weight1_ref = (
              weight1.detach().clone().float().requires_grad_()
          )
          bias1_pt = (
              bias1.detach().clone().requires_grad_()
              if bias1 is not None
              else None
          )
          bias1_ref = (
              bias1.detach().clone().float().requires_grad_()
              if bias1 is not None
              else None
          )
      else:
          weight1, bias1 = None, None
      epsilon = 1e-5
      residual_in_fp32 = (
          not has_residual
      ) and residual_dtype == torch.float32
  
      out0, out1, residual, dmask0, dmask1 = our_layer_norm_func(
          x0,
          x1,
          res,
          weight0,
          bias0,
          weight1,
          bias1,
          dropout_p,
          epsilon,
          prenorm=True,
          residual_in_fp32=residual_in_fp32,
          return_dropout_mask=True,
      )
      assert out0.dtype == input_dtype
      if not tied_norm:
          assert out1.dtype == input_dtype
      print(
          f"Actual dropout fraction: {1 - dmask0.float().mean().item()}"
      )
      if has_residual:
          if has_x1:
              residual_pt = (
                  (x0_pt.float() * dmask0.float()) / (1 - dropout_p)
                  + (x1_pt.float() * dmask1.float()) / (1 - dropout_p)
                  + res_pt.float()
              ).to(dtype=residual_dtype)
              residual_ref = (
                  (x0_ref * dmask0.float()) / (1 - dropout_p)
                  + (x1_ref * dmask1.float()) / (1 - dropout_p)
              ) + res_ref
          else:
              residual_pt = (
                  (x0_pt.float() * dmask0.float()) / (1 - dropout_p)
                  + res_pt.float()
              ).to(dtype=residual_dtype)
              residual_ref = (x0_ref * dmask0.float()) / (
                  1 - dropout_p
              ) + res_ref
      else:
          if has_x1:
              residual_pt = (
                  (x0_pt.float() * dmask0.float()) / (1 - dropout_p)
                  + (x1_pt.float() * dmask1.float()) / (1 - dropout_p)
              ).to(dtype=residual_dtype)
              residual_ref = (x0_ref * dmask0.float()) / (
                  1 - dropout_p
              ) + (x1_ref * dmask1.float()) / (1 - dropout_p)
          else:
              residual_pt = (
                  (x0_pt.float() * dmask0.float()) / (1 - dropout_p)
              ).to(dtype=residual_dtype)
              residual_ref = (x0_ref * dmask0.float()) / (1 - dropout_p)
      if not is_rms_norm:
          out0_pt = F.layer_norm(
              residual_pt.to(dtype=weight_dtype),
              (hidden_size,),
              weight0_pt,
              bias0_pt,
              eps=epsilon,
          ).to(dtype=input_dtype)
          out0_ref = F.layer_norm(
              residual_ref,
              (hidden_size,),
              weight0_ref,
              bias0_ref,
              eps=epsilon,
          )
          if not tied_norm:
              out1_pt = F.layer_norm(
                  residual_pt.to(dtype=weight_dtype),
                  (hidden_size,),
                  weight1_pt,
                  bias1_pt,
                  eps=epsilon,
              ).to(dtype=input_dtype)
              out1_ref = F.layer_norm(
                  residual_ref,
                  (hidden_size,),
                  weight1_ref,
                  bias1_ref,
                  eps=epsilon,
              )
      else:
          out0_pt = fused_rms_norm_affine(
              residual_pt.to(dtype=weight_dtype),
              weight0_pt,
              (hidden_size,),
              eps=epsilon,
          ).to(dtype=input_dtype)
          out0_ref = fused_rms_norm_affine(
              residual_ref, weight0_ref, (hidden_size,), eps=epsilon
          )
          if not tied_norm:
              out1_pt = fused_rms_norm_affine(
                  residual_pt.to(dtype=weight_dtype),
                  weight1_pt,
                  (hidden_size,),
                  eps=epsilon,
              ).to(dtype=input_dtype)
              out1_ref = fused_rms_norm_affine(
                  residual_ref, weight1_ref, (hidden_size,), eps=epsilon
              )
  
      assert (out0 - out0_ref).abs().max() <= 4 * (
          out0_pt - out0_ref
      ).abs().max() + 1e-4
      if not tied_norm:
          assert (out1 - out1_ref).abs().max() <= 4 * (
              out1_pt - out1_ref
          ).abs().max() + 1e-4
      assert (residual - residual_ref).abs().max() <= 4 * (
          residual_pt - residual_ref
      ).abs().max() + 1e-4
  
      g0 = torch.randn_like(out0) / batch_size
      if tied_norm:
          (out0 * F.sigmoid(residual)).backward(g0)
          (out0_pt * F.sigmoid(residual_pt)).backward(g0)
          (out0_ref * F.sigmoid(residual_ref)).backward(g0)
      else:
          g1 = torch.randn_like(out1) / batch_size
          (out0 * F.sigmoid(residual) * g0 + out1 * g1).sum().backward()
          (
              out0_pt * F.sigmoid(residual_pt) * g0 + out1_pt * g1
          ).sum().backward()
          (
              out0_ref * F.sigmoid(residual_ref) * g0 + out1_ref * g1
          ).sum().backward()
      assert (x0.grad - x0_ref.grad).abs().max() <= 4 * (
          x0_pt.grad - x0_ref.grad
      ).abs().max() + 1e-4
      if has_x1:
          assert (x1.grad - x1_ref.grad).abs().max() <= 4 * (
              x1_pt.grad - x1_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 (weight0.grad - weight0_ref.grad).abs().max() <= 3 * (
          weight0_pt.grad - weight0_ref.grad
      ).abs().max() + 3e-5
      if not is_rms_norm:
          assert (bias0.grad - bias0_ref.grad).abs().max() <= 2 * (
              bias0_pt.grad - bias0_ref.grad
          ).abs().max() + 3e-5
      if not tied_norm:
          assert (weight1.grad - weight1_ref.grad).abs().max() <= 3 * (
              weight1_pt.grad - weight1_ref.grad
          ).abs().max() + 3e-5
          if not is_rms_norm:
              assert (bias1.grad - bias1_ref.grad).abs().max() <= 2 * (
                  bias1_pt.grad - bias1_ref.grad
              ).abs().max() + 3e-5
  
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 8 for calling a global function with example usage: test_dropout_layer_norm_parallel_residual_prenorm_training(hidden_size, input_dtype, residual_dtype, weight_dtype, dropout_p, has_x1, has_residual, tied_norm, is_rms_norm) without return types