code: def test_bert_optimized(model_name):
      """Check that our implementation of BERT (with all optimizations enabled) matches the
      HF implementation: the output of our forward pass in fp16 should be around the same as the HF
      forward pass in fp16, when compared to the HF forward pass in fp32.
      """
      dtype = torch.float16
      config = BertConfig.from_pretrained(model_name)
      # Our implementation of fused_mlp assumes the activation is
      # nn.GELU(approximate='tanh'). Huggingface calls it "gelu_new", "gelu_fast", or "gelu_pytorch_tanh".
      # If you just want "gelu", disable fused_mlp.
      config.hidden_act = "gelu_new"
      config.use_flash_attn = True
      config.fused_bias_fc = True
      config.fused_mlp = True
      config.fused_dropout_add_ln = True
  
      model = BertForPreTraining.from_pretrained(model_name, config)
      model = model.cuda().to(dtype=dtype)
  
      model_ref = get_hf_models(model_name, config, torch.float32)
      model_hf = get_hf_models(model_name, config, dtype)
  
      model.eval()
      model_ref.eval()
      model_hf.eval()
  
      torch.manual_seed(0)
      batch_size = 4
      max_seqlen = 512
      seqlens = torch.randint(
          max_seqlen // 2, max_seqlen + 1, (batch_size,), device="cuda"
      )
      attention_mask = (
          torch.arange(max_seqlen, device="cuda")[None, :]
          < seqlens[:, None]
      )
      input_ids = torch.randint(
          0,
          config.vocab_size,
          (batch_size, max_seqlen),
          dtype=torch.long,
          device="cuda",
      )
      out = model.bert(input_ids, attention_mask=attention_mask)
      sequence_output, pooled_output = (
          out.last_hidden_state,
          out.pooler_output,
      )
      out_hf = model_hf.bert(input_ids, attention_mask=attention_mask)
      sequence_output_hf, pooled_output_hf = (
          out_hf.last_hidden_state,
          out_hf.pooler_output,
      )
      # Need to zero out the padded tokens in the sequence before comparison.
      sequence_output_hf[~attention_mask, :] = 0.0
      out_ref = model_ref.bert(input_ids, attention_mask=attention_mask)
      sequence_output_ref, pooled_output_ref = (
          out_ref.last_hidden_state,
          out_ref.pooler_output,
      )
      sequence_output_ref[~attention_mask, :] = 0.0
  
      print(
          f"BertModel output max diff: {(sequence_output - sequence_output_ref).abs().max().item()}"
      )
      print(
          f"BertModel output mean diff: {(sequence_output - sequence_output_ref).abs().mean().item()}"
      )
      print(
          f"HF fp16 BertModel max diff: {(sequence_output_hf - sequence_output_ref).abs().max().item()}"
      )
      print(
          f"HF fp16 BertModel mean diff: {(sequence_output_hf - sequence_output_ref).abs().mean().item()}"
      )
      assert (
          sequence_output - sequence_output_ref
      ).abs().max().item() < 4 * (
          sequence_output_hf - sequence_output_ref
      ).abs().max().item()
      assert (
          pooled_output - pooled_output_ref
      ).abs().max().item() < 4 * (
          pooled_output_hf - pooled_output_ref
      ).abs().max().item()
  
      out = model(input_ids, attention_mask=attention_mask)
      prediction_scores, seq_relationship_scores = (
          out.prediction_logits,
          out.seq_relationship_logits,
      )
      # Need to zero out the padded tokens in the sequence before comparison.
      prediction_scores = prediction_scores.clone()
      prediction_scores[~attention_mask, :] = 0.0
      out_hf = model_hf(input_ids, attention_mask=attention_mask)
      prediction_scores_hf, seq_relationship_scores_hf = (
          out_hf.prediction_logits,
          out_hf.seq_relationship_logits,
      )
      prediction_scores_hf[~attention_mask, :] = 0.0
      out_ref = model_ref(input_ids, attention_mask=attention_mask)
      prediction_scores_ref, seq_relationship_scores_ref = (
          out_ref.prediction_logits,
          out_ref.seq_relationship_logits,
      )
      prediction_scores_ref[~attention_mask, :] = 0.0
  
      print(
          f"prediction_scores max diff: {(prediction_scores - prediction_scores_ref).abs().max().item()}"
      )
      print(
          f"prediction_scores mean diff: {(prediction_scores - prediction_scores_ref).abs().mean().item()}"
      )
      print(
          f"HF fp16 prediction_scoresff: {(prediction_scores_hf - prediction_scores_ref).abs().max().item()}"
      )
      print(
          f"HF fp16 prediction_scoresiff: {(prediction_scores_hf - prediction_scores_ref).abs().mean().item()}"
      )
      assert (
          prediction_scores - prediction_scores_ref
      ).abs().max().item() < 2 * (
          prediction_scores_hf - prediction_scores_ref
      ).abs().max().item()
      assert (
          seq_relationship_scores - seq_relationship_scores_ref
      ).abs().max().item() < 2 * (
          seq_relationship_scores_hf - seq_relationship_scores_ref
      ).abs().max().item()
  
file_path: flash-attention/tests/models/test_bert.py
imports: import re
  from collections import OrderedDict
  import pytest
  import torch
  import torch.nn.functional
  from einops import rearrange
  from transformers import BertConfig
  from transformers.models.bert.modeling_bert import BertForPreTraining
  from transformers.models.bert.modeling_bert import BertModel
  from flash_attn.models.bert import BertForPreTraining, BertModel, inv_remap_state_dict, remap_state_dict
  from flash_attn.utils.pretrained import state_dict_from_pretrained
label_desc: Use image node_id 4 for calling a global function with example usage: test_bert_optimized(model_name) without return types