# Copyright (c) 2023 Alibaba PAI Team. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import types import numpy as np from typing import List, Optional, Union import torch import torch.nn.functional as F import transformers from tqdm import tqdm from megatron import get_args from megatron.checkpointing import load_checkpoint from megatron.core.enums import ModelType from megatron.core import mpu from megatron.core.tensor_parallel.mappings import gather_from_tensor_model_parallel_region from megatron.core.pipeline_parallel.p2p_communication import recv_forward from megatron.core.pipeline_parallel.p2p_communication import send_forward from megatron.utils import get_ltor_masks_and_position_ids, unwrap_model from megatron.arguments import core_transformer_config_from_args from lm_eval import utils from lm_eval.api.instance import Instance from lm_eval.models.huggingface import HFLM, eval_logger from megatron_patch.training import get_model from megatron_patch.tokenizer import build_tokenizer, get_tokenizer class EvalHarnessAdaptor(HFLM): def __init__( self, pretrained: Optional[Union[str, transformers.PreTrainedModel]] = "gpt2", max_length: Optional[int] = None, batch_size: Optional[Union[int, str]] = 1, trust_remote_code: Optional[bool] = False, **kwargs, ) -> None: self.args = get_args() build_tokenizer(self.args) self.tokenizer = get_tokenizer() self.is_main = torch.distributed.get_rank() == 0 self.adaptive_seq_len = self.args.adaptive_seq_len self.model_provider = kwargs['model_provider'] super().__init__(pretrained=pretrained, batch_size=batch_size, trust_remote_code=trust_remote_code, max_length=max_length, tokenizer=self.tokenizer) def _create_model( self, pretrained: str, **kwargs, ) -> None: model_list = get_model(self.model_provider, model_type=ModelType.encoder_or_decoder, wrap_with_ddp=False) if pretrained is not None: load_checkpoint(model_list, None, None) self._model = model_list[0] def tie_weights(self): pass self._model.tie_weights = types.MethodType(tie_weights, self._model) return None def create_model_inputs(self, tokens): attention_mask, loss_mask, position_ids = get_ltor_masks_and_position_ids( tokens, self.eot_token_id, self.args.reset_position_ids, self.args.reset_attention_mask, self.args.eod_mask_loss) return (tokens, position_ids, attention_mask), (tokens, loss_mask) def _model_call(self, inps, attn_mask=None, labels=None): args = get_args() # Since the shape of the micro-batch will change # We need set the correct shapes here # So that latter pipeline stages knows which shapes to expect. # Otherwise we will deadlock. args.micro_batch_size = len(inps) args.seq_length = len(inps[0]) config = core_transformer_config_from_args(args) tensor_shape = (args.seq_length, args.micro_batch_size, args.hidden_size) input_tensor = recv_forward(tensor_shape, config) # Forward pass through the model. unwrapped_model = unwrap_model(self.model) unwrapped_model.set_input_tensor(input_tensor) output = self.model(*self.create_model_inputs(inps)[0]) send_forward(output, config) if mpu.is_pipeline_last_stage(): return gather_from_tensor_model_parallel_region(output)[..., :self.tokenizer.vocab_size] else: return None def loglikelihood_rolling(self, requests: List[Instance]) -> List[float]: # TODO: Implement caching once we've confirmed the perplexity implementation # TODO: automatic batch size detection for vectorization loglikelihoods = [] with torch.no_grad(): for string, in tqdm(requests): rolling_token_windows = list(map(utils.make_disjoint_window, utils.get_rolling_token_windows( token_list=self.tokenizer_encode(string), prefix_token=self.eot_token_id, max_seq_len=self.max_length, context_len=1, ))) rolling_token_windows = [(None,) + x for x in rolling_token_windows] # TODO: extract out this call so it only gets called once and also somehow figure out partial caching for that string_nll = self._loglikelihood_tokens(rolling_token_windows, disable_tqdm=True) # discard is_greedy string_nll = [x[0] for x in string_nll] string_nll = sum(string_nll) loglikelihoods.append(string_nll) return loglikelihoods def _loglikelihood_tokens(self, requests, disable_tqdm=False): disable_tqdm = disable_tqdm if self.is_main else True res = [] res_len = 0 # storing the result length for later self.model.eval() with torch.no_grad(): def _collate(x): toks = x[1] + x[2] return (-len(toks), tuple(toks)) reord = utils.Reorderer(requests, _collate) for chunk in utils.chunks(tqdm(reord.get_reordered(), disable=disable_tqdm), self.batch_size): inps, contlens, inplens, padding_length = [], [], [], None for _, context_enc, continuation_enc in chunk: # when too long to fit in context, truncate from the left inp = torch.tensor( (context_enc + continuation_enc)[-(self.max_length + 1):][:-1] , dtype=torch.long).to(self.device) inplen, = inp.shape cont = continuation_enc # since in _collate we make sure length is descending, the longest is always the first one. padding_length = padding_length if padding_length is not None else inplen if not self.adaptive_seq_len: padding_length = self.max_length # pad to length inp = torch.cat([ inp, # [seq] torch.zeros(padding_length - inplen, dtype=torch.long).to(inp.device) # [padding_length - seq] ], dim=0) inps.append(inp.unsqueeze(0)) contlens.append(cont) inplens.append(inplen) logits = self._model_call(torch.cat(inps, dim=0)) res_len += len(chunk) if logits is not None: multi_logits = F.log_softmax(logits, dim=-1).cpu() # [batch, seq, vocab] for (cache_key, _, _), logits, inp, inplen, cont_toks in zip(chunk, multi_logits, inps, inplens, contlens): contlen = len(cont_toks) logits = logits[inplen - contlen:inplen].unsqueeze(0) # [1, seq, vocab] greedy_tokens = logits.argmax(dim=-1) # cont_toks :: [1, seq] cont_toks = torch.tensor(cont_toks, dtype=torch.long).unsqueeze(0) max_equal = (greedy_tokens == cont_toks).all() # last_token_slice = logits[:, -1, :].squeeze(0).tolist() logits = torch.gather(logits, 2, cont_toks.unsqueeze(-1)).squeeze(-1) # [1, seq] answer = (float(logits.sum()), bool(max_equal)) # partial caching res.append(answer) if not mpu.is_pipeline_last_stage(): # @HACK: To make the eval harness happy on threads that don't have access to the results. # We just randomly generate some data. res = [(np.random.rand(), np.random.rand()>0.5) for _ in requests] return reord.get_original(res)