| import logging |
| from typing import Union, List, Optional, Dict, Any, Literal |
| import torch |
| import torch.nn.functional as F |
| from transformers import AutoTokenizer |
| from transformers_neuronx import MistralForSampling, GQA, NeuronConfig |
| import time |
| import math |
|
|
|
|
| model_name = './checkpoint-3000' |
| amp = 'bf16' |
| batch_size = 1 |
| tp_degree = 8 |
| n_positions = 8192 |
| neuron_config = NeuronConfig(group_query_attention=GQA.SHARD_OVER_HEADS) |
|
|
|
|
| model = MistralForSampling.from_pretrained( |
| model_name, |
| amp=amp, |
| batch_size=batch_size, |
| tp_degree=tp_degree, |
| n_positions=n_positions, |
| neuron_config=neuron_config |
| ) |
| model.to_neuron() |
