| 58 | return torch.sum(loss).item(), num_tokens |
| 59 | |
| 60 | def load_model_tokenizer_config(): |
| 61 | from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig |
| 62 | from transformers.generation import GenerationConfig |
| 63 | config = AutoConfig.from_pretrained(args.model_path, trust_remote_code=True) |
| 64 | tokenizer = AutoTokenizer.from_pretrained(args.model_path, use_fast=False, trust_remote_code=True) |
| 65 | model = AutoModelForCausalLM.from_pretrained(args.model_path, device_map="auto", config=config, trust_remote_code=True).eval() |
| 66 | while config.num_attention_heads % args.n_gpus != 0: |
| 67 | args.n_gpus //= 2 |
| 68 | args.batch_size //= 2 |
| 69 | return model, tokenizer, config |
| 70 | |
| 71 | def main(): |
| 72 | |