code: def test_tensor_parallel(model_name, rotary, world_size):
      """Check that our implementation of GPT2 generation matches the HF implementation:
      the scores in fp16 should be around the same as the HF scores in fp16, when compared to
      the HF scores in fp32.
      """
      dtype = torch.float16
      rtol, atol = 3e-3, 3e-1
      config = GPT2Config.from_pretrained(model_name)
      if rotary:
          config.n_positions = 0
          config.rotary_emb_dim = 64
      config.residual_in_fp32 = True
      config.use_flash_attn = True
      config.fused_bias_fc = True
      config.fused_mlp = True
      config.fused_dropout_add_ln = True
      config.pad_vocab_size_multiple = 8 * world_size
      config.sequence_parallel = (
          False  # Need to set this to False for generation
      )
  
      os.environ["NCCL_ASYNC_ERROR_HANDLING"] = "0"
      if not torch.distributed.is_initialized():
          torch.distributed.init_process_group(
              backend="nccl", init_method="env://"
          )
      device = f"cuda:{torch.distributed.get_rank()}"
      assert world_size <= torch.distributed.get_world_size()
      # Need this, otherwise when we capture the graph the process for GPU 1 would run on both
      # GPU0 and GPU1 and things would hang
      torch.cuda.set_device(device)
  
      from apex.transformer import parallel_state
  
      parallel_state.initialize_model_parallel(
          tensor_model_parallel_size_=world_size
      )
      rank = parallel_state.get_tensor_model_parallel_rank()
      process_group = parallel_state.get_tensor_model_parallel_group()
  
      # if not rotary, we load the weight from HF but ignore the position embeddings.
      # The model would be nonsense but it doesn't matter for the test.
      model = GPTLMHeadModel.from_pretrained(
          model_name,
          config,
          strict=not rotary,
          device=device,
          dtype=dtype,
          process_group=process_group,
          world_size=world_size,
          rank=rank,
      )
      model.eval()
  
      if not rotary:
          model_ref = GPT2LMHeadModelHF.from_pretrained(model_name).to(
              device=device
          )
          model_hf = GPT2LMHeadModelHF.from_pretrained(model_name).to(
              device=device, dtype=dtype
          )
          model_ref.eval()
          model_hf.eval()
  
      torch.manual_seed(0)
      tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
      input_ids = tokenizer(
          "Hello, my dog is cute and ", return_tensors="pt"
      ).input_ids.to(device=device)
      max_length = 30
      # input_ids = torch.randint(0, 100, (1, 10), dtype=torch.long, device='cuda')
      # max_length = input_ids.shape[1] + 40
  
      # Slow generation for reference
      sequences = []
      scores = []
      cur_input_ids = input_ids
      with torch.inference_mode():
          logits, _ = all_gather_raw(
              model(cur_input_ids).logits[:, -1], process_group
          )
          logits = rearrange(
              logits, "(n b) d -> b (n d)", b=input_ids.shape[0]
          )[..., : config.vocab_size]
          scores.append(logits)
          sequences.append(scores[-1].argmax(dim=-1))
          for _ in range(input_ids.shape[1] + 1, max_length):
              cur_input_ids = torch.cat(
                  [cur_input_ids, rearrange(sequences[-1], "b -> b 1")],
                  dim=-1,
              )
              logits, _ = all_gather_raw(
                  model(cur_input_ids).logits[:, -1], process_group
              )
              logits = rearrange(
                  logits, "(n b) d -> b (n d)", b=input_ids.shape[0]
              )[..., : config.vocab_size]
              scores.append(logits)
              sequences.append(scores[-1].argmax(dim=-1))
      sequences = torch.cat(
          [input_ids, torch.stack(sequences, dim=1)], dim=1
      )
      scores = tuple(scores)
      print(sequences)
  
      out = model.generate(
          input_ids=input_ids,
          max_length=max_length,
          tensor_parallel=world_size,
          vocab_size=config.vocab_size,
          return_dict_in_generate=True,
          output_scores=True,
          enable_timing=True,
      )
      print(out.sequences)
      if getattr(config, "use_flash_attn", False):
          out_cg = model.generate(
              input_ids=input_ids,
              max_length=max_length,
              tensor_parallel=world_size,
              vocab_size=config.vocab_size,
              cg=True,
              return_dict_in_generate=True,
              output_scores=True,
              enable_timing=True,
          )
          print(out_cg.sequences)
  
      parallel_state.destroy_model_parallel()
  
      if not rotary:
          out_hf = model_hf.generate(
              input_ids=input_ids,
              max_length=max_length,
              return_dict_in_generate=True,
              output_scores=True,
          )
          out_ref = model_ref.generate(
              input_ids=input_ids,
              max_length=max_length,
              return_dict_in_generate=True,
              output_scores=True,
          )
  
          print(
              f"Scores max diff: {(torch.stack(out.scores, 1) - torch.stack(out_ref.scores, 1)).abs().max().item()}"
          )
          print(
              f"Scores mean diff: {(torch.stack(out.scores, 1) - torch.stack(out_ref.scores, 1)).abs().mean().item()}"
          )
          print(
              f"HF fp16 max diff: {(torch.stack(out_hf.scores, 1) - torch.stack(out_ref.scores, 1)).abs().max().item()}"
          )
          print(
              f"HF fp16 mean diff: {(torch.stack(out_hf.scores, 1) - torch.stack(out_ref.scores, 1)).abs().mean().item()}"
          )
  
      assert torch.all(out.sequences == sequences)
      assert torch.allclose(
          torch.stack(out.scores, dim=1),
          torch.stack(scores, dim=1),
          rtol=rtol,
          atol=atol,
      )
      assert torch.equal(
          torch.stack(out.scores, dim=1),
          torch.stack(out_cg.scores, dim=1),
      )
      if not rotary:
          assert torch.all(out.sequences == out_ref.sequences)
          assert torch.all(out.sequences == out_hf.sequences)
  
          assert (
              torch.stack(out.scores, 1)
              - torch.stack(out_ref.scores, 1)
          ).abs().max().item() < 3 * (
              torch.stack(out_hf.scores, 1)
              - torch.stack(out_ref.scores, 1)
          ).abs().max().item()
  
file_path: flash-attention/tests/models/test_gpt_generation_parallel.py
imports: import os
  import re
  import pytest
  import torch
  from einops import rearrange
  from flash_attn.models.gpt import GPTLMHeadModel, remap_state_dict_hf_gpt2
  from flash_attn.utils.distributed import all_gather_raw
  from flash_attn.utils.pretrained import state_dict_from_pretrained
  from transformers import GPT2Config, GPT2Tokenizer
  from transformers.models.gpt2.modeling_gpt2 import GPT2LMHeadModel
label_desc: Use image node_id 1 for calling a global function with example usage: test_tensor_parallel(model_name, rotary, world_size) without return types