bumblebee-testing/tiny-random-NomicBertModel
0166
1# Copyright (c) 2022, Tri Dao.2# This BERT implementation is based on our MLPerf 2.0 and MLPerf 2.1 BERT implementation.3# https://github.com/mlcommons/training_results_v2.0/blob/main/HazyResearch/benchmarks/bert/implementations/pytorch/modeling.py4# https://github.com/mlcommons/training_results_v2.1/blob/main/Azure-HazyResearch/benchmarks/bert/implementations/ND96amsr_A100_v4/modeling.py5 6# Inspired by https://github.com/huggingface/transformers/blob/main/src/transformers/models/bert/modeling_bert.py7 8import collections9import inspect10import logging11import math12import os13import re14import warnings15from collections import OrderedDict16from functools import partial17from typing import List, Optional, Tuple, Union18 19import numpy as np20import torch21import torch.nn as nn22import torch.nn.functional as F23from einops import rearrange, repeat24from safetensors.torch import load_file as safe_load_file25from torch.nn.modules.utils import _pair26from transformers import GPT2Config, PreTrainedModel, ViTConfig, ViTModel27from transformers.models.bert.modeling_bert import (28 BaseModelOutputWithPoolingAndCrossAttentions,29 MaskedLMOutput,30 SequenceClassifierOutput,31)32from transformers.modeling_outputs import (33 BaseModelOutput,34 BaseModelOutputWithPast,35 BaseModelOutputWithPooling,36 MaskedLMOutput,37 MultipleChoiceModelOutput,38 QuestionAnsweringModelOutput,39 SequenceClassifierOutput,40 ModelOutput,41 TokenClassifierOutput,42)43from transformers.utils import SAFE_WEIGHTS_INDEX_NAME, SAFE_WEIGHTS_NAME, WEIGHTS_INDEX_NAME, WEIGHTS_NAME44from transformers.utils.hub import cached_file, get_checkpoint_shard_files45 46from .configuration_hf_nomic_bert import NomicBertConfig47logger = logging.getLogger(__name__)48 49try:50 from torch.nn.functional import scaled_dot_product_attention51except ImportError:52 logger.warning("scaled_dot_product_attention not available, using torch.matmul instead")53 scaled_dot_product_attention = None54 55try:56 from megablocks.layers import dmoe57 from megablocks.layers.arguments import Arguments58except ImportError:59 dmoe = None60else:61 dmoe_is_nomic = 'attention_mask' in inspect.signature(dmoe.dMoE.forward).parameters62 63 64 65# adapted from flash attention, added safe serialization option for hf models66def state_dict_from_pretrained(model_name, safe_serialization=False, device=None, dtype=None):67 # If not fp32, then we don't want to load directly to the GPU68 mapped_device = "cpu" if dtype not in [torch.float32, None] else device69 is_sharded = False70 load_safe = False71 resolved_archive_file = None72 73 weights_path = os.path.join(model_name, WEIGHTS_NAME)74 weights_index_path = os.path.join(model_name, WEIGHTS_INDEX_NAME)75 safe_weights_path = os.path.join(model_name, SAFE_WEIGHTS_NAME)76 safe_weights_index_path = os.path.join(model_name, SAFE_WEIGHTS_INDEX_NAME)77 78 if os.path.isfile(weights_path):79 resolved_archive_file = cached_file(model_name, WEIGHTS_NAME, _raise_exceptions_for_missing_entries=False)80 elif os.path.isfile(weights_index_path):81 resolved_archive_file = cached_file(model_name, WEIGHTS_INDEX_NAME, _raise_exceptions_for_missing_entries=False)82 is_sharded = True83 elif os.path.isfile(safe_weights_path):84 resolved_archive_file = cached_file(model_name, SAFE_WEIGHTS_NAME, _raise_exceptions_for_missing_entries=False)85 load_safe = True86 elif os.path.isfile(safe_weights_index_path):87 resolved_archive_file = cached_file(88 model_name, SAFE_WEIGHTS_INDEX_NAME, _raise_exceptions_for_missing_entries=False89 )90 is_sharded = True91 load_safe = True92 else: # Try loading from HF hub instead of from local files93 resolved_archive_file = None94 for weight_name in [WEIGHTS_NAME, SAFE_WEIGHTS_NAME, WEIGHTS_INDEX_NAME, SAFE_WEIGHTS_INDEX_NAME]:95 resolved_archive_file = cached_file(model_name, weight_name, _raise_exceptions_for_missing_entries=False)96 if resolved_archive_file is not None:97 if weight_name in [SAFE_WEIGHTS_NAME, SAFE_WEIGHTS_INDEX_NAME]:98 load_safe = True99 if weight_name in [WEIGHTS_INDEX_NAME, SAFE_WEIGHTS_INDEX_NAME]:100 is_sharded = True101 break102 103 if resolved_archive_file is None:104 raise EnvironmentError(f"Model name {model_name} was not found.")105 106 if load_safe:107 loader = partial(safe_load_file, device=mapped_device)108 else:109 loader = partial(torch.load, map_location=mapped_device)110 111 if is_sharded:112 # resolved_archive_file becomes a list of files that point to the different113 # checkpoint shards in this case.114 resolved_archive_file, sharded_metadata = get_checkpoint_shard_files(model_name, resolved_archive_file)115 state_dict = {}116 for sharded_file in resolved_archive_file:117 state_dict.update(loader(sharded_file))118 else:119 state_dict = loader(resolved_archive_file)120 # Convert dtype before moving to GPU to save memory121 if dtype is not None:122 state_dict = {k: v.to(dtype=dtype) for k, v in state_dict.items()}123 state_dict = {k: v.to(device=device) for k, v in state_dict.items()}124 return state_dict125 126 127def filter_shapes(state_dict, model):128 """129 Filters the state dict to match the current model shape.130 """131 filtered_state_dict = {}132 for key, value in state_dict.items():133 if key in model.state_dict():134 if value.shape == model.state_dict()[key].shape:135 filtered_state_dict[key] = value136 return filtered_state_dict137 138 139def remap_bert_state_dict(140 state_dict,141 config,142 remove_bert=False,143 remove_cls_weights=False,144 add_pooling_layer=False,145):146 """147 Map the state_dict of a Huggingface BERT model to be flash_attn compatible.148 """149 150 def add_bert_prefix(key):151 # prepend bert. to the key152 if key.startswith("bert.") or key.startswith("cls."):153 return key154 return f"bert.{key}"155 156 state_dict = OrderedDict((add_bert_prefix(k), v) for k, v in state_dict.items())157 158 # LayerNorm159 def key_mapping_ln_gamma_beta(key):160 key = re.sub(r"LayerNorm.gamma$", "LayerNorm.weight", key)161 key = re.sub(r"LayerNorm.beta$", "LayerNorm.bias", key)162 return key163 164 state_dict = OrderedDict((key_mapping_ln_gamma_beta(k), v) for k, v in state_dict.items())165 166 # Layers167 def key_mapping_layers(key):168 return re.sub(r"^bert.encoder.layer\.", "bert.encoder.layers.", key)169 170 state_dict = OrderedDict((key_mapping_layers(k), v) for k, v in state_dict.items())171 172 # LayerNorm173 def key_mapping_ln(key):174 key = re.sub(r"^bert.embeddings.LayerNorm.", "bert.emb_ln.", key)175 key = re.sub(176 r"^bert.encoder.layers.(\d+).attention.output.LayerNorm.(weight|bias)",177 r"bert.encoder.layers.\1.norm1.\2",178 key,179 )180 key = re.sub(181 r"^bert.encoder.layers.(\d+).output.LayerNorm.(weight|bias)",182 r"bert.encoder.layers.\1.norm2.\2",183 key,184 )185 key = re.sub(186 r"^cls.predictions.transform.LayerNorm.(weight|bias)",187 r"cls.predictions.transform.layer_norm.\1",188 key,189 )190 return key191 192 state_dict = OrderedDict((key_mapping_ln(k), v) for k, v in state_dict.items())193 194 # MLP195 def key_mapping_mlp(key):196 key = re.sub(197 r"^bert.encoder.layers.(\d+).intermediate.dense.(weight|bias)",198 r"bert.encoder.layers.\1.mlp.fc1.\2",199 key,200 )201 key = re.sub(202 r"^bert.encoder.layers.(\d+).output.dense.(weight|bias)",203 r"bert.encoder.layers.\1.mlp.fc2.\2",204 key,205 )206 return key207 208 state_dict = OrderedDict((key_mapping_mlp(k), v) for k, v in state_dict.items())209 210 # Attention211 last_layer_subset = getattr(config, "last_layer_subset", False)212 for d in range(config.num_hidden_layers):213 if f"bert.encoder.layers.{d}.attention.self.query.weight" not in state_dict:214 continue215 Wq = state_dict.pop(f"bert.encoder.layers.{d}.attention.self.query.weight")216 Wk = state_dict.pop(f"bert.encoder.layers.{d}.attention.self.key.weight")217 Wv = state_dict.pop(f"bert.encoder.layers.{d}.attention.self.value.weight")218 bq = state_dict.pop(f"bert.encoder.layers.{d}.attention.self.query.bias")219 bk = state_dict.pop(f"bert.encoder.layers.{d}.attention.self.key.bias")220 bv = state_dict.pop(f"bert.encoder.layers.{d}.attention.self.value.bias")221 if not (last_layer_subset and d == config.num_hidden_layers - 1):222 state_dict[f"bert.encoder.layers.{d}.attn.Wqkv.weight"] = torch.cat([Wq, Wk, Wv], dim=0)223 state_dict[f"bert.encoder.layers.{d}.attn.Wqkv.bias"] = torch.cat([bq, bk, bv], dim=0)224 else:225 state_dict[f"bert.encoder.layers.{d}.attn.Wq.weight"] = Wq226 state_dict[f"bert.encoder.layers.{d}.attn.Wkv.weight"] = torch.cat([Wk, Wv], dim=0)227 state_dict[f"bert.encoder.layers.{d}.attn.Wq.bias"] = bq228 state_dict[f"bert.encoder.layers.{d}.attn.Wkv.bias"] = torch.cat([bk, bv], dim=0)229 230 def key_mapping_attn(key):231 return re.sub(232 r"^bert.encoder.layers.(\d+).attention.output.dense.(weight|bias)",233 r"bert.encoder.layers.\1.attn.out_proj.\2",234 key,235 )236 237 state_dict = OrderedDict((key_mapping_attn(k), v) for k, v in state_dict.items())238 239 def key_mapping_decoder_bias(key):240 return re.sub(r"^cls.predictions.bias", "cls.predictions.decoder.bias", key)241 242 # remove nsp weights, we don't use243 state_dict.pop("cls.seq_relationship.weight", None)244 state_dict.pop("cls.seq_relationship.bias", None)245 state_dict.pop("bert.embeddings.position_ids", None)246 247 state_dict = OrderedDict((key_mapping_decoder_bias(k), v) for k, v in state_dict.items())248 249 if remove_cls_weights:250 cls_weights = [251 "cls.predictions.decoder.bias",252 "cls.predictions.transform.dense.weight",253 "cls.predictions.transform.dense.bias",254 "cls.predictions.transform.layer_norm.weight",255 "cls.predictions.transform.layer_norm.bias",256 "cls.predictions.decoder.weight",257 ]258 for weight in cls_weights:259 state_dict.pop(weight, None)260 261 # Word embedding262 pad_vocab_size_multiple = getattr(config, "pad_vocab_size_multiple", 1)263 if pad_vocab_size_multiple > 1:264 word_embeddings = state_dict["bert.embeddings.word_embeddings.weight"]265 state_dict["bert.embeddings.word_embeddings.weight"] = F.pad(266 word_embeddings, (0, 0, 0, config.vocab_size - word_embeddings.shape[0])267 )268 if not remove_cls_weights:269 decoder_weight = state_dict["cls.predictions.decoder.weight"]270 state_dict["cls.predictions.decoder.weight"] = F.pad(271 decoder_weight, (0, 0, 0, config.vocab_size - decoder_weight.shape[0])272 )273 # If the vocab was padded, we want to set the decoder bias for those padded indices to be274 # strongly negative (i.e. the decoder shouldn't predict those indices).275 # TD [2022-05-09]: I don't think it affects the MLPerf training.276 if "cls.predictions.decoder.bias" in state_dict:277 decoder_bias = state_dict["cls.predictions.decoder.bias"]278 state_dict["cls.predictions.decoder.bias"] = F.pad(279 decoder_bias, (0, config.vocab_size - decoder_bias.shape[0]), value=-100.0280 )281 282 if add_pooling_layer is False:283 pooler_weights = [284 "bert.pooler.dense.weight",285 "bert.pooler.dense.bias",286 ]287 for key in pooler_weights:288 state_dict.pop(key, None)289 290 if remove_bert:291 292 def remove_bert_prefix(key):293 key = re.sub(r"^bert.", "", key)294 return key295 296 state_dict = OrderedDict((remove_bert_prefix(k), v) for k, v in state_dict.items())297 298 return state_dict299 300 301def _trunc_normal_(tensor, mean, std, a, b):302 # Cut & paste from PyTorch official master until it's in a few official releases - RW303 # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf304 def norm_cdf(x):305 # Computes standard normal cumulative distribution function306 return (1.0 + math.erf(x / math.sqrt(2.0))) / 2.0307 308 if (mean < a - 2 * std) or (mean > b + 2 * std):309 print(310 "mean is more than 2 std from [a, b] in nn.init.trunc_normal_. "311 "The distribution of values may be incorrect.",312 stacklevel=2,313 )314 315 # Values are generated by using a truncated uniform distribution and316 # then using the inverse CDF for the normal distribution.317 # Get upper and lower cdf values318 l = norm_cdf((a - mean) / std)319 u = norm_cdf((b - mean) / std)320 321 # Uniformly fill tensor with values from [l, u], then translate to322 # [2l-1, 2u-1].323 tensor.uniform_(2 * l - 1, 2 * u - 1)324 325 # Use inverse cdf transform for normal distribution to get truncated326 # standard normal327 tensor.erfinv_()328 329 # Transform to proper mean, std330 tensor.mul_(std * math.sqrt(2.0))331 tensor.add_(mean)332 333 # Clamp to ensure it's in the proper range334 tensor.clamp_(min=a, max=b)335 return tensor336 337 338def trunc_normal_tf_(tensor, mean=0.0, std=1.0, a=-2.0, b=2.0):339 r"""Fills the input Tensor with values drawn from a truncated340 normal distribution. The values are effectively drawn from the341 normal distribution :math:`\mathcal{N}(\text{mean}, \text{std}^2)`342 with values outside :math:`[a, b]` redrawn until they are within343 the bounds. The method used for generating the random values works344 best when :math:`a \leq \text{mean} \leq b`.345 346 NOTE: this 'tf' variant behaves closer to Tensorflow / JAX impl where the347 bounds [a, b] are applied when sampling the normal distribution with mean=0, std=1.0348 and the result is subsquently scaled and shifted by the mean and std args.349 350 Args:351 tensor: an n-dimensional `torch.Tensor`352 mean: the mean of the normal distribution353 std: the standard deviation of the normal distribution354 a: the minimum cutoff value355 b: the maximum cutoff value356 Examples:357 >>> w = torch.empty(3, 5)358 >>> nn.init.trunc_normal_(w)359 """360 with torch.no_grad():361 _trunc_normal_(tensor, 0, 1.0, a, b)362 tensor.mul_(std).add_(mean)363 return tensor364 365 366class NomicBertPreTrainedModel(PreTrainedModel):367 """An abstract class to handle weights initialization and368 a simple interface for dowloading and loading pretrained models.369 """370 371 config_class = NomicBertConfig372 base_model_prefix = "model"373 supports_gradient_checkpointing = True374 _no_split_modules = ["Block"]375 _skip_keys_device_placement = "past_key_values"376 377 def __init__(self, config, *inputs, **kwargs):378 super().__init__(config)379 if not isinstance(config, GPT2Config):380 raise ValueError(381 "Parameter config in `{}(config)` should be an instance of class `GPT2Config`. "382 "To create a model from a Google pretrained model use "383 "`model = {}.from_pretrained(PRETRAINED_MODEL_NAME)`".format(384 self.__class__.__name__, self.__class__.__name__385 )386 )387 self.config = config388 389 @classmethod390 def from_pretrained(cls, model_name, config=None, *inputs, **kwargs):391 """392 Instantiate a NomicBertPreTrainedModel from a pre-trained model file or a pytorch state dict.393 Download and cache the pre-trained model file if needed.394 395 Params:396 pretrained_model_name_or_path: either:397 - a path or url to a pretrained model archive containing:398 . `bert_config.json` a configuration file for the model399 . `pytorch_model.bin` a PyTorch dump of a NomicBertForPretraining instance400 - a path or url to a pretrained model archive containing:401 . `bert_config.json` a configuration file for the model402 . `model.chkpt` a TensorFlow checkpoint403 *inputs, **kwargs: additional input for the specific NomicBert class404 (ex: num_labels for NomicBertForSequenceClassification)405 """406 # Instantiate model.407 if config is None:408 config = cls.config_class.from_pretrained(model_name)409 remove_cls = cls != NomicBertForPreTraining410 remove_bert_prefix = cls not in [NomicBertForPreTraining, NomicBertForSequenceClassification, NomicBertForTokenClassification, NomicBertForMultipleChoice, NomicBertForQuestionAnswering]411 ignore_mismatched_shapes = kwargs.pop("ignore_mismatched_sizes", False)412 num_labels = kwargs.pop("num_labels", None)413 rotary_scaling_factor = kwargs.pop("rotary_scaling_factor", None)414 strict = kwargs.pop("strict", True)415 dtype = kwargs.pop("torch_dtype", None)416 if rotary_scaling_factor:417 config.rotary_scaling_factor = rotary_scaling_factor418 419 if config.n_positions <= 0 and config.rotary_emb_fraction > 0:420 config.n_positions = 2048421 if num_labels:422 config.num_labels = num_labels423 424 if "add_pooling_layer" in kwargs:425 model = cls(config, *inputs, add_pooling_layer=kwargs.pop("add_pooling_layer"))426 else:427 if cls == NomicBertModel:428 model = cls(config, *inputs, add_pooling_layer=False)429 else:430 model = cls(config, *inputs)431 432 if dtype is not None:433 model = model.to(dtype=dtype)434 # TODO: fix this435 # Assuming we know what we're doing when loading from disk436 # Prob a bad assumption but i'm tired and want to train this asap437 if os.path.exists(model_name):438 model_path = f"{model_name}/pytorch_model.bin"439 if os.path.exists(model_path):440 state_dict = torch.load(f"{model_name}/pytorch_model.bin")441 else:442 model_path = f"{model_name}/model.safetensors"443 if not os.path.exists(model_path):444 raise ValueError(f"Model path {model_path} not found")445 state_dict = safe_load_file(model_path)446 447 if ignore_mismatched_shapes:448 state_dict = filter_shapes(state_dict, model)449 load_return = model.load_state_dict(state_dict, strict=False)450 else:451 # TODO: can probably check config class and see if we need to remap from a bert model452 state_dict = state_dict_from_pretrained(model_name, dtype=dtype)453 state_dict = remap_bert_state_dict(454 state_dict,455 config,456 remove_bert=remove_bert_prefix,457 remove_cls_weights=remove_cls,458 add_pooling_layer=getattr(config, "add_pooling_layer", False),459 )460 if ignore_mismatched_shapes:461 state_dict = filter_shapes(state_dict, model)462 463 load_return = model.load_state_dict(state_dict, strict=strict)464 # moe models load new weights 465 if getattr(config, "moe_top_k", 0) == 0:466 logger.warning(load_return)467 return model468 469 def _set_gradient_checkpointing(self, module, value=False):470 if isinstance(module, NomicBertEncoder):471 module.gradient_checkpointing = value472 473 474# https://github.com/huggingface/transformers/blob/7032e0203262ebb2ebf55da8d2e01f873973e835/src/transformers/models/bert/modeling_bert.py#L748475def _init_weights(module, initializer_range=0.02):476 if isinstance(module, nn.Linear):477 nn.init.normal_(module.weight, std=initializer_range)478 if module.bias is not None:479 nn.init.zeros_(module.bias)480 elif isinstance(module, nn.Embedding):481 nn.init.normal_(module.weight, std=initializer_range)482 if module.padding_idx is not None:483 nn.init.zeros_(module.weight[module.padding_idx])484 485 486def _ntuple(n):487 def parse(x):488 if isinstance(x, collections.abc.Iterable) and not isinstance(x, str):489 return tuple(x)490 return tuple(repeat(x, n))491 492 return parse493 494 495to_1tuple = _ntuple(1)496to_2tuple = _ntuple(2)497to_3tuple = _ntuple(3)498to_4tuple = _ntuple(4)499to_ntuple = _ntuple500 501 502def get_2d_sincos_pos_embed(embed_dim, grid_size, add_cls_token=False):503 """504 Create 2D sin/cos positional embeddings.505 506 Args:507 embed_dim (`int`):508 Embedding dimension.509 grid_size (`int`):510 The grid height and width.511 add_cls_token (`bool`, *optional*, defaults to `False`):512 Whether or not to add a classification (CLS) token.513 514 Returns:515 (`torch.FloatTensor` of shape (grid_size*grid_size, embed_dim) or (1+grid_size*grid_size, embed_dim): the516 position embeddings (with or without classification token)517 """518 grid_h = np.arange(grid_size, dtype=np.float32)519 520 grid_w = np.arange(grid_size, dtype=np.float32)521 grid = np.meshgrid(grid_w, grid_h) # here w goes first522 grid = np.stack(grid, axis=0)523 524 grid = grid.reshape([2, 1, grid_size, grid_size])525 pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)526 if add_cls_token:527 pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0)528 return pos_embed529 530 531def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):532 if embed_dim % 2 != 0:533 raise ValueError("embed_dim must be even")534 535 # use half of dimensions to encode grid_h536 emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0]) # (H*W, D/2)537 emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1]) # (H*W, D/2)538 539 emb = np.concatenate([emb_h, emb_w], axis=1) # (H*W, D)540 return emb541 542 543def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):544 """545 embed_dim: output dimension for each position pos: a list of positions to be encoded: size (M,) out: (M, D)546 """547 if embed_dim % 2 != 0:548 raise ValueError("embed_dim must be even")549 550 omega = np.arange(embed_dim // 2, dtype=float)551 omega /= embed_dim / 2.0552 omega = 1.0 / 10000**omega # (D/2,)553 554 pos = pos.reshape(-1) # (M,)555 out = np.einsum("m,d->md", pos, omega) # (M, D/2), outer product556 557 emb_sin = np.sin(out) # (M, D/2)558 emb_cos = np.cos(out) # (M, D/2)559 560 emb = np.concatenate([emb_sin, emb_cos], axis=1) # (M, D)561 return emb562 563 564def ndgrid(*tensors) -> Tuple[torch.Tensor, ...]:565 """generate N-D grid in dimension order.566 567 The ndgrid function is like meshgrid except that the order of the first two input arguments are switched.568 569 That is, the statement570 [X1,X2,X3] = ndgrid(x1,x2,x3)571 572 produces the same result as573 574 [X2,X1,X3] = meshgrid(x2,x1,x3)575 576 This naming is based on MATLAB, the purpose is to avoid confusion due to torch's change to make577 torch.meshgrid behaviour move from matching ndgrid ('ij') indexing to numpy meshgrid defaults of ('xy').578 579 """580 try:581 return torch.meshgrid(*tensors, indexing='ij')582 except TypeError:583 # old PyTorch < 1.10 will follow this path as it does not have indexing arg,584 # the old behaviour of meshgrid was 'ij'585 return torch.meshgrid(*tensors)586 587 588def build_fourier_pos_embed(589 feat_shape: List[int],590 bands: Optional[torch.Tensor] = None,591 num_bands: int = 64,592 max_res: int = 224,593 temperature: float = 10000.0,594 linear_bands: bool = False,595 include_grid: bool = False,596 in_pixels: bool = True,597 ref_feat_shape: Optional[List[int]] = None,598 dtype: torch.dtype = torch.float32,599 device: Optional[torch.device] = None,600) -> List[torch.Tensor]:601 """602 603 Args:604 feat_shape: Feature shape for embedding.605 bands: Pre-calculated frequency bands.606 num_bands: Number of frequency bands (determines output dim).607 max_res: Maximum resolution for pixel based freq.608 temperature: Temperature for non-pixel freq.609 linear_bands: Linear band spacing for pixel based freq.610 include_grid: Include the spatial grid in output.611 in_pixels: Output in pixel freq.612 ref_feat_shape: Reference feature shape for resize / fine-tune.613 dtype: Output dtype.614 device: Output device.615 616 Returns:617 618 """619 if bands is None:620 if in_pixels:621 bands = pixel_freq_bands(622 num_bands,623 float(max_res),624 linear_bands=linear_bands,625 device=device,626 )627 else:628 bands = freq_bands(629 num_bands,630 temperature=temperature,631 step=1,632 device=device,633 )634 else:635 if device is None:636 device = bands.device637 if dtype is None:638 dtype = bands.dtype639 640 if in_pixels:641 t = [torch.linspace(-1.0, 1.0, steps=s, device=device, dtype=torch.float32) for s in feat_shape]642 else:643 t = [torch.arange(s, device=device, dtype=torch.int64).to(torch.float32) for s in feat_shape]644 645 if ref_feat_shape is not None:646 # eva's scheme for resizing rope embeddings (ref shape = pretrain)647 t = [x / f * r for x, f, r in zip(t, feat_shape, ref_feat_shape)]648 649 grid = torch.stack(ndgrid(t), dim=-1)650 grid = grid.unsqueeze(-1)651 pos = grid * bands652 653 pos_sin, pos_cos = pos.sin().to(dtype=dtype), pos.cos().to(dtype)654 out = [grid, pos_sin, pos_cos] if include_grid else [pos_sin, pos_cos]655 return out656 657 658def build_rotary_pos_embed(659 feat_shape: List[int],660 bands: Optional[torch.Tensor] = None,661 dim: int = 64,662 max_res: int = 224,663 temperature: float = 10000.0,664 linear_bands: bool = False,665 in_pixels: bool = True,666 ref_feat_shape: Optional[List[int]] = None,667 dtype: torch.dtype = torch.float32,668 device: Optional[torch.device] = None,669):670 """671 672 Args:673 feat_shape: Spatial shape of the target tensor for embedding.674 bands: Optional pre-generated frequency bands675 dim: Output dimension of embedding tensor.676 max_res: Maximum resolution for pixel mode.677 temperature: Temperature (inv freq) for non-pixel mode678 linear_bands: Linearly (instead of log) spaced bands for pixel mode679 in_pixels: Pixel vs language (inv freq) mode.680 dtype: Output dtype.681 device: Output device.682 683 Returns:684 685 """686 sin_emb, cos_emb = build_fourier_pos_embed(687 feat_shape,688 bands=bands,689 num_bands=dim // 4,690 max_res=max_res,691 temperature=temperature,692 linear_bands=linear_bands,693 in_pixels=in_pixels,694 ref_feat_shape=ref_feat_shape,695 device=device,696 dtype=dtype,697 )698 num_spatial_dim = 1699 # this would be much nicer as a .numel() call to torch.Size(), but torchscript sucks700 for x in feat_shape:701 num_spatial_dim *= x702 sin_emb = sin_emb.reshape(num_spatial_dim, -1).repeat_interleave(2, -1)703 cos_emb = cos_emb.reshape(num_spatial_dim, -1).repeat_interleave(2, -1)704 return sin_emb, cos_emb705 706 707def freq_bands(708 num_bands: int,709 temperature: float = 10000.0,710 step: int = 2,711 device: Optional[torch.device] = None,712) -> torch.Tensor:713 exp = torch.arange(0, num_bands, step, dtype=torch.int64, device=device).to(torch.float32) / num_bands714 bands = 1.0 / (temperature**exp)715 return bands716 717 718def pixel_freq_bands(719 num_bands: int,720 max_freq: float = 224.0,721 linear_bands: bool = True,722 device: Optional[torch.device] = None,723):724 if linear_bands:725 bands = torch.linspace(1.0, max_freq / 2, num_bands, dtype=torch.float32, device=device)726 else:727 bands = 2 ** torch.linspace(0, math.log(max_freq, 2) - 1, num_bands, dtype=torch.float32, device=device)728 return bands * torch.pi729 730 731def rot(x):732 return torch.stack([-x[..., 1::2], x[..., ::2]], -1).reshape(x.shape)733 734 735def apply_rot_embed_cat(x: torch.Tensor, emb):736 sin_emb, cos_emb = emb.tensor_split(2, -1)737 if sin_emb.ndim == 3:738 return x * cos_emb.unsqueeze(1).expand_as(x) + rot(x) * sin_emb.unsqueeze(1).expand_as(x)739 return x * cos_emb + rot(x) * sin_emb740 741 742# taken from https://github.com/huggingface/pytorch-image-models/blob/cb0e4391beedcc5ac3ae4bce16561b95c326f32c/timm/layers/pos_embed_sincos.py#L363743class NomicVisionRotaryEmbeddingCat(nn.Module):744 """Rotary position embedding w/ concatenatd sin & cos745 746 The following impl/resources were referenced for this impl:747 * https://github.com/lucidrains/vit-pytorch/blob/6f3a5fcf0bca1c5ec33a35ef48d97213709df4ba/vit_pytorch/rvt.py748 * https://blog.eleuther.ai/rotary-embeddings/749 """750 751 def __init__(752 self,753 dim,754 max_res=224,755 temperature=10000,756 in_pixels=True,757 linear_bands: bool = False,758 feat_shape: Optional[List[int]] = None,759 ref_feat_shape: Optional[List[int]] = None,760 ):761 super().__init__()762 self.dim = dim763 self.max_res = max_res764 self.temperature = temperature765 self.in_pixels = in_pixels766 self.feat_shape = feat_shape767 self.ref_feat_shape = ref_feat_shape768 769 if feat_shape is None:770 # only cache bands771 if in_pixels:772 bands = pixel_freq_bands(773 dim // 4,774 float(max_res),775 linear_bands=linear_bands,776 )777 else:778 bands = freq_bands(779 dim // 4,780 temperature=temperature,781 step=1,782 )783 self.register_buffer(784 'bands',785 bands,786 persistent=False,787 )788 self.pos_embed = None789 else:790 # cache full sin/cos embeddings if shape provided up front791 embeds = build_rotary_pos_embed(792 feat_shape=feat_shape,793 dim=dim,794 max_res=max_res,795 linear_bands=linear_bands,796 in_pixels=in_pixels,797 ref_feat_shape=self.ref_feat_shape,798 )799 self.bands = None800 self.register_buffer(801 'pos_embed',802 torch.cat(embeds, -1),803 persistent=False,804 )805 806 def get_embed(self, shape: Optional[List[int]] = None):807 if self.bands is not None and shape is not None:808 # rebuild embeddings every call, use if target shape changes809 embeds = build_rotary_pos_embed(810 shape,811 self.bands,812 in_pixels=self.in_pixels,813 ref_feat_shape=self.ref_feat_shape,814 )815 return torch.cat(embeds, -1)816 elif self.pos_embed is not None:817 return self.pos_embed818 else:819 assert False, "get_embed() requires pre-computed pos_embed or valid shape w/ pre-computed bands"820 821 def forward(self, x):822 # assuming channel-first tensor where spatial dim are >= 2823 pos_embed = self.get_embed(x.shape[2:])824 return apply_rot_embed_cat(x, pos_embed)825 826 827class NomicVisionPatchEmbeddings(nn.Module):828 def __init__(829 self,830 config,831 ):832 super().__init__()833 img_size = _pair(config.img_size)834 patch_size = _pair(config.patch_size)835 self.img_size = img_size836 self.patch_size = patch_size837 self.grid_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1])838 self.num_patches = self.grid_size[0] * self.grid_size[1]839 840 self.proj = nn.Linear(841 config.num_channels * patch_size[0] * patch_size[1], config.n_embd, bias=config.patch_embed_bias842 )843 844 self.learned_pos_embedding = False845 self.sinusoidal_pos_embedding = False846 self.no_embed_class = getattr(config, "no_embed_class", False)847 848 self.cls_token = (849 nn.Parameter(torch.zeros(1, 1, config.n_embd)) if not getattr(config, "no_cls_token", False) else None850 )851 if config.learned_pos_embedding:852 # this is the default in DINO853 self.learned_pos_embedding = True854 # hack for timm dinov2 with registers855 num_patches = self.num_patches if getattr(config, "register_tokens", 0) > 0 else self.num_patches + 1856 self.pos_embed = (857 nn.Parameter(torch.randn(1, num_patches, config.n_embd) * 0.02)858 if getattr(config, "use_pos_embed", True)859 else None860 )861 elif getattr(config, "sinusoidal_pos_embedding", False):862 self.sinusoidal_pos_embedding = True863 if getattr(config, "use_pos_embed", True):864 self.pos_embed = nn.Parameter(torch.zeros(1, self.num_patches + 1, config.n_embd), requires_grad=False)865 pos_embed = get_2d_sincos_pos_embed(config.n_embd, self.grid_size[0], add_cls_token=True)866 self.pos_embed.data.copy_(torch.from_numpy(pos_embed).to(self.pos_embed))867 else:868 self.pos_embed = None869 else:870 self.pos_embed = (871 nn.Parameter(torch.randn(1, self.num_patches + 1, config.n_embd) * 0.02)872 if getattr(config, "use_pos_embed", True)873 else None874 )875 876 if getattr(config, "register_tokens", 0) > 0:877 self.reg_token = nn.Parameter(torch.randn(1, config.register_tokens, config.n_embd) * 0.02)878 else:879 self.reg_token = None880 881 if config.mask_token:882 self.mask_token = nn.Parameter(torch.zeros(1, config.n_embd))883 884 self.patch_dropout = nn.Identity()885 886 if getattr(config, "use_rotary_pos_emb", False):887 ref_feat_shape = getattr(config, "ref_feat_shape", None)888 ref_feat_shape = to_2tuple(ref_feat_shape) if ref_feat_shape is not None else None889 self.rope = NomicVisionRotaryEmbeddingCat(890 config.n_embd // config.n_head,891 in_pixels=False,892 feat_shape=self.grid_size,893 ref_feat_shape=ref_feat_shape,894 )895 else:896 self.rope = None897 898 def interpolate_pos_encoding(self, embeddings: torch.Tensor, height: int, width: int) -> torch.Tensor:899 """900 This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher901 resolution images.902 903 Source:904 https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py#L174905 """906 num_patches = embeddings.shape[1] - 1907 num_positions = self.pos_embed.shape[1] - 1908 if num_patches == num_positions and height == width:909 return self.pos_embed910 class_pos_embed = self.pos_embed[:, 0]911 patch_pos_embed = self.pos_embed[:, 1:]912 dim = embeddings.shape[-1]913 height = height // self.patch_size[0]914 width = width // self.patch_size[1]915 # we add a small number to avoid floating point error in the interpolation916 # see discussion at https://github.com/facebookresearch/dino/issues/8917 height, width = height + 0.1, width + 0.1918 patch_pos_embed = patch_pos_embed.reshape(1, int(math.sqrt(num_positions)), int(math.sqrt(num_positions)), dim)919 patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2)920 patch_pos_embed = nn.functional.interpolate(921 patch_pos_embed,922 scale_factor=(height / math.sqrt(num_positions), width / math.sqrt(num_positions)),923 mode="bicubic",924 align_corners=False,925 )926 if int(height) != patch_pos_embed.shape[-2] or int(width) != patch_pos_embed.shape[-1]:927 raise ValueError("Width or height does not match with the interpolated position embeddings")928 patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)929 return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1)930 931 def forward(self, x):932 # deepspeed case where the input is in fp32933 if x.dtype != self.proj.weight.dtype:934 x = x.to(dtype=self.proj.weight.dtype)935 936 _, _, height, width = x.shape937 x = self.proj(938 rearrange(939 x,940 "b c (h p1) (w p2) -> b h w (c p1 p2)",941 p1=self.patch_size[0],942 p2=self.patch_size[1],943 )944 )945 embeddings = rearrange(x, "b h w c -> b (h w) c")946 947 to_cat = []948 if self.cls_token is not None:949 if self.sinusoidal_pos_embedding:950 cls_token = self.cls_token + self.pos_embed[:, 0]951 cls_token = cls_token.expand(embeddings.shape[0], -1, -1)952 to_cat += [cls_token]953 else:954 cls_token = self.cls_token.expand(embeddings.shape[0], 1, -1)955 to_cat += [cls_token]956 957 if self.reg_token is not None:958 to_cat += [self.reg_token.expand(embeddings.shape[0], -1, -1)]959 960 rot_pos_embed = self.rope.get_embed() if self.rope is not None else None961 962 if self.no_embed_class:963 if self.learned_pos_embedding:964 embeddings = embeddings + self.interpolate_pos_encoding(embeddings, height, width)965 else:966 if self.pos_embed is not None:967 embeddings = embeddings + self.pos_embed968 if to_cat:969 embeddings = torch.cat(to_cat + [embeddings], dim=1)970 else:971 if to_cat:972 embeddings = torch.cat(to_cat + [embeddings], dim=1)973 if self.learned_pos_embedding:974 if self.pos_embed is not None:975 embeddings = embeddings + self.interpolate_pos_encoding(embeddings, height, width)976 else:977 if self.pos_embed is not None:978 embeddings = embeddings + self.pos_embed979 980 embeddings = self.patch_dropout(embeddings)981 982 return embeddings, rot_pos_embed983 984 985class NomicBertEmbeddings(nn.Module):986 def __init__(self, config):987 """988 If max_position_embeddings <= 0, there's no position embeddings989 If type_vocab_size <= 0, there's no token type embeddings990 """991 super().__init__()992 self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)993 self.max_position_embeddings = config.max_position_embeddings if config.rotary_emb_fraction <= 0 else 0994 self.type_vocab_size = config.type_vocab_size995 if self.max_position_embeddings > 0 and config.rotary_emb_fraction <= 0:996 self.position_embeddings = nn.Embedding(997 config.max_position_embeddings,998 config.hidden_size,999 )1000 if self.type_vocab_size > 0:1001 self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)1002 1003 def forward(self, input_ids=None, position_ids=None, token_type_ids=None, inputs_embeds=None):1004 """1005 input_ids: (batch, seqlen)1006 position_ids: (batch, seqlen)1007 token_type_ids: (batch, seqlen)1008 """1009 if inputs_embeds is None:1010 embeddings = self.word_embeddings(input_ids)1011 else:1012 embeddings = inputs_embeds1013 batch_size, seqlen, _ = embeddings.shape1014 1015 if self.type_vocab_size > 0:1016 if token_type_ids is None:1017 token_type_ids = torch.zeros(seqlen, dtype=torch.long, device=embeddings.device)1018 token_type_embeddings = self.token_type_embeddings(token_type_ids)1019 embeddings = embeddings + token_type_embeddings1020 1021 if self.max_position_embeddings > 0:1022 if position_ids is None:1023 position_ids = torch.arange(seqlen, dtype=torch.long, device=embeddings.device)1024 position_embeddings = self.position_embeddings(position_ids)1025 embeddings = embeddings + position_embeddings1026 return embeddings1027 1028 1029class NomicBertMLP(nn.Module):1030 def __init__(1031 self,1032 in_features,1033 hidden_features=None,1034 out_features=None,1035 activation=F.gelu,1036 bias1=True,1037 bias2=True,1038 return_residual=False,1039 fused_bias_fc=False,1040 ):1041 super().__init__()1042 out_features = out_features if out_features is not None else in_features1043 hidden_features = hidden_features if hidden_features is not None else in_features * 41044 self.return_residual = return_residual1045 self.fc1 = nn.Linear(in_features, hidden_features, bias=bias1)1046 approximate = "tanh" if activation in ["gelu_new", "gelu_fast", "gelu_pytorch_tanh"] else "none"1047 self.activation = nn.GELU(approximate=approximate) if activation == "gelu" else activation1048 self.fc2 = nn.Linear(hidden_features, out_features, bias=bias2)1049 1050 def forward(self, x):1051 y = self.fc1(x)1052 y = self.activation(y)1053 y = self.fc2(y)1054 return y if not self.return_residual else (y, x)1055 1056 1057class NomciBertGatedMLP(nn.Module):1058 def __init__(1059 self,1060 in_features,1061 hidden_features=None,1062 out_features=None,1063 activation=F.sigmoid,1064 bias1=True,1065 bias2=True,1066 multiple_of=256,1067 return_residual=False,1068 fused_bias_fc=True,1069 device=None,1070 dtype=None,1071 norm_layer=False,1072 ):1073 super().__init__()1074 out_features = out_features if out_features is not None else in_features1075 hidden_features = hidden_features if hidden_features is not None else int(8 * in_features / 3)1076 hidden_features = int((hidden_features + multiple_of - 1) // multiple_of * multiple_of)1077 self.return_residual = return_residual1078 1079 self.fc11 = nn.Linear(in_features, hidden_features, bias=bias1)1080 self.fc12 = nn.Linear(in_features, hidden_features, bias=bias1)1081 self.activation = activation1082 self.fc2 = nn.Linear(hidden_features, out_features, bias=bias2)1083 self.norm = nn.LayerNorm(hidden_features) if norm_layer else nn.Identity()1084 1085 def forward(self, x):1086 y = self.fc11(x)1087 gate = self.fc12(x)1088 if self.activation == F.sigmoid: # Special case for GLU1089 y = F.glu(torch.cat([y, gate], dim=-1), dim=-1)1090 else:1091 y = y * self.activation(gate)1092 1093 # eva uses layer norm after the activation1094 y = self.norm(y)1095 1096 y = self.fc2(y)1097 return y if not self.return_residual else (y, x)1098 1099class NomicRouter(nn.Module):1100 def __init__(self, hidden_size: int, moe_num_experts: int, moe_top_k: int,1101 moe_jitter_eps: Optional[float] = None,1102 moe_normalize_expert_weights: Optional[float] = None,1103 uniform_expert_assignment: bool = False):1104 super().__init__()1105 self.hidden_size = hidden_size1106 self.moe_num_experts = moe_num_experts1107 self.moe_top_k = moe_top_k1108 self.moe_jitter_eps = moe_jitter_eps1109 self.moe_normalize_expert_weights = moe_normalize_expert_weights1110 self.uniform_expert_assignment = uniform_expert_assignment1111 1112 self.layer = nn.Linear(self.hidden_size,1113 self.moe_num_experts,1114 bias=False)1115 1116 def jitter(self, x: torch.Tensor) -> torch.Tensor:1117 if self.moe_jitter_eps is None:1118 raise RuntimeError('The router does not have moe_jitter_eps set.')1119 low = 1.0 - self.moe_jitter_eps1120 high = 1.0 + self.moe_jitter_eps1121 noise = torch.rand(x.size(), dtype=x.dtype, device=x.device)1122 return low + noise * (high - low)1123 1124 def forward(1125 self, x: torch.Tensor1126 ) -> Tuple[torch.Tensor, torch.Tensor, torch.LongTensor]:1127 if self.training and self.moe_jitter_eps is not None:1128 x = x * self.jitter(x)1129 1130 weights = self.layer(x.view(-1,1131 x.shape[-1])).softmax(dim=-1,1132 dtype=torch.float32)1133 top_weights, top_experts = torch.topk(weights, self.moe_top_k, dim=-1)1134 1135 if self.moe_normalize_expert_weights:1136 top_weights = top_weights / torch.norm(1137 top_weights,1138 p=self.moe_normalize_expert_weights,1139 dim=-1,1140 keepdim=True)1141 1142 if self.uniform_expert_assignment:1143 with torch.no_grad():1144 uniform_tensor = torch.arange(1145 0,1146 top_experts.numel(),1147 device=top_experts.device,1148 dtype=top_experts.dtype) % self.moe_num_experts1149 top_experts = uniform_tensor.reshape(top_experts.shape)1150 # Note, weights and top_weights are not changed1151 1152 weights = weights.to(x.dtype)1153 top_weights = top_weights.to(x.dtype)1154 return weights, top_weights, top_experts # type: ignore1155 1156 1157class NomicExpertMLP(nn.Module):1158 1159 def __init__(self, hidden_size: int, ffn_hidden_size: int,1160 moe_num_experts: int, ffn_act_fn: dict):1161 super().__init__()1162 self.hidden_size = hidden_size1163 self.ffn_hidden_size = ffn_hidden_size1164 self.moe_num_experts = moe_num_experts1165 1166 self.w1 = nn.Parameter(1167 torch.empty(moe_num_experts * ffn_hidden_size, hidden_size))1168 self.w2 = nn.Parameter(1169 torch.empty(moe_num_experts * ffn_hidden_size, hidden_size))1170 self.activation_fn = ffn_act_fn1171 1172 def forward(self, x: torch.Tensor, expert_idx: int) -> torch.Tensor:1173 expert_w1 = self.w1.view(self.moe_num_experts, self.ffn_hidden_size,1174 self.hidden_size)[expert_idx]1175 expert_w2 = self.w2.view(self.moe_num_experts, self.ffn_hidden_size,1176 self.hidden_size)[expert_idx]1177 1178 x1 = x.matmul(expert_w1.t())1179 act_out = self.activation_fn(x1)1180 x2 = act_out.matmul(expert_w2)1181 return x21182 1183class NomicExperts(nn.Module):1184 def __init__(self, config, hidden_size: int, ffn_hidden_size: int,1185 moe_num_experts: int):1186 super().__init__()1187 self.moe_num_experts = moe_num_experts1188 activation = (1189 F.sigmoid1190 if config.activation_function == "glu"1191 else (F.silu if config.activation_function == "swiglu" else F.gelu)1192 )1193 self.mlp = NomicExpertMLP(1194 hidden_size=config.n_embd,1195 ffn_hidden_size=config.n_inner,1196 moe_num_experts=moe_num_experts,1197 ffn_act_fn=activation,1198 )1199 self.bias = nn.Parameter(torch.zeros(config.n_embd))1200