File size: 1,322 Bytes
c92e699 | 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 | # SPDX-FileCopyrightText: © 2023 Tenstorrent USA, Inc.
# SPDX-License-Identifier: Apache-2.0
import os
import torch.nn.functional as F
def load_partial_weights(auto_model, weights_path, layer_prefix):
partial_state_dict = {}
model = auto_model.from_pretrained(weights_path, torch_dtype="auto", local_files_only=os.getenv("CI") == "true")
weights = model.state_dict()
keys = weights.keys()
for key in keys:
if layer_prefix in key:
# Caution it may cause potential failures. In future versions and different formats the below prefix may change
key_name = key[len(layer_prefix) :]
partial_state_dict.update({key_name: weights[key]})
return partial_state_dict
def expand_num_tokens_to_mult8(tensor):
num_padding_patches = (8 - (tensor.shape[-2] % 8)) % 8
# Compute padding tuple for pad function
padding = (0, 0, 0, num_padding_patches) # (pad_left, pad_right, pad_left for dim -2, pad_right for dim -2)
# Pad the tensor
tensor = F.pad(tensor, padding, mode="constant", value=0)
slice_index = -num_padding_patches if num_padding_patches > 0 else 0
return tensor, slice_index
def contract_num_tokens_from_mult8(tensor, slice_index):
if slice_index == 0:
return tensor
return tensor[:, :, :slice_index, :]
|