File size: 4,669 Bytes
b3a3b15 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 | from transformers.models.opt import OPTConfig
from transformers import AutoTokenizer
from flash_attn.models.opt import opt_config_to_gpt2_config
import os
import torch
import argparse
from apex.transformer import parallel_state
from HybridTensor.utils.utils import arg_parser, _get_device
from HybridTensor.utils.activations import OPT_MODELS
from HybridTensor.models.opt import SparseConfig, build_sparse_opt
def update_router_config(model, num_layers, mlp_act_th, attn_topk, layer_config = None):
for i in range(num_layers):
model.transformer.layers[i].mlp_router.act_th = mlp_act_th
model.transformer.layers[i].mha_router.topk = attn_topk
def initialize_distributed_environment():
# Set environment variables for NCCL
os.environ["TORCH_NCCL_ASYNC_ERROR_HANDLING"] = "0"
os.environ["NCCL_GRAPH_MIXING_SUPPORT"] = "0"
# Initialize the distributed process group
torch.distributed.init_process_group(backend="nccl", init_method="env://")
# Set the device based on the rank of the current process
device = f"cuda:{torch.distributed.get_rank()}"
world_size = torch.distributed.get_world_size()
# Set the current CUDA device to avoid operations being executed on the wrong GPU
torch.cuda.set_device(device)
# You can return device, world_size, and any other relevant information
return device, world_size
def arg_parser():
parser = argparse.ArgumentParser(description='Inference benchmarking')
parser.add_argument('--batch_size', type=int, default=128)
parser.add_argument('--model_index', type=int, default=5)
parser.add_argument('--seq_len', type=int, default=28)
parser.add_argument('--index_size', type=int, default=8192)
parser.add_argument('--head_density', type=float, default=0.25)
parser.add_argument('--print_results', type=bool, default=True)
parser.add_argument('--iterations', type=int, default=100)
parser.add_argument('--check_results', type=bool, default=False)
parser.add_argument('--results_dir', type=str, default='results')
parser.add_argument('--gpu', type=int, default=0)
parser.add_argument('--mlp_ckpt_dir', type=str, default='<PATH_TO_MLP_ROUTER_CHECKPOINTS>')
parser.add_argument('--attn_topk', type=float, default=0.5, help='Attention topk for sparse model')
parser.add_argument('--attn_ckpt_dir', type=str, default='<PATH_TO_ATTENTION_CHECKPOINTS>')
return parser.parse_args()
if __name__ == "__main__":
args = arg_parser()
model_name = OPT_MODELS[args.model_index-1]
device, world_size = initialize_distributed_environment()
dtype = torch.float16
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()
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = build_sparse_opt(model_name, args.mlp_ckpt_dir, args.attn_ckpt_dir, device = device, dtype=dtype, process_group = process_group, world_size = world_size, rank = rank)
model.eval()
print("Model loaded with sparse routers")
mlp_act_th = 0.5
attn_topk = 0.5
update_router_config(model, model.config.n_layer, mlp_act_th, attn_topk)
print("Router config updated")
# print router configs from all layers
# for i in range(model.config.n_layer):
# print(f"Layer {i}: mlp_act_th = {model.transformer.layers[i].mlp_router.act_th}, attn_topk = {model.transformer.layers[i].mha_router.topk}")
input_texts = ["Hello, my dog is cute and", "The future of AI is", "In a distant galaxy, a spaceship", "The cat is sleeping on the "]
# input_texts = ["Hello, my dog is cute and", "Hello, my rat is cute and"]
tokenized_inputs = tokenizer(input_texts, return_tensors="pt", padding=True, truncation=True).to(device)
input_ids=tokenized_inputs["input_ids"]
# input_ids = tokenizer("Hello, my dog is cute and", return_tensors="pt").input_ids.to(device=device)
max_length = args.seq_len
position_ids = None
eos_token_id = tokenizer.eos_token_id
num_layers = model.config.n_layer
# print all the model weights and check the accuracy
# if rank == 0:
# print(model.state_dict())
# out = model(input_ids)
# print(out)
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,
)
if rank == 0:
print(tokenizer.batch_decode(out.sequences.tolist()))
|