# Copyright (c) 2023, Tri Dao. import os import time from pathlib import Path current_dir = Path(__file__).parent.absolute() import pytest import torch from einops import rearrange from flash_attn.models.falcon import falcon_config_to_gpt2_config, remap_state_dict_hf_falcon from flash_attn.models.gpt import GPTLMHeadModel, combine_state_dicts_tp, shard_state_dict_tp from flash_attn.utils.distributed import all_gather_raw from flash_attn.utils.generation import update_graph_cache from flash_attn.utils.pretrained import state_dict_from_pretrained from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer @pytest.mark.parametrize("model_name", ["tiiuae/falcon-7b", "tiiuae/falcon-40b"]) def test_falcon_state_dict(model_name): config = falcon_config_to_gpt2_config( AutoConfig.from_pretrained(model_name, trust_remote_code=True) ) pretrained_state_dict = remap_state_dict_hf_falcon( state_dict_from_pretrained(model_name), config ) model = GPTLMHeadModel(config, device="meta") # Without device='meta' init is very slow state_dict = model.state_dict() assert state_dict.keys() == pretrained_state_dict.keys() for k in state_dict.keys(): assert state_dict[k].shape == pretrained_state_dict[k].shape @pytest.mark.parametrize("model_name", ["tiiuae/falcon-7b"]) def test_falcon_optimized(model_name): """Check that our implementation (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 device = "cuda" config = falcon_config_to_gpt2_config( AutoConfig.from_pretrained(model_name, trust_remote_code=True) ) config.use_flash_attn = True config.fused_bias_fc = True config.fused_mlp = False # We don't have fused MLP for "gelu" activation config.fused_dropout_add_ln = True config.residual_in_fp32 = True model = GPTLMHeadModel.from_pretrained(model_name, config, device=device, dtype=dtype) model.eval() torch.manual_seed(0) batch_size = 2 max_seqlen = 256 input_ids = torch.randint( 0, config.vocab_size, (batch_size, max_seqlen), dtype=torch.long, device=device ) with torch.no_grad(): out = model.transformer(input_ids) logits = model(input_ids).logits del model # Without device_map, the model is loaded on the CPU, which is very slow model_ref = AutoModelForCausalLM.from_pretrained( model_name, device_map={"": device}, trust_remote_code=True ) model_ref.eval() with torch.no_grad(): out_ref = model_ref.transformer(input_ids).last_hidden_state.to(device=device) logits_ref = model_ref(input_ids).logits.to(device=device) del model_ref model_hf = AutoModelForCausalLM.from_pretrained( model_name, torch_dtype=dtype, device_map={"": device}, trust_remote_code=True ) model_hf.eval() out_hf = model_hf.transformer(input_ids).last_hidden_state logits_hf = model_hf(input_ids).logits del model_hf print(f"Output max diff: {(out - out_ref).abs().max().item()}") print(f"Output mean diff: {(out - out_ref).abs().mean().item()}") print(f"HF fp16 max diff: {(out_hf - out_ref).abs().max().item()}") print(f"HF fp16 mean diff: {(out_hf - out_ref).abs().mean().item()}") assert (out - out_ref).abs().max().item() < 3 * (out_hf - out_ref).abs().max().item() print(f"Logits max diff: {(logits - logits_ref).abs().max().item()}") print(f"Logits mean diff: {(logits - logits_ref).abs().mean().item()}") print(f"HF fp16 max diff: {(logits_hf - logits_ref).abs().max().item()}") print(f"HF fp16 mean diff: {(logits_hf - logits_ref).abs().mean().item()}") assert (logits - logits_ref).abs().max().item() < 3 * ( logits_hf - logits_ref ).abs().max().item() # torchrun --no_python --nproc_per_node=4 pytest -q -s tests/models/test_falcon.py -k "falcon_parallel_forward" # We want to run this on a machine with 4 x A100 80GB or 8 x A100 40GB so we have enough # memory to run the model in fp32. @pytest.mark.parametrize("world_size", [4]) @pytest.mark.parametrize("model_name", ["tiiuae/falcon-40b"]) def test_falcon_parallel_forward(model_name, world_size): from apex.transformer import parallel_state dtype = torch.float16 config = falcon_config_to_gpt2_config( AutoConfig.from_pretrained(model_name, trust_remote_code=True) ) config.use_flash_attn = False config.fused_bias_fc = True config.fused_mlp = False # We don't have fused MLP for "gelu" activation config.fused_dropout_add_ln = False config.residual_in_fp32 = True 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() 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() pretrained_state_dict = remap_state_dict_hf_falcon( state_dict_from_pretrained(model_name), config ) model = GPTLMHeadModel(config, process_group=process_group, device=device, dtype=dtype) model.load_state_dict(shard_state_dict_tp(pretrained_state_dict, config, world_size, rank)) model.eval() torch.manual_seed(0) batch_size = 2 max_seqlen = 256 seqlens = torch.randint(max_seqlen // 2, max_seqlen + 1, (batch_size,), device=device) input_ids = torch.randint( 0, config.vocab_size, (batch_size, max_seqlen), dtype=torch.long, device=device ) with torch.no_grad(): out = model.transformer(input_ids) out, _ = all_gather_raw(out, process_group=process_group) out = rearrange(out, "(b s) d -> b s d", b=batch_size) logits = model(input_ids).logits logits = rearrange(logits, "(b s) d -> b s d", b=batch_size) logits, _ = all_gather_raw(logits, process_group) logits = rearrange(logits, "(n b) ... d -> b ... (n d)", b=batch_size) del model parallel_state.destroy_model_parallel() if rank == 0: model_hf = AutoModelForCausalLM.from_pretrained( model_name, torch_dtype=dtype, device_map="auto", trust_remote_code=True ) model_hf.eval() out_hf = model_hf.transformer(input_ids).last_hidden_state.to(device=device) logits_hf = model_hf(input_ids).logits.to(device=device) del model_hf # Without device_map, the model is loaded on the CPU, which is very slow model_ref = AutoModelForCausalLM.from_pretrained( model_name, device_map="auto", trust_remote_code=True ) model_ref.eval() with torch.no_grad(): out_ref = model_ref.transformer(input_ids).last_hidden_state.to(device=device) logits_ref = model_ref(input_ids).logits.to(device=device) del model_ref print(f"Output max diff: {(out - out_ref).abs().max().item()}") print(f"Output mean diff: {(out - out_ref).abs().mean().item()}") print(f"HF fp16 max diff: {(out_hf - out_ref).abs().max().item()}") print(f"HF fp16 mean diff: {(out_hf - out_ref).abs().mean().item()}") assert (out - out_ref).abs().max().item() < 2 * (out_hf - out_ref).abs().max().item() print(f"Logits max diff: {(logits - logits_ref).abs().max().item()}") print(f"Logits mean diff: {(logits - logits_ref).abs().mean().item()}") print(f"HF fp16 max diff: {(logits_hf - logits_ref).abs().max().item()}") print(f"HF fp16 mean diff: {(logits_hf - logits_ref).abs().mean().item()}") assert (logits - logits_ref).abs().max().item() < 2 * ( logits_hf - logits_ref ).abs().max().item() @pytest.mark.parametrize("model_name", ["tiiuae/falcon-7b"]) def test_falcon_generation(model_name): """Check that our implementation (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 device = "cuda" config = falcon_config_to_gpt2_config( AutoConfig.from_pretrained(model_name, trust_remote_code=True) ) config.use_flash_attn = True config.fused_bias_fc = True config.fused_mlp = False # We don't have fused MLP for "gelu" activation config.fused_dropout_add_ln = True config.residual_in_fp32 = True tokenizer = AutoTokenizer.from_pretrained(model_name) eos_token_id = tokenizer.eos_token_id torch.manual_seed(0) batch_size = 1 seqlen = 100 max_length = 150 input_ids = torch.randint( 0, config.vocab_size, (batch_size, seqlen), dtype=torch.long, device=device ) model_hf = AutoModelForCausalLM.from_pretrained( model_name, torch_dtype=dtype, device_map={"": device}, trust_remote_code=True ) model_hf.eval() print("HF fp16") torch.cuda.synchronize() start = time.time() out_hf = model_hf.generate( input_ids=input_ids, max_length=max_length, return_dict_in_generate=True, output_scores=True ) torch.cuda.synchronize() print(f"Prompt processing + decoding time: {(time.time() - start) * 1000:.0f}ms") del model_hf model_ref = AutoModelForCausalLM.from_pretrained( model_name, device_map={"": device}, trust_remote_code=True ) model_ref.eval() with torch.no_grad(): logits_ref = model_ref(out_hf.sequences).logits[:, (seqlen - 1) : -1] del model_ref model = GPTLMHeadModel.from_pretrained(model_name, config, device=device, dtype=dtype) model.eval() print("Without CUDA graph") torch.cuda.synchronize() start = time.time() out = model.generate( input_ids=input_ids, max_length=max_length, eos_token_id=eos_token_id, return_dict_in_generate=True, output_scores=True, enable_timing=True, teacher_outputs=out_hf.sequences, ) torch.cuda.synchronize() print(f"Prompt processing + decoding time: {(time.time() - start) * 1000:.0f}ms") # Capture graph outside the timing loop batch_size, seqlen_og = input_ids.shape model._decoding_cache = update_graph_cache(model, None, batch_size, seqlen_og, max_length) print("With CUDA graph") torch.cuda.synchronize() start = time.time() out_cg = model.generate( input_ids=input_ids, max_length=max_length, cg=True, return_dict_in_generate=True, output_scores=True, enable_timing=True, teacher_outputs=out_hf.sequences, ) torch.cuda.synchronize() print(f"Prompt processing + decoding time: {(time.time() - start) * 1000:.0f}ms") with torch.no_grad(): logits_parallel = model(out_hf.sequences).logits[:, (seqlen - 1) : -1] logits_hf = torch.stack(out_hf.scores, dim=1) logits = torch.stack(out.scores, dim=1) logits_cg = torch.stack(out_cg.scores, dim=1) del model hf_error = (logits_hf - logits_ref).abs().max().item() assert (logits_parallel - logits_ref).abs().max().item() < 2 * hf_error print(f"HF fp16 logits max diff: {hf_error}") print(f"Logits max diff: {(logits - logits_ref).abs().max().item() }") assert (logits - logits_ref).abs().max().item() < 2 * hf_error print(f"Logits CG max diff: {(logits_cg - logits_ref).abs().max().item() }") assert torch.equal(logits_cg, logits) # torchrun --no_python --nproc_per_node=4 pytest -q -s tests/models/test_falcon.py -k "falcon_parallel_generation" # We want to run this on a machine with 4 x A100 80GB or 8 x A100 40GB so we have enough # memory to run the model in fp32. @pytest.mark.parametrize("world_size", [4]) @pytest.mark.parametrize("model_name", ["tiiuae/falcon-40b"]) def test_falcon_parallel_generation(model_name, world_size): """Check that our implementation 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. """ from apex.transformer import parallel_state dtype = torch.float16 config = falcon_config_to_gpt2_config( AutoConfig.from_pretrained(model_name, trust_remote_code=True) ) config.use_flash_attn = False config.fused_bias_fc = True config.fused_mlp = False # We don't have fused MLP for "gelu" activation config.fused_dropout_add_ln = False config.residual_in_fp32 = 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() 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() torch.manual_seed(0) batch_size = 1 seqlen = 100 max_length = 150 input_ids = torch.randint( 0, config.vocab_size, (batch_size, seqlen), dtype=torch.long, device=device ) # 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) pretrained_state_dict = remap_state_dict_hf_falcon( state_dict_from_pretrained(model_name), config ) model = GPTLMHeadModel(config, process_group=process_group, device=device, dtype=dtype) model.load_state_dict(shard_state_dict_tp(pretrained_state_dict, config, world_size, rank)) model.eval() print("Without CUDA graph") out = model.generate( input_ids=input_ids, max_length=max_length, tensor_parallel=world_size, vocab_size=config.vocab_size, # teacher_outputs=out_hf.sequences, return_dict_in_generate=True, output_scores=True, enable_timing=True, ) # Capture graph outside the timing loop batch_size, seqlen_og = input_ids.shape model._decoding_cache = update_graph_cache(model, None, batch_size, seqlen_og, max_length) print("With CUDA graph") out_cg = model.generate( input_ids=input_ids, max_length=max_length, tensor_parallel=world_size, vocab_size=config.vocab_size, cg=True, # teacher_outputs=out_hf.sequences, return_dict_in_generate=True, output_scores=True, enable_timing=True, ) del model parallel_state.destroy_model_parallel() if rank == 0: model_hf = AutoModelForCausalLM.from_pretrained( model_name, torch_dtype=dtype, device_map="auto", trust_remote_code=True ) model_hf.eval() print("HF fp16") torch.cuda.synchronize() start = time.time() with torch.inference_mode(): out_hf = model_hf.generate( input_ids=input_ids, max_length=max_length, return_dict_in_generate=True, output_scores=True, ) torch.cuda.synchronize() print(f"Prompt processing + decoding time: {(time.time() - start) * 1000:.0f}ms") del model_hf model_ref = AutoModelForCausalLM.from_pretrained( model_name, device_map="auto", trust_remote_code=True ) model_ref.eval() with torch.inference_mode(): logits_ref = model_ref(out_hf.sequences).logits[:, (seqlen - 1) : -1] del model_ref logits_hf = torch.stack(out_hf.scores, dim=1) logits = torch.stack(out.scores, dim=1) logits_cg = torch.stack(out_cg.scores, dim=1) hf_error = (logits_hf - logits_ref).abs().max().item() print(f"HF fp16 logits max diff: {hf_error}") print(f"Logits max diff: {(logits - logits_ref).abs().max().item() }") assert (logits - logits_ref).abs().max().item() < 2 * hf_error print(f"Logits CG max diff: {(logits_cg - logits_ref).abs().max().item() }") assert torch.equal(logits_cg, logits)